Traffic Light Bounding Box using Facebook Hugging Face Transformer

traffic_light_bbox.py

# source: https://github.com/christianversloot/machine-learning-articles/blob/main/easy-object-detection-with-python-huggingface-transformers-and-machine-learning.md

import os
from os import listdir
from transformers import pipeline
from PIL import Image, ImageDraw, ImageFont
import numpy as np
import cv2


# Load font
font = ImageFont.load_default() #truetype("./arial.ttf", 40)
# Store Hue, Saturation, Value of color lower and upper bounding_boxes
green_lb = np.array([35,0,0])
green_ub = np.array([90,255,255])

yellow_lb = np.array([20,100,100])
yellow_ub = np.array([35,255,255])

red_lb = np.array([0,75,70])
red_ub = np.array([25,255,255])

pink_lb = np.array([160,50,50])
pink_ub = np.array([180,255,255])

#Initialize the object detection pipeline
object_detector = pipeline(model="facebook/detr-resnet-50")

# Draw bounding box definition
def draw_bounding_box(im, score, label, xmin, ymin, xmax, ymax, index, num_boxes):
	""" Draw a bounding box. """

	print(f"Drawing bounding box {index} of {num_boxes}...")

	# Draw the actual bounding box
	im_with_rectangle = ImageDraw.Draw(im)
	im_with_rectangle.rectangle((xmin, ymin, xmax, ymax), outline = "blue", width = 3)

	# Draw the label
	im_with_rectangle.text((xmin+35, ymin-25), label, fill="white", stroke_fill = "red", font = font)

	# Return the intermediate result
	return im

def tl_filter(image, all_bbox):
	# filter to only bounding boxes with red, yellow or green color in them
	out = []
	color_detect = []
	full_image = cv2.imread(image)
	traffic_light_bboxes = [tl_box for tl_box in all_bbox if tl_box["label"] == "traffic light"]

	for i in range(len(traffic_light_bboxes)):
		tl_bbox = traffic_light_bboxes[i]["box"]
		offset = 0
		xmin = tl_bbox["xmin"] - offset
		ymin = tl_bbox["ymin"] - offset
		xmax = tl_bbox["xmax"] - offset
		ymax = tl_bbox["ymax"] - offset

		clip_tl_bbox = full_image[ymin:ymax, xmin:xmax].copy()
		hsv_clip = cv2.cvtColor(clip_tl_bbox, cv2.COLOR_BGR2HSV)

		green_mask = cv2.inRange(hsv_clip, green_lb, green_ub)
		yellow_mask = cv2.inRange(hsv_clip, yellow_lb, yellow_ub)
		red_mask = cv2.inRange(hsv_clip, red_lb, red_ub)
		pink_mask = cv2.inRange(hsv_clip, pink_lb, pink_ub)

		signal_colors = (green_mask, yellow_mask, red_mask, pink_mask)
		print(np.unique(yellow_mask, return_counts=True)[0].size)
		signal_light = lambda x:True if (np.unique(x, return_counts=True)[0].size > 1) else False
		print(signal_light(signal_colors[1]))
		sig_light = signal_light(signal_colors[0]) or \
					signal_light(signal_colors[1]) or \
					signal_light(signal_colors[2]) or \
					signal_light(signal_colors[3])
		print(f"{sig_light}")

		if sig_light:
			out.append(traffic_light_bboxes[i])

	print(f"{len(traffic_light_bboxes)}:{len(out)}")
	#print(out)
	return out

# Open the image
folder_dir = "./pics/"
for picture in os.listdir(folder_dir):
	with Image.open(folder_dir + picture) as im:

		# Perform object detection
		print(f"Current Image {picture}")
		all_bounding_boxes = object_detector(im)
		#print(all_bounding_boxes)
		bounding_boxes = tl_filter(folder_dir + picture, all_bounding_boxes)
		#traffic_light_boxes = [ts_box for ts_box in all_bounding_boxes if ts_box["label"] == "traffic light"]


		#print(bounding_boxes)


		# Iteration elements
		num_boxes = len(bounding_boxes)
		index = 0

		# Draw bounding box for each result
		for bounding_box in bounding_boxes:

			# Get actual box
			box = bounding_box["box"]

			# Draw the bounding box
			im = draw_bounding_box(im, bounding_box["score"], bounding_box["label"],\
				box["xmin"], box["ymin"], box["xmax"], box["ymax"], index, num_boxes)

			# Increase index by one
			index += 1

		# Save image
		im.save("./out/" + picture.replace(".jpg","") + "_bboxes.jpg")

		# Done
		print("Done!")