Commit d8bc6cb1 authored by Jeeken's avatar Jeeken
Browse files

debug mode improved; drawing simplified 

# On process:
rewrite detection: connected_cpn -> minAreaRec

# Possible issue:
performance may decline
parent 96615efc

src/car_target.py

+22 −18
Original line number Diff line number Diff line
@@ -80,6 +80,7 @@ class CarTargetApp: 
    def __init__(self, color: BarColor, debug: bool = False):

        self.detector = Detector(color=color, debug=debug)

        self.debug = debug

        self.color = color



    def __del__(self):

        cv2.destroyAllWindows()
@@ -88,6 +89,15 @@ class CarTargetApp: 
    def run(self):

        pass



    def draw_target(self, img, target):

        if img.ndim == 2:

            aim_color = 255

        else:

            aim_color = (0, 0, 255) if self.color == BarColor.BLUE else (0, 255, 0)

        cv2.circle(img=img, center=target, radius=20, color=aim_color, thickness=2)

        cv2.drawMarker(img=img, position=target, color=aim_color, markerType=cv2.MARKER_CROSS, markerSize=80,

                       thickness=2)





class ImgCarTargetApp(CarTargetApp):

    def __init__(self, folder: str, color: BarColor, frame_size: tuple, ext_name: str = "jpg", debug: bool = False):
@@ -105,31 +115,28 @@ class ImgCarTargetApp(CarTargetApp): 


            plt.figure("Armor Detection")



            ROW, COL = 3, 3

            count = 1

            plt.subplot(2, 3, count)

            plt.subplot(ROW, COL, count)

            plt.title("Original")

            plt.imshow(cv2.cvtColor(mat, cv2.COLOR_BGR2RGB))

            plt.axis("off")



            target, debug_imgs = self.detector.target(mat)

            target = self.detector.target(mat)

            print("target:", target)



            if self.debug:

                count = 2

                for title, img in debug_imgs:

                    plt.subplot(2, 3, count)

                for title, img in self.detector.debug_imgs:

                    plt.subplot(ROW, COL, count)

                    plt.title(title)

                    plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB) if img.ndim == 3 else img)

                    plt.axis("off")

                    count += 1



            if target is not None:

                x, y = target

                aim_color = (0, 255, 0)  # green

                cv2.circle(img=mat, center=target, radius=18, color=aim_color, thickness=2)

                cv2.line(img=mat, pt1=(x - 40, y), pt2=(x + 40, y), color=aim_color, thickness=2)

                cv2.line(img=mat, pt1=(x, y - 40), pt2=(x, y + 40), color=aim_color, thickness=2)

            plt.subplot(2, 3, count)

                self.draw_target(img=mat, target=target)

            plt.subplot(ROW, COL, count)

            plt.title("Aimed")

            plt.imshow(cv2.cvtColor(mat, cv2.COLOR_BGR2RGB))

            plt.axis("off")
@@ -208,7 +215,8 @@ class VideoCarTargetApp(CarTargetApp): 
                break

            frame = cv2.resize(frame, self.frame_size)



            target, debug_imgs = self.detector.target(frame)

            target = self.detector.target(frame)

            # target, debug_imgs = self.detector.target(frame)

            target_smoothed = self.smoother.smoothed(target)



            print("target:  ", target)
@@ -218,16 +226,12 @@ class VideoCarTargetApp(CarTargetApp): 
            cv2.imshow("Original", frame)



            if self.debug:

                for title, img in debug_imgs:

                for title, img in self.detector.debug_imgs:

                    cv2.imshow(title, img)



            if target is not None:

                cv2.circle(img=frame, center=target, radius=4, color=(0, 255, 0), thickness=-1) # green, -1: filled

                aim_color = (0, 0, 255) if self.color == BarColor.BLUE else (0, 255, 0)

                x, y = target_smoothed

                cv2.circle(img=frame, center=target_smoothed, radius=20, color=aim_color, thickness=2) # red circle

                cv2.line(img=frame, pt1=(x-40, y), pt2=(x+40, y), color=aim_color, thickness=1)

                cv2.line(img=frame, pt1=(x, y-40), pt2=(x, y+40), color=aim_color, thickness=1)

                self.draw_target(img=frame, target=target_smoothed)

            cv2.imshow("Smoothed", frame)



            # out.write(frame)
@@ -250,7 +254,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=False)

                            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)

    app.run()

src/detector.py

+76 −32
Original line number Diff line number Diff line
@@ -5,8 +5,16 @@ 
import cv2

import math

import numpy as np

from matplotlib import pyplot as plt

from enum import Enum

from random import randint





# class BgrColor(Enum):

#     RED = (0, 0, 255)

#     BLUE = (255, 0, 0)

#     GREEN = (0, 255, 0)

#     YELLOW = (0, 255, 255)

#     PURPLE = (255, 0, 255)





class BarColor(Enum):
@@ -23,12 +31,24 @@ class TargetDis(Enum): 
    ULTRA = 3





# TODO: Detector for red light bars



class Detector:

    # TODO: remove shape from self.__init__

    def __init__(self, color: BarColor, shape: tuple = (480, 360), debug = False):

        self.color = color

        self.debug_imgs = [] if debug else None

        self.refresh(frame=None, debug=debug)



    def refresh(self, frame, debug=None):

        self.frame = frame

        if debug is not None:

            self.debug = debug

        self.debug_imgs = []



    def add_debug_img(self, title: str, img):

        if self.debug:

            if callable(img):

                self.debug_imgs.append((title, img()))

            else:

                self.debug_imgs.append((title, img))



    def hsv(self, frame):

        hsv_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
@@ -38,12 +58,11 @@ class Detector: 
            h += 107

            h %= 180

            hsv_frame = cv2.merge((h, s, v)).astype(np.uint8)



            if self.debug_imgs is not None:

                self.debug_imgs.append(("Red to Blue", cv2.cvtColor(hsv_frame, cv2.COLOR_HSV2BGR)))

            self.add_debug_img("Red to Blue", cv2.cvtColor(hsv_frame, cv2.COLOR_HSV2BGR))



        # TODO: decide the range of lightness based on histogram

        lower_light = np.array([90, 0, 215])

        # lower_light = np.array([90, 0, 215])

        lower_light = np.array([90, 0, 210])

        upper_light = np.array([120, 100, 255])

        lower_halo  = np.array([90, 100, 185])

        upper_halo  = np.array([120, 255, 255])
@@ -54,6 +73,7 @@ class Detector: 


    @staticmethod

    def get_connected_components_info(frame):

        # cpn: component

        cpn_count, label_map, cpn_info, centroids = \

            cv2.connectedComponentsWithStats(image=frame, connectivity=4, ltype=cv2.CV_32S)
@@ -198,13 +218,42 @@ class Detector: 
        i1, i2 = winner["indices"]

        return bars[i1], bars[i2]



    # NEW!

    def get_lights(self, binary_img):

        height, 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(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 get_img_lights_recs_as_elps():

            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,

                            startAngle=0, endAngle=360, color=rand_color, thickness=2)

            return show_lights

        self.add_debug_img("Light_rec", get_img_lights_recs_as_elps)



        return rectangles



    def halo_circle(self, binary_img):

        # TODO

        pass



    def target(self, frame):

        if self.debug_imgs is not None:

            self.debug_imgs = []

        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)

        selected_bars = Detector.select_valid_bars(lights_in_halo)

        selected_pair = Detector.select_pair(selected_bars)
@@ -218,28 +267,23 @@ class Detector: 
        else:

            target = None



        if self.debug_imgs is not None:

            if selected_area is not None:

        def get_img_selected_halo():

            cv2.rectangle(halo, pt1=selected_area["northwest"], pt2=selected_area["southeast"], color=255, thickness=1)

                self.debug_imgs.append(("Halo", halo))



            return halo

        def get_img_selected_light():

            cv2.rectangle(light, pt1=selected_area["northwest"], pt2=selected_area["southeast"], color=255, thickness=1)

                self.debug_imgs.append(("Light", light))

            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)



            selected = np.copy(frame)

        def get_img_selected():

            selected = np.copy(self.frame)

            for each in selected_bars:

                cv2.circle(img=selected, center=each["centroid"], radius=7, color=(0, 255, 0), thickness=2)

                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)

            self.debug_imgs.append(("Seleted", selected))



            # aimed = np.copy(frame)

            # if target is not None:

            #     x, y = target

            #     aim_color = (0, 255, 0)  # green

            #     cv2.circle(img=aimed, center=target, radius=18, color=aim_color, thickness=2)

            #     cv2.line(img=aimed, pt1=(x - 40, y), pt2=(x + 40, y), color=aim_color, thickness=2)

            #     cv2.line(img=aimed, pt1=(x, y - 40), pt2=(x, y + 40), color=aim_color, thickness=2)

            # self.debug_imgs.append(("Aimed", aimed))



        return target, self.debug_imgs
 
            return selected

        self.add_debug_img("Selected", get_img_selected)



        return target

CRA Git | Maintained and supported by SUSTech CRA and CCSE