courtline detector

court_line_detector/__init__.py

@@ -0,0 +1 @@
+from .court_line_detector import CourtLineDetector

court_line_detector/court_line_detector.py

@@ -0,0 +1,48 @@
+import torch
+import torchvision.transforms as transforms
+import cv2
+from torchvision import models
+import numpy as np
+
+class CourtLineDetector:
+    def __init__(self, model_path):
+        self.model = models.resnet50(pretrained=True)
+        self.model.fc = torch.nn.Linear(self.model.fc.in_features, 14*2) 
+        self.model.load_state_dict(torch.load(model_path, map_location='cpu'))
+        self.transform = transforms.Compose([
+            transforms.ToPILImage(),
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        ])
+
+    def predict(self, image):
+
+    
+        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        image_tensor = self.transform(image_rgb).unsqueeze(0)
+        with torch.no_grad():
+            outputs = self.model(image_tensor)
+        keypoints = outputs.squeeze().cpu().numpy()
+        original_h, original_w = image.shape[:2]
+        keypoints[::2] *= original_w / 224.0
+        keypoints[1::2] *= original_h / 224.0
+
+        return keypoints
+
+    def draw_keypoints(self, image, keypoints):
+        if isinstance(image, np.ndarray):
+            for i in range(0, len(keypoints), 2):
+                x, y = int(keypoints[i]), int(keypoints[i+1])
+                cv2.circle(image, (x, y), 5, (0, 255, 0), -1)
+                cv2.putText(image, str(i//2), (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
+        else:
+            print("image is not a numpy array")
+        return image
+    
+    def draw_keypoints_on_video(self, video_frames, keypoints):
+        output_video_frames = []
+        for frame in video_frames:
+            frame = self.draw_keypoints(frame, keypoints)
+            output_video_frames.append(frame)
+        return output_video_frames

yolov8x.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3df4ada6b4dad6d657868f2fdf7faecfb34dcfccf3a25c4b82079064718524c8
+size 136890692