Delete lbw_detector.py

lbw_detector.py

@@ -1,127 +0,0 @@
-# lbw_detector.py
-import torch
-import torch.nn as nn
-import numpy as np
-from torchvision import transforms
-import cv2
-from utils import extract_frames
-from trajectory_predictor import predict_trajectory
-from visualizer import draw_visuals
-
-# -----------------------------
-# Model Definition (UNet-lite)
-# -----------------------------
-class DoubleConv(nn.Module):
-    def __init__(self, in_channels, out_channels):
-        super(DoubleConv, self).__init__()
-        self.double_conv = nn.Sequential(
-            nn.Conv2d(in_channels, out_channels, 3, padding=1), nn.ReLU(inplace=True),
-            nn.Conv2d(out_channels, out_channels, 3, padding=1), nn.ReLU(inplace=True)
-        )
-
-    def forward(self, x):
-        return self.double_conv(x)
-
-class UNet(nn.Module):
-    def __init__(self, in_channels=3, out_channels=1):  # ✅ Match model file
-        super(UNet, self).__init__()
-        self.down1 = DoubleConv(in_channels, 64)
-        self.down2 = DoubleConv(64, 128)
-        self.middle = DoubleConv(128, 256)
-
-        self.up2 = nn.ConvTranspose2d(256, 128, 2, stride=2)
-        self.upconv2 = DoubleConv(256, 128)
-
-        self.up1 = nn.ConvTranspose2d(128, 64, 2, stride=2)
-        self.upconv1 = DoubleConv(128, 64)
-
-        self.final = nn.Conv2d(64, out_channels, 1)
-
-    def forward(self, x):
-        d1 = self.down1(x)
-        d2 = self.down2(nn.MaxPool2d(2)(d1))
-        m = self.middle(nn.MaxPool2d(2)(d2))
-
-        u2 = self.up2(m)
-        u2 = self.upconv2(torch.cat([u2, d2], dim=1))
-
-        u1 = self.up1(u2)
-        u1 = self.upconv1(torch.cat([u1, d1], dim=1))
-
-        out = self.final(u1)
-        return out
-
-# -----------------------------
-# Load Model
-# -----------------------------
-model_path = "models/lbw_drs_unet_model.pth"
-device = "cpu"
-
-model = UNet(in_channels=3, out_channels=1)  # ✅ Must match trained model
-state_dict = torch.load(model_path, map_location=device)
-model.load_state_dict(state_dict)
-model.to(device)
-model.eval()
-
-# -----------------------------
-# Frame Preprocessing
-# -----------------------------
-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)
-    mask = torch.sigmoid(output).squeeze().cpu().numpy()
-    return mask  # shape: (H, W)
-
-# -----------------------------
-# Main Analysis Function
-# -----------------------------
-def analyze_video(video_path):
-    frames = extract_frames(video_path)
-    ball_positions = []
-    impact_frame_idx = None
-    impact_zone = "unknown"
-
-    for i, frame in enumerate(frames):
-        mask = detect_objects_with_model(frame)
-
-        ball_mask = mask > 0.5
-        pad_mask = None  # Not used with single-channel model
-
-        # Ball center detection
-        contours, _ = cv2.findContours(ball_mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        if contours:
-            largest = max(contours, key=cv2.contourArea)
-            M = cv2.moments(largest)
-            if M['m00'] != 0:
-                cx = int(M['m10']/M['m00'])
-                cy = int(M['m01']/M['m00'])
-                ball_positions.append((i, cx, cy))
-
-        # Optional: skip impact detection for now (no pad_mask)
-        # If you later upgrade to multi-class model, enable below
-        # if pad_mask is not None and contours:
-        #     overlap = np.logical_and(ball_mask, pad_mask).sum()
-        #     if overlap > 10:
-        #         impact_frame_idx = i
-        #         impact_zone = "pad"
-        #         break
-
-    trajectory = predict_trajectory(ball_positions)
-    decision = "OUT" if trajectory_hits_stumps(trajectory) else "NOT OUT"
-    result_path = draw_visuals(frames, ball_positions, trajectory, impact_frame_idx, decision)
-    return result_path, decision
-
-# -----------------------------
-# Basic Rule: Hits Stumps
-# -----------------------------
-def trajectory_hits_stumps(trajectory):
-    for (x, y) in trajectory:
-        if 300 < x < 340 and y < 480:
-            return True
-    return False