Update app.py

app.py

@@ -4,40 +4,59 @@ from PIL import Image
 import numpy as np
 from ultralytics import YOLO
 from huggingface_hub import hf_hub_download
+import os
+
+# Verify that Hugging Face repo and file paths are correct
+REPO_ID = "StephanST/WALDO30"  # Update if the repository ID is different
+MODEL_FILENAME = "WALDO30_yolov8m_640x640.pt"
 
 # Download the model from Hugging Face
-model_path = hf_hub_download(repo_id="StephanST/WALDO30", filename="WALDO30_yolov8m_640x640.pt")
-model = YOLO(model_path)  # Load YOLOv8 model
+try:
+    model_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILENAME)
+except Exception as e:
+    raise RuntimeError(f"Failed to download model from Hugging Face. Verify `repo_id` and `filename`. Error: {e}")
+
+# Load the YOLOv8 model
+try:
+    model = YOLO(model_path)  # Ensure the model path is correct
+except Exception as e:
+    raise RuntimeError(f"Failed to load the YOLO model. Verify the model file at `{model_path}`. Error: {e}")
 
 # Detection function for images
 def detect_on_image(image):
-    results = model(image)  # Perform detection
-    annotated_frame = results[0].plot()  # Get annotated image
-    return Image.fromarray(annotated_frame)
+    try:
+        results = model(image)  # Perform detection
+        annotated_frame = results[0].plot()  # Get annotated image
+        return Image.fromarray(annotated_frame)
+    except Exception as e:
+        raise RuntimeError(f"Error during image processing: {e}")
 
 # 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),
-                          (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)  # Perform detection
-        annotated_frame = results[0].plot()  # Get annotated frame
-        out.write(annotated_frame)
-
-    cap.release()
-    out.release()
-    return temp_video_path
+    try:
+        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),
+                              (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)  # Perform detection
+            annotated_frame = results[0].plot()  # Get annotated frame
+            out.write(annotated_frame)
+
+        cap.release()
+        out.release()
+        return temp_video_path
+    except Exception as e:
+        raise RuntimeError(f"Error during video processing: {e}")
 
 # Gradio Interface
 image_input = gr.Image(type="pil", label="Upload Image")
-video_input = gr.Video(type="file", label="Upload Video")
+video_input = gr.Video(label="Upload Video")  # Removed invalid `type` argument
 image_output = gr.Image(type="pil", label="Detected Image")
 video_output = gr.Video(label="Detected Video")