Update app.py

app.py

@@ -5,49 +5,65 @@ from ultralytics import YOLO
 import torch
 from scipy.interpolate import interp1d
 import tempfile
-from typing import Dict, Tuple
+import os
+from typing import Dict, Tuple, Optional
 
-# Initialize models
-BALL_MODEL = YOLO('yolov8n.pt')  # Auto-downloads if not present
-STUMP_MODEL = YOLO('yolov8m.pt')
+# Initialize models with error handling
+def safe_load_model(model_name: str):
+    try:
+        model = YOLO(model_name)
+        # Verify model works
+        dummy = model(np.zeros((640,640,3)), verbose=False)
+        if dummy[0].boxes is not None:
+            print(f"✅ {model_name} loaded successfully")
+            return model
+    except Exception as e:
+        print(f"❌ Error loading {model_name}: {str(e)}")
+    return None
 
-# Pitch constants (in pixels)
-PITCH_LENGTH = 2000  # From crease to crease
-STUMPS_WIDTH = 71    # Standard cricket stump width (9 inches)
+BALL_MODEL = safe_load_model("yolov8n.pt")  # Ball detection
+STUMP_MODEL = safe_load_model("yolov8m.pt")  # Stump detection
+
+# Constants (adjust based on your video)
+PITCH_LENGTH_PX = 1800  # Pixels between creases
+STUMPS_WIDTH_PX = 60    # Width of stumps in pixels
+FRAME_SKIP = 2          # Process every 2nd frame for speed
 
 def predict_trajectory(positions: list) -> Tuple[list, float]:
-    """Predict ball trajectory using cubic interpolation"""
+    """Predict ball path with cubic spline interpolation"""
     if len(positions) < 3:
         return positions, 0.0
     
     x = [p[0] for p in positions]
     y = [p[1] for p in positions]
-    
-    # Cubic spline interpolation
     t = np.linspace(0, 1, len(positions))
-    fx = interp1d(t, x, kind='cubic', fill_value="extrapolate")
-    fy = interp1d(t, y, kind='cubic', fill_value="extrapolate")
-    
-    # Predict next 10 frames
-    new_t = np.linspace(0, 1.5, len(positions)+10)
-    new_x = fx(new_t)
-    new_y = fy(new_t)
     
-    # Calculate speed (pixels/frame to km/h)
-    dx = new_x[-1] - new_x[-2]
-    dy = new_y[-1] - new_y[-2]
-    speed = np.sqrt(dx**2 + dy**2) * 25 * 3.6 / PITCH_LENGTH  # Convert to km/h
-    
-    return list(zip(new_x, new_y)), speed
+    try:
+        fx = interp1d(t, x, kind='cubic', fill_value="extrapolate")
+        fy = interp1d(t, y, kind='cubic', fill_value="extrapolate")
+        new_t = np.linspace(0, 1.5, len(positions)+10)
+        new_x = fx(new_t)
+        new_y = fy(new_t)
+        
+        # Calculate speed (km/h)
+        dx = new_x[-1] - new_x[-2]
+        dy = new_y[-1] - new_y[-2]
+        speed = np.sqrt(dx**2 + dy**2) * 25 * 3.6 / PITCH_LENGTH_PX
+        return list(zip(new_x, new_y)), speed
+    except:
+        return positions, 0.0
 
-def check_lbw(ball_pos: tuple, stump_pos: tuple, impact: tuple) -> Dict:
-    """LBW decision system"""
-    # Simplified decision logic
-    hitting = "Hitting" if abs(ball_pos[0] - stump_pos[0]) < STUMPS_WIDTH else "Missing"
-    in_line = "In Line" if impact[0] < stump_pos[0] + STUMPS_WIDTH else "Not in Line"
-    pitching = "In Line" if impact[1] < stump_pos[1] + 100 else "Outside"
-    
-    decision = "Out" if all([hitting == "Hitting", in_line == "In Line", pitching == "In Line"]) else "Not Out"
+def check_lbw(ball_pos: tuple, stump_pos: tuple) -> Dict:
+    """Make LBW decision with all parameters"""
+    hitting = "Hitting" if abs(ball_pos[0] - stump_pos[0]) < STUMPS_WIDTH_PX else "Missing"
+    in_line = "In Line" if ball_pos[0] < stump_pos[0] + STUMPS_WIDTH_PX else "Not in Line"
+    pitching = "In Line" if ball_pos[1] < stump_pos[1] + 100 else "Outside"
+    
+    decision = "Out" if all([
+        hitting == "Hitting",
+        in_line == "In Line",
+        pitching == "In Line"
+    ]) else "Not Out"
     
     return {
         "decision": decision,
@@ -57,79 +73,127 @@ def check_lbw(ball_pos: tuple, stump_pos: tuple, impact: tuple) -> Dict:
         "pitching": pitching
     }
 
-def process_video(video_path: str) -> Dict:
-    """Main processing function"""
-    cap = cv2.VideoCapture(video_path)
-    fps = cap.get(cv2.CAP_PROP_FPS)
-    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    
-    # Video writer setup
-    temp_file = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False)
-    out = cv2.VideoWriter(temp_file.name, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
-    
-    ball_positions = []
-    lbw_data = None
-    max_speed = 0.0
-    
-    while cap.isOpened():
-        ret, frame = cap.read()
-        if not ret:
-            break
+def process_video(video_input) -> Dict:
+    """Main processing function with full error handling"""
+    try:
+        # Handle Gradio file input
+        if isinstance(video_input, dict):
+            video_path = video_input["name"]
+        else:
+            video_path = video_input
+        
+        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")
         
-        # Ball detection
-        ball_results = BALL_MODEL(frame, classes=32, verbose=False)
-        boxes = ball_results[0].boxes.xyxy.cpu().numpy()
+        # Get video properties
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
         
-        if len(boxes) > 0:
-            x1, y1, x2, y2 = boxes[0]
-            x, y = (x1 + x2) // 2, (y1 + y2) // 2
-            ball_positions.append((x, y))
+        # Create temp output file
+        temp_path = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False).name
+        out = cv2.VideoWriter(
+            temp_path,
+            cv2.VideoWriter_fourcc(*'mp4v'),
+            fps/FRAME_SKIP,  # Adjusted for frame skipping
+            (width, height)
+        )
+        
+        ball_positions = []
+        lbw_data = None
+        max_speed = 0.0
+        frame_count = 0
+        
+        while True:
+            ret, frame = cap.read()
+            if not ret:
+                break
             
-            # Predict trajectory
-            trajectory, speed = predict_trajectory(ball_positions[-10:])
-            max_speed = max(max_speed, speed)
+            frame_count += 1
+            if frame_count % FRAME_SKIP != 0:
+                continue
             
-            # Draw trajectory
-            for i in range(1, len(trajectory)):
-                cv2.line(frame, tuple(map(int, trajectory[i-1])), tuple(map(int, trajectory[i])), (0, 255, 255), 2)
+            # Ball detection
+            if BALL_MODEL:
+                results = BALL_MODEL(frame, classes=32, verbose=False)
+                boxes = results[0].boxes.xyxy.cpu().numpy()
+                
+                if len(boxes) > 0:
+                    x1, y1, x2, y2 = boxes[0]
+                    x, y = (x1 + x2) // 2, (y1 + y2) // 2
+                    ball_positions.append((x, y))
+                    
+                    # Predict trajectory and speed
+                    trajectory, speed = predict_trajectory(ball_positions[-10:])
+                    max_speed = max(max_speed, speed)
+                    
+                    # Draw trajectory
+                    for i in range(1, len(trajectory)):
+                        cv2.line(
+                            frame,
+                            tuple(map(int, trajectory[i-1])),
+                            tuple(map(int, trajectory[i])),
+                            (0, 255, 255), 2
+                        )
+                    
+                    # LBW check (every 5 processed frames)
+                    if len(ball_positions) % 5 == 0 and STUMP_MODEL:
+                        stump_results = STUMP_MODEL(frame, classes=33, verbose=False)
+                        if len(stump_results[0].boxes) > 0:
+                            sx1, sy1, sx2, sy2 = stump_results[0].boxes.xyxy[0].cpu().numpy()
+                            stump_pos = ((sx1 + sx2) // 2, (sy1 + sy2) // 2)
+                            lbw_data = check_lbw((x, y), stump_pos)
             
-            # LBW check (every 5 frames)
-            if len(ball_positions) % 5 == 0:
-                stump_results = STUMP_MODEL(frame, classes=33, verbose=False)
-                if len(stump_results[0].boxes) > 0:
-                    sx1, sy1, sx2, sy2 = stump_results[0].boxes.xyxy[0].cpu().numpy()
-                    stump_pos = ((sx1 + sx2) // 2, (sy1 + sy2) // 2)
-                    lbw_data = check_lbw((x, y), stump_pos, ball_positions[-1])
-        
-        # Draw DRS overlay
-        if lbw_data:
-            cv2.putText(frame, f"Final Decision: {lbw_data['decision']}", (50, 50), 
-                       cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255) if lbw_data['decision'] == "Out" else (0, 255, 0), 3)
-            cv2.putText(frame, f"Hitting: {lbw_data['hitting']}", (50, 100), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2)
-            cv2.putText(frame, f"Impact: {lbw_data['impact']}", (50, 140), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2)
-            cv2.putText(frame, f"In Line: {lbw_data['in_line']}", (50, 180), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2)
-            cv2.putText(frame, f"Pitching: {lbw_data['pitching']}", (50, 220), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2)
+            # Draw DRS overlay
+            if lbw_data:
+                cv2.putText(
+                    frame, f"Final Decision: {lbw_data['decision']}", 
+                    (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1,
+                    (0, 0, 255) if lbw_data['decision'] == "Out" else (0, 255, 0), 3
+                )
+                cv2.putText(frame, f"Hitting: {lbw_data['hitting']}", (50, 100), 
+                          cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2)
+                cv2.putText(frame, f"Impact: {lbw_data['impact']}", (50, 140), 
+                          cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2)
+                cv2.putText(frame, f"In Line: {lbw_data['in_line']}", (50, 180), 
+                          cv2.FERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2)
+                cv2.putText(frame, f"Pitching: {lbw_data['pitching']}", (50, 220), 
+                          cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2)
+            
+            # Draw speed
+            cv2.putText(
+                frame, f"Speed: {max_speed:.1f} km/h", 
+                (width-300, 50), cv2.FONT_HERSHEY_SIMPLEX, 
+                1, (255, 255, 0), 2
+            )
+            
+            out.write(frame)
         
-        # Draw speed
-        cv2.putText(frame, f"Speed: {max_speed:.1f} km/h", (width-300, 50), 
-                   cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 0), 2)
+        cap.release()
+        out.release()
         
-        out.write(frame)
-    
-    cap.release()
-    out.release()
+        return {
+            "output_video": temp_path,
+            "lbw_decision": lbw_data["decision"] if lbw_data else "No Decision",
+            "max_speed": max_speed,
+            "analytics": lbw_data if lbw_data else {}
+        }
     
-    return {
-        "output_video": temp_file.name,
-        "lbw_decision": lbw_data["decision"] if lbw_data else "No Decision",
-        "max_speed": max_speed,
-        "analytics": lbw_data if lbw_data else {}
-    }
+    except Exception as e:
+        return {
+            "output_video": None,
+            "lbw_decision": f"Error: {str(e)}",
+            "max_speed": 0.0,
+            "analytics": {}
+        }
 
 # Gradio Interface
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🏏 Professional DRS System")
+    gr.Markdown("# 🏏 Professional Cricket DRS System")
     
     with gr.Row():
         input_video = gr.Video(label="Input Match Footage", format="mp4")
@@ -137,7 +201,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     
     with gr.Row():
         with gr.Column():
-            gr.Markdown("### 📊 Decision Review System")
+            gr.Markdown("### 📊 Decision Review")
             decision = gr.Textbox(label="Final Decision")
             hitting = gr.Textbox(label="Hitting")
             impact = gr.Textbox(label="Impact")
@@ -150,7 +214,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     
     analyze_btn = gr.Button("Run DRS Analysis", variant="primary")
     
-    def analyze_wrapper(video):
+    def wrapper(video):
         result = process_video(video)
         return {
             output_video: result["output_video"],
@@ -163,9 +227,10 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         }
     
     analyze_btn.click(
-        fn=analyze_wrapper,
+        fn=wrapper,
         inputs=input_video,
         outputs=[output_video, decision, max_speed, hitting, impact, in_line, pitching]
     )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch(debug=True)