144 lines
4.8 KiB
Python
144 lines
4.8 KiB
Python
"""Asset: Calibrate camera using court corners and net position"""
|
||
|
||
import cv2
|
||
import numpy as np
|
||
from typing import Dict
|
||
from dagster import asset, AssetExecutionContext
|
||
|
||
|
||
@asset(
|
||
io_manager_key="json_io_manager",
|
||
compute_kind="opencv",
|
||
description="Calibrate camera using cv2.solvePnP with court corners and net"
|
||
)
|
||
def calibrate_camera_3d(
|
||
context: AssetExecutionContext,
|
||
detect_court_keypoints: Dict,
|
||
detect_net: Dict
|
||
) -> Dict:
|
||
"""
|
||
Calibrate camera to get 3D pose using known 3D↔2D point correspondences
|
||
|
||
Inputs:
|
||
- detect_court_keypoints: 4 court corners in pixels
|
||
- detect_net: 4 net corners in pixels
|
||
|
||
Known 3D points:
|
||
- Court: 13.4m × 6.1m rectangle at Z=0
|
||
- Net: height 0.914m at middle of court (Y=3.05m)
|
||
|
||
Outputs:
|
||
Camera calibration parameters
|
||
|
||
Returns:
|
||
Dict with:
|
||
- camera_matrix: [[fx, 0, cx], [0, fy, cy], [0, 0, 1]]
|
||
- rotation_vector: [rx, ry, rz]
|
||
- translation_vector: [tx, ty, tz]
|
||
- rotation_matrix: 3x3 matrix
|
||
- reprojection_error: RMS error in pixels
|
||
- calibrated: success flag
|
||
"""
|
||
# Get 2D points (pixels)
|
||
court_corners = np.array(detect_court_keypoints['corners_pixel'], dtype=np.float32)
|
||
net_corners = np.array(detect_net['net_corners_pixel'], dtype=np.float32)
|
||
|
||
# Define 3D points (meters) in world coordinates
|
||
# Court corners (Z=0, on ground)
|
||
court_3d = np.array([
|
||
[0, 0, 0], # TL
|
||
[13.4, 0, 0], # TR
|
||
[13.4, 6.1, 0], # BR
|
||
[0, 6.1, 0] # BL
|
||
], dtype=np.float32)
|
||
|
||
# Net endpoints (2 points)
|
||
# Net is at Y=3.05m (middle of 6.1m court width)
|
||
# We have 2 endpoints: left and right side at top of net
|
||
net_3d = np.array([
|
||
[0, 3.05, 0.914], # Left endpoint (top of net)
|
||
[13.4, 3.05, 0.914], # Right endpoint (top of net)
|
||
], dtype=np.float32)
|
||
|
||
# Combine all 3D and 2D points
|
||
object_points = np.vstack([court_3d, net_3d]) # 6 points total (4 court + 2 net)
|
||
image_points = np.vstack([court_corners, net_corners])
|
||
|
||
context.log.info(f"Calibrating with {len(object_points)} point correspondences")
|
||
context.log.info(f"3D points shape: {object_points.shape}")
|
||
context.log.info(f"2D points shape: {image_points.shape}")
|
||
|
||
# Initial camera matrix estimate
|
||
# Assume principal point at image center
|
||
w = detect_court_keypoints['frame_width']
|
||
h = detect_court_keypoints['frame_height']
|
||
cx = w / 2
|
||
cy = h / 2
|
||
# Estimate focal length (typical for drone/action cameras)
|
||
focal_length = max(w, h) # Initial guess
|
||
|
||
camera_matrix = np.array([
|
||
[focal_length, 0, cx],
|
||
[0, focal_length, cy],
|
||
[0, 0, 1]
|
||
], dtype=np.float32)
|
||
|
||
# No lens distortion (assume corrected or minimal)
|
||
dist_coeffs = None
|
||
|
||
# Solve PnP to get camera pose
|
||
try:
|
||
success, rotation_vec, translation_vec = cv2.solvePnP(
|
||
object_points,
|
||
image_points,
|
||
camera_matrix,
|
||
dist_coeffs,
|
||
flags=cv2.SOLVEPNP_ITERATIVE
|
||
)
|
||
|
||
if not success:
|
||
context.log.error("cv2.solvePnP failed")
|
||
return {
|
||
"calibrated": False,
|
||
"error": "solvePnP failed"
|
||
}
|
||
|
||
# Convert rotation vector to rotation matrix
|
||
rotation_matrix, _ = cv2.Rodrigues(rotation_vec)
|
||
|
||
# Calculate reprojection error
|
||
projected_points, _ = cv2.projectPoints(
|
||
object_points,
|
||
rotation_vec,
|
||
translation_vec,
|
||
camera_matrix,
|
||
dist_coeffs
|
||
)
|
||
projected_points = projected_points.reshape(-1, 2)
|
||
reprojection_error = np.sqrt(np.mean((image_points - projected_points) ** 2))
|
||
|
||
context.log.info(f"✓ Camera calibration successful")
|
||
context.log.info(f" Reprojection error: {reprojection_error:.2f} pixels")
|
||
context.log.info(f" Focal length: {focal_length:.1f}")
|
||
context.log.info(f" Rotation vector: {rotation_vec.flatten()}")
|
||
context.log.info(f" Translation vector: {translation_vec.flatten()}")
|
||
|
||
return {
|
||
"camera_matrix": camera_matrix.tolist(),
|
||
"rotation_vector": rotation_vec.flatten().tolist(),
|
||
"translation_vector": translation_vec.flatten().tolist(),
|
||
"rotation_matrix": rotation_matrix.tolist(),
|
||
"reprojection_error": float(reprojection_error),
|
||
"focal_length": float(focal_length),
|
||
"principal_point": [float(cx), float(cy)],
|
||
"image_size": [w, h],
|
||
"calibrated": True
|
||
}
|
||
|
||
except Exception as e:
|
||
context.log.error(f"Error during camera calibration: {e}")
|
||
return {
|
||
"calibrated": False,
|
||
"error": str(e)
|
||
}
|