Update app.py

app.py

@@ -1,26 +1,70 @@
 from ultralytics import YOLO
 from PIL import Image
 import gradio as gr
+import cv2
+import tempfile
 
 # Load YOLOv8 model
-model = YOLO("best.pt")  # Ensure best.pt is in the same directory
+model = YOLO("best.pt")  # Ensure best.pt is in the same directory or provide the correct path
 
-# Preprocess and run inference
-def predict(image):
+# Preprocess and run inference for images
+def predict_image(image):
     # Perform prediction
     results = model.predict(source=image, conf=0.5)
-    
+
     # Annotate the image with bounding boxes
     annotated_image = results[0].plot()
-    
+
     # Convert to PIL Image
     return Image.fromarray(annotated_image)
 
+# Preprocess and run inference for videos
+def predict_video(video):
+    # Save video to a temporary file
+    temp_video_path = tempfile.NamedTemporaryFile(suffix=".mp4").name
+    with open(temp_video_path, "wb") as f:
+        f.write(video.read())
+
+    # Open the video file
+    cap = cv2.VideoCapture(temp_video_path)
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Codec for .mp4
+    output_path = tempfile.NamedTemporaryFile(suffix=".mp4").name
+
+    # Get video properties
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+    # Create video writer for output
+    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break  # Exit when video ends
+
+        # Perform predictions on the frame
+        results = model.predict(source=frame, conf=0.5)
+        annotated_frame = results[0].plot()  # Annotate frame
+
+        # Write the frame to the output video
+        out.write(annotated_frame)
+
+    # Release resources
+    cap.release()
+    out.release()
+
+    # Return the annotated video path
+    return output_path
+
 # Gradio interface
 gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil"),
-    outputs="image",
+    fn={"Image Detection": predict_image, "Video Detection": predict_video},
+    inputs=[
+        gr.Image(type="pil", label="Upload an Image"),
+        gr.Video(label="Upload a Video")
+    ],
+    outputs=["image", "video"],
     title="Hippo or Rhino Detection",
-    description="Upload an image for object detection with YOLOv8."
-).launch()
+    description="Upload an image or video for object detection using YOLOv8."
+).launch()