Update lbw_detector.py

lbw_detector.py

@@ -1,79 +1,28 @@
-from ultralytics import YOLO
-import cv2
+# lbw_detector.py
+import torch
 import numpy as np
-import os
-
-# Load YOLO model (custom-trained or pretrained with compatible classes)
-model = YOLO("yolov8n.pt")  # Replace with "lbw_yolov8.pt" if custom-trained
-
-#model_path = os.path.join("models", "yolov8n.pt")
-#model = YOLO(model_path)
-
-# Target class IDs — update based on your custom model class mapping
-CLASS_NAMES = {
-    0: "ball",
-    1: "bat",
-    2: "pad",
-    3: "stump",
-    4: "player"
-}
-
-def detect_lbw_event(frames):
-    """
-    Detects ball, bat, stump, and pad in each frame.
-    Identifies impact and prepares coordinates for trajectory modeling.
-
-    Returns:
-        dict: {
-            "ball_positions": [x, y] list per frame,
-            "impact_frame": int,
-            "impact_type": str,
-            "objects_per_frame": [
-                {"ball": (x, y), "pad": (x, y), ...}
-            ]
-        }
-    """
-    ball_positions = []
-    impact_frame = -1
-    impact_type = None
-    objects_per_frame = []
-
-    for idx, frame in enumerate(frames):
-        results = model(frame)[0]
-        frame_objects = {}
-
-        for det in results.boxes.data:
-            x1, y1, x2, y2, conf, cls = det.cpu().numpy()
-            class_id = int(cls)
-            class_name = CLASS_NAMES.get(class_id, "unknown")
-            center_x = int((x1 + x2) / 2)
-            center_y = int((y1 + y2) / 2)
-            frame_objects[class_name] = (center_x, center_y)
-
-            if class_name == "ball":
-                ball_positions.append((idx, center_x, center_y))
-
-        objects_per_frame.append(frame_objects)
-
-        # Basic impact logic: ball overlaps pad or bat
-        if "ball" in frame_objects and ("pad" in frame_objects or "bat" in frame_objects):
-            bx, by = frame_objects["ball"]
-            if "pad" in frame_objects:
-                px, py = frame_objects["pad"]
-                if abs(bx - px) < 30 and abs(by - py) < 30:
-                    impact_frame = idx
-                    impact_type = "pad"
-                    break
-            if "bat" in frame_objects:
-                tx, ty = frame_objects["bat"]
-                if abs(bx - tx) < 30 and abs(by - ty) < 30:
-                    impact_frame = idx
-                    impact_type = "bat"
-                    break
-
-    return {
-        "ball_positions": ball_positions,
-        "impact_frame": impact_frame,
-        "impact_type": impact_type,
-        "objects_per_frame": objects_per_frame
-    }
+from torchvision import transforms
+import cv2
+from utils import extract_frames
+from trajectory_predictor import predict_trajectory
+from visualizer import draw_visuals
+
+# Load the custom LBW model
+model_path = "models/lbw_drs_unet_model.pth"
+device = "cpu"  # Hugging Face Free Tier
+
+model = torch.load(model_path, map_location=device)
+model.eval()
+
+transform = transforms.Compose([
+    transforms.ToTensor(),
+])
+
+def detect_objects_with_model(frame):
+    """Run segmentation on a frame using the custom model"""
+    input_tensor = transform(frame).unsqueeze(0).to(device)
+    with torch.no_grad():
+        output = model(input_tensor)
+    # Convert output to mask
+    mask = torch.sigmoid(output).squeeze().cpu().numpy()
+    return mask  # Assumed to be binary mask (ball/pad/stump segmentation)