Add video output and put in comment cv2 window

detectvideo.py

@@ -23,6 +23,8 @@
 flags.DEFINE_string('video', './data/road.mp4', 'path to input video')
 flags.DEFINE_float('iou', 0.45, 'iou threshold')
 flags.DEFINE_float('score', 0.25, 'score threshold')
+flags.DEFINE_string('output', None, 'path to output video')
+flags.DEFINE_string('output_format', 'XVID', 'codec used in VideoWriter when saving video to file')
 
 def main(_argv):
     config = ConfigProto()
@@ -45,6 +47,14 @@ def main(_argv):
     else:
         saved_model_loaded = tf.saved_model.load(FLAGS.weights, tags=[tag_constants.SERVING])
         infer = saved_model_loaded.signatures['serving_default']
+
+    if FLAGS.output:
+        # by default VideoCapture returns float instead of int
+        width = int(vid.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        fps = int(vid.get(cv2.CAP_PROP_FPS))
+        codec = cv2.VideoWriter_fourcc(*FLAGS.output_format)
+        out = cv2.VideoWriter(FLAGS.output, codec, fps, (width, height))
 
     while True:
         return_value, frame = vid.read()
@@ -92,10 +102,14 @@ def main(_argv):
         result = np.asarray(image)
         info = "time: %.2f ms" %(1000*exec_time)
         print(info)
-        cv2.namedWindow("result", cv2.WINDOW_AUTOSIZE)
-        result = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
-        cv2.imshow("result", result)
-        if cv2.waitKey(1) & 0xFF == ord('q'): break
+        # cv2.namedWindow("result", cv2.WINDOW_AUTOSIZE)
+        # result = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+        # cv2.imshow("result", result)
+        # if cv2.waitKey(1) & 0xFF == ord('q'): break
+
+        if FLAGS.output:
+            result = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
+            out.write(result)
 
 if __name__ == '__main__':
     try: