import numpy as np import cv2 from .math_functions import euclidean_dist_2_pts def extract_img_markers(img, workspace_ratio=1.0): """ Extract working area from an image thanks to 4 Niryo's markers :param img: OpenCV image which contain 4 Niryo's markers :param workspace_ratio: Ratio between the width and the height of the area represented by the markers :return: extracted and warped working area image """ gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) img_thresh = cv2.adaptiveThreshold(gray, maxValue=255, adaptiveMethod=cv2.ADAPTIVE_THRESH_MEAN_C, thresholdType=cv2.THRESH_BINARY, blockSize=15, C=25) list_good_candidates = find_markers_from_img_thresh(img_thresh) if not list_good_candidates or len(list_good_candidates) > 6: return None if len(list_good_candidates) == 4: list_good_candidates = sort_markers_detection(list_good_candidates) else: list_good_candidates = complicated_sort_markers(list_good_candidates, workspace_ratio=workspace_ratio) if list_good_candidates is None: return None im_cut = extract_sub_img(img, list_good_candidates, ratio_w_h=workspace_ratio) return im_cut def extract_sub_img(img, list_corners, ratio_w_h=1.0): """ Extract an small image from a big one using a Perspective Warp :param img: Big image from which the small one will be extracted :param list_corners: corners list of the small image :param ratio_w_h: Width over Height ratio of the area. It helps to not stretch the working area image :return: extracted and warped image """ if list_corners is None or len(list_corners) != 4: return None if ratio_w_h >= 1.0: target_w_area = int(round(ratio_w_h * 200)) target_h_area = 200 else: ratio_w_h = 1.0 / ratio_w_h target_h_area = int(round(ratio_w_h * 200)) target_w_area = 200 points_grid = [] for marker in list_corners: points_grid.append(marker.get_center()) points_grid = np.array(points_grid, dtype=np.float32) final_pts = np.array( [[0, 0], [target_w_area - 1, 0], [target_w_area - 1, target_h_area - 1], [0, target_h_area - 1]], dtype=np.float32) transfo_matrix = cv2.getPerspectiveTransform(points_grid, final_pts) # print transfo_matrix # print np.linalg.det(transfo_matrix) area_im = cv2.warpPerspective(img, transfo_matrix, (target_w_area, target_h_area)) return area_im def draw_markers(img, workspace_ratio=1.0): gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) img_thresh = cv2.adaptiveThreshold(gray, maxValue=255, adaptiveMethod=cv2.ADAPTIVE_THRESH_MEAN_C, thresholdType=cv2.THRESH_BINARY, blockSize=15, C=32) list_good_candidates = find_markers_from_img_thresh(img_thresh) if not list_good_candidates: return False, img im_draw = img.copy() for marker in list_good_candidates: cx, cy = marker.get_center() radius = marker.get_radius() cv2.circle(im_draw, (cx, cy), radius, (0, 0, 255), 2) if len(list_good_candidates) > 6: return False, im_draw if len(list_good_candidates) == 4: list_good_candidates = sort_markers_detection(list_good_candidates) else: list_good_candidates = complicated_sort_markers(list_good_candidates, workspace_ratio=workspace_ratio) if list_good_candidates is None: return False, im_draw for i, marker in enumerate(list_good_candidates[:4]): cx, cy = marker.get_center() radius = marker.get_radius() cv2.circle(im_draw, (cx, cy), radius, (0, 200, 0), 2) cv2.putText(im_draw, "{}".format(i + 1), (cx + 5, cy - 15), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 0, 0), 3) cv2.putText(im_draw, "{}".format(i + 1), (cx + 5, cy - 15), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 200, 0), 2) return True, im_draw class PotentialMarker: def __init__(self, center, radius, cnt): self.center = center self.x = center[0] self.y = center[1] self.radius = radius self.contour = cnt self.is_merged = False def get_center(self): return self.center def __str__(self): return "{} - {} - {}".format(self.x, self.y, self.radius) def __repr__(self): return self.__str__() class Marker: def __init__(self, potential_marker): self.list_centers = [potential_marker.get_center()] self.list_radius = [potential_marker.radius] self.list_contours = [potential_marker.contour] self.cx = self.list_centers[0][0] self.cy = self.list_centers[0][1] self.radius = potential_marker.radius self.identifiant = None self.value_for_id = None def get_radius(self): return self.radius def get_center(self): return self.cx, self.cy def add_circle(self, obj_potential_marker): self.list_centers.append(obj_potential_marker.get_center()) self.list_radius.append(obj_potential_marker.radius) obj_potential_marker.is_merged = True (x, y) = np.mean(self.list_centers, axis=0) self.cx, self.cy = int(round(x)), int(round(y)) self.radius = int(round(max(self.list_radius))) def nb_circles(self): return len(self.list_centers) def get_id_from_slice(self, img_thresh): x, y, w, h = self.cx - 1, self.cy - 1, 3, 3 self.value_for_id = np.mean(img_thresh[y:y + h, x:x + w]) # return value_for_id if self.value_for_id > 200: self.identifiant = "A" else: self.identifiant = "B" return self.identifiant # return value_for_id def __str__(self): return "{} - {}".format(self.nb_circles(), self.list_centers) def __repr__(self): return self.__str__() def sort_markers_detection(list_markers): def rotate(l, n): return l[n:] + l[:n] list_sort_y = sorted(list_markers, key=lambda m: m.cy) top1, top2, bottom1, bottom2 = list_sort_y if top1.cx < top2.cx: top_left = top1 top_right = top2 else: top_left = top2 top_right = top1 if bottom1.cx < bottom2.cx: bottom_left = bottom1 bottom_right = bottom2 else: bottom_left = bottom2 bottom_right = bottom1 list_markers_unsorted = [top_left, top_right, bottom_right, bottom_left] list_id = [marker.identifiant for marker in list_markers_unsorted] if list_id.count("A") == 1: list_corners_sorted = rotate(list_markers_unsorted, n=list_id.index("A")) elif list_id.count("B") == 1: list_corners_sorted = rotate(list_markers_unsorted, n=list_id.index("B")) else: return list_markers_unsorted return list_corners_sorted def complicated_sort_markers(list_markers, workspace_ratio): import itertools if workspace_ratio >= 1.0: target_w_area = int(round(workspace_ratio * 200)) target_h_area = 200 else: ratio_w_h = 1.0 / workspace_ratio target_h_area = int(round(ratio_w_h * 200)) target_w_area = 200 list_id = [marker.identifiant for marker in list_markers] count_A = list_id.count("A") count_B = list_id.count("B") if count_A < 3 > count_B: return None if count_A < count_B: id_first_marker = "A" id_second_marker = "B" else: id_first_marker = "B" id_second_marker = "A" list_combinaisons = [] list_marker_1 = [marker for marker in list_markers if marker.identifiant == id_first_marker] list_marker_2 = [marker for marker in list_markers if marker.identifiant == id_second_marker] if list_marker_1: list_combinaisons_marker_2 = itertools.combinations(list_marker_2, 3) for marker1 in list_marker_1: for combi_markers2 in list_combinaisons_marker_2: combin = [marker1] + list(combi_markers2) list_combinaisons.append(sort_markers_detection(combin)) else: for combinaison in itertools.combinations(list_marker_2, 4): list_combinaisons.append(combinaison) if not list_combinaisons: return None final_pts = np.array( [[0, 0], [target_w_area - 1, 0], [target_w_area - 1, target_h_area - 1], [0, target_h_area - 1]], dtype=np.float32) list_det_transfo_matrix = [] for combin in list_combinaisons: points_grid = np.array([[mark.cx, mark.cy] for mark in combin], dtype=np.float32) transfo_matrix = cv2.getPerspectiveTransform(points_grid, final_pts) list_det_transfo_matrix.append(np.linalg.det(transfo_matrix)) best_combin_ind = np.argmin(abs(np.array(list_det_transfo_matrix) - 1)) best_markers = list_combinaisons[best_combin_ind] return best_markers def find_markers_from_img_thresh(img_thresh, max_dist_between_centers=3, min_radius_circle=4, max_radius_circle=35, min_radius_marker=7): contours = cv2.findContours(img_thresh, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)[-2] list_potential_markers = [] for cnt in contours: (x, y), radius = cv2.minEnclosingCircle(cnt) if not min_radius_circle < radius < max_radius_circle: continue center = (int(round(x)), int(round(y))) radius = int(radius) list_potential_markers.append(PotentialMarker(center, radius, cnt)) list_potential_markers = sorted(list_potential_markers, key=lambda m: m.x) list_good_candidates = [] for i, potential_marker in enumerate(list_potential_markers): if potential_marker.is_merged: continue marker1 = Marker(potential_marker) center_marker = marker1.get_center() for potential_marker2 in list_potential_markers[i + 1:]: if potential_marker.is_merged: continue center_potential = potential_marker2.get_center() if center_potential[0] - center_marker[0] > max_dist_between_centers: break dist = euclidean_dist_2_pts(center_marker, center_potential) if dist <= max_dist_between_centers: marker1.add_circle(potential_marker2) center_marker = marker1.get_center() if marker1.nb_circles() > 2 and marker1.radius >= min_radius_marker: list_good_candidates.append(marker1) marker1.get_id_from_slice(img_thresh) return list_good_candidates