Update app.py

app.py

@@ -8,15 +8,16 @@ import plotly.graph_objects as go
 import uuid
 import os
 
-# Load the trained YOLOv8n model
+# Load the trained YOLOv8n model with optimizations
 model = YOLO("best.pt")
+model.to('cuda' if torch.cuda.is_available() else 'cpu')  # Use GPU if available
 
 # Constants for LBW decision and video processing
 STUMPS_WIDTH = 0.2286  # meters (width of stumps)
 BALL_DIAMETER = 0.073  # meters (approx. cricket ball diameter)
-FRAME_RATE = 30  # Input video frame rate
+FRAME_RATE = 30  # Input video frame rate (adjust if known)
 SLOW_MOTION_FACTOR = 6  # For very slow motion (6x slower)
-CONF_THRESHOLD = 0.2  # Lowered confidence threshold to improve detection
+CONF_THRESHOLD = 0.2  # Confidence threshold
 IMPACT_ZONE_Y = 0.85  # Fraction of frame height where impact is likely
 IMPACT_DELTA_Y = 50  # Pixels for detecting sudden y-position change
 PITCH_LENGTH = 20.12  # meters (standard cricket pitch length)
@@ -28,6 +29,9 @@ def process_video(video_path):
     if not os.path.exists(video_path):
         return [], [], [], "Error: Video file not found"
     cap = cv2.VideoCapture(video_path)
+    # Get native video resolution
+    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
     frames = []
     ball_positions = []
     detection_frames = []
@@ -40,7 +44,8 @@ def process_video(video_path):
             break
         frame_count += 1
         frames.append(frame.copy())
-        results = model.predict(frame, conf=CONF_THRESHOLD, imgsz=640)
+        # Use native resolution for inference
+        results = model.predict(frame, conf=CONF_THRESHOLD, imgsz=(frame_height, frame_width), iou=0.5, max_det=1)
         detections = 0
         for detection in results[0].boxes:
             if detection.cls == 0:  # Class 0 is the ball
@@ -57,6 +62,7 @@ def process_video(video_path):
         debug_log.append("No balls detected in any frame")
     else:
         debug_log.append(f"Total ball detections: {len(ball_positions)}")
+        debug_log.append(f"Video resolution: {frame_width}x{frame_height}")
 
     return frames, ball_positions, detection_frames, "\n".join(debug_log)
 
@@ -91,7 +97,6 @@ def estimate_trajectory(ball_positions, frames, detection_frames):
             impact_frame = detection_frames[i]
             break
         elif y_coords[i] > frame_height * IMPACT_ZONE_Y:
-            # Fallback to y-position if no significant y-change
             impact_idx = i
             impact_frame = detection_frames[i]
             break
@@ -206,7 +211,7 @@ def create_3d_plot(detections_3d, trajectory_3d, pitch_point_3d, impact_point_3d
         impact_scatter = go.Scatter3d(
             x=[impact_point_3d[0]] if impact_point_3d else [], 
             y=[impact_point_3d[1]] if impact_point_3d else [], 
-            z=[impact_point_3d[2]] if pitch_point_3d else [],
+            z=[impact_point_3d[2]] if impact_point_3d else [],
             mode='markers', marker=dict(size=8, color='yellow'), name='Impact Point'
         )
         data = [trajectory_line, pitch_scatter, impact_scatter] + stump_traces + bail_traces
@@ -228,8 +233,9 @@ def create_3d_plot(detections_3d, trajectory_3d, pitch_point_3d, impact_point_3d
 def generate_slow_motion(frames, trajectory, pitch_point, impact_point, detection_frames, pitch_frame, impact_frame, output_path):
     if not frames:
         return None
+    frame_height, frame_width = frames[0].shape[:2]
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter(output_path, fourcc, FRAME_RATE / SLOW_MOTION_FACTOR, (frames[0].shape[1], frames[0].shape[0]))
+    out = cv2.VideoWriter(output_path, fourcc, FRAME_RATE / SLOW_MOTION_FACTOR, (frame_width, frame_height))
 
     if trajectory and detection_frames:
         trajectory_points = np.array(trajectory[:len(detection_frames)], dtype=np.int32).reshape((-1, 1, 2))