"""Asset 2: Detect court keypoints using Roboflow Hosted API""" import os import cv2 import numpy as np from pathlib import Path from typing import Dict, List from dagster import asset, AssetExecutionContext from inference_sdk import InferenceHTTPClient @asset( io_manager_key="json_io_manager", compute_kind="roboflow", description="Detect pickleball court corners using Roboflow keypoint detection model" ) def detect_court_keypoints( context: AssetExecutionContext, extract_video_frames: Dict ) -> Dict: """ Detect court keypoints from first frame using Roboflow model Inputs: - extract_video_frames: metadata from frame extraction - data/frames/frame_0000.jpg: first frame Outputs: - data/detect_court_keypoints.json Returns: Dict with: - corners_pixel: list of 4 corner coordinates [[x,y], ...] - court_width_m: court width in meters (6.1) - court_length_m: court length in meters (13.4) - keypoints: all detected keypoints """ from inference import get_model frames_dir = Path(extract_video_frames['frames_dir']) first_frame_path = frames_dir / "frame_0000.jpg" context.log.info(f"Loading first frame: {first_frame_path}") if not first_frame_path.exists(): raise FileNotFoundError(f"First frame not found: {first_frame_path}") # Load frame frame = cv2.imread(str(first_frame_path)) h, w = frame.shape[:2] context.log.info(f"Frame dimensions: {w}x{h}") # Get API key api_key = os.getenv("ROBOFLOW_API_KEY") if not api_key: context.log.warning("ROBOFLOW_API_KEY not set, using estimated corners") corners = _estimate_court_corners(w, h) else: # Try to detect court using Roboflow Hosted API try: context.log.info("Detecting court using Roboflow Hosted API...") client = InferenceHTTPClient( api_url="https://serverless.roboflow.com", api_key=api_key ) result = client.infer(str(first_frame_path), model_id="pickleball-court-cfyv4/1") # Extract keypoints from result all_points = [] if result and 'predictions' in result and len(result['predictions']) > 0: pred = result['predictions'][0] if 'points' in pred and len(pred['points']) >= 4: # Модель возвращает много points (линии корта) all_points = [[p['x'], p['y']] for p in pred['points']] context.log.info(f"✓ Detected {len(all_points)} keypoints from court lines") # Находим 4 угла для калибровки (но не для визуализации) corners = _extract_court_corners_from_points(all_points, w, h) context.log.info(f"✓ Extracted 4 corners from keypoints") else: context.log.warning("No keypoints in prediction, using estimated corners") corners = _estimate_court_corners(w, h) else: context.log.warning("No predictions from model, using estimated corners") corners = _estimate_court_corners(w, h) except Exception as e: context.log.warning(f"Court detection failed: {e}. Using estimated corners.") corners = _estimate_court_corners(w, h) context.log.info(f"Court corners: {corners}") # Save visualization - рисуем ВСЕ точки и линии от модели vis_frame = frame.copy() # Рисуем все точки от модели if len(all_points) > 0: context.log.info(f"Drawing {len(all_points)} keypoints on visualization") # Рисуем все точки for i, point in enumerate(all_points): x, y = int(point[0]), int(point[1]) cv2.circle(vis_frame, (x, y), 5, (0, 255, 0), -1) # Подписываем каждую точку номером cv2.putText( vis_frame, str(i), (x + 8, y), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 0), 1 ) # Соединяем все соседние точки линиями for i in range(len(all_points) - 1): p1 = tuple(map(int, all_points[i])) p2 = tuple(map(int, all_points[i + 1])) cv2.line(vis_frame, p1, p2, (0, 255, 0), 2) # Save visualization with run_id run_id = context.run_id vis_path = Path(f"data/{run_id}/court_detection_preview.jpg") cv2.imwrite(str(vis_path), vis_frame) context.log.info(f"Saved court visualization to {vis_path}") return { "corners_pixel": corners, "court_width_m": 6.1, "court_length_m": 13.4, "frame_width": w, "frame_height": h } def _estimate_court_corners(width: int, height: int) -> List[List[float]]: """ Estimate court corners based on typical DJI camera position (camera in corner at angle) Returns corners in order: [TL, TR, BR, BL] """ # Assume court takes up ~80% of frame with perspective margin_x = width * 0.05 margin_y = height * 0.1 # Perspective: far edge narrower than near edge return [ [margin_x + width * 0.1, margin_y], # Top-left (far) [width - margin_x - width * 0.1, margin_y], # Top-right (far) [width - margin_x, height - margin_y], # Bottom-right (near) [margin_x, height - margin_y] # Bottom-left (near) ] def _extract_court_corners_from_points(points: List[List[float]], width: int, height: int) -> List[List[float]]: """ Extract 4 court corners from many detected points (court lines) Strategy: 1. Build convex hull from all points 2. Classify hull points into 4 sides (left, right, top, bottom) 3. Fit line for each side using linear regression 4. Find 4 corners as intersections of fitted lines This works even if one corner is not visible on frame (extrapolation) """ if len(points) < 4: return _estimate_court_corners(width, height) # Build convex hull from all points points_array = np.array(points, dtype=np.float32) hull = cv2.convexHull(points_array) hull_points = np.array([p[0] for p in hull], dtype=np.float32) # Classify hull points into 4 sides # Strategy: sort hull points by angle from centroid, then split into 4 groups center = hull_points.mean(axis=0) # Calculate angle for each point relative to center angles = np.arctan2(hull_points[:, 1] - center[1], hull_points[:, 0] - center[0]) # Sort points by angle sorted_indices = np.argsort(angles) sorted_points = hull_points[sorted_indices] # Split into 4 groups (4 sides) n = len(sorted_points) quarter = n // 4 side1 = sorted_points[0:quarter] side2 = sorted_points[quarter:2*quarter] side3 = sorted_points[2*quarter:3*quarter] side4 = sorted_points[3*quarter:] # Fit lines for each side using cv2.fitLine def fit_line_coefficients(pts): if len(pts) < 2: return None # cv2.fitLine returns (vx, vy, x0, y0) - direction vector and point on line line = cv2.fitLine(pts, cv2.DIST_L2, 0, 0.01, 0.01) vx, vy, x0, y0 = line[0][0], line[1][0], line[2][0], line[3][0] # Convert to line equation: y = mx + b or vertical line x = c if abs(vx) < 1e-6: # Vertical line return ('vertical', x0) m = vy / vx b = y0 - m * x0 return ('normal', m, b) line1 = fit_line_coefficients(side1) line2 = fit_line_coefficients(side2) line3 = fit_line_coefficients(side3) line4 = fit_line_coefficients(side4) lines = [line1, line2, line3, line4] # Find intersections between adjacent sides def line_intersection(line_a, line_b): if line_a is None or line_b is None: return None # Handle vertical lines if line_a[0] == 'vertical' and line_b[0] == 'vertical': return None elif line_a[0] == 'vertical': x = line_a[1] m2, b2 = line_b[1], line_b[2] y = m2 * x + b2 return [float(x), float(y)] elif line_b[0] == 'vertical': x = line_b[1] m1, b1 = line_a[1], line_a[2] y = m1 * x + b1 return [float(x), float(y)] else: m1, b1 = line_a[1], line_a[2] m2, b2 = line_b[1], line_b[2] if abs(m1 - m2) < 1e-6: # Parallel lines return None x = (b2 - b1) / (m1 - m2) y = m1 * x + b1 return [float(x), float(y)] # Find 4 corners as intersections corners = [] for i in range(4): next_i = (i + 1) % 4 corner = line_intersection(lines[i], lines[next_i]) if corner: corners.append(corner) # If we got 4 corners, return them if len(corners) == 4: return corners # Fallback: use convex hull extreme points tl = hull_points[np.argmin(hull_points[:, 0] + hull_points[:, 1])].tolist() tr = hull_points[np.argmax(hull_points[:, 0] - hull_points[:, 1])].tolist() br = hull_points[np.argmax(hull_points[:, 0] + hull_points[:, 1])].tolist() bl = hull_points[np.argmin(hull_points[:, 0] - hull_points[:, 1])].tolist() return [tl, tr, br, bl] def _extract_court_corners(keypoints: List[Dict], width: int, height: int) -> List[List[float]]: """ Extract 4 court corners from detected keypoints (old function for compatibility) """ if len(keypoints) < 4: return _estimate_court_corners(width, height) points = [[kp['x'], kp['y']] for kp in keypoints] return _extract_court_corners_from_points(points, width, height)