added video writerr

.gitignore

@@ -1,2 +1,4 @@
 tf_models/
-*.csv
+*.csv
+*.mp4
+__pycache__/

analytics/tracking.py

@@ -34,7 +34,7 @@ def update_person_status(self, class_names):
 
 
     def write_to_report(self, frame_number):
-        self.writer.writerow({'frame': frame_number, 'detections': self.class_names})
+        self.writer.writerow({'frame': frame_number, 'detections': self.class_counts})
 
 
     def __call__(self, context):

detect_object.py

@@ -0,0 +1,42 @@
+import os
+import numpy as np
+from utils import label_map_util
+from utils.webcam import draw_boxes_and_labels
+
+
+CWD_PATH = os.getcwd()
+PATH_TO_LABELS = os.path.join(CWD_PATH, 'detection', 'data', 'mscoco_label_map.pbtxt')
+
+NUM_CLASSES = 90
+# 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)
+
+# pass in image_np, returns 
+def detect_objects(image_np, sess, detection_graph):
+
+    image_np_expanded = np.expand_dims(image_np, axis=0)
+    image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
+
+    boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
+
+    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')
+
+    # Do the detection/model prediction here
+    (boxes, scores, classes, num_detections) = sess.run(
+        [boxes, scores, classes, num_detections],
+        feed_dict={image_tensor: image_np_expanded})
+
+    rect_points, class_names, class_colors = draw_boxes_and_labels(
+        boxes=np.squeeze(boxes),
+        classes=np.squeeze(classes).astype(np.int32),
+        scores=np.squeeze(scores),
+        category_index=category_index,
+        min_score_thresh=.5
+    )
+
+    return dict(rect_points=rect_points, class_names=class_names, class_colors=class_colors)

objection_detection_app.py

@@ -5,61 +5,19 @@
 import numpy as np
 import tensorflow as tf
 
-from utils.webcam import FPS, WebcamVideoStream, draw_boxes_and_labels
+from utils.webcam import FPS, WebcamVideoStream
 from queue import Queue
 from threading import Thread
-from utils import label_map_util
 from analytics.tracking import ObjectTracker
+from video_writer import VideoWriter
+from detect_object import detect_objects
 
 CWD_PATH = os.getcwd()
 
 MODEL_NAME = 'ssd_mobilenet_v1_coco_2017_11_17'
 PATH_TO_MODEL = os.path.join(CWD_PATH, 'detection', 'tf_models', MODEL_NAME, 'frozen_inference_graph.pb')
 PATH_TO_VIDEO = os.path.join(CWD_PATH, 'input.mp4')
 
-print('cv2', cv2.__version__)
-
-PATH_TO_LABELS = os.path.join(CWD_PATH, 'detection', 'data', 'mscoco_label_map.pbtxt')
-
-NUM_CLASSES = 90
-
-# 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)
-
-# pass in image_np, returns 
-def detect_objects(image_np, sess, detection_graph):
-
-    image_np_expanded = np.expand_dims(image_np, axis=0)
-    # print('image_np_expanded')
-    # print(image_np)
-    image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
-
-    boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
-
-    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')
-
-    # Do the detection/model prediction here
-    (boxes, scores, classes, num_detections) = sess.run(
-        [boxes, scores, classes, num_detections],
-        feed_dict={image_tensor: image_np_expanded})
-
-    # print('num_detections:', num_detections)
-    # Visualization of the results of a detection.
-    # print('classes:', classes)
-    rect_points, class_names, class_colors = draw_boxes_and_labels(
-        boxes=np.squeeze(boxes),
-        classes=np.squeeze(classes).astype(np.int32),
-        scores=np.squeeze(scores),
-        category_index=category_index,
-        min_score_thresh=.5
-    )
-
-    return dict(rect_points=rect_points, class_names=class_names, class_colors=class_colors)
 
 def worker(input_q, output_q):
     # load the frozen tensorflow model into memory
@@ -104,6 +62,7 @@ def worker(input_q, output_q):
     video_capture = WebcamVideoStream(src=args.video_source,
                                       width=args.width,
                                       height=args.height).start()
+    writer = VideoWriter('output.mp4', (args.width, args.height))
 
     '''
     stream = cv2.VideoCapture(0)
@@ -131,8 +90,9 @@ def worker(input_q, output_q):
         else:
             data = output_q.get()
             context = {'frame': frame, 'class_names': data['class_names'], 'rec_points': data['rect_points'], 'class_colors': data['class_colors'],
-                        'width': args.width, 'height': args.height, 'frame_numer': fps.get_numFrames()}
+                        'width': args.width, 'height': args.height, 'frame_number': fps.get_numFrames()}
             new_frame = object_tracker(context)
+            writer(new_frame)
             cv2.imshow('Video', new_frame)
 
         print('[INFO] elapsed time: {:.2f}'.format(time.time() - t))
@@ -145,4 +105,5 @@ def worker(input_q, output_q):
     print('[INFO] approx. FPS: {:.2f}'.format(fps.fps()))
 
     video_capture.stop()
+    writer.close()
     cv2.destroyAllWindows()

utils/webcam.py

@@ -78,6 +78,7 @@ def read(self):
     def stop(self):
         # stopping the thread
         self.stopped = True
+        self.stream.release()
 
 
 def standard_colors():

video_writer.py

@@ -0,0 +1,19 @@
+import cv2
+
+class VideoWriter(object):
+
+    def __init__(self, path, size):
+        self.path = path
+        self.size = size
+        self.writer = cv2.VideoWriter(self.path, 
+                         cv2.VideoWriter_fourcc('M','J','P','G'), 
+                         20.0, self.size, True)
+
+    def __call__(self, frame):
+        self.writer.write(frame)
+
+    def close():
+        self.writer.release()
+
+
+