Upload app.py

app.py

@@ -48,9 +48,49 @@ def load_model2():
 threshold = 0.50
 
 def predict(pilimg,video_in_filepath,threshold):
+    if pilimg != "":
+        image_np = pil_image_as_numpy_array(pilimg)
+        return predict2(image_np,threshold),None
+    else:
+        video_reader = cv2.VideoCapture(video_in_filepath)
+        
+        nb_frames = int(video_reader.get(cv2.CAP_PROP_FRAME_COUNT))
+        frame_h = int(video_reader.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        frame_w = int(video_reader.get(cv2.CAP_PROP_FRAME_WIDTH))
+        fps = video_reader.get(cv2.CAP_PROP_FPS)
+
+       
+        video_out_filepath = 'detected.mp4'
+        video_writer = cv2.VideoWriter(video_out_filepath,
+                                cv2.VideoWriter_fourcc(*'mp4v'), 
+                                fps, 
+                                (frame_w, frame_h))
+
+        for i in tqdm(range(nb_frames)):
+            ret, image_np = video_reader.read()
+            input_tensor = tf.convert_to_tensor(np.expand_dims(image_np, 0), dtype=tf.uint8)
+            results = detection_model(input_tensor)
+            viz_utils.visualize_boxes_and_labels_on_image_array(
+                    image_np,
+                    results['detection_boxes'][0].numpy(),
+                    (results['detection_classes'][0].numpy()+ label_id_offset).astype(int),
+                    results['detection_scores'][0].numpy(),
+                    category_index,
+                    use_normalized_coordinates=True,
+                    max_boxes_to_draw=200,
+                    min_score_thresh=.50,
+                    agnostic_mode=False,
+                    line_thickness=2)
+
+            video_writer.write(np.uint8(image_np))
+                
+        # Release camera and close windows
+        video_reader.release()
+        video_writer.release() 
+        cv2.destroyAllWindows() 
+        cv2.waitKey(1)
+        return None,video_out_filepath
 
-    image_np = pil_image_as_numpy_array(pilimg)
-    return predict2(image_np,threshold),None
 
 def predict2(image_np,threshold):
 
@@ -130,20 +170,20 @@ detection_model = load_model()
 # predicted_img = predict(image_arr)
 # predicted_img.save('predicted.jpg')
 
-# gr.Interface(fn=predict,
-#              inputs=[gr.Image(type="pil",label="Input Image"),gr.Video(label="Input Video"),gr.Textbox(placeholder="0.50",label="Set the confidence threshold (0.00-1.00)")],
-#              outputs=[gr.Image(type="pil",label="Output Image"),gr.Video(label="Output Video")], 
-#              title="Facemask & Glasses",
-#              description="Model: ssd_mobilenet_v2_320x320",
-#              theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"),
-#              examples=[["test_samples/image489.png","",0.55], ["test_samples/image825.png","",0.55], ["test_samples/image833.png","",0.55], ["test_samples/image846.png","",0.55]]
-#              ).launch(share=True)
-
-gr.Interface(fn=video_fn,
-             inputs=gr.Video(label="Input Video"),
-             outputs=gr.Video(label="Output Video"), 
+gr.Interface(fn=predict,
+             inputs=[gr.Image(type="pil",label="Input Image"),gr.Video(label="Input Video"),gr.Textbox(placeholder="0.50",label="Set the confidence threshold (0.00-1.00)")],
+             outputs=[gr.Image(type="pil",label="Output Image"),gr.Video(label="Output Video")], 
              title="Facemask & Glasses",
              description="Model: ssd_mobilenet_v2_320x320",
              theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"),
-             #examples="test_samples/test_video.mp4"
-             ).launch(share=True)
+             examples=[["test_samples/image489.png","",0.55], ["test_samples/image825.png","",0.55], ["test_samples/image833.png","",0.55], ["test_samples/image846.png","",0.55]]
+             ).launch(share=True)
+
+# gr.Interface(fn=video_fn,
+#              inputs=gr.Video(label="Input Video"),
+#              outputs=gr.Video(label="Output Video"), 
+#              title="Facemask & Glasses",
+#              description="Model: ssd_mobilenet_v2_320x320",
+#              theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"),
+#              #examples="test_samples/test_video.mp4"
+#              ).launch(share=True)