Update app.py

app.py

@@ -3,177 +3,169 @@ import numpy as np
 import gradio as gr
 from ultralytics import YOLO
 import torch
-import os
+from scipy.interpolate import interp1d
 import tempfile
-from typing import Dict, List
+from typing import Dict, Tuple
 
-# Initialize models safely
-def load_model(model_name: str):
-    try:
-        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
-        raise RuntimeError("Model test failed")
-    except Exception as e:
-        print(f"❌ Model load error: {str(e)}")
-        return None
+# Initialize models
+BALL_MODEL = YOLO('yolov8n.pt')  # Auto-downloads if not present
+STUMP_MODEL = YOLO('yolov8m.pt')
 
-BALL_MODEL = load_model("yolov8n.pt")
+# Pitch constants (in pixels)
+PITCH_LENGTH = 2000  # From crease to crease
+STUMPS_WIDTH = 71    # Standard cricket stump width (9 inches)
+
+def predict_trajectory(positions: list) -> Tuple[list, float]:
+    """Predict ball trajectory using cubic 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
+
+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"
+    
+    return {
+        "decision": decision,
+        "hitting": hitting,
+        "impact": "Impact" if decision == "Out" else "No Impact",
+        "in_line": in_line,
+        "pitching": pitching
+    }
 
 def process_video(video_path: str) -> Dict:
-    """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)
-        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.0,
-            "events": [],
-            "status": "Processing...",
-            "fps": fps,
-            "resolution": f"{width}x{height}"
-        }
+    """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
         
-        prev_pos = None
-        frame_count = 0
+        # Ball detection
+        ball_results = BALL_MODEL(frame, classes=32, verbose=False)
+        boxes = ball_results[0].boxes.xyxy.cpu().numpy()
         
-        while True:
-            ret, frame = cap.read()
-            if not ret:
-                break
+        if len(boxes) > 0:
+            x1, y1, x2, y2 = boxes[0]
+            x, y = (x1 + x2) // 2, (y1 + y2) // 2
+            ball_positions.append((x, y))
             
-            frame_count += 1
+            # Predict trajectory
+            trajectory, speed = predict_trajectory(ball_positions[-10:])
+            max_speed = max(max_speed, speed)
             
-            # Resize for consistent processing (optional)
-            frame = cv2.resize(frame, (1280, 720))
+            # 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 = map(int, boxes[0])
-                    x, y = (x1 + x2) // 2, (y1 + y2) // 2
-                    
-                    # Speed calculation
-                    if prev_pos:
-                        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)
-            
-            # 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"
+            # 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])
         
-        cap.release()
+        # 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)
         
-        # Handle empty output
-        if not output_frames:
-            raise ValueError("No frames processed - check video format")
+        # Draw speed
+        cv2.putText(frame, f"Speed: {max_speed:.1f} km/h", (width-300, 50), 
+                   cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 0), 2)
         
-        # 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
-            }
+        out.write(frame)
     
-    except Exception as e:
-        return {
-            "output_video": None,
-            "analytics": {
-                "status": f"❌ Error: {str(e)}",
-                "max_speed": 0.0,
-                "events": [],
-                "fps": 0,
-                "resolution": "0x0"
-            }
-        }
+    cap.release()
+    out.release()
+    
+    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 {}
+    }
 
 # Gradio Interface
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("""
-    # 🏏 Professional Cricket Tracker
-    *Ball Tracking • Speed Analysis • Event Detection*
-    """)
+    gr.Markdown("# 🏏 Professional DRS System")
     
     with gr.Row():
-        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")
+        input_video = gr.Video(label="Input Match Footage", format="mp4")
+        output_video = gr.Video(label="DRS Analysis", format="mp4")
     
     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")
+            gr.Markdown("### 📊 Decision Review System")
+            decision = gr.Textbox(label="Final Decision")
+            hitting = gr.Textbox(label="Hitting")
+            impact = gr.Textbox(label="Impact")
         
         with gr.Column():
-            gr.Markdown("### 📝 Processing Info")
-            status = gr.Textbox(label="Status")
-            fps = gr.Number(label="Video FPS")
+            gr.Markdown("### 📝 Ball Tracking")
+            max_speed = gr.Number(label="Max Speed (km/h)", precision=1)
+            in_line = gr.Textbox(label="In Line")
+            pitching = gr.Textbox(label="Pitching")
     
-    analyze_btn = gr.Button("Analyze Video", variant="primary")
+    analyze_btn = gr.Button("Run DRS Analysis", variant="primary")
     
     def analyze_wrapper(video):
         result = process_video(video)
         return {
             output_video: result["output_video"],
-            max_speed: result["analytics"]["max_speed"],
-            resolution: result["analytics"]["resolution"],
-            status: result["analytics"]["status"],
-            fps: result["analytics"]["fps"]
+            decision: result["lbw_decision"],
+            max_speed: result["max_speed"],
+            hitting: result["analytics"].get("hitting", "N/A"),
+            impact: result["analytics"].get("impact", "N/A"),
+            in_line: result["analytics"].get("in_line", "N/A"),
+            pitching: result["analytics"].get("pitching", "N/A")
         }
     
     analyze_btn.click(
         fn=analyze_wrapper,
         inputs=input_video,
-        outputs=[output_video, max_speed, resolution, status, fps]
+        outputs=[output_video, decision, max_speed, hitting, impact, in_line, pitching]
     )
 
-if __name__ == "__main__":
-    demo.launch(debug=True)
+demo.launch()