1. red->blue; 2. improve debug mode (96615efc) · Commits · Ziqin Wang / ArmorDetection

src/car_target.py

+61 −14

Original line number	Diff line number	Diff line
		@@ -79,6 +79,7 @@ class Smoother:
		class CarTargetApp:
		def __init__(self, color: BarColor, debug: bool = False):
		self.detector = Detector(color=color, debug=debug)
		self.debug = debug

		def __del__(self):
		cv2.destroyAllWindows()
		@@ -89,23 +90,53 @@ class CarTargetApp:


		class ImgCarTargetApp(CarTargetApp):
		def __init__(self, color: BarColor, frame_size: tuple, debug: bool = False):
		CarTargetApp.__init__(self, color, debug)
		def __init__(self, folder: str, color: BarColor, frame_size: tuple, ext_name: str = "jpg", debug: bool = False):
		CarTargetApp.__init__(self, color, debug=debug)
		self.frame_size = frame_size
		self.folder = folder
		self.ext_name = ext_name

		@quitable
		def run(self): # FIXME: move folder path to constructor
		file_list = folder("/home/jeeken/Pictures/DMovie", "jpg")
		file_list = folder(self.folder, self.ext_name)
		for i in file_list:
		print("img_file:", i)
		mat = cv2.resize(cv2.imread(i), self.frame_size)
		target = self.detector.target(mat)

		plt.figure("Armor Detection")

		count = 1
		plt.subplot(2, 3, count)
		plt.title("Original")
		plt.imshow(cv2.cvtColor(mat, cv2.COLOR_BGR2RGB))
		plt.axis("off")

		target, debug_imgs = self.detector.target(mat)
		print("target:", target)

		if self.debug:
		count = 2
		for title, img in debug_imgs:
		plt.subplot(2, 3, 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)
		plt.title("Aimed")
		plt.imshow(cv2.cvtColor(mat, cv2.COLOR_BGR2RGB))
		plt.axis("off")

		if plt.waitforbuttonpress():
		plt.close()
		break
		# if cv2.waitKey(0) & 0xFF == ord('q'):
		# break


		class CamCarTargetApp(CarTargetApp):
		@@ -157,8 +188,9 @@ class CamCarTargetApp(CarTargetApp):


		class VideoCarTargetApp(CarTargetApp):
		def __init__(self, color: BarColor, file: str, frame_size: tuple = (640, 480), debug: bool = False):
		def __init__(self, file: str, color: BarColor, frame_size: tuple = (640, 480), debug: bool = False):
		CarTargetApp.__init__(self, color, debug)
		self.color = color
		self.file = file
		self.debug = debug
		self.frame_size = frame_size
		@@ -168,30 +200,41 @@ class VideoCarTargetApp(CarTargetApp):
		def run(self):
		# out = cv2.VideoWriter(filename="/home/jeeken/Videos/180_out.mp4", fourcc=cv2.VideoWriter_fourcc(*"XVID"), fps=30.0, frameSize=(640, 480))
		cap = cv2.VideoCapture(self.file)
		# debug:
		# cap = cv2.VideoCapture(1)
		while cap.isOpened():
		ok, frame = cap.read()
		if not ok:
		break
		frame = cv2.resize(frame, self.frame_size)
		target = self.detector.target(frame)

		target, debug_imgs = self.detector.target(frame)
		target_smoothed = self.smoother.smoothed(target)

		print("target: ", target)
		print("smoothed:", target_smoothed)
		print("--------------------")

		cv2.imshow("Original", frame)

		if self.debug:
		for title, img in 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)
		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)
		cv2.imshow("Aimed", frame)
		cv2.imshow("Smoothed", frame)

		# out.write(frame)

		if cv2.waitKey(10) & 0xFF == ord('q'):
		break
		# out.release()
		cap.release()


		@@ -201,9 +244,13 @@ if __name__ == "__main__":

		# app = CamCarTargetApp(color=BarColor.BLUE, msg=ser, cam_idx=1, frame_size=(480, 360), debug=True)

		# app = ImgCarTargetApp(color=BarColor.BLUE, frame_size=(480, 360), debug=True)
		# app = ImgCarTargetApp(color=BarColor.RED, frame_size=(480, 360), debug=True)
		# app = ImgCarTargetApp(color=BarColor.BLUE, folder="/home/jeeken/Pictures/blue",
		# frame_size=(640, 480), ext_name="jpg", debug=True)
		# app = ImgCarTargetApp(color=BarColor.RED, folder="/home/jeeken/Pictures/red",
		# frame_size=(640, 360), ext_name="jpg", debug=True)

		# edit commented code if you want to debug Video mode
		app = VideoCarTargetApp(color=BarColor.BLUE, file="/home/jeeken/Videos/live_blue.avi", frame_size=(640, 480), debug=False)
		app = VideoCarTargetApp(file="/home/jeeken/Videos/live_blue.avi",
		color=BarColor.BLUE, frame_size=(640, 480), debug=False)
		# app = VideoCarTargetApp(file="/home/jeeken/Videos/live_red.avi",
		# color=BarColor.RED, frame_size=(640, 480), debug=False)
		app.run()

src/detector.py

+52 −64

Original line number	Diff line number	Diff line
		@@ -26,30 +26,27 @@ class TargetDis(Enum):
		# TODO: Detector for red light bars

		class Detector:
		def __init__(self, color: BarColor, shape: tuple = (480, 360), debug: bool = False):
		def __init__(self, color: BarColor, shape: tuple = (480, 360), debug = False):
		self.color = color
		self.debug = debug
		self.distance = TargetDis.NEAR # for target_old
		self.debug_imgs = [] if debug else None

		def hsv(self, frame):
		hsv_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)

		if self.color == BarColor.RED:
		h, s, v = cv2.split(hsv_frame.astype(np.uint16))
		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)))

		# TODO: decide the range of lightness based on histogram
		if self.color == BarColor.BLUE:
		lower_light = np.array([90, 0, 215])
		upper_light = np.array([120, 100, 255])
		lower_halo = np.array([90, 100, 185])
		upper_halo = np.array([120, 255, 255])
		else: # self.color == BarColor.RED:
		lower_light = np.array([0, 0, 215])
		upper_light = np.array([25, 100, 255])
		# TODO
		lower_halo1 = np.array([170, 100, 185])
		upper_halo1 = np.array([180, 255, 255])
		lower_halo2 = np.array([0, 100, 185])
		upper_halo2 = np.array([25, 255, 255])
		lower_halo = cv2.bitwise_or(lower_halo1, lower_halo2)
		upper_halo = cv2.bitwise_or(upper_halo1, upper_halo2)

		light_area = cv2.inRange(hsv_frame, lower_light, upper_light)
		halo_area = cv2.inRange(hsv_frame, lower_halo, upper_halo)
		@@ -119,8 +116,9 @@ class Detector:
		for i in insides:
		nw_x, nw_y = i["northwest"]
		se_x, se_y = i["southeast"]
		tx, ty = nw_x + np.argmax(light_map[nw_y, nw_x:se_x]), nw_y # top peak
		bx, by = nw_x + np.argmax(light_map[se_y-1, nw_x:se_x]), se_y-1 # bottom peak
		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

		@@ -139,13 +137,13 @@ class Detector:
		width, height = bar["size"]
		square_ratio = height / width
		return 1.5 <= square_ratio < 8
		# def is_bar_far(bar):
		# width, height = bar["size"]
		# square_ratio = height / width
		# return 1.3 < square_ratio < 4
		def is_bar_far(bar):
		width, height = bar["size"]
		square_ratio = height / width
		return 0.8 < square_ratio < 2
		selected = list(filter(is_bar_near, big_enough))
		# if len(selected) < 2:
		# selected = list(filter(is_bar_far, big_enough))
		if len(selected) < 2:
		selected = list(filter(is_bar_far, big_enough))

		vertical_bars = filter(lambda x: abs(x["angle"] - 0.5) < 0.2, selected)

		@@ -169,6 +167,9 @@ class Detector:
		height_base = min((bar1["size"][1], bar2["size"][1]))
		return abs(y1 - y2) / height_base

		def square_ratio(bar1, bar2): # TODO
		pass

		def area(bar1, bar2):
		a1 = bar1["pixels"]
		a2 = bar2["pixels"]
		@@ -198,6 +199,9 @@ class Detector:
		return bars[i1], bars[i2]

		def target(self, frame):
		if self.debug_imgs is not None:
		self.debug_imgs = []

		# frame = cv2.blur(frame, ksize=(4, 4))
		frame = cv2.pyrUp(cv2.pyrDown(frame)) # down-sample
		light, halo = self.hsv(frame)
		@@ -214,44 +218,28 @@ class Detector:
		else:
		target = None

		# # debug
		# for each in selected_pair:
		# cv2.circle(img=frame, center=each["centroid"], radius=7, color=(0,255,0), thickness=2)
		# cv2.imshow('Debug', frame)

		if self.debug:
		show = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

		plt.figure("Car Detection Debug")
		if self.debug_imgs is not None:
		if selected_area is not None:
		cv2.rectangle(halo, pt1=selected_area["northwest"], pt2=selected_area["southeast"], color=255, thickness=1)
		self.debug_imgs.append(("Halo", halo))

		plt.subplot(2, 2, 1)
		plt.title("Original")
		plt.imshow(show)
		plt.axis('off')
		cv2.rectangle(light, pt1=selected_area["northwest"], pt2=selected_area["southeast"], color=255, thickness=1)
		self.debug_imgs.append(("Light", light))

		if selected_area is not None:
		cv2.rectangle(halo, pt1=selected_area["northwest"], pt2=selected_area["southeast"], color=255,
		thickness=1)
		cv2.rectangle(light, pt1=selected_area["northwest"], pt2=selected_area["southeast"], color=255,
		thickness=1)
		selected = np.copy(frame)
		for each in selected_bars:
		cv2.circle(img=show, center=each["centroid"], radius=7, color=(0, 255, 0), thickness=2)
		cv2.circle(img=selected, center=each["centroid"], radius=7, color=(0, 255, 0), thickness=2)
		for each in selected_pair:
		cv2.circle(img=show, center=each["centroid"], radius=5, color=(255, 0, 0), thickness=2)

		plt.subplot(2, 2, 2)
		plt.title("halo")
		plt.imshow(halo)
		plt.axis('off')

		plt.subplot(2, 2, 3)
		plt.title("light")
		plt.imshow(light)
		plt.axis('off')

		plt.subplot(2, 2, 4)
		plt.title("selected")
		plt.imshow(show)
		plt.axis('off')

		return target
		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

Admin message