Update app.py

app.py

@@ -1,57 +1,48 @@
 import gradio as gr
-import torch
-from PIL import Image
 import cv2
+from PIL import Image
 import numpy as np
+from ultralytics import YOLO
 from huggingface_hub import hf_hub_download
 
-# Load the model from Hugging Face
+# Download the model from Hugging Face
 model_path = hf_hub_download(repo_id="StephanST/WALDO30", filename="WALDO30_yolov8m_640x640.pt")
-model = torch.hub.load('ultralytics/yolov8', 'custom', path=model_path)
+model = YOLO(model_path)  # Load YOLOv8 model
 
 # Detection function for images
 def detect_on_image(image):
-    results = model(image)
-    results.render()  # Render the bounding boxes on the image
-    detected_img = Image.fromarray(results.imgs[0])  # Convert to PIL format
-    return detected_img
+    results = model(image)  # Perform detection
+    annotated_frame = results[0].plot()  # Get annotated image
+    return Image.fromarray(annotated_frame)
 
 # Detection function for videos
 def detect_on_video(video):
     temp_video_path = "processed_video.mp4"
     cap = cv2.VideoCapture(video)
     fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-    out = cv2.VideoWriter(temp_video_path, fourcc, cap.get(cv2.CAP_PROP_FPS), 
+    out = cv2.VideoWriter(temp_video_path, fourcc, cap.get(cv2.CAP_PROP_FPS),
                           (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))))
 
     while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
             break
-        results = model(frame)  # Run detection
-        results.render()
-        frame = np.squeeze(results.imgs[0])  # Extract processed frame
-        out.write(frame)  # Write frame to output video
+        results = model(frame)  # Perform detection
+        annotated_frame = results[0].plot()  # Get annotated frame
+        out.write(annotated_frame)
 
     cap.release()
     out.release()
     return temp_video_path
 
-# Create Gradio Interface
-image_input = gr.inputs.Image(type="pil", label="Upload Image")
-video_input = gr.inputs.Video(type="file", label="Upload Video")
-
-image_output = gr.outputs.Image(type="pil", label="Detected Image")
-video_output = gr.outputs.Video(label="Detected Video")
-
+# Gradio Interface
 app = gr.Interface(
     fn=[detect_on_image, detect_on_video],
-    inputs=[image_input, video_input],
-    outputs=[image_output, video_output],
+    inputs=[gr.inputs.Image(type="pil", label="Upload Image"), gr.inputs.Video(type="file", label="Upload Video")],
+    outputs=[gr.outputs.Image(type="pil", label="Detected Image"), gr.outputs.Video(label="Detected Video")],
     title="WALDO30 YOLOv8 Object Detection",
-    description="Upload an image or video to see object detection results using WALDO30 YOLOv8 model."
+    description="Upload an image or video to see object detection results using the WALDO30 YOLOv8 model."
 )
 
-# Launch the app
 if __name__ == "__main__":
     app.launch()