Initial commit

dagster_project/assets/camera_calibration.py

@@ -0,0 +1,143 @@
+"""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)
+        }