Update app.py

app.py

@@ -3,69 +3,72 @@ import numpy as np
 import gradio as gr
 from ultralytics import YOLO
 import torch
-from typing import List, Dict
 import os
+import tempfile
+from typing import Dict, List
 
-# Verify PyTorch and CUDA
-print(f"PyTorch version: {torch.__version__}")
-print(f"CUDA available: {torch.cuda.is_available()}")
-
-# Initialize models with safe loading
-def load_model(model_path: str, model_type: str = "ball"):
-    """Safely load YOLO model with verification"""
+# Initialize models safely
+def load_model(model_name: str):
     try:
-        # For Hugging Face Spaces - use cached models
-        if not os.path.exists(model_path):
-            model = YOLO(f"yolov8{model_type}.pt")  # Auto-download
-            torch.save(model.model.state_dict(), model_path)
-        else:
-            model = YOLO(model_path)
-        
-        # Verify model
-        if model.predict(np.zeros((640, 640, 3)), verbose=False)[0].boxes is not None:
-            print(f"{model_type.upper()} model loaded successfully!")
+        model = YOLO(model_name)
+        # Test model with dummy data
+        dummy_result = model(np.zeros((640, 640, 3)), verbose=False)
+        if dummy_result[0].boxes is not None:
+            print(f"✅ {model_name} loaded successfully!")
             return model
-        else:
-            raise RuntimeError("Model verification failed")
+        raise RuntimeError("Model test failed")
     except Exception as e:
-        raise RuntimeError(f"Model loading error: {str(e)}")
+        print(f"❌ Model load error: {str(e)}")
+        return None
 
-# Load models (will auto-download if not present)
-try:
-    BALL_MODEL = load_model("yolov8n.pt", "n")
-    STUMP_MODEL = load_model("yolov8m.pt", "m")
-except Exception as e:
-    print(f"Critical error: {str(e)}")
-    # Fallback to CPU-only basic detection
-    BALL_MODEL = None
+BALL_MODEL = load_model("yolov8n.pt")
 
 def process_video(video_path: str) -> Dict:
-    """Robust video processing with error handling"""
+    """Robust video processing with full error handling"""
     try:
+        # Handle Gradio's temporary file path
+        if isinstance(video_path, dict):
+            video_path = video_path["name"]
+        
+        # Verify file exists
+        if not os.path.exists(video_path):
+            raise FileNotFoundError(f"Video file not found: {video_path}")
+        
         cap = cv2.VideoCapture(video_path)
         if not cap.isOpened():
             raise ValueError("Could not open video file")
         
+        # Get video properties
         fps = cap.get(cv2.CAP_PROP_FPS)
-        frames = []
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        
+        # Prepare output
+        output_frames = []
         analytics = {
-            "max_speed": 0,
+            "max_speed": 0.0,
             "events": [],
-            "status": "Success"
+            "status": "Processing...",
+            "fps": fps,
+            "resolution": f"{width}x{height}"
         }
         
         prev_pos = None
+        frame_count = 0
         
-        while cap.isOpened():
+        while True:
             ret, frame = cap.read()
             if not ret:
                 break
             
+            frame_count += 1
+            
+            # Resize for consistent processing (optional)
             frame = cv2.resize(frame, (1280, 720))
             
-            # Ball detection (only if model loaded)
+            # Ball detection
             if BALL_MODEL:
-                results = BALL_MODEL(frame, classes=32, verbose=False)  # Class 32 = sports ball
+                results = BALL_MODEL(frame, classes=32, verbose=False)
                 boxes = results[0].boxes.xyxy.cpu().numpy()
                 
                 if len(boxes) > 0:
@@ -74,55 +77,102 @@ def process_video(video_path: str) -> Dict:
                     
                     # Speed calculation
                     if prev_pos:
-                        speed = np.sqrt((x - prev_pos[0])**2 + (y - prev_pos[1])**2) * fps * 3.6 / 100
+                        px_per_meter = 100  # Calibration needed
+                        speed = np.sqrt((x - prev_pos[0])**2 + (y - prev_pos[1])**2) * fps * 3.6 / px_per_meter
                         analytics["max_speed"] = max(analytics["max_speed"], speed)
+                        
+                        # Visualize
+                        cv2.circle(frame, (x, y), 10, (0, 255, 0), -1)
                         cv2.putText(frame, f"{speed:.1f} km/h", (x+15, y), 
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2)
                     
                     prev_pos = (x, y)
-                    cv2.circle(frame, (x, y), 10, (0, 255, 0), -1)
             
-            frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+            # Convert frame to RGB for Gradio
+            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            output_frames.append(frame_rgb)
+            
+            # Update status every 10 frames
+            if frame_count % 10 == 0:
+                analytics["status"] = f"Processed {frame_count} frames"
         
         cap.release()
-        return {"frames": frames, "analytics": analytics}
+        
+        # Handle empty output
+        if not output_frames:
+            raise ValueError("No frames processed - check video format")
+        
+        # Save output as temporary video file
+        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp:
+            out = cv2.VideoWriter(tmp.name, cv2.VideoWriter_fourcc(*'mp4v'), fps, (1280, 720))
+            for frame in output_frames:
+                out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
+            out.release()
+            
+            analytics["status"] = "✅ Processing complete"
+            return {
+                "output_video": tmp.name,
+                "analytics": analytics
+            }
     
     except Exception as e:
         return {
-            "frames": [],
+            "output_video": None,
             "analytics": {
-                "status": f"Error: {str(e)}",
-                "max_speed": 0,
-                "events": []
+                "status": f"❌ Error: {str(e)}",
+                "max_speed": 0.0,
+                "events": [],
+                "fps": 0,
+                "resolution": "0x0"
             }
         }
 
 # Gradio Interface
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🏏 Cricket Ball Tracker (Fixed Version)")
+    gr.Markdown("""
+    # 🏏 Professional Cricket Tracker
+    *Ball Tracking • Speed Analysis • Event Detection*
+    """)
     
     with gr.Row():
-        input_video = gr.Video(label="Input Video", format="mp4")
-        output_video = gr.Video(label="Tracking Results")
+        with gr.Column():
+            input_video = gr.Video(label="Input Match Footage", format="mp4")
+            gr.Examples(
+                examples=["sample.mp4"],
+                inputs=input_video,
+                label="Try Sample Video"
+            )
+        
+        with gr.Column():
+            output_video = gr.Video(label="Tracking Results", format="mp4")
     
-    with gr.Accordion("Advanced Info", open=False):
-        status = gr.Textbox(label="Processing Status")
-        max_speed = gr.Number(label="Max Ball Speed (km/h)")
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("### 📊 Match Analytics")
+            max_speed = gr.Number(label="Max Ball Speed (km/h)", precision=1)
+            resolution = gr.Textbox(label="Video Resolution")
+        
+        with gr.Column():
+            gr.Markdown("### 📝 Processing Info")
+            status = gr.Textbox(label="Status")
+            fps = gr.Number(label="Video FPS")
     
-    analyze_btn = gr.Button("Analyze", variant="primary")
+    analyze_btn = gr.Button("Analyze Video", variant="primary")
     
     def analyze_wrapper(video):
         result = process_video(video)
         return {
-            output_video: result["frames"],
+            output_video: result["output_video"],
+            max_speed: result["analytics"]["max_speed"],
+            resolution: result["analytics"]["resolution"],
             status: result["analytics"]["status"],
-            max_speed: result["analytics"]["max_speed"]
+            fps: result["analytics"]["fps"]
         }
     
     analyze_btn.click(
         fn=analyze_wrapper,
         inputs=input_video,
-        outputs=[output_video, status, max_speed]
+        outputs=[output_video, max_speed, resolution, status, fps]
     )
 
 if __name__ == "__main__":