Vehicle detection main

vehicle_detection_main.py

+# Imports
+import os
+import tensorflow as tf
+import cv2
+import numpy as np
+import csv
+
+# Object detection imports
+from utils import label_map_util
+from utils import visualization_utils as vis_util
+
+# initialize .csv
+with open('traffic_measurement.csv', 'w') as f:
+    writer = csv.writer(f)
+    csv_line = 'Vehicle Type/Size, Vehicle Color, Vehicle Movement Direction, Vehicle Speed (km/h)'
+    writer.writerows([csv_line.split(',')])
+
+# input video
+# source_video = 'input_videos/k.mp4'
+
+
+total_passed_vehicle = 0  # using it to count vehicles
+
+MODEL_NAME = 'ssd_mobilenet_v1_coco_2018_01_28'
+MODEL_FILE = MODEL_NAME + '.tar.gz'
+DOWNLOAD_BASE = 'http://download.tensorflow.org/models/object_detection/'
+
+# Path to frozen detection graph. This is the actual model that is used for the object detection.
+PATH_TO_CKPT = MODEL_NAME + '/frozen_inference_graph.pb'
+
+# List of the strings that is used to add correct label for each box.
+PATH_TO_LABELS = os.path.join('data', 'mscoco_label_map.pbtxt')
+
+NUM_CLASSES = 90
+
+detection_graph = tf.Graph()
+with detection_graph.as_default():
+    od_graph_def = tf.compat.v1.GraphDef()
+    with tf.compat.v2.io.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
+        serialized_graph = fid.read()
+        od_graph_def.ParseFromString(serialized_graph)
+        tf.import_graph_def(od_graph_def, name='')
+
+label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
+categories = label_map_util.convert_label_map_to_categories(label_map,
+                                                            max_num_classes=NUM_CLASSES, use_display_name=True)
+category_index = label_map_util.create_category_index(categories)
+
+
+# Helper code
+def load_image_into_numpy_array(image):
+    (im_width, im_height) = image.size
+    return np.array(image.getdata()).reshape((im_height, im_width,
+                                              3)).astype(np.uint8)
+
+
+# Detection
+def object_detection_function(source_video):
+    cap = cv2.VideoCapture(source_video)
+
+    # Variables
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    fps = int(cap.get(cv2.CAP_PROP_FPS))
+    total_passed_vehicle = 0
+    speed = 'waiting...'
+    direction = 'waiting...'
+    size = 'waiting...'
+    color = 'waiting...'
+
+    # fourcc = cv2.VideoWriter_fourcc(*'XVID')
+    # output_movie = cv2.VideoWriter(source_video.split(".")[0] + '_output.avi', fourcc, fps, (width, height))
+
+    with detection_graph.as_default():
+        with tf.compat.v1.Session(graph=detection_graph) as sess:
+            # Definite input and output Tensors for detection_graph
+            image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
+
+            # Each box represents a part of the image where a particular object was detected.
+            detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
+
+            # Each score represent how level of confidence for each of the objects.
+            # Score is shown on the result image, together with the class label.
+            detection_scores = detection_graph.get_tensor_by_name('detection_scores:0')
+            detection_classes = detection_graph.get_tensor_by_name('detection_classes:0')
+            num_detections = detection_graph.get_tensor_by_name('num_detections:0')
+
+            # for all the frames that are extracted from input video
+            while cap.isOpened():
+                (ret, frame) = cap.read()
+
+                if not ret:
+                    print('end of the video file...')
+                    break
+
+                input_frame = frame
+
+                # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
+                image_np_expanded = np.expand_dims(input_frame, axis=0)
+
+                # Actual detection.
+                (boxes, scores, classes, num) = \
+                    sess.run([detection_boxes, detection_scores,
+                              detection_classes, num_detections],
+                             feed_dict={image_tensor: image_np_expanded})
+
+                # Visualization of the results of a detection.
+                (counter, csv_line) = vis_util.visualize_boxes_and_labels_on_image_array(
+                    cap.get(1),
+                    input_frame,
+                    np.squeeze(boxes),
+                    np.squeeze(classes).astype(np.int32),
+                    np.squeeze(scores),
+                    category_index,
+                    use_normalized_coordinates=True,
+                    line_thickness=4,
+                )
+
+                total_passed_vehicle += counter
+
+                # insert information text to video frame
+                font = cv2.FONT_HERSHEY_SIMPLEX
+                cv2.putText(
+                    input_frame,
+                    'Detected Vehicles: ' + str(total_passed_vehicle),
+                    (10, 35),
+                    font,
+                    0.8,
+                    (0, 0xFF, 0xFF),
+                    2,
+                    cv2.FONT_HERSHEY_SIMPLEX,
+                )
+
+                # when the vehicle passed over line and counted, make the color line green
+                # if counter == 1:
+                #     cv2.line(input_frame, (0, 200), (width, 200), (0, 0xFF, 0), 5)
+                # else:
+                #     cv2.line(input_frame, (0, 200), (width, 200), (0, 0, 0xFF), 5)
+
+                # insert information text to video frame
+                cv2.rectangle(input_frame, (10, 275), (230, 337), (180, 132, 109), -1)
+                cv2.putText(
+                    input_frame,
+                    'LAST PASSED VEHICLE INFO',
+                    (11, 290),
+                    font,
+                    0.5,
+                    (0xFF, 0xFF, 0xFF),
+                    1,
+                    cv2.FONT_HERSHEY_SIMPLEX,
+                )
+                cv2.putText(
+                    input_frame,
+                    '-Movement Direction: ' + direction,
+                    (14, 302),
+                    font,
+                    0.4,
+                    (0xFF, 0xFF, 0xFF),
+                    1,
+                    cv2.FONT_HERSHEY_COMPLEX_SMALL,
+                )
+                cv2.putText(
+                    input_frame,
+                    '-Speed(km/h): ' + str(speed).split(".")[0],
+                    (14, 312),
+                    font,
+                    0.4,
+                    (0xFF, 0xFF, 0xFF),
+                    1,
+                    cv2.FONT_HERSHEY_COMPLEX_SMALL,
+                )
+                cv2.putText(
+                    input_frame,
+                    '-Color: ' + color,
+                    (14, 322),
+                    font,
+                    0.4,
+                    (0xFF, 0xFF, 0xFF),
+                    1,
+                    cv2.FONT_HERSHEY_COMPLEX_SMALL,
+                )
+                cv2.putText(
+                    input_frame,
+                    '-Vehicle Size/Type: ' + size,
+                    (14, 332),
+                    font,
+                    0.4,
+                    (0xFF, 0xFF, 0xFF),
+                    1,
+                    cv2.FONT_HERSHEY_COMPLEX_SMALL,
+                )
+
+                cv2.imshow('vehicle detection', input_frame)
+
+                if cv2.waitKey(1) & 0xFF == ord('q'):
+                    break
+
+                if csv_line != 'not_available':
+                    with open('traffic_measurement.csv', 'a') as f:
+                        writer = csv.writer(f)
+                        (size, color, direction, speed) = \
+                            csv_line.split(',')
+                        writer.writerows([csv_line.split(',')])
+            cap.release()
+            cv2.destroyAllWindows()
+
+
+# object_detection_function()