Update app.py

app.py

@@ -6,6 +6,7 @@ import plotly.graph_objects as go
 import torch
 import gradio as gr
 import os
+import time
 from scipy.optimize import curve_fit
 import sys
 
@@ -24,6 +25,9 @@ STUMP_WIDTH = 0.2286  # Stump width (including bails)
 
 # Model input size (adjust if best.pt was trained with a different size)
 MODEL_INPUT_SIZE = (640, 640)  # (height, width)
+FRAME_SKIP = 2  # Process every 2nd frame
+MIN_DETECTIONS = 10  # Stop after 10 detections
+BATCH_SIZE = 4  # Process 4 frames at a time
 
 # Load model
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -40,45 +44,70 @@ def process_video(video_path):
     frame_numbers = []
     bounce_frame = None
     bounce_point = None
+    batch_frames = []
+    batch_frame_nums = []
+    frame_count = 0
 
+    start_time = time.time()
     while cap.isOpened():
         frame_num = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
         ret, frame = cap.read()
         if not ret:
             break
 
+        # Skip frames
+        if frame_count % FRAME_SKIP != 0:
+            frame_count += 1
+            continue
+
         # Resize frame to model input size
         frame = cv2.resize(frame, MODEL_INPUT_SIZE, interpolation=cv2.INTER_AREA)
-
-        # Preprocess frame for YOLOv5
-        img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        img = torch.from_numpy(img).to(device).float() / 255.0
-        img = img.permute(2, 0, 1).unsqueeze(0)  # [1, 3, H, W]
-
-        # Run inference
-        with torch.no_grad():
-            pred = model(img)[0]
-        pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45)
-
-        # Process detections
-        for det in pred:
-            if det is not None and len(det):
-                det = xywh2xyxy(det)  # Convert to [x1, y1, x2, y2]
-                for *xyxy, conf, cls in det:
-                    x_center = (xyxy[0] + xyxy[2]) / 2
-                    y_center = (xyxy[1] + xyxy[3]) / 2
-                    # Scale coordinates back to original frame size
-                    x_center = x_center * frame_width / MODEL_INPUT_SIZE[1]
-                    y_center = y_center * frame_height / MODEL_INPUT_SIZE[0]
-                    positions.append((x_center.item(), y_center.item()))
-                    frame_numbers.append(frame_num)
-
-                    # Detect bounce (lowest y_center point)
-                    if bounce_frame is None or y_center > positions[bounce_frame][1]:
-                        bounce_frame = len(frame_numbers) - 1
-                        bounce_point = (x_center.item(), y_center.item())
+        batch_frames.append(frame)
+        batch_frame_nums.append(frame_num)
+        frame_count += 1
+
+        # Process batch when full or at end
+        if len(batch_frames) == BATCH_SIZE or not ret:
+            # Preprocess batch
+            batch = [cv2.cvtColor(f, cv2.COLOR_BGR2RGB) for f in batch_frames]
+            batch = np.stack(batch)  # [batch_size, H, W, 3]
+            batch = torch.from_numpy(batch).to(device).float() / 255.0
+            batch = batch.permute(0, 3, 1, 2)  # [batch_size, 3, H, W]
+
+            # Run inference
+            frame_start_time = time.time()
+            with torch.no_grad():
+                pred = model(batch)[0]
+            pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45)
+            print(f"Batch inference time: {time.time() - frame_start_time:.2f}s for {len(batch_frames)} frames")
+
+            # Process detections
+            for i, det in enumerate(pred):
+                if det is not None and len(det):
+                    det = xywh2xyxy(det)  # Convert to [x1, y1, x2, y2]
+                    for *xyxy, conf, cls in det:
+                        x_center = (xyxy[0] + xyxy[2]) / 2
+                        y_center = (xyxy[1] + xyxy[3]) / 2
+                        # Scale coordinates back to original frame size
+                        x_center = x_center * frame_width / MODEL_INPUT_SIZE[1]
+                        y_center = y_center * frame_height / MODEL_INPUT_SIZE[0]
+                        positions.append((x_center.item(), y_center.item()))
+                        frame_numbers.append(batch_frame_nums[i])
+
+                        # Detect bounce (lowest y_center point)
+                        if bounce_frame is None or y_center > positions[bounce_frame][1]:
+                            bounce_frame = len(frame_numbers) - 1
+                            bounce_point = (x_center.item(), y_center.item())
+
+            batch_frames = []
+            batch_frame_nums = []
+
+            # Early termination
+            if len(positions) >= MIN_DETECTIONS:
+                break
 
     cap.release()
+    print(f"Total video processing time: {time.time() - start_time:.2f}s")
     return positions, frame_numbers, bounce_point, frame_rate, frame_width, frame_height
 
 # Polynomial function for trajectory fitting