Add files via upload

6 years ago · 34bb59fc1a
1 changed files with 182 additions and 0 deletions
--- a/detectRecognizeLight.py
+++ b/detectRecognizeLight.py
@ -0,0 +1,182 @@
+import tensorflow as tf
+import numpy as np
+from PIL import Image
+import cv2
+import os
+
+
+def recognize_color(im):
+    '''
+    Given a traffic light image, returns class of the image (red, yellow, green)
+
+    Args:
+        im: traffic light image (detected bbox of traffic light)
+
+    Returns:
+        predicted_label: output label of the light color
+                        ([1, 0, 0]: red, [0, 1, 0]: yellow, [0, 0, 1]: green)
+    '''
+
+    # Standardize input and convert it from RGB to HSV
+    standard_im = cv2.resize(im, (32, 32))
+    hsv = cv2.cvtColor(standard_im, cv2.COLOR_RGB2HSV)
+    h = hsv[:, :, 0]
+    s = hsv[:, :, 1]
+    v = hsv[:, :, 2]
+
+    # Define the mask & its boundaries, based on value channel
+    # The limits of the masks are deduced by trial & error
+    lower_hsv = np.array([0, 0, 0])
+    upper_hsv = np.array([255, 255, 40])
+    mask_hsv = cv2.inRange(hsv, lower_hsv, upper_hsv)
+    masked_image=np.copy(hsv)
+    masked_image[mask_hsv != 0] = [0, 0, 0]
+
+    # Grayscale is used to enhance the quality of brightness
+    grayimage = cv2.cvtColor(standard_im, cv2.COLOR_RGB2GRAY)
+
+    # Using image slicing, crop the image into 3 horizontal parts
+    upper_part_v = masked_image[3:12, 5:27, 2]
+    mid_part_v = masked_image[12:21, 5:27, 2]
+    lower_part_v = masked_image[21:30, 5:27, 2]
+
+    upper_part_g = grayimage[3:12, 5:27]
+    mid_part_g = grayimage[12:21, 5:27]
+    lower_part_g = grayimage[21:30, 5:27]
+
+    # Average brightness of all the pixels on value channel & grayscale for the masked image
+    area = 9.0*22
+    avg_brightness_upper = (np.sum(upper_part_g[:,:]) + np.sum(upper_part_v[:,:]))/area
+    avg_brightness_mid =  (np.sum(mid_part_g[:, :]) + np.sum(mid_part_v[:, :]))/area
+    avg_brightness_lower =  (np.sum(lower_part_g[:, :]) + np.sum(lower_part_v[:, :]))/area
+
+    # Feature vector
+    feature = []
+    feature = [avg_brightness_upper, avg_brightness_mid, avg_brightness_lower]
+
+    # Convert feature vector to output label
+    predicted_label = [0, 0, 0]
+    predicted_label[np.argmax(feature)] = 1
+    return predicted_label
+
+
+def draw_boxes(image_fed, best_boxes_roi, best_boxes_classes, best_boxes_scores, frame_w, frame_h, num_pred, video_writer):
+    '''
+    Draws boxes for detected objects
+
+    Args:
+        image_fed: image batch in which object detection performed
+        best_boxes_roi: roi of detected objects in images
+        best_boxes_classes: classes of detected objects in images
+        best_boxes_scores: scores of detected objects in images
+        frame_w: output frame width
+        frame_h: output frame height
+        num_pred: number of predictions
+        video_writer: video writer object to write drawn frame to video
+    '''
+
+    for i in range(best_boxes_roi.shape[0]):
+        im = np.reshape(image_fed[i], (frame_w, frame_h, 3))
+        im = cv2.cvtColor(im, cv2.COLOR_RGB2BGR)
+
+        for j in range(num_pred):
+
+            # Choose traffic lights from detected objects (label number of traffic light category in COCO dataset is 10)
+            if best_boxes_scores[i][j] > 0.35 and best_boxes_classes[i][j] == 10.:
+                x = best_boxes_roi[i][j][1]
+                y = best_boxes_roi[i][j][0]
+                x_max = best_boxes_roi[i][j][3]
+                y_max = best_boxes_roi[i][j][2]
+
+                # Classify color of the traffic light detected
+                color = recognize_color(im[int(x):int(x_max), int(y):int(y_max), :])
+
+                # Use different colors according to classified color of traffic light
+                if color == [1, 0, 0]:
+                    cv2.rectangle(im, (x,y), (x_max,y_max), (255, 0, 0), 2)
+                    font = cv2.FONT_HERSHEY_SIMPLEX
+                    cv2.putText(im, 'Red', (x,y), font, 1e-3*frame_h, (0, 0, 255), 2)
+                elif color == [0, 1, 0]:
+                    cv2.rectangle(im, (x,y), (x_max,y_max), (255, 0, 0), 2)
+                    font = cv2.FONT_HERSHEY_SIMPLEX
+                    cv2.putText(im, 'Yellow', (x,y), font, 1e-3*frame_h, (0, 255, 255), 2)
+                elif color == [0, 0, 1]:
+                    cv2.rectangle(im, (x,y), (x_max,y_max), (255, 0, 0), 2)
+                    font = cv2.FONT_HERSHEY_SIMPLEX
+                    cv2.putText(im, 'Green', (x,y), font, 1e-3*frame_h, (0, 255, 0), 2)
+
+        video_writer.write(im)
+
+
+def main():
+    im_size = 512
+    pb_dir = './model/frozen_inference_graph.pb'  # Pretrained model path
+    img_dir = './object-dataset'  # Dataset folder path
+    vid_out = './out'  # Output video path
+
+    # Load pretrained model
+    graph = tf.Graph()
+    with graph.as_default():
+        with tf.gfile.FastGFile(pb_dir, 'rb') as file:
+            graph_def = tf.GraphDef()
+            graph_def.ParseFromString(file.read())
+            tf.import_graph_def(graph_def, name='')
+
+            img = graph.get_tensor_by_name('image_tensor:0')
+            detection_boxes = graph.get_tensor_by_name('detection_boxes:0')
+            detection_scores = graph.get_tensor_by_name('detection_scores:0')
+            num_detections = graph.get_tensor_by_name('num_detections:0')
+            detection_classes = graph.get_tensor_by_name('detection_classes:0')
+            sess = tf.Session(graph=graph)
+
+    vid_out = vid_out + '/outputVideo.mp4'
+    num_iter = 67
+    batch_size = 32
+    num_pred = 30
+    frame_w = 512
+    frame_h = 512
+    video_writer = cv2.VideoWriter(vid_out, cv2.VideoWriter_fourcc(*'MP4V'), 5.0, (frame_w, frame_h))
+
+    # Read image name list from the directory and sort it (os.listdir generates output in random order)
+    images = sorted(os.listdir(img_dir))
+    k = 2
+    for i in range((len(images)-2)//batch_size):
+        image_bat = []
+        for j in range(batch_size):
+            # Read images from dataset directory
+            image = cv2.imread(img_dir + '/' + images[k])
+            # Resize images and make them batch
+            image = cv2.resize(image, (im_size, im_size))
+            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+            image_bat.append(image)
+            image_batch = np.asarray(image_bat)
+            k = k + 1
+
+        # Give image batch to the model as input
+        feed_dict = {img:image_batch}
+        y_p_boxes, y_p_scores, y_p_num_detections, y_p_classes = sess.run([detection_boxes,
+                                                                           detection_scores,
+                                                                           num_detections,
+                                                                           detection_classes], feed_dict=feed_dict)
+
+        # Process detection outputs
+        best_boxes_roi = []
+        best_boxes_scores = []
+        best_boxes_classes = []
+        for i in range(y_p_boxes.shape[0]):
+            temp = y_p_boxes[i, :num_pred] * frame_w
+            best_boxes_roi.append(temp)
+            best_boxes_scores.append(y_p_scores[i, :num_pred])
+            best_boxes_classes.append(y_p_classes[i, :num_pred])
+        best_boxes_roi = np.asarray(best_boxes_roi)
+        best_boxes_scores = np.asarray(best_boxes_scores)
+        best_boxes_classes = np.asarray(best_boxes_classes)
+
+        # Draw boxes for detected objects
+        draw_boxes(image_batch, best_boxes_roi, best_boxes_classes, best_boxes_scores, frame_w, frame_h, num_pred, video_writer)
+
+    video_writer.release()
+
+
+if __name__ == "__main__":
+    main()