Update app.py

app.py

@@ -1,20 +1,24 @@
 import cv2
 import gradio as gr
 import numpy as np
+from sklearn.preprocessing import PolynomialFeatures
 from sklearn.linear_model import LinearRegression
 import tempfile
 
 def process_video(video_path):
     cap = cv2.VideoCapture(video_path)
-    frame_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    frame_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
 
     out_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4").name
-    out = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*"mp4v"), 20.0, (frame_w, frame_h))
+    out = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*"mp4v"), 20.0, (width, height))
 
-    ball_color_lower = np.array([5, 100, 100])    # HSV orange-red
+    # Color range for ball (adjust if needed)
+    ball_color_lower = np.array([5, 100, 100])
     ball_color_upper = np.array([20, 255, 255])
+
     trajectory = []
+    predicted_points = []
 
     while True:
         ret, frame = cap.read()
@@ -37,51 +41,64 @@ def process_video(video_path):
                 trajectory.append(center)
                 cv2.circle(frame, center, int(radius), (0, 0, 255), 2)
 
-        # Draw stumps area
-        stump_box = (frame_w // 2 - 30, frame_h - 120, frame_w // 2 + 30, frame_h - 50)
-        cv2.rectangle(frame, stump_box[:2], stump_box[2:], (0, 255, 255), 2)
-
-        # Draw trajectory
+        # Draw trajectory line
         for i in range(1, len(trajectory)):
             cv2.line(frame, trajectory[i - 1], trajectory[i], (255, 0, 0), 2)
 
-        out.write(frame)
+        # Draw stumps box
+        stump_box = (width // 2 - 30, height - 120, width // 2 + 30, height - 50)
+        cv2.rectangle(frame, stump_box[:2], stump_box[2:], (0, 255, 255), 2)
 
-    cap.release()
-    out.release()
+        # Draw projected path if enough points collected
+        if len(trajectory) >= 5 and not predicted_points:
+            X = np.array([x for x, y in trajectory]).reshape(-1, 1)
+            Y = np.array([y for x, y in trajectory])
 
-    decision = "NOT OUT"
-    projected_hit = False
+            poly = PolynomialFeatures(degree=2)
+            X_poly = poly.fit_transform(X)
 
-    # Predict trajectory
-    if len(trajectory) >= 5:
-        X = np.array([x for x, y in trajectory]).reshape(-1, 1)
-        y = np.array([y for x, y in trajectory])
+            model = LinearRegression()
+            model.fit(X_poly, Y)
 
-        model = LinearRegression()
-        model.fit(X, y)
+            x_future = np.linspace(min(X)[0], max(X)[0] + 150, num=20).reshape(-1, 1)
+            y_future = model.predict(poly.transform(x_future))
 
-        # Predict at stump x-position
-        stump_x = frame_w // 2
-        pred_y = int(model.predict(np.array([[stump_x]]))[0])
+            predicted_points = list(zip(x_future.flatten().astype(int), y_future.astype(int)))
 
-        # Check if predicted Y is in stump area
-        if stump_box[1] <= pred_y <= stump_box[3]:
-            projected_hit = True
+        # Draw projected path (dotted yellow)
+        for pt in predicted_points:
+            if 0 <= pt[0] < width and 0 <= pt[1] < height:
+                cv2.circle(frame, pt, 3, (0, 255, 255), -1)
+
+        out.write(frame)
+
+    cap.release()
+    out.release()
+
+    # Determine OUT/NOT OUT
+    decision = "NOT OUT"
+    for x, y in predicted_points:
+        if stump_box[0] <= x <= stump_box[2] and stump_box[1] <= y <= stump_box[3]:
             decision = "OUT"
+            break
 
-    # Write final frame with decision
-    final_frame = cv2.imread(out_path)
-    cv2.putText(final_frame, f"DECISION: {decision}", (50, 80), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 255, 0) if decision == "NOT OUT" else (0, 0, 255), 4)
+    # Final frame update
+    final_frame = cv2.VideoCapture(out_path)
+    ret, last_frame = final_frame.read()
+    if ret:
+        cv2.putText(last_frame, f"DECISION: {decision}", (50, 100), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255) if decision == "OUT" else (0, 255, 0), 4)
+        out_final = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*"mp4v"), 20.0, (width, height))
+        out_final.write(last_frame)
+        out_final.release()
 
     return out_path
 
 iface = gr.Interface(
     fn=process_video,
-    inputs=gr.Video(label="Upload Bowling Video (Camera View)"),
-    outputs=gr.Video(label="Replay with Trajectory and Decision"),
-    title="Cricket DRS LBW Tracker",
-    description="Track ball trajectory and check if it's LBW OUT based on impact prediction."
+    inputs=gr.Video(label="Upload Bowling Video"),
+    outputs=gr.Video(label="LBW Tracker Output"),
+    title="DRS LBW Review System",
+    description="Detect ball trajectory, project path, and decide OUT/NOT OUT using AI"
 )
 
 iface.launch()