Merge pull request #176 from bessszilard/add_video_output_and_dont_show_flag

Add video output and dont show flag

Author: hunglc007

detectvideo.py

@@ -23,6 +23,9 @@
 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')
+flags.DEFINE_boolean('dis_cv2_window', False, 'disable cv2 window during the process') # this is good for the .ipynb
 
 def main(_argv):
     config = ConfigProto()
@@ -45,14 +48,27 @@ 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))
 
+    frame_id = 0
     while True:
         return_value, frame = vid.read()
         if return_value:
             frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
             image = Image.fromarray(frame)
         else:
+            if frame_id == vid.get(cv2.CAP_PROP_FRAME_COUNT):
+                print("Video processing complete")
+                break
             raise ValueError("No image! Try with another video format")
+        
         frame_size = frame.shape[:2]
         image_data = cv2.resize(frame, (input_size, input_size))
         image_data = image_data / 255.
@@ -92,10 +108,18 @@ 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
+
+        if not FLAGS.dis_cv2_window:
+            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)
+        
+        frame_id += 1
 
 if __name__ == '__main__':
     try: