Initial commit
This commit is contained in:
Ruslan Bakiev
284
src/court_calibrator.py
Normal file
284
src/court_calibrator.py
Normal file
@@ -0,0 +1,284 @@
|
||||
"""
|
||||
Court calibration module for mapping pixel coordinates to real-world coordinates
|
||||
Uses homography transformation based on court keypoints
|
||||
"""
|
||||
import numpy as np
|
||||
import cv2
|
||||
from typing import List, Tuple, Optional, Dict
|
||||
import json
|
||||
|
||||
|
||||
class CourtCalibrator:
|
||||
"""
|
||||
Calibrates camera perspective to map pixel coordinates to real-world court coordinates
|
||||
Uses homography transformation
|
||||
"""
|
||||
|
||||
def __init__(self, court_width_m: float = 6.1, court_length_m: float = 13.4):
|
||||
"""
|
||||
Initialize court calibrator
|
||||
|
||||
Args:
|
||||
court_width_m: Width of pickleball court in meters (default: 6.1m)
|
||||
court_length_m: Length of pickleball court in meters (default: 13.4m)
|
||||
"""
|
||||
self.court_width = court_width_m
|
||||
self.court_length = court_length_m
|
||||
self.homography_matrix = None
|
||||
self.court_corners_pixel = None
|
||||
self.court_corners_real = np.array([
|
||||
[0, 0], # Top-left
|
||||
[court_length_m, 0], # Top-right
|
||||
[court_length_m, court_width_m], # Bottom-right
|
||||
[0, court_width_m] # Bottom-left
|
||||
], dtype=np.float32)
|
||||
|
||||
def calibrate_manual(self, corner_points: List[Tuple[float, float]]) -> bool:
|
||||
"""
|
||||
Manually calibrate using 4 corner points of the court
|
||||
|
||||
Args:
|
||||
corner_points: List of 4 (x, y) tuples representing court corners in pixels
|
||||
Order: [top-left, top-right, bottom-right, bottom-left]
|
||||
|
||||
Returns:
|
||||
True if calibration successful, False otherwise
|
||||
"""
|
||||
if len(corner_points) != 4:
|
||||
print("Error: Need exactly 4 corner points")
|
||||
return False
|
||||
|
||||
self.court_corners_pixel = np.array(corner_points, dtype=np.float32)
|
||||
|
||||
# Calculate homography matrix
|
||||
self.homography_matrix, status = cv2.findHomography(
|
||||
self.court_corners_pixel,
|
||||
self.court_corners_real,
|
||||
method=cv2.RANSAC
|
||||
)
|
||||
|
||||
if self.homography_matrix is None:
|
||||
print("Error: Failed to calculate homography matrix")
|
||||
return False
|
||||
|
||||
print("✓ Court calibration successful")
|
||||
return True
|
||||
|
||||
def calibrate_auto(self, frame: np.ndarray) -> bool:
|
||||
"""
|
||||
Automatically detect court corners using computer vision
|
||||
This is a simplified version - can be enhanced with YOLO keypoint detection
|
||||
|
||||
Args:
|
||||
frame: Video frame to detect court corners from
|
||||
|
||||
Returns:
|
||||
True if calibration successful, False otherwise
|
||||
"""
|
||||
# Convert to grayscale
|
||||
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
|
||||
|
||||
# Apply edge detection
|
||||
edges = cv2.Canny(gray, 50, 150, apertureSize=3)
|
||||
|
||||
# Detect lines using Hough Transform
|
||||
lines = cv2.HoughLinesP(edges, 1, np.pi/180, threshold=100,
|
||||
minLineLength=100, maxLineGap=10)
|
||||
|
||||
if lines is None or len(lines) < 4:
|
||||
print("Error: Could not detect enough lines for court corners")
|
||||
return False
|
||||
|
||||
# This is a placeholder - in production, you'd use:
|
||||
# 1. YOLO keypoint detection model trained on court corners
|
||||
# 2. More sophisticated line intersection detection
|
||||
# 3. Court line template matching
|
||||
|
||||
print("Warning: Auto-calibration not fully implemented")
|
||||
print("Please use calibrate_manual() with corner points")
|
||||
return False
|
||||
|
||||
def pixel_to_real(self, pixel_coords: Tuple[float, float]) -> Optional[Tuple[float, float]]:
|
||||
"""
|
||||
Transform pixel coordinates to real-world court coordinates
|
||||
|
||||
Args:
|
||||
pixel_coords: (x, y) tuple in pixel space
|
||||
|
||||
Returns:
|
||||
(x, y) tuple in real-world meters, or None if not calibrated
|
||||
"""
|
||||
if self.homography_matrix is None:
|
||||
print("Error: Court not calibrated. Call calibrate_manual() first")
|
||||
return None
|
||||
|
||||
# Convert to homogeneous coordinates
|
||||
pixel_point = np.array([[pixel_coords[0], pixel_coords[1]]], dtype=np.float32)
|
||||
pixel_point = pixel_point.reshape(-1, 1, 2)
|
||||
|
||||
# Apply homography transformation
|
||||
real_point = cv2.perspectiveTransform(pixel_point, self.homography_matrix)
|
||||
|
||||
x, y = real_point[0][0]
|
||||
return (float(x), float(y))
|
||||
|
||||
def real_to_pixel(self, real_coords: Tuple[float, float]) -> Optional[Tuple[float, float]]:
|
||||
"""
|
||||
Transform real-world court coordinates to pixel coordinates
|
||||
|
||||
Args:
|
||||
real_coords: (x, y) tuple in real-world meters
|
||||
|
||||
Returns:
|
||||
(x, y) tuple in pixel space, or None if not calibrated
|
||||
"""
|
||||
if self.homography_matrix is None:
|
||||
print("Error: Court not calibrated")
|
||||
return None
|
||||
|
||||
# Use inverse homography
|
||||
inv_homography = np.linalg.inv(self.homography_matrix)
|
||||
|
||||
real_point = np.array([[real_coords[0], real_coords[1]]], dtype=np.float32)
|
||||
real_point = real_point.reshape(-1, 1, 2)
|
||||
|
||||
pixel_point = cv2.perspectiveTransform(real_point, inv_homography)
|
||||
|
||||
x, y = pixel_point[0][0]
|
||||
return (float(x), float(y))
|
||||
|
||||
def is_calibrated(self) -> bool:
|
||||
"""Check if court is calibrated"""
|
||||
return self.homography_matrix is not None
|
||||
|
||||
def draw_court_overlay(self, frame: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Draw court boundaries on frame for visualization
|
||||
|
||||
Args:
|
||||
frame: Input frame
|
||||
|
||||
Returns:
|
||||
Frame with court overlay
|
||||
"""
|
||||
if not self.is_calibrated():
|
||||
return frame
|
||||
|
||||
overlay = frame.copy()
|
||||
|
||||
# Draw court corners
|
||||
if self.court_corners_pixel is not None:
|
||||
for point in self.court_corners_pixel:
|
||||
cv2.circle(overlay, (int(point[0]), int(point[1])), 10, (0, 255, 0), -1)
|
||||
|
||||
# Draw court boundary lines
|
||||
pts = self.court_corners_pixel.astype(np.int32).reshape((-1, 1, 2))
|
||||
cv2.polylines(overlay, [pts], True, (0, 255, 0), 2)
|
||||
|
||||
return overlay
|
||||
|
||||
def save_calibration(self, filepath: str):
|
||||
"""Save calibration data to file"""
|
||||
if not self.is_calibrated():
|
||||
print("Error: No calibration to save")
|
||||
return
|
||||
|
||||
data = {
|
||||
'court_width_m': self.court_width,
|
||||
'court_length_m': self.court_length,
|
||||
'homography_matrix': self.homography_matrix.tolist(),
|
||||
'court_corners_pixel': self.court_corners_pixel.tolist()
|
||||
}
|
||||
|
||||
with open(filepath, 'w') as f:
|
||||
json.dump(data, f, indent=2)
|
||||
|
||||
print(f"✓ Calibration saved to {filepath}")
|
||||
|
||||
def load_calibration(self, filepath: str) -> bool:
|
||||
"""Load calibration data from file"""
|
||||
try:
|
||||
with open(filepath, 'r') as f:
|
||||
data = json.load(f)
|
||||
|
||||
self.court_width = data['court_width_m']
|
||||
self.court_length = data['court_length_m']
|
||||
self.homography_matrix = np.array(data['homography_matrix'])
|
||||
self.court_corners_pixel = np.array(data['court_corners_pixel'])
|
||||
|
||||
print(f"✓ Calibration loaded from {filepath}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error loading calibration: {e}")
|
||||
return False
|
||||
|
||||
|
||||
class InteractiveCalibrator:
|
||||
"""
|
||||
Interactive tool for manual court calibration
|
||||
Click on 4 corners of the court in order: top-left, top-right, bottom-right, bottom-left
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.points = []
|
||||
self.window_name = "Court Calibration - Click 4 corners (TL, TR, BR, BL)"
|
||||
|
||||
def _mouse_callback(self, event, x, y, flags, param):
|
||||
"""Handle mouse clicks"""
|
||||
if event == cv2.EVENT_LBUTTONDOWN and len(self.points) < 4:
|
||||
self.points.append((x, y))
|
||||
print(f"Point {len(self.points)}: ({x}, {y})")
|
||||
|
||||
def calibrate_interactive(self, frame: np.ndarray) -> Optional[List[Tuple[float, float]]]:
|
||||
"""
|
||||
Interactive calibration - user clicks 4 corners
|
||||
|
||||
Args:
|
||||
frame: First frame of video to calibrate on
|
||||
|
||||
Returns:
|
||||
List of 4 corner points, or None if cancelled
|
||||
"""
|
||||
display = frame.copy()
|
||||
cv2.namedWindow(self.window_name)
|
||||
cv2.setMouseCallback(self.window_name, self._mouse_callback)
|
||||
|
||||
print("\nClick on 4 court corners in this order:")
|
||||
print("1. Top-left")
|
||||
print("2. Top-right")
|
||||
print("3. Bottom-right")
|
||||
print("4. Bottom-left")
|
||||
print("Press 'q' to cancel, 'r' to reset")
|
||||
|
||||
while True:
|
||||
# Draw points
|
||||
temp = display.copy()
|
||||
for i, point in enumerate(self.points):
|
||||
cv2.circle(temp, point, 5, (0, 255, 0), -1)
|
||||
cv2.putText(temp, str(i+1), (point[0]+10, point[1]),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
|
||||
|
||||
# Draw lines between points
|
||||
if len(self.points) > 1:
|
||||
for i in range(len(self.points) - 1):
|
||||
cv2.line(temp, self.points[i], self.points[i+1], (0, 255, 0), 2)
|
||||
|
||||
cv2.imshow(self.window_name, temp)
|
||||
key = cv2.waitKey(1) & 0xFF
|
||||
|
||||
if key == ord('q'):
|
||||
print("Calibration cancelled")
|
||||
cv2.destroyAllWindows()
|
||||
return None
|
||||
|
||||
if key == ord('r'):
|
||||
print("Reset points")
|
||||
self.points = []
|
||||
|
||||
if len(self.points) == 4:
|
||||
print("\n✓ All 4 points selected")
|
||||
cv2.destroyAllWindows()
|
||||
return self.points
|
||||
|
||||
return None
|
||||
Reference in New Issue
Block a user