Update app.py

app.py

@@ -1,63 +1,56 @@
 import gradio as gr
 import torch
-from ultralytics import YOLO
 import cv2
-import tempfile
+import numpy as np
+from PIL import Image
 
-# Load the trained YOLOv8 model
-model = YOLO('best.pt')
+# Load the YOLOv8 model
+model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt', force_reload=True)  # Adjust to yolov8 if needed
 
-def predict(image):
+def process_image(image):
     results = model(image)
-    # You might want to process results to return bounding boxes, class labels, etc.
-    annotated_image = results[0].plot()  # plot the results on the image
-    return annotated_image
+    return results.render()[0]  # Returns an image with boxes drawn
 
-def predict_video(video):
-    # Read the video file
+def process_video(video):
     cap = cv2.VideoCapture(video)
-    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    fps = int(cap.get(cv2.CAP_PROP_FPS))
-    
-    # Create a temporary file to save the output video
-    out_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
-    out_path = out_file.name
-    
-    # Define the codec and create VideoWriter object
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter(out_path, fourcc, fps, (width, height))
+    frames = []
 
-    while cap.isOpened():
+    while(cap.isOpened()):
         ret, frame = cap.read()
         if not ret:
             break
         results = model(frame)
-        annotated_frame = results[0].plot()  # plot the results on the frame
-        out.write(annotated_frame)
-    
+        frames.append(results.render()[0])
+
     cap.release()
-    out.release()
     
-    return out_path
+    # Convert frames back to a video format
+    height, width, layers = frames[0].shape
+    video_out = cv2.VideoWriter('output.mp4', cv2.VideoWriter_fourcc(*'mp4v'), 30, (width, height))
+
+    for frame in frames:
+        video_out.write(frame)
+    
+    video_out.release()
+    
+    return 'output.mp4'
 
 # Create Gradio interface
-interface = gr.Interface(
-    fn=lambda img, vid: (predict(img), predict_video(vid)),
-    inputs=[
-        gr.inputs.Image(type="numpy", label="Input Image"),
-        gr.inputs.Video(label="Input Video")
-    ],
-    outputs=[
-        gr.outputs.Image(type="numpy", label="Output Image"),
-        gr.outputs.Video(label="Output Video")
-    ],
-    title="YOLOv8 Object Detection",
-    description="Upload an image or a video and get the object detection results using a YOLOv8 model."
-)
+image_input = gr.inputs.Image(type="numpy", label="Upload an image")
+video_input = gr.inputs.Video(type="mp4", label="Upload a video")
+
+image_output = gr.outputs.Image(type="numpy", label="Detected image")
+video_output = gr.outputs.Video(type="mp4", label="Detected video")
+
+iface = gr.Interface(fn={'image': process_image, 'video': process_video},
+                     inputs=[image_input, video_input],
+                     outputs=[image_output, video_output],
+                     live=True,
+                     title="YOLOv8 Object Detection",
+                     description="Upload an image or video to detect objects using YOLOv8.")
 
 if __name__ == "__main__":
-    interface.launch()
+    iface.launch()