Update app.py

app.py

@@ -83,6 +83,40 @@ def predict2(image_np):
     
     return result_pil_img
 
+import cv2
+
+def predict_video(video_path):
+    cap = cv2.VideoCapture(video_path)
+    frame_width = int(cap.get(3))
+    frame_height = int(cap.get(4))
+
+    # Define the codec and create a video writer object
+    out = cv2.VideoWriter('output.avi', cv2.VideoWriter_fourcc('M','J','P','G'), 10, (frame_width, frame_height))
+
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        # Convert the frame to PIL image
+        pil_image = Image.fromarray(frame)
+        
+        # Perform object detection on the frame
+        result_pil_img = predict(pil_image)
+
+        # Convert the result back to a NumPy array
+        result_np_img = tf.keras.utils.img_to_array(result_pil_img)
+
+        # Write the frame with detected objects to the video output
+        out.write(result_np_img.astype('uint8'))
+
+    # Release the video capture and writer objects
+    cap.release()
+    out.release()
+
+    return "output.avi"
+
+
 
 REPO_ID = "Louisw3399/burgerorfriesdetector"
 detection_model = load_model()
@@ -96,3 +130,10 @@ gr.Interface(fn=predict,
              inputs=gr.Image(type="pil"),
              outputs=gr.Image(type="pil")
              ).launch(share=True)
+
+gr.Interface(
+    fn=predict_video,
+    inputs=gr.Video(type="file", label="Upload a video"),
+    outputs=gr.Video(type="file", label="Download the processed video")
+).launch(share=True)
+