Update pages/live_match.py

pages/live_match.py

@@ -1,74 +1,122 @@
-import gradio as gr
-import cv2
-import time
-import re
-from gully_drs_core.model_utils import load_yolov5_model, detect_objects
-from gully_drs_core.ball_detection import track_ball_and_estimate_speed, detect_lbw
-from gully_drs_core.replay_utils import generate_replay
-
-def validate_url(url):
-    """Validate live video URL."""
+def process_live_video(url, state, button_click):
+    """
+    Process a live video stream for LBW decision and ball speed.
+    url: String URL of the live video stream (RTSP/HTTP).
+    state: Gradio state for maintaining context.
+    button_click: Button state (triggered on DRS LBW click).
+    """
     try:
-        if not url:
-            return False, "No URL provided."
-        if not re.match(r'^https?://[a-zA-Z0-9:/._-]*$|^rtsp://[a-zA-Z0-9:/._-]*$', url):
-            return False, "Invalid URL format. Use HTTP/HTTPS or RTSP URLs."
-        return True, ""
-    except Exception as e:
-        return False, f"Error validating URL: {str(e)}"
-
-def process_live_video(live_url, state, drs_button):
-    """Process live video stream with DRS LBW functionality."""
-    try:
-        is_valid, error_msg = validate_url(live_url)
+        # Validate URL
+        is_valid, message = validate_url(url)
         if not is_valid:
-            yield None, error_msg, None, state
-            return
-        
-        cap = cv2.VideoCapture(live_url)
+            return None, message, "N/A", state
+
+        # Open the live stream
+        cap = cv2.VideoCapture(url)
         if not cap.isOpened():
-            yield None, "Unable to open live stream.", None, state
-            return
-        
-        model = load_yolov5_model()
+            return None, "Error: Could not open video stream", "N/A", state
+
+        # Get video properties
         frame_rate = cap.get(cv2.CAP_PROP_FPS) or 30
-        frames, detections, speeds, lbw_verdicts = [], [], [], []
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        slow_motion_factor = 4  # Maintain slow motion
+        out = cv2.VideoWriter("processed_video.mp4", fourcc, frame_rate / slow_motion_factor, (width, height))
+
+        # Initialize variables
         ball_positions = state.get("ball_positions", [])
-        start_time = state.get("start_time", None)
-        
-        if drs_button:
-            state["drs_triggered"] = True
-            state["start_time"] = time.time()
-            start_time = state["start_time"]
-        
-        while cap.isOpened():
+        ball_speed = 0.0
+        lbw_detected = False
+        frame_count = 0
+        prev_position = None
+        start_time = time.time() if button_click else state.get("start_time", None)
+        drs_review = "DRS Review: N/A"
+
+        while cap.isOpened() and (not button_click or time.time() - start_time < 5):
             ret, frame = cap.read()
             if not ret:
                 break
-            
-            if start_time and time.time() - start_time > 5:  # Trim to 5 seconds
-                break
-            
-            frames.append(frame)
-            results = detect_objects(frame, model)
-            detections.append(results.xyxy[0])
-            frame, speed, pos = track_ball_and_estimate_speed(frame, results.xyxy[0], frame_rate)
-            speeds.append(speed)
-            if pos:
-                ball_positions.append(pos)
-            lbw_verdict, ball_positions = detect_lbw(results.xyxy[0], frame.shape, ball_positions)
-            lbw_verdicts.append(lbw_verdict)
-            state["ball_positions"] = ball_positions
-            
-            yield frame, None, None, state
-        
+
+            # Refined ball detection (tighter HSV, size, aspect ratio, and height filters)
+            hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
+            lower_red = np.array([0, 150, 120])  # Tighter red range to exclude shoes
+            upper_red = np.array([10, 255, 255])
+            mask = cv2.inRange(hsv, lower_red, upper_red)
+            contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+            ball_detected = False
+            if contours:
+                largest_contour = max(contours, key=cv2.contourArea)
+                x, y, w, h = cv2.boundingRect(largest_contour)
+                aspect_ratio = float(w) / h
+                if 20 < cv2.contourArea(largest_contour) < 500 and 0.5 < aspect_ratio < 1.5 and y < height - 200:  # Exclude elongated shoes
+                    ball_position = (x + w // 2, y + h // 2)
+                    if not prev_position and frame_count < 10:  # Initial position
+                        prev_position = ball_position
+                    if prev_position and abs(ball_position[1] - prev_position[1]) > 2:  # Motion filter
+                        ball_positions.append(ball_position)
+                        cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 4)  # Thicker green
+                        cv2.putText(frame, "Ball", (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+                        ball_detected = True
+
+            # Refined wickets detection
+            stump_region = frame[height - 250:height, width // 2 - 120:width // 2 + 120]
+            gray_stump = cv2.cvtColor(stump_region, cv2.COLOR_BGR2GRAY)
+            _, stump_mask = cv2.threshold(gray_stump, 200, 255, cv2.THRESH_BINARY)
+            if np.sum(stump_mask) > 1200:
+                cv2.rectangle(frame, (width // 2 - 120, height - 250), (width // 2 + 120, height - 50), (0, 0, 255), 3)
+                cv2.putText(frame, "Wickets", (width // 2 - 120, height - 260), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
+
+            # Refined batsman leg pads detection
+            leg_pads_region = frame[height // 2 + 50:height, width // 3:2 * width // 3]
+            gray_leg_pads = cv2.cvtColor(leg_pads_region, cv2.COLOR_BGR2GRAY)
+            _, leg_pads_mask = cv2.threshold(gray_leg_pads, 90, 255, cv2.THRESH_BINARY_INV)
+            if np.sum(leg_pads_mask) > 4000:
+                cv2.rectangle(frame, (width // 3, height // 2 + 50), (2 * width // 3, height), (255, 0, 0), 3)
+                cv2.putText(frame, "Leg Pads", (width // 3, height // 2 + 40), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
+
+            # Ball speed estimation
+            if ball_detected and prev_position and frame_count > 1:
+                dx = ball_position[0] - prev_position[0]
+                dy = ball_position[1] - prev_position[1]
+                distance = np.sqrt(dx**2 + dy**2)
+                time_diff = 1 / frame_rate
+                ball_speed = min(200, max(20, distance / time_diff * 0.012))
+
+            # Enhanced LBW detection
+            if ball_detected and len(ball_positions) > 5:
+                last_five = ball_positions[-5:]
+                impact_wickets = any(width // 2 - 120 < pos[0] < width // 2 + 120 and height - 250 < pos[1] < height - 50 for pos in last_five)
+                impact_pads = any(width // 3 < pos[0] < 2 * width // 3 and height // 2 + 50 < pos[1] < height for pos in last_five)
+                if impact_wickets or (impact_pads and ball_speed > 40):
+                    lbw_detected = True
+                    drs_review = f"DRS Review: Ball impacting {'wickets' if impact_wickets else 'leg pads'} at {ball_position[0]}px, {ball_position[1]}px. Speed: {ball_speed:.2f} km/h. Decision: Out."
+                    cv2.putText(frame, drs_review[:50], (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1)
+                else:
+                    drs_review = "DRS Review: No LBW. Ball missed wickets and pads."
+
+            # Write frame to slow-motion video
+            if button_click and start_time:
+                for _ in range(slow_motion_factor):
+                    out.write(frame)
+
+            prev_position = ball_position
+            frame_count += 1
+
         cap.release()
-        
-        if state.get("drs_triggered", False):
-            replay_path = generate_replay(frames, detections, speeds, lbw_verdicts[-1])
-            yield replay_path, lbw_verdicts[-1], f"{speeds[-1]:.2f} km/h", state
-        else:
-            yield None, "Press DRS LBW to analyze.", None, state
-            
+        out.release()
+
+        # Check if any ball was detected
+        if not any(ball_positions):
+            return "processed_video.mp4", "No red ball detected", "N/A", state
+
+        lbw_decision = "Out" if lbw_detected else "Not Out"
+        lbw_decision = f"{lbw_decision}\n{drs_review}"
+        ball_speed_str = f"{ball_speed:.2f} km/h" if ball_speed > 0 else "N/A"
+
+        state["ball_positions"] = ball_positions
+        state["start_time"] = start_time
+
+        return "processed_video.mp4", lbw_decision, ball_speed_str, state
     except Exception as e:
-        yield None, f"Error: {str(e)}", None, state
+        return None, f"Error processing video: {str(e)}. Ensure video is MP4 with H.264 codec and contains a visible red ball.", "N/A", state