detect_circle_and_rects.py

import sys
import cv2
import numpy as np
from skimage.transform import hough_circle, hough_circle_peaks
from skimage.feature import canny


def detect_circles(image, max_circles=15, single_radius=False):
    edges = canny(image, sigma=1)

    hough_radii = np.arange(10, 80, 2)
    hough_res = hough_circle(edges, hough_radii)
    _, cxs, cys, radii = hough_circle_peaks(hough_res, hough_radii, total_num_peaks=max_circles)
    if single_radius:
        circle_rects = np.array([[cx-r, cy-r, 2*r, 2*r, 1] for r, cx, cy in zip(radii, cxs, cys) if r == radii[0]])
    else:
        circle_rects = np.array([[cx-r, cy-r, 2*r, 2*r, 1] for r, cx, cy in zip(radii, cxs, cys)])
    circle_rects = nms(circle_rects, 0.5)
    return circle_rects[:, :4]


def detect_rectangles(image, min_rect_length=15, single_size=False):
    square_rects = []

    thresh = 255 - cv2.threshold(image, 128, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1]
    _, cnts, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    for cnt in cnts:
        peri = cv2.arcLength(cnt, True)
        approx = np.squeeze(np.array(cv2.approxPolyDP(cnt, 0.04 * peri, True)), axis=1)
        min_area = min_rect_length*min_rect_length
        if len(approx) == 4  and cv2.contourArea(cnt) > min_area:
            x, y, w, h = cv2.boundingRect(approx)
            square_rects.append([x, y, w, h])

    square_rects = np.array(square_rects)
    if single_size and square_rects.size:
        areas = list(square_rects[:,2] * square_rects[:,3])
        most_common_area = max(areas, key=areas.count)
        return square_rects[np.where(areas == most_common_area)[0]]
    return square_rects