Update app.py

app.py

@@ -5,43 +5,45 @@ import cv2
 import os
 import torch
 
-# ✅ Auto-select device
-device = "cuda" if torch.cuda.is_available() else "cpu"
+
 model = YOLO("best.pt").to(device)
 
-# ✅ Image detection
+# Image prediction
 def predict_image(image):
     results = model.predict(image, imgsz=640, conf=0.3, verbose=False)
     return results[0].plot()
 
-# ✅ Video detection
+# Video prediction
 def predict_video(video_path):
     cap = cv2.VideoCapture(video_path)
     fps = cap.get(cv2.CAP_PROP_FPS) or 25
-    frame_width, frame_height = 640, 360  # ✅ Fixed size for performance
+    width, height = 640, 360
 
-    # ✅ Output file path
     fd, output_path = tempfile.mkstemp(suffix=".mp4")
     os.close(fd)
+
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height))
+    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
 
-    # ✅ Fast streaming detection
     for result in model.track(source=video_path, stream=True, imgsz=640, conf=0.3, verbose=False):
         if result.orig_img is None:
             continue
-        resized = cv2.resize(result.orig_img, (frame_width, frame_height))
+        frame = cv2.resize(result.orig_img, (width, height))
         annotated = result.plot()
         out.write(annotated)
 
     cap.release()
     out.release()
 
+    # Check if file is valid
+    if not os.path.exists(output_path) or os.path.getsize(output_path) == 0:
+        raise RuntimeError("Output video file is missing or corrupted.")
+
     return output_path
 
-# ✅ Gradio Interface
+# Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# 🚀 YOLOv8 Detection\nOptimized for Speed and Accuracy")
+    gr.Markdown("# YOLOv8 Object Detection\nFast & Accurate for Images and Videos")
 
     with gr.Tab("Image"):
         img_input = gr.Image(type="pil")