Loading src/car_target.py +6 −5 Original line number Diff line number Diff line Loading @@ -50,7 +50,7 @@ class Smoother: width, height = shape self.last_target = None self.invalid_count = 0 self.far_enough = math.sqrt(width*width + height*height) / 12 self.far_enough = math.sqrt(width*width + height*height) / 12 # diagonal length / 12 self.count_threshold = 5 def smoothed(self, new_target): Loading Loading @@ -115,7 +115,7 @@ class ImgCarTargetApp(CarTargetApp): plt.figure("Armor Detection") ROW, COL = 3, 3 ROW, COL = 2, 4 count = 1 plt.subplot(ROW, COL, count) plt.title("Original") Loading @@ -124,6 +124,7 @@ class ImgCarTargetApp(CarTargetApp): target = self.detector.target(mat) print("target:", target) print("------------------------------------") if self.debug: count = 2 Loading Loading @@ -254,7 +255,7 @@ if __name__ == "__main__": # frame_size=(640, 360), ext_name="jpg", debug=True) app = VideoCarTargetApp(file="/home/jeeken/Videos/live_blue.avi", color=BarColor.BLUE, frame_size=(640, 480), debug=True) # app = VideoCarTargetApp(file="/home/jeeken/Videos/live_red.avi", # color=BarColor.RED, frame_size=(640, 480), debug=False) color=BarColor.BLUE, frame_size=(640, 480), debug=False) # app = VideoCarTargetApp(file="/home/jeeken/Videos/red.avi", # color=BarColor.RED, frame_size=(640, 480), debug=True) app.run() src/detector.py +215 −92 Original line number Diff line number Diff line Loading @@ -31,13 +31,24 @@ class TargetDis(Enum): ULTRA = 3 def round_point(pt): x, y = pt return (round(x), round(y)) def mid_point(pt1, pt2): x1, y1 = pt1 x2, y2 = pt2 return (x1 + x2) // 2, (y1 + y2) // 2 class Detector: # TODO: remove shape from self.__init__ def __init__(self, color: BarColor, shape: tuple = (480, 360), debug = False): self.color = color self.refresh(frame=None, debug=debug) def refresh(self, frame, debug=None): def refresh(self, frame: np.ndarray, debug=None): self.frame = frame if debug is not None: self.debug = debug Loading @@ -50,7 +61,12 @@ class Detector: else: self.debug_imgs.append((title, img)) def hsv(self, frame): def debug_print(self, *msg): if self.debug: print("Debug>>>", *msg) # print("Debug>>>", *msg, file=sys.stderr) def hsv(self, frame: np.ndarray): hsv_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV) if self.color == BarColor.RED: Loading Loading @@ -103,46 +119,46 @@ class Detector: "width": se_x - nw_x, "height": se_y - se_x } @staticmethod def select_lights_in_halo(light, halo): light_components, light_map = Detector.get_connected_components_info(light) halo_components= Detector.get_connected_components_info(halo)[0] selected_halo = list(filter(lambda c: c["pixels"] > 20, halo_components)) # TODO: eliminate magic number # find halo area try: northwest_x = min([x["northwest"][0] for x in selected_halo]) northwest_y = min([x["northwest"][1] for x in selected_halo]) southeast_x = max([x["southeast"][0] for x in selected_halo]) southeast_y = max([x["southeast"][1] for x in selected_halo]) except ValueError: # selected_halo is [] return [], None width = southeast_x - northwest_x height = southeast_y - northwest_y # zoom and shift northwest_x -= width // 4 southeast_x += width // 4 northwest_y -= height // 2 # select inside ones def is_fully_inside(c): c_nw_x, c_nw_y = c["northwest"] c_se_x, c_se_y = c["southeast"] return northwest_x <= c_nw_x and northwest_y <= c_nw_y and c_se_x <= southeast_x and c_se_y <= southeast_y insides = list(filter(is_fully_inside, light_components)) # FIXME: why doesn't work if don't convert it to a list? # add angle info for i in insides: nw_x, nw_y = i["northwest"] se_x, se_y = i["southeast"] top_line, bottom_line = light_map[nw_y, nw_x:se_x], light_map[se_y-1, nw_x:se_x] tx, ty = nw_x + np.argmax(top_line), nw_y # top peak bx, by = nw_x + np.argmax(bottom_line), se_y-1 # bottom peak angle = math.atan2(by-ty, tx-bx) / math.pi # unit: pi rad i["angle"] = angle return insides, {"northwest": (northwest_x, northwest_y), "southeast": (southeast_x, southeast_y)} # @staticmethod # def select_lights_in_halo(light, halo): # light_components, light_map = Detector.get_connected_components_info(light) # halo_components= Detector.get_connected_components_info(halo)[0] # selected_halo = list(filter(lambda c: c["pixels"] > 20, halo_components)) # TODO: eliminate magic number # # # find halo area # try: # northwest_x = min([x["northwest"][0] for x in selected_halo]) # northwest_y = min([x["northwest"][1] for x in selected_halo]) # southeast_x = max([x["southeast"][0] for x in selected_halo]) # southeast_y = max([x["southeast"][1] for x in selected_halo]) # except ValueError: # selected_halo is [] # return [], None # width = southeast_x - northwest_x # height = southeast_y - northwest_y # # # zoom and shift # northwest_x -= width // 4 # southeast_x += width // 4 # northwest_y -= height // 2 # # # select inside ones # def is_fully_inside(c): # c_nw_x, c_nw_y = c["northwest"] # c_se_x, c_se_y = c["southeast"] # return northwest_x <= c_nw_x and northwest_y <= c_nw_y and c_se_x <= southeast_x and c_se_y <= southeast_y # insides = list(filter(is_fully_inside, light_components)) # FIXME: why doesn't work if don't convert it to a list? # # # add angle info # for i in insides: # nw_x, nw_y = i["northwest"] # se_x, se_y = i["southeast"] # top_line, bottom_line = light_map[nw_y, nw_x:se_x], light_map[se_y-1, nw_x:se_x] # tx, ty = nw_x + np.argmax(top_line), nw_y # top peak # bx, by = nw_x + np.argmax(bottom_line), se_y-1 # bottom peak # angle = math.atan2(by-ty, tx-bx) / math.pi # unit: pi rad # i["angle"] = angle # # return insides, {"northwest": (northwest_x, northwest_y), "southeast": (southeast_x, southeast_y)} @staticmethod def select_valid_bars(lights): Loading Loading @@ -218,72 +234,179 @@ class Detector: i1, i2 = winner["indices"] return bars[i1], bars[i2] # NEW! def get_lights(self, binary_img): def get_lights(self, binary_img: np.ndarray): _, width = binary_img.shape # https://docs.opencv.org/2.4/modules/imgproc/doc/structural_analysis_and_shape_descriptors.html?highlight=findcontours#findcontours img, contours, hierarchy = cv2.findContours(binary_img, mode=cv2.RETR_EXTERNAL, method=cv2.CHAIN_APPROX_SIMPLE) rectangles_info = [cv2.minAreaRect(i) for i in contours if len(i) * 60 > width] # filter (omit small ones) and map # rectangles_info = [cv2.fitEllipse(i) for i in contours if len(i) * 60 > width] return rectangles_info def halo_circle(self, binary_img: np.ndarray): height, width = binary_img.shape img, contours, hierarchy = cv2.findContours(binary_img, mode=cv2.RETR_EXTERNAL, method=cv2.CHAIN_APPROX_SIMPLE) selected_contours = [i for i in contours if len(i) * 60 > width] # small contours omitted if len(selected_contours) > 0: combined_halo_points = np.concatenate(selected_contours) original_center, original_radius = cv2.minEnclosingCircle(combined_halo_points) x, y = original_center center = round(x), round(y - original_radius * 0.6) # shift up radius = round(original_radius * 1.8 if original_radius < (width / 4) else width * 0.375) # zoom else: # no valid halo, consider the central area center = width // 2, height // 2 radius = width // 3 return center, radius def select_lights_in_halo(self, lights, halo_center: tuple, halo_radius): x, y = halo_center def is_inside(light): (xl, yl), shape, angle = light return (xl-x)**2 + (yl-y)**2 < halo_radius**2 return filter(is_inside, lights) def select_vertical_lights(self, lights): def normalize_angle(rec): center, shape, angle = rec w, h = shape if w > h: return center, (h, w), angle+90 else: return rec # https://docs.opencv.org/2.4/modules/imgproc/doc/structural_analysis_and_shape_descriptors.html?highlight=findcontours#findcontours img, contours, hierarchy = cv2.findContours(image=binary_img, mode=cv2.RETR_EXTERNAL, method=cv2.CHAIN_APPROX_SIMPLE) # rectangle: (centroid, shape, angle) rectangles = [cv2.minAreaRect(i) for i in contours if len(i) * 60 > width] # omit small ones # rectangles = [cv2.fitEllipse(i) for i in contours if len(i) * 40 > width] def is_vertical(rec): center, shape, angle = rec return -20 < angle < 20 angle_normalized_lights = map(normalize_angle, lights) vertical_lights = [i for i in angle_normalized_lights if is_vertical(i)] def get_img_lights_recs_as_elps(): # debug def get_img_vertical_lights(): show_lights = np.copy(self.frame) for centroid, shape, angle in rectangles: to_int_point = lambda x, y: (round(x), round(y)) rand_color = (randint(40, 230), randint(40, 230), randint(40, 230)) cv2.ellipse(img=show_lights, center=to_int_point(*centroid), axes=to_int_point(*shape), angle=angle, for i in vertical_lights: center, shape, angle = i rand_color = (randint(40, 230), randint(40, 255), randint(40, 255)) cv2.ellipse(show_lights, center=round_point(center), axes=round_point(shape), angle=angle, startAngle=0, endAngle=360, color=rand_color, thickness=2) return show_lights self.add_debug_img("Light_rec", get_img_lights_recs_as_elps) self.add_debug_img("Vertical Lights", get_img_vertical_lights) return vertical_lights return rectangles def get_armor(self, lights): lights = list(lights) lights_quantity = len(lights) def halo_circle(self, binary_img): # TODO pass if lights_quantity < 2: return None pairs = [(i, j) for i in range(0, lights_quantity - 1) for j in range(i + 1, lights_quantity)] def parallel(light1, light2): angle1, angle2 = light1[2], light2[2] return abs(angle1 - angle2) / 20 def size_similarity(light1, light2): (w1, h1), (w2, h2) = light1[1], light2[1] area1, area2 = w1 * h1, w2 * h2 min_size = area1 if area1 < area2 else area2 return abs(area1 - area2) / min_size def square_ratio(light1, light2): (x1, y1), (w1, h1), angle1 = light1 (x2, y2), (w2, h2), angle2 = light2 ratio = 2 * math.sqrt((x1-x2)**2 + (y1-y2)**2) / (h1+h2) return abs(ratio - 2.4) ** 2 if ratio > 0.85 else 1e9 def target(self, frame): def y_dis(light1, light2): (x1, y1), shape1, angle1 = light1 (x2, y2), shape2, angle2 = light2 min_y = y1 if y1 < y2 else y2 return abs(y1 - y2) / min_y def judge(pair): i1, i2 = pair bar1, bar2 = lights[i1], lights[i2] policy = [(square_ratio, 5), (y_dis, 15), (size_similarity, 1), (parallel, 1)] # [(func, coefficient),...] return sum((func(bar1, bar2) * coefficient for func, coefficient in policy)) # weighted sum judge_results = [(pair, judge(pair)) for pair in pairs] (i1, i2), winner_score = min(judge_results, key=lambda x: x[1]) if winner_score > 5: return None def get_img_selected_pair(): selected_pair_show = np.copy(self.frame) l1, l2 = lights[i1], lights[i2] self.debug_print("l1:", l1) self.debug_print("l2:", l2) self.debug_print("score:", winner_score) self.debug_print("parallel:", parallel(l1, l2)) self.debug_print("size_similarity:", size_similarity(l1, l2)) self.debug_print("square_ratio:", 5 * square_ratio(l1, l2)) self.debug_print("y_dis:", 15 * y_dis(l1, l2)) c1, s1, a1 = l1 c2, s2, a2 = l2 cv2.drawMarker(selected_pair_show, position=round_point(c1), color=(0, 255, 255), markerSize=15, thickness=2) cv2.drawMarker(selected_pair_show, position=round_point(c2), color=(0, 255, 255), markerSize=15, thickness=2) return selected_pair_show self.add_debug_img("Selected Pair", get_img_selected_pair) x1, y1 = lights[i1][0] x2, y2 = lights[i2][0] return round((x1+x2) * 0.5), round((y1+y2) * 0.5) def target(self, frame: np.ndarray): self.refresh(frame) # frame = cv2.blur(frame, ksize=(4, 4)) frame = cv2.pyrUp(cv2.pyrDown(frame)) # down-sample light, halo = self.hsv(frame) lights = self.get_lights(light) # lights_in_halo, selected_area = Detector.select_lights_in_halo(light, halo) lights_in_halo, selected_area = Detector.select_lights_in_halo(light, halo) selected_bars = Detector.select_valid_bars(lights_in_halo) selected_pair = Detector.select_pair(selected_bars) if selected_pair != (): bar1, bar2 = selected_pair def mid_point(pt1, pt2): x1, y1 = pt1 x2, y2 = pt2 return (x1 + x2) // 2, (y1 + y2) // 2 target = mid_point(bar1["centroid"], bar2["centroid"]) else: target = None halo_center, halo_radius = self.halo_circle(halo) lights_info = self.get_lights(light) def get_img_selected_halo(): cv2.rectangle(halo, pt1=selected_area["northwest"], pt2=selected_area["southeast"], color=255, thickness=1) # only draw on light/halo after info got, or draw on a copied mat def get_img_halo(): cv2.circle(img=halo, center=halo_center, radius=halo_radius, color=255, thickness=2) return halo def get_img_selected_light(): cv2.rectangle(light, pt1=selected_area["northwest"], pt2=selected_area["southeast"], color=255, thickness=1) def get_img_light(): cv2.circle(img=light, center=halo_center, radius=halo_radius, color=255, thickness=2) return light if selected_area is not None: self.add_debug_img("Halo", get_img_selected_halo) self.add_debug_img("Light", get_img_selected_light) def get_img_selected(): selected = np.copy(self.frame) for each in selected_bars: cv2.drawMarker(img=selected, position=each["centroid"], color=(0, 255, 255), markerSize=25, thickness=2) for each in selected_pair: cv2.circle(img=selected, center=each["centroid"], radius=5, color=(0, 0, 255), thickness=-1) return selected self.add_debug_img("Selected", get_img_selected) self.add_debug_img("Halo", get_img_halo) self.add_debug_img("Light", get_img_light) lights_inside = self.select_lights_in_halo(lights_info, halo_center, halo_radius) vertical_lights = self.select_vertical_lights(lights_inside) target = self.get_armor(vertical_lights) # selected_bars = Detector.select_valid_bars(lights_in_halo) # selected_pair = Detector.select_pair(selected_bars) # if selected_pair != (): # bar1, bar2 = selected_pair # target = mid_point(bar1["centroid"], bar2["centroid"]) # else: # target = None # # # def get_img_selected_halo(): # # cv2.rectangle(halo, pt1=selected_area["northwest"], pt2=selected_area["southeast"], color=255, thickness=1) # # return halo # # def get_img_selected_light(): # # cv2.rectangle(light, pt1=selected_area["northwest"], pt2=selected_area["southeast"], color=255, thickness=1) # # return light # # if selected_area is not None: # # self.add_debug_img("Halo", get_img_selected_halo) # # self.add_debug_img("Light", get_img_selected_light) # def get_img_selected(): # selected = np.copy(self.frame) # for each in selected_bars: # cv2.drawMarker(img=selected, position=each["centroid"], color=(0, 255, 255), markerSize=25, thickness=2) # for each in selected_pair: # cv2.circle(img=selected, center=each["centroid"], radius=5, color=(0, 0, 255), thickness=-1) # return selected # self.add_debug_img("Selected", get_img_selected) return target Loading
src/car_target.py +6 −5 Original line number Diff line number Diff line Loading @@ -50,7 +50,7 @@ class Smoother: width, height = shape self.last_target = None self.invalid_count = 0 self.far_enough = math.sqrt(width*width + height*height) / 12 self.far_enough = math.sqrt(width*width + height*height) / 12 # diagonal length / 12 self.count_threshold = 5 def smoothed(self, new_target): Loading Loading @@ -115,7 +115,7 @@ class ImgCarTargetApp(CarTargetApp): plt.figure("Armor Detection") ROW, COL = 3, 3 ROW, COL = 2, 4 count = 1 plt.subplot(ROW, COL, count) plt.title("Original") Loading @@ -124,6 +124,7 @@ class ImgCarTargetApp(CarTargetApp): target = self.detector.target(mat) print("target:", target) print("------------------------------------") if self.debug: count = 2 Loading Loading @@ -254,7 +255,7 @@ if __name__ == "__main__": # frame_size=(640, 360), ext_name="jpg", debug=True) app = VideoCarTargetApp(file="/home/jeeken/Videos/live_blue.avi", color=BarColor.BLUE, frame_size=(640, 480), debug=True) # app = VideoCarTargetApp(file="/home/jeeken/Videos/live_red.avi", # color=BarColor.RED, frame_size=(640, 480), debug=False) color=BarColor.BLUE, frame_size=(640, 480), debug=False) # app = VideoCarTargetApp(file="/home/jeeken/Videos/red.avi", # color=BarColor.RED, frame_size=(640, 480), debug=True) app.run()
src/detector.py +215 −92 Original line number Diff line number Diff line Loading @@ -31,13 +31,24 @@ class TargetDis(Enum): ULTRA = 3 def round_point(pt): x, y = pt return (round(x), round(y)) def mid_point(pt1, pt2): x1, y1 = pt1 x2, y2 = pt2 return (x1 + x2) // 2, (y1 + y2) // 2 class Detector: # TODO: remove shape from self.__init__ def __init__(self, color: BarColor, shape: tuple = (480, 360), debug = False): self.color = color self.refresh(frame=None, debug=debug) def refresh(self, frame, debug=None): def refresh(self, frame: np.ndarray, debug=None): self.frame = frame if debug is not None: self.debug = debug Loading @@ -50,7 +61,12 @@ class Detector: else: self.debug_imgs.append((title, img)) def hsv(self, frame): def debug_print(self, *msg): if self.debug: print("Debug>>>", *msg) # print("Debug>>>", *msg, file=sys.stderr) def hsv(self, frame: np.ndarray): hsv_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV) if self.color == BarColor.RED: Loading Loading @@ -103,46 +119,46 @@ class Detector: "width": se_x - nw_x, "height": se_y - se_x } @staticmethod def select_lights_in_halo(light, halo): light_components, light_map = Detector.get_connected_components_info(light) halo_components= Detector.get_connected_components_info(halo)[0] selected_halo = list(filter(lambda c: c["pixels"] > 20, halo_components)) # TODO: eliminate magic number # find halo area try: northwest_x = min([x["northwest"][0] for x in selected_halo]) northwest_y = min([x["northwest"][1] for x in selected_halo]) southeast_x = max([x["southeast"][0] for x in selected_halo]) southeast_y = max([x["southeast"][1] for x in selected_halo]) except ValueError: # selected_halo is [] return [], None width = southeast_x - northwest_x height = southeast_y - northwest_y # zoom and shift northwest_x -= width // 4 southeast_x += width // 4 northwest_y -= height // 2 # select inside ones def is_fully_inside(c): c_nw_x, c_nw_y = c["northwest"] c_se_x, c_se_y = c["southeast"] return northwest_x <= c_nw_x and northwest_y <= c_nw_y and c_se_x <= southeast_x and c_se_y <= southeast_y insides = list(filter(is_fully_inside, light_components)) # FIXME: why doesn't work if don't convert it to a list? # add angle info for i in insides: nw_x, nw_y = i["northwest"] se_x, se_y = i["southeast"] top_line, bottom_line = light_map[nw_y, nw_x:se_x], light_map[se_y-1, nw_x:se_x] tx, ty = nw_x + np.argmax(top_line), nw_y # top peak bx, by = nw_x + np.argmax(bottom_line), se_y-1 # bottom peak angle = math.atan2(by-ty, tx-bx) / math.pi # unit: pi rad i["angle"] = angle return insides, {"northwest": (northwest_x, northwest_y), "southeast": (southeast_x, southeast_y)} # @staticmethod # def select_lights_in_halo(light, halo): # light_components, light_map = Detector.get_connected_components_info(light) # halo_components= Detector.get_connected_components_info(halo)[0] # selected_halo = list(filter(lambda c: c["pixels"] > 20, halo_components)) # TODO: eliminate magic number # # # find halo area # try: # northwest_x = min([x["northwest"][0] for x in selected_halo]) # northwest_y = min([x["northwest"][1] for x in selected_halo]) # southeast_x = max([x["southeast"][0] for x in selected_halo]) # southeast_y = max([x["southeast"][1] for x in selected_halo]) # except ValueError: # selected_halo is [] # return [], None # width = southeast_x - northwest_x # height = southeast_y - northwest_y # # # zoom and shift # northwest_x -= width // 4 # southeast_x += width // 4 # northwest_y -= height // 2 # # # select inside ones # def is_fully_inside(c): # c_nw_x, c_nw_y = c["northwest"] # c_se_x, c_se_y = c["southeast"] # return northwest_x <= c_nw_x and northwest_y <= c_nw_y and c_se_x <= southeast_x and c_se_y <= southeast_y # insides = list(filter(is_fully_inside, light_components)) # FIXME: why doesn't work if don't convert it to a list? # # # add angle info # for i in insides: # nw_x, nw_y = i["northwest"] # se_x, se_y = i["southeast"] # top_line, bottom_line = light_map[nw_y, nw_x:se_x], light_map[se_y-1, nw_x:se_x] # tx, ty = nw_x + np.argmax(top_line), nw_y # top peak # bx, by = nw_x + np.argmax(bottom_line), se_y-1 # bottom peak # angle = math.atan2(by-ty, tx-bx) / math.pi # unit: pi rad # i["angle"] = angle # # return insides, {"northwest": (northwest_x, northwest_y), "southeast": (southeast_x, southeast_y)} @staticmethod def select_valid_bars(lights): Loading Loading @@ -218,72 +234,179 @@ class Detector: i1, i2 = winner["indices"] return bars[i1], bars[i2] # NEW! def get_lights(self, binary_img): def get_lights(self, binary_img: np.ndarray): _, width = binary_img.shape # https://docs.opencv.org/2.4/modules/imgproc/doc/structural_analysis_and_shape_descriptors.html?highlight=findcontours#findcontours img, contours, hierarchy = cv2.findContours(binary_img, mode=cv2.RETR_EXTERNAL, method=cv2.CHAIN_APPROX_SIMPLE) rectangles_info = [cv2.minAreaRect(i) for i in contours if len(i) * 60 > width] # filter (omit small ones) and map # rectangles_info = [cv2.fitEllipse(i) for i in contours if len(i) * 60 > width] return rectangles_info def halo_circle(self, binary_img: np.ndarray): height, width = binary_img.shape img, contours, hierarchy = cv2.findContours(binary_img, mode=cv2.RETR_EXTERNAL, method=cv2.CHAIN_APPROX_SIMPLE) selected_contours = [i for i in contours if len(i) * 60 > width] # small contours omitted if len(selected_contours) > 0: combined_halo_points = np.concatenate(selected_contours) original_center, original_radius = cv2.minEnclosingCircle(combined_halo_points) x, y = original_center center = round(x), round(y - original_radius * 0.6) # shift up radius = round(original_radius * 1.8 if original_radius < (width / 4) else width * 0.375) # zoom else: # no valid halo, consider the central area center = width // 2, height // 2 radius = width // 3 return center, radius def select_lights_in_halo(self, lights, halo_center: tuple, halo_radius): x, y = halo_center def is_inside(light): (xl, yl), shape, angle = light return (xl-x)**2 + (yl-y)**2 < halo_radius**2 return filter(is_inside, lights) def select_vertical_lights(self, lights): def normalize_angle(rec): center, shape, angle = rec w, h = shape if w > h: return center, (h, w), angle+90 else: return rec # https://docs.opencv.org/2.4/modules/imgproc/doc/structural_analysis_and_shape_descriptors.html?highlight=findcontours#findcontours img, contours, hierarchy = cv2.findContours(image=binary_img, mode=cv2.RETR_EXTERNAL, method=cv2.CHAIN_APPROX_SIMPLE) # rectangle: (centroid, shape, angle) rectangles = [cv2.minAreaRect(i) for i in contours if len(i) * 60 > width] # omit small ones # rectangles = [cv2.fitEllipse(i) for i in contours if len(i) * 40 > width] def is_vertical(rec): center, shape, angle = rec return -20 < angle < 20 angle_normalized_lights = map(normalize_angle, lights) vertical_lights = [i for i in angle_normalized_lights if is_vertical(i)] def get_img_lights_recs_as_elps(): # debug def get_img_vertical_lights(): show_lights = np.copy(self.frame) for centroid, shape, angle in rectangles: to_int_point = lambda x, y: (round(x), round(y)) rand_color = (randint(40, 230), randint(40, 230), randint(40, 230)) cv2.ellipse(img=show_lights, center=to_int_point(*centroid), axes=to_int_point(*shape), angle=angle, for i in vertical_lights: center, shape, angle = i rand_color = (randint(40, 230), randint(40, 255), randint(40, 255)) cv2.ellipse(show_lights, center=round_point(center), axes=round_point(shape), angle=angle, startAngle=0, endAngle=360, color=rand_color, thickness=2) return show_lights self.add_debug_img("Light_rec", get_img_lights_recs_as_elps) self.add_debug_img("Vertical Lights", get_img_vertical_lights) return vertical_lights return rectangles def get_armor(self, lights): lights = list(lights) lights_quantity = len(lights) def halo_circle(self, binary_img): # TODO pass if lights_quantity < 2: return None pairs = [(i, j) for i in range(0, lights_quantity - 1) for j in range(i + 1, lights_quantity)] def parallel(light1, light2): angle1, angle2 = light1[2], light2[2] return abs(angle1 - angle2) / 20 def size_similarity(light1, light2): (w1, h1), (w2, h2) = light1[1], light2[1] area1, area2 = w1 * h1, w2 * h2 min_size = area1 if area1 < area2 else area2 return abs(area1 - area2) / min_size def square_ratio(light1, light2): (x1, y1), (w1, h1), angle1 = light1 (x2, y2), (w2, h2), angle2 = light2 ratio = 2 * math.sqrt((x1-x2)**2 + (y1-y2)**2) / (h1+h2) return abs(ratio - 2.4) ** 2 if ratio > 0.85 else 1e9 def target(self, frame): def y_dis(light1, light2): (x1, y1), shape1, angle1 = light1 (x2, y2), shape2, angle2 = light2 min_y = y1 if y1 < y2 else y2 return abs(y1 - y2) / min_y def judge(pair): i1, i2 = pair bar1, bar2 = lights[i1], lights[i2] policy = [(square_ratio, 5), (y_dis, 15), (size_similarity, 1), (parallel, 1)] # [(func, coefficient),...] return sum((func(bar1, bar2) * coefficient for func, coefficient in policy)) # weighted sum judge_results = [(pair, judge(pair)) for pair in pairs] (i1, i2), winner_score = min(judge_results, key=lambda x: x[1]) if winner_score > 5: return None def get_img_selected_pair(): selected_pair_show = np.copy(self.frame) l1, l2 = lights[i1], lights[i2] self.debug_print("l1:", l1) self.debug_print("l2:", l2) self.debug_print("score:", winner_score) self.debug_print("parallel:", parallel(l1, l2)) self.debug_print("size_similarity:", size_similarity(l1, l2)) self.debug_print("square_ratio:", 5 * square_ratio(l1, l2)) self.debug_print("y_dis:", 15 * y_dis(l1, l2)) c1, s1, a1 = l1 c2, s2, a2 = l2 cv2.drawMarker(selected_pair_show, position=round_point(c1), color=(0, 255, 255), markerSize=15, thickness=2) cv2.drawMarker(selected_pair_show, position=round_point(c2), color=(0, 255, 255), markerSize=15, thickness=2) return selected_pair_show self.add_debug_img("Selected Pair", get_img_selected_pair) x1, y1 = lights[i1][0] x2, y2 = lights[i2][0] return round((x1+x2) * 0.5), round((y1+y2) * 0.5) def target(self, frame: np.ndarray): self.refresh(frame) # frame = cv2.blur(frame, ksize=(4, 4)) frame = cv2.pyrUp(cv2.pyrDown(frame)) # down-sample light, halo = self.hsv(frame) lights = self.get_lights(light) # lights_in_halo, selected_area = Detector.select_lights_in_halo(light, halo) lights_in_halo, selected_area = Detector.select_lights_in_halo(light, halo) selected_bars = Detector.select_valid_bars(lights_in_halo) selected_pair = Detector.select_pair(selected_bars) if selected_pair != (): bar1, bar2 = selected_pair def mid_point(pt1, pt2): x1, y1 = pt1 x2, y2 = pt2 return (x1 + x2) // 2, (y1 + y2) // 2 target = mid_point(bar1["centroid"], bar2["centroid"]) else: target = None halo_center, halo_radius = self.halo_circle(halo) lights_info = self.get_lights(light) def get_img_selected_halo(): cv2.rectangle(halo, pt1=selected_area["northwest"], pt2=selected_area["southeast"], color=255, thickness=1) # only draw on light/halo after info got, or draw on a copied mat def get_img_halo(): cv2.circle(img=halo, center=halo_center, radius=halo_radius, color=255, thickness=2) return halo def get_img_selected_light(): cv2.rectangle(light, pt1=selected_area["northwest"], pt2=selected_area["southeast"], color=255, thickness=1) def get_img_light(): cv2.circle(img=light, center=halo_center, radius=halo_radius, color=255, thickness=2) return light if selected_area is not None: self.add_debug_img("Halo", get_img_selected_halo) self.add_debug_img("Light", get_img_selected_light) def get_img_selected(): selected = np.copy(self.frame) for each in selected_bars: cv2.drawMarker(img=selected, position=each["centroid"], color=(0, 255, 255), markerSize=25, thickness=2) for each in selected_pair: cv2.circle(img=selected, center=each["centroid"], radius=5, color=(0, 0, 255), thickness=-1) return selected self.add_debug_img("Selected", get_img_selected) self.add_debug_img("Halo", get_img_halo) self.add_debug_img("Light", get_img_light) lights_inside = self.select_lights_in_halo(lights_info, halo_center, halo_radius) vertical_lights = self.select_vertical_lights(lights_inside) target = self.get_armor(vertical_lights) # selected_bars = Detector.select_valid_bars(lights_in_halo) # selected_pair = Detector.select_pair(selected_bars) # if selected_pair != (): # bar1, bar2 = selected_pair # target = mid_point(bar1["centroid"], bar2["centroid"]) # else: # target = None # # # def get_img_selected_halo(): # # cv2.rectangle(halo, pt1=selected_area["northwest"], pt2=selected_area["southeast"], color=255, thickness=1) # # return halo # # def get_img_selected_light(): # # cv2.rectangle(light, pt1=selected_area["northwest"], pt2=selected_area["southeast"], color=255, thickness=1) # # return light # # if selected_area is not None: # # self.add_debug_img("Halo", get_img_selected_halo) # # self.add_debug_img("Light", get_img_selected_light) # def get_img_selected(): # selected = np.copy(self.frame) # for each in selected_bars: # cv2.drawMarker(img=selected, position=each["centroid"], color=(0, 255, 255), markerSize=25, thickness=2) # for each in selected_pair: # cv2.circle(img=selected, center=each["centroid"], radius=5, color=(0, 0, 255), thickness=-1) # return selected # self.add_debug_img("Selected", get_img_selected) return target
