Update app.py

app.py

@@ -5,96 +5,192 @@ import cv2
 import numpy as np
 from PIL import Image
 import os
+import matplotlib.pyplot as plt
+from scipy.interpolate import interp1d
 
-# Load the YOLOv5 model
+# Load YOLOv5 model
 model = YOLO("best.pt")
 
-def detect_ball(input_media, conf_threshold=0.5, iou_threshold=0.5):
+class CentroidTracker:
+    def __init__(self, max_disappeared=50):
+        self.next_object_id = 0
+        self.objects = {}
+        self.disappeared = {}
+        self.max_disappeared = max_disappeared
+
+    def register(self, centroid):
+        self.objects[self.next_object_id] = centroid
+        self.disappeared[self.next_object_id] = 0
+        self.next_object_id += 1
+
+    def deregister(self, object_id):
+        del self.objects[object_id]
+        del self.disappeared[object_id]
+
+    def update(self, rects):
+        if len(rects) == 0:
+            for object_id in list(self.disappeared.keys()):
+                self.disappeared[object_id] += 1
+                if self.disappeared[object_id] > self.max_disappeared:
+                    self.deregister(object_id)
+            return self.objects
+
+        input_centroids = np.zeros((len(rects), 2), dtype="int")
+        for (i, (x1, y1, x2, y2)) in enumerate(rects):
+            cX = int((x1 + x2) / 2.0)
+            cY = int((y1 + y2) / 2.0)
+            input_centroids[i] = (cX, cY)
+
+        if len(self.objects) == 0:
+            for i in range(len(input_centroids)):
+                self.register(input_centroids[i])
+        else:
+            object_ids = list(self.objects.keys())
+            object_centroids = list(self.objects.values())
+            D = np.sqrt(((input_centroids[:, None] - object_centroids) ** 2).sum(axis=2))
+            rows = D.min(axis=1).argsort()
+            cols = D.argmin(axis=1)[rows]
+            used_rows = set()
+            used_cols = set()
+            for (row, col) in zip(rows, cols):
+                if row in used_rows or col in used_cols:
+                    continue
+                object_id = object_ids[col]
+                self.objects[object_id] = input_centroids[row]
+                self.disappeared[object_id] = 0
+                used_rows.add(row)
+                used_cols.add(col)
+            unused_rows = set(range(0, D.shape[0])).difference(used_rows)
+            unused_cols = set(range(0, D.shape[1])).difference(used_cols)
+            if D.shape[0] >= D.shape[1]:
+                for row in unused_rows:
+                    self.register(input_centroids[row])
+            else:
+                for col in unused_cols:
+                    object_id = object_ids[col]
+                    self.disappeared[object_id] += 1
+                    if self.disappeared[object_id] > self.max_disappeared:
+                        self.deregister(object_id)
+        return self.objects
+
+def detect_and_track_ball(video_path, conf_threshold=0.5, iou_threshold=0.5):
     """
-    Perform ball detection on image or video input.
+    Detect and track ball in video, generate pitch map, and predict LBW outcome.
     
     Args:
-        input_media: Uploaded image or video file
+        video_path: Path to uploaded video
         conf_threshold: Confidence threshold for detection
         iou_threshold: IoU threshold for non-max suppression
     
     Returns:
-        Annotated image or video path
+        Tuple of (annotated video path, pitch map image path, LBW decision)
     """
-    # Check if input is image or video based on file extension
-    file_extension = os.path.splitext(input_media)[1].lower()
+    # Initialize tracker
+    tracker = CentroidTracker(max_disappeared=10)
+    cap = cv2.VideoCapture(video_path)
+    output_path = "output_video.mp4"
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter(output_path, fourcc, 30.0, 
+                         (int(cap.get(3)), int(cap.get(4))))
     
-    if file_extension in ['.jpg', '.jpeg', '.png']:
-        # Process image
-        img = cv2.imread(input_media)
-        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-        
-        # Perform detection
-        results = model.predict(img, conf=conf_threshold, iou=iou_threshold)
+    # Store ball centroids for trajectory
+    centroids = []
+    pitch_points = []
+    
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+            
+        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        results = model.predict(frame_rgb, conf=conf_threshold, iou=iou_threshold)
         
-        # Draw bounding boxes
+        rects = []
         for box in results[0].boxes:
             x1, y1, x2, y2 = map(int, box.xyxy[0])
             conf = box.conf[0]
             label = f"Ball: {conf:.2f}"
-            cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
-            cv2.putText(img, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
+            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
+            cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
+            rects.append((x1, y1, x2, y2))
         
-        # Convert to PIL Image for Gradio output
-        output_img = Image.fromarray(img)
-        return output_img
+        # Update tracker
+        objects = tracker.update(rects)
+        for object_id, centroid in objects.items():
+            cv2.circle(frame, centroid, 5, (0, 0, 255), -1)
+            centroids.append(centroid)
+        
+        out.write(frame)
+    
+    cap.release()
+    out.release()
+    
+    # Generate pitch map
+    pitch_map_path = "pitch_map.png"
+    fig, ax = plt.subplots(figsize=(8, 4))
+    ax.set_xlim(0, 22)  # Pitch length in meters (approx)
+    ax.set_ylim(-1.5, 1.5)  # Pitch width (approx)
+    ax.set_xlabel("Length (m)")
+    ax.set_ylabel("Width (m)")
+    ax.set_title("Pitch Map with Ball Trajectory")
+    
+    # Plot stumps
+    ax.plot([20.12, 20.12], [-0.135, 0.135], 'k-', lw=5)  # Stumps at bowling end
+    ax.plot([0, 0], [-0.135, 0.135], 'k-', lw=5)  # Stumps at batting end
+    ax.plot([0, 20.12], [0, 0], 'k--')  # Pitch center line
+    
+    # Map centroids to pitch coordinates (simplified scaling)
+    if centroids:
+        x_coords = [20.12 - (c[1] / cap.get(4)) * 20.12 for c in centroids]  # Scale y to pitch length
+        y_coords = [(c[0] / cap.get(3)) * 2.7 - 1.35 for c in centroids]  # Scale x to pitch width
+        ax.plot(x_coords, y_coords, 'ro-', label="Ball Trajectory")
+        pitch_points = list(zip(x_coords, y_coords))
     
-    elif file_extension in ['.mp4', '.avi', '.mov']:
-        # Process video
-        cap = cv2.VideoCapture(input_media)
-        output_path = "output_video.mp4"
-        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-        out = cv2.VideoWriter(output_path, fourcc, 30.0, 
-                             (int(cap.get(3)), int(cap.get(4))))
+    ax.legend()
+    plt.savefig(pitch_map_path)
+    plt.close()
+    
+    # LBW Decision (simplified physics-based model)
+    lbw_decision = "Not Out"
+    if pitch_points:
+        # Check pitching, impact, and wickets
+        pitching = any(0 <= x <= 20.12 and -0.135 <= y <= 0.135 for x, y in pitch_points[:len(pitch_points)//2])
+        impact = any(18 <= x <= 20.12 for x, y in pitch_points[len(pitch_points)//2:])
         
-        while cap.isOpened():
-            ret, frame = cap.read()
-            if not ret:
-                break
-                
-            # Perform detection
-            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-            results = model.predict(frame_rgb, conf=conf_threshold, iou=iou_threshold)
+        # Fit a quadratic curve to predict trajectory post-impact
+        if len(x_coords) > 2:
+            t = np.linspace(0, 1, len(x_coords))
+            f_x = interp1d(t, x_coords, kind='quadratic', fill_value="extrapolate")
+            f_y = interp1d(t, y_coords, kind='quadratic', fill_value="extrapolate")
+            t_future = np.array([1.5])  # Predict beyond impact
+            x_future = f_x(t_future)[0]
+            y_future = f_y(t_future)[0]
+            wickets = (18 <= x_future <= 20.12) and (-0.135 <= y_future <= 0.135)
             
-            # Draw bounding boxes
-            for box in results[0].boxes:
-                x1, y1, x2, y2 = map(int, box.xyxy[0])
-                conf = box.conf[0]
-                label = f"Ball: {conf:.2f}"
-                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
-                cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
+            if pitching and impact and wickets:
+                lbw_decision = "Out"
+            elif pitching and impact:
+                lbw_decision = "Umpire's Call"  # Marginal case
             
-            out.write(frame)
-        
-        cap.release()
-        out.release()
-        return output_path
-    
-    else:
-        return "Unsupported file format. Please upload an image (.jpg, .png) or video (.mp4, .avi, .mov)."
+    return output_path, pitch_map_path, lbw_decision
 
 # Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Decision Review System (DRS) for Ball Detection")
-    gr.Markdown("Upload an image or video to detect the ball using a trained YOLOv5 model. Adjust confidence and IoU thresholds for detection.")
+    gr.Markdown("# DRS Review System for Cricket")
+    gr.Markdown("Upload a cricket video to analyze ball tracking, pitch mapping, and LBW review. Adjust thresholds for detection accuracy.")
     
- --
-
-    input_media = gr.File(label="Upload Image or Video")
+    video_input = gr.Video(label="Upload Cricket Video")
     conf_slider = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.5, label="Confidence Threshold")
     iou_slider = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.5, label="IoU Threshold")
-    output = gr.Image(label="Output (Image or Video)")
-    submit_button = gr.Button("Detect Ball")
+    output_video = gr.Video(label="Annotated Video with Ball Tracking")
+    output_image = gr.Image(label="Pitch Map")
+    output_text = gr.Textbox(label="LBW Decision")
+    submit_button = gr.Button("Analyze DRS")
     
     submit_button.click(
-        fn=detect_ball,
-        inputs=[input_media, conf_slider, iou_slider],
-        outputs=output
+        fn=detect_and_track_ball,
+        inputs=[video_input, conf_slider, iou_slider],
+        outputs=[output_video, output_image, output_text]
     )
 
 demo.launch()