Update gully_drs_core/ball_detection.py

gully_drs_core/ball_detection.py

@@ -1,91 +1,36 @@
+# gully_drs_core/ball_detection.py
 import cv2
 import numpy as np
-import torch
-import logging
-import os
-from gully_drs_core.replay_utils import generate_replay
-from gully_drs_core.video_utils import get_video_properties
-from gully_drs_core.model_utils import load_yolo_model
+from .model_utils import load_model
 
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+def analyze_video(file_path):
+    model = load_model()
+    cap = cv2.VideoCapture(file_path)
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    width  = int(cap.get(3))
+    height = int(cap.get(4))
 
-# Load YOLOv5 model
-model = load_yolo_model()
+    ball_path = []
+    frames = []
 
-# Stump zone coordinates (example, adjust based on video resolution)
-STUMP_ZONE = [(200, 400), (300, 400), (300, 600), (200, 600)]  # [x1,y1, x2,y2, x3,y3, x4,y4]
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        results = model(frame)
+        for r in results:
+            for box in r.boxes:
+                cls = int(box.cls[0])
+                if cls == 32:  # Assuming class 32 = cricket ball
+                    x1, y1, x2, y2 = map(int, box.xyxy[0])
+                    cx, cy = (x1 + x2) // 2, (y1 + y2) // 2
+                    ball_path.append((cx, cy))
+                    cv2.circle(frame, (cx, cy), 6, (0, 255, 0), -1)
+        frames.append(frame)
 
-def process_video(video_path):
-    try:
-        # Validate video file
-        if not os.path.exists(video_path):
-            return {"status": "error", "error": "Video file not found"}
-        
-        # Get video properties
-        fps, width, height = get_video_properties(video_path)
-        if fps == 0:
-            return {"status": "error", "error": "Invalid video file"}
-        
-        # Initialize video capture
-        cap = cv2.VideoCapture(video_path)
-        ball_positions = []
-        bounce_point = None
-        frame_count = 0
-        
-        while cap.isOpened():
-            ret, frame = cap.read()
-            if not ret:
-                break
-            
-            # Detect ball using YOLOv5
-            results = model(frame)
-            detections = results.xyxy[0].cpu().numpy()  # [x1, y1, x2, y2, conf, class]
-            
-            ball_center = None
-            for det in detections:
-                if det[5] == 0:  # Assuming class 0 is the ball
-                    x1, y1, x2, y2 = map(int, det[:4])
-                    ball_center = ((x1 + x2) // 2, (y1 + y2) // 2)
-                    ball_positions.append(ball_center)
-                    # Detect bounce point (simplified: assume bounce near ground)
-                    if y1 > height * 0.8 and bounce_point is None:
-                        bounce_point = ball_center
-                    break
-            
-            frame_count += 1
-        
-        cap.release()
-        
-        # Check LBW decision
-        decision = "Not Out"
-        for pos in ball_positions:
-            if is_ball_in_stump_zone(pos):
-                decision = "Out"
-                break
-        
-        # Calculate speed (pixels per frame to km/h)
-        speed_kmh = 0
-        if len(ball_positions) >= 2:
-            pixel_dist = np.sqrt((ball_positions[-1][0] - ball_positions[-2][0])**2 +
-                                 (ball_positions[-1][1] - ball_positions[-2][1])**2)
-            speed_kmh = (pixel_dist / (1/fps)) * 0.036  # Simplified conversion, adjust scale factor
-        
-        # Generate replay video
-        replay_path = generate_replay(video_path, ball_positions, STUMP_ZONE, decision, speed_kmh, bounce_point)
-        
-        return {
-            "status": "success",
-            "decision": decision,
-            "speed_kmh": speed_kmh,
-            "replay_path": replay_path
-        }
-    except Exception as e:
-        logger.error(f"Error processing video: {str(e)}")
-        return {"status": "error", "error": str(e)}
-
-def is_ball_in_stump_zone(ball_center):
-    x, y = ball_center
-    x1, y1, x2, y2, x3, y3, x4, y4 = STUMP_ZONE[0] + STUMP_ZONE[1] + STUMP_ZONE[2] + STUMP_ZONE[3]
-    return (x1 <= x <= x2) and (y1 <= y <= y3)
+    cap.release()
+    return {
+        "trajectory": ball_path,
+        "fps": fps,
+        "frames": frames
+    }