WIP: Add multithreading example with lag; need to draw boxes only instead of entire predicted image.

Author: datitran

object_detection_multithreading.py

@@ -0,0 +1,132 @@
+import os
+import cv2
+import time
+import argparse
+import numpy as np
+import tensorflow as tf
+
+from queue import Queue
+from threading import Thread
+from utils import FPS, WebcamVideoStream
+from object_detection.utils import label_map_util
+from object_detection.utils import visualization_utils as vis_util
+
+CWD_PATH = os.getcwd()
+
+# Path to frozen detection graph. This is the actual model that is used for the object detection.
+MODEL_NAME = 'ssd_mobilenet_v1_coco_11_06_2017'
+PATH_TO_CKPT = os.path.join(CWD_PATH, 'object_detection', 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(CWD_PATH, 'object_detection', 'data', 'mscoco_label_map.pbtxt')
+
+NUM_CLASSES = 90
+
+# Loading label map
+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)
+
+
+def detect_objects(image_np, sess, detection_graph):
+ # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
+ image_np_expanded = np.expand_dims(image_np, axis=0)
+ 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.
+ 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.
+ scores = detection_graph.get_tensor_by_name('detection_scores:0')
+ classes = detection_graph.get_tensor_by_name('detection_classes:0')
+ num_detections = detection_graph.get_tensor_by_name('num_detections:0')
+
+ # Actual detection.
+ (boxes, scores, classes, num_detections) = sess.run(
+ [boxes, scores, classes, num_detections],
+ feed_dict={image_tensor: image_np_expanded})
+
+ # Visualization of the results of a detection.
+ vis_util.visualize_boxes_and_labels_on_image_array(
+ image_np,
+ np.squeeze(boxes),
+ np.squeeze(classes).astype(np.int32),
+ np.squeeze(scores),
+ category_index,
+ use_normalized_coordinates=True,
+ line_thickness=8)
+ return image_np
+
+
+def worker(input_q, output_q):
+ # Load a (frozen) Tensorflow model into memory.
+ detection_graph = tf.Graph()
+ with detection_graph.as_default():
+ od_graph_def = tf.GraphDef()
+ with tf.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='')
+
+ sess = tf.Session(graph=detection_graph)
+
+ fps = FPS().start()
+ while True:
+ fps.update()
+ frame = input_q.get()
+ output_q.put(detect_objects(frame, sess, detection_graph))
+
+ fps.stop()
+ sess.close()
+
+
+if __name__ == '__main__':
+ parser = argparse.ArgumentParser()
+ parser.add_argument('-src', '--source', dest='video_source', type=int,
+ default=0, help='Device index of the camera.')
+ parser.add_argument('-wd', '--width', dest='width', type=int,
+ default=480, help='Width of the frames in the video stream.')
+ parser.add_argument('-ht', '--height', dest='height', type=int,
+ default=360, help='Height of the frames in the video stream.')
+ args = parser.parse_args()
+
+ input_q = Queue(30)
+ output_q = Queue()
+ for i in range(1):
+ t = Thread(target=worker, args=(input_q, output_q))
+ t.daemon = True
+ t.start()
+
+ video_capture = WebcamVideoStream(src=args.video_source,
+ width=args.width,
+ height=args.height).start()
+ fps = FPS().start()
+
+ while True:
+ frame = video_capture.read()
+ input_q.put(frame)
+
+ t = time.time()
+
+ if output_q.empty():
+ cv2.imshow('Video', frame)
+ else:
+ # TO-DO need to draw the boxes here
+ cv2.imshow('Video', frame)
+ print(output_q.get())
+
+ fps.update()
+
+ print('[INFO] elapsed time: {:.2f}'.format(time.time() - t))
+
+ if cv2.waitKey(1) & 0xFF == ord('q'):
+ break
+
+ fps.stop()
+ print('[INFO] elapsed time (total): {:.2f}'.format(fps.elapsed()))
+ print('[INFO] approx. FPS: {:.2f}'.format(fps.fps()))
+
+ video_capture.stop()
+ cv2.destroyAllWindows()