diff --git a/jetson/main.py b/jetson/main.py index 4cb4989..2e1ca02 100644 --- a/jetson/main.py +++ b/jetson/main.py @@ -35,12 +35,14 @@ _args = None def auto_calibrate(): - """Detect all visible court lines, identify them, compute camera 3D position. + """Detect the green court surface, find its boundary quadrilateral, + map to known court dimensions, compute camera 3D position via solvePnP. - Each camera sees one half of the pickleball court from the net. - We detect all white lines, classify them by angle and position, - try to match them to known court lines, then use solvePnP. - Returns debug images showing every detected line + camera position. + Debug images show: + - Green mask overlay (what the system sees as court) + - Detected quadrilateral (court boundary) + - White line segments found on court + - Computed camera position """ results = {} @@ -54,357 +56,188 @@ def auto_calibrate(): side = 'left' if sensor_id == 0 else 'right' debug_frame = frame.copy() - # Detect all court lines - detection = _detect_court_lines(frame) - all_segments = detection['segments'] - grouped = detection['grouped'] + # Step 1: Detect green court surface + court = _detect_court_surface(frame) - # Draw ALL detected line segments on debug frame - for seg in all_segments: - x1, y1, x2, y2 = seg - cv2.line(debug_frame, (x1, y1), (x2, y2), (80, 80, 80), 1) - - # Draw grouped/merged lines with colors - # "across" lines (roughly horizontal in image = baseline, kitchen, net) - across_colors = [(0, 255, 255), (0, 200, 200), (0, 150, 150), - (0, 100, 200), (0, 80, 160)] - for i, line in enumerate(grouped.get('across', [])): - x1, y1, x2, y2 = line['endpoints'] - color = across_colors[i % len(across_colors)] - cv2.line(debug_frame, (x1, y1), (x2, y2), color, 3) - cv2.putText(debug_frame, f"across#{i} ({line['count']}seg)", - (x1, max(y1 - 8, 15)), - cv2.FONT_HERSHEY_SIMPLEX, 0.4, color, 1) - - # "along" lines (roughly vertical in image = sidelines, center service) - along_colors = [(255, 0, 255), (200, 0, 200), (150, 0, 150), - (200, 0, 100), (160, 0, 80)] - for i, line in enumerate(grouped.get('along', [])): - x1, y1, x2, y2 = line['endpoints'] - color = along_colors[i % len(along_colors)] - cv2.line(debug_frame, (x1, y1), (x2, y2), color, 3) - cv2.putText(debug_frame, f"along#{i} ({line['count']}seg)", - (x1 + 5, (y1 + y2) // 2), - cv2.FONT_HERSHEY_SIMPLEX, 0.4, color, 1) - - # Summary text - n_across = len(grouped.get('across', [])) - n_along = len(grouped.get('along', [])) - cv2.putText(debug_frame, - f"Lines: {len(all_segments)} raw, {n_across} across, {n_along} along", - (10, h - 15), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1) - - # Try to match lines to known court geometry and calibrate - match = _match_court_lines(grouped, side, w, h) - - if match['points_2d'] is not None and len(match['points_2d']) >= 4: - # We have enough correspondences — calibrate - cal = CameraCalibrator() - cal.calibrate( - np.array(match['points_2d'], dtype=np.float32), - np.array(match['points_3d'], dtype=np.float32), - w, h - ) - - cam_pos = (-cal.rotation_matrix.T @ cal.translation_vec).flatten() - - # Draw matched points - for i, (px, py) in enumerate(match['points_2d']): - cv2.circle(debug_frame, (int(px), int(py)), 6, (0, 255, 0), -1) - - cv2.putText(debug_frame, - f"CAM{sensor_id}: X={cam_pos[0]:.2f} Y={cam_pos[1]:.2f} Z={cam_pos[2]:.2f}m", - (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) - cv2.putText(debug_frame, - f"Matched {len(match['points_2d'])} points, {side} half", - (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1) - - # Save - cal_path = os.path.join(_args.calibration_dir, - f'cam{sensor_id}_calibration.json') - os.makedirs(os.path.dirname(cal_path), exist_ok=True) - cal.save(cal_path) - state['calibrators'][sensor_id] = cal - - _, jpeg = cv2.imencode('.jpg', debug_frame, [cv2.IMWRITE_JPEG_QUALITY, 85]) - results[str(sensor_id)] = { - 'ok': True, - 'camera_position': cam_pos.tolist(), - 'debug_image': base64.b64encode(jpeg.tobytes()).decode('ascii'), - 'lines_detected': {'across': n_across, 'along': n_along}, - 'points_matched': len(match['points_2d']), - 'matched_lines_3d': match.get('matched_lines_3d', []), - } - print(f"[CAM {sensor_id}] Calibrated! Camera at " - f"({cam_pos[0]:.2f}, {cam_pos[1]:.2f}, {cam_pos[2]:.2f}), " - f"{len(match['points_2d'])} correspondences") - else: - cv2.putText(debug_frame, - f"Not enough correspondences: {match.get('error', 'unknown')}", - (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2) + # Draw green mask as semi-transparent overlay + if court['green_mask'] is not None: + green_overlay = np.zeros_like(debug_frame) + green_overlay[court['green_mask'] > 0] = (0, 80, 0) + debug_frame = cv2.addWeighted(debug_frame, 1.0, green_overlay, 0.3, 0) + if court['quad'] is None: + cv2.putText(debug_frame, f"FAILED: {court['error']}", (10, 30), + cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2) _, jpeg = cv2.imencode('.jpg', debug_frame, [cv2.IMWRITE_JPEG_QUALITY, 85]) results[str(sensor_id)] = { 'ok': False, - 'error': f"CAM {sensor_id}: {match.get('error', 'Not enough line correspondences')}", + 'error': f"CAM {sensor_id}: {court['error']}", 'debug_image': base64.b64encode(jpeg.tobytes()).decode('ascii'), - 'lines_detected': {'across': n_across, 'along': n_along}, } + continue + + quad = court['quad'] + + # Draw detected court quadrilateral + pts = quad.astype(int) + for i in range(4): + p1 = tuple(pts[i]) + p2 = tuple(pts[(i + 1) % 4]) + cv2.line(debug_frame, p1, p2, (0, 255, 0), 3) + cv2.circle(debug_frame, p1, 8, (0, 0, 255), -1) + cv2.putText(debug_frame, f"C{i}", (p1[0] + 10, p1[1] - 5), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2) + + # Step 2: Detect white lines on court surface + white_lines = _detect_white_lines_on_court(frame, court['green_mask']) + for seg in white_lines: + x1, y1, x2, y2 = seg + cv2.line(debug_frame, (x1, y1), (x2, y2), (255, 255, 0), 1) + + # Step 3: Map quad corners to 3D court coordinates + # Quad corners are ordered: TL, TR, BR, BL + # For camera at net looking at one half-court: + # TL = far-left, TR = far-right (baseline) + # BL = near-left, BR = near-right (near net/camera) + corners_3d = get_half_court_3d_points(side) + + # Calibrate + cal = CameraCalibrator() + cal.calibrate(quad, corners_3d, w, h) + + cam_pos = (-cal.rotation_matrix.T @ cal.translation_vec).flatten() + + cv2.putText(debug_frame, + f"CAM{sensor_id}: X={cam_pos[0]:.2f} Y={cam_pos[1]:.2f} Z={cam_pos[2]:.2f}m", + (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) + cv2.putText(debug_frame, + f"{side} half | {len(white_lines)} white line segments", + (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1) + + # Save + cal_path = os.path.join(_args.calibration_dir, + f'cam{sensor_id}_calibration.json') + os.makedirs(os.path.dirname(cal_path), exist_ok=True) + cal.save(cal_path) + state['calibrators'][sensor_id] = cal + + _, jpeg = cv2.imencode('.jpg', debug_frame, [cv2.IMWRITE_JPEG_QUALITY, 85]) + + # Court lines in 3D for visualization + matched_lines_3d = _get_half_court_lines_3d(side) + + results[str(sensor_id)] = { + 'ok': True, + 'camera_position': cam_pos.tolist(), + 'debug_image': base64.b64encode(jpeg.tobytes()).decode('ascii'), + 'matched_lines_3d': matched_lines_3d, + } + print(f"[CAM {sensor_id}] Calibrated! Camera at " + f"({cam_pos[0]:.2f}, {cam_pos[1]:.2f}, {cam_pos[2]:.2f})") return results -def _detect_court_lines(frame): - """Detect all white line segments in the frame. +def _get_half_court_lines_3d(side): + """Return 3D line segments for all lines of a half-court.""" + if side == 'left': + bx, kx = 0, 4.57 # baseline, kitchen + else: + bx, kx = 13.4, 8.83 + lines = [ + {'name': 'baseline', 'from': [bx, 0, 0], 'to': [bx, 6.1, 0]}, + {'name': 'kitchen', 'from': [kx, 0, 0], 'to': [kx, 6.1, 0]}, + {'name': 'sideline_near', 'from': [bx, 0, 0], 'to': [kx, 0, 0]}, + {'name': 'sideline_far', 'from': [bx, 6.1, 0], 'to': [kx, 6.1, 0]}, + {'name': 'center_service', 'from': [bx, 3.05, 0], 'to': [kx, 3.05, 0]}, + ] + return lines - Returns: - segments: list of [x1, y1, x2, y2] raw segments - grouped: dict with 'across' and 'along' merged line groups + +def _detect_court_surface(frame): + """Detect the green court surface and find its boundary as a quadrilateral. + + The court surface is distinctly green against pink/gray surroundings. + We segment by color, find the largest contour, and approximate with 4 corners. + + Returns dict with: + quad: 4x2 numpy array of court boundary corners in image, or None + green_mask: binary mask of detected green surface + error: description if detection failed + """ + hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV) + + # Green court surface — relatively broad range to handle lighting + lower_green = np.array([30, 30, 50]) + upper_green = np.array([90, 255, 220]) + green_mask = cv2.inRange(hsv, lower_green, upper_green) + + # Clean up noise + kernel = np.ones((7, 7), np.uint8) + green_mask = cv2.morphologyEx(green_mask, cv2.MORPH_CLOSE, kernel) + green_mask = cv2.morphologyEx(green_mask, cv2.MORPH_OPEN, kernel) + + # Find contours + contours, _ = cv2.findContours(green_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + + if not contours: + return {'quad': None, 'green_mask': green_mask, 'error': 'No green surface detected'} + + # Take largest contour (should be the court) + largest = max(contours, key=cv2.contourArea) + area = cv2.contourArea(largest) + frame_area = frame.shape[0] * frame.shape[1] + + if area < frame_area * 0.05: + return {'quad': None, 'green_mask': green_mask, + 'error': f'Green area too small ({area / frame_area * 100:.0f}% of frame)'} + + # Approximate contour with polygon + epsilon = 0.02 * cv2.arcLength(largest, True) + approx = cv2.approxPolyDP(largest, epsilon, True) + + # We want 4 corners — if we got more, use convex hull + min area rect + if len(approx) == 4: + quad = approx.reshape(4, 2).astype(np.float32) + else: + # Use minimum area rectangle as fallback + rect = cv2.minAreaRect(largest) + quad = cv2.boxPoints(rect).astype(np.float32) + + # Order corners: top-left, top-right, bottom-right, bottom-left + # Sort by y first to get top/bottom pairs, then by x within each pair + sorted_by_y = quad[np.argsort(quad[:, 1])] + top_pair = sorted_by_y[:2] + bottom_pair = sorted_by_y[2:] + top_pair = top_pair[np.argsort(top_pair[:, 0])] + bottom_pair = bottom_pair[np.argsort(bottom_pair[:, 0])] + quad = np.array([top_pair[0], top_pair[1], bottom_pair[1], bottom_pair[0]], dtype=np.float32) + + return {'quad': quad, 'green_mask': green_mask, 'error': None} + + +def _detect_white_lines_on_court(frame, green_mask): + """Detect white lines within the green court area. + + Returns list of line segments [x1, y1, x2, y2] that are on the court surface. """ - # White line detection via color thresholding + edges gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) - # White lines are bright — threshold to isolate them - _, white_mask = cv2.threshold(gray, 180, 255, cv2.THRESH_BINARY) + # White lines are bright pixels on the green surface + _, white = cv2.threshold(gray, 180, 255, cv2.THRESH_BINARY) - # Also use edge detection - blur = cv2.GaussianBlur(gray, (5, 5), 0) - edges = cv2.Canny(blur, 50, 150) + # Only keep white pixels within/near the court + dilated_court = cv2.dilate(green_mask, np.ones((15, 15), np.uint8)) + white_on_court = cv2.bitwise_and(white, dilated_court) - # Combine: edges that are also near white regions - dilated_white = cv2.dilate(white_mask, np.ones((5, 5), np.uint8)) - combined = cv2.bitwise_and(edges, dilated_white) - - # Hough line detection on combined mask - lines = cv2.HoughLinesP(combined, 1, np.pi / 180, threshold=50, - minLineLength=60, maxLineGap=30) + # Hough line detection + lines = cv2.HoughLinesP(white_on_court, 1, np.pi / 180, threshold=40, + minLineLength=50, maxLineGap=25) if lines is None: - return {'segments': [], 'grouped': {'across': [], 'along': []}} + return [] - segments = [line[0].tolist() for line in lines] - - # Classify by angle and group nearby parallel lines - across_lines = [] # roughly horizontal (court width direction) - along_lines = [] # roughly vertical (court length direction) - - for seg in segments: - x1, y1, x2, y2 = seg - angle = np.arctan2(y2 - y1, x2 - x1) * 180 / np.pi - length = np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2) - - if length < 40: - continue - - abs_angle = abs(angle) - if abs_angle < 40 or abs_angle > 140: - across_lines.append(seg) - elif 50 < abs_angle < 130: - along_lines.append(seg) - - def merge_nearby(segs, position_key, threshold=30): - """Merge line segments that are close together (same court line).""" - if not segs: - return [] - - # Sort by mid-position perpendicular to the line direction - def get_pos(s): - if position_key == 'y': - return (s[1] + s[3]) / 2 - return (s[0] + s[2]) / 2 - - sorted_segs = sorted(segs, key=get_pos) - groups = [] - current_group = [sorted_segs[0]] - - for seg in sorted_segs[1:]: - if abs(get_pos(seg) - get_pos(current_group[-1])) < threshold: - current_group.append(seg) - else: - groups.append(current_group) - current_group = [seg] - groups.append(current_group) - - # Merge each group into one representative line - merged = [] - for group in groups: - all_pts = [] - for s in group: - all_pts.extend([(s[0], s[1]), (s[2], s[3])]) - all_pts = np.array(all_pts) - - # Fit a line through all points - if len(all_pts) >= 2: - vx, vy, cx, cy = cv2.fitLine(all_pts, cv2.DIST_L2, 0, 0.01, 0.01).flatten() - # Extend line across frame - t_max = max(np.sqrt((p[0] - cx) ** 2 + (p[1] - cy) ** 2) for p in all_pts) - x1, y1 = int(cx - vx * t_max), int(cy - vy * t_max) - x2, y2 = int(cx + vx * t_max), int(cy + vy * t_max) - mid_pos = get_pos(group[0]) - merged.append({ - 'endpoints': [x1, y1, x2, y2], - 'mid_pos': mid_pos, - 'count': len(group), - }) - - return merged - - grouped = { - 'across': merge_nearby(across_lines, 'y'), - 'along': merge_nearby(along_lines, 'x'), - } - - return {'segments': segments, 'grouped': grouped} - - -def _match_court_lines(grouped, side, frame_w, frame_h): - """Try to match detected line groups to known court lines. - - For a camera at the net looking at one half: - - 'across' lines map to: baseline, kitchen line (and possibly net line) - - 'along' lines map to: left sideline, right sideline, center service line - - Returns dict with points_2d, points_3d for solvePnP, or error. - """ - across = grouped.get('across', []) - along = grouped.get('along', []) - - if len(across) < 2 or len(along) < 2: - return { - 'points_2d': None, 'points_3d': None, - 'error': f'{len(across)} across + {len(along)} along lines (need 2+ each)', - } - - # Sort across lines by vertical position (top of frame = far from camera = baseline) - across_sorted = sorted(across, key=lambda l: l['mid_pos']) - - # Sort along lines by horizontal position - along_sorted = sorted(along, key=lambda l: l['mid_pos']) - - # Known court geometry for this half - # Court: 13.4m x 6.1m, half = 6.7m - # Kitchen (NVZ) = 2.13m from net - if side == 'left': - # Camera at net (X=6.7) looking toward baseline (X=0) - # "across" lines (perpendicular to camera view): - # - baseline at X=0 - # - kitchen at X=6.7-2.13=4.57 - # "along" lines (parallel to camera view): - # - near sideline at Y=0 - # - far sideline at Y=6.1 - # - center service at Y=3.05 - across_3d_x = [0, 4.57] # baseline, kitchen - along_3d_y = [0, 6.1] # near sideline, far sideline - center_service_y = 3.05 - else: - # Camera at net (X=6.7) looking toward baseline (X=13.4) - across_3d_x = [13.4, 6.7 + 2.13] # baseline (far), kitchen - along_3d_y = [0, 6.1] - center_service_y = 3.05 - - # Match: take outermost across lines as baseline (farthest) and kitchen (closest) - # In image: top across line = far = baseline, bottom = near = kitchen - baseline_px = across_sorted[0] # topmost = farthest = baseline - kitchen_px = across_sorted[-1] # bottommost = nearest = kitchen - - # Take outermost along lines as sidelines - left_sideline_px = along_sorted[0] - right_sideline_px = along_sorted[-1] - - # Find intersection points as 2D-3D correspondences - def line_intersect(l1_ep, l2_ep): - x1, y1, x2, y2 = l1_ep - x3, y3, x4, y4 = l2_ep - denom = (x1 - x2) * (y3 - y4) - (y1 - y2) * (x3 - x4) - if abs(denom) < 1e-6: - return None - t = ((x1 - x3) * (y3 - y4) - (y1 - y3) * (x3 - x4)) / denom - return [x1 + t * (x2 - x1), y1 + t * (y2 - y1)] - - points_2d = [] - points_3d = [] - - # Baseline x left sideline - pt = line_intersect(baseline_px['endpoints'], left_sideline_px['endpoints']) - if pt: - points_2d.append(pt) - points_3d.append([across_3d_x[0], along_3d_y[0], 0]) - - # Baseline x right sideline - pt = line_intersect(baseline_px['endpoints'], right_sideline_px['endpoints']) - if pt: - points_2d.append(pt) - points_3d.append([across_3d_x[0], along_3d_y[1], 0]) - - # Kitchen x right sideline - pt = line_intersect(kitchen_px['endpoints'], right_sideline_px['endpoints']) - if pt: - points_2d.append(pt) - points_3d.append([across_3d_x[1], along_3d_y[1], 0]) - - # Kitchen x left sideline - pt = line_intersect(kitchen_px['endpoints'], left_sideline_px['endpoints']) - if pt: - points_2d.append(pt) - points_3d.append([across_3d_x[1], along_3d_y[0], 0]) - - # If we have a center service line, add more correspondences - if len(along_sorted) >= 3: - center_px = along_sorted[len(along_sorted) // 2] - pt = line_intersect(baseline_px['endpoints'], center_px['endpoints']) - if pt: - points_2d.append(pt) - points_3d.append([across_3d_x[0], center_service_y, 0]) - pt = line_intersect(kitchen_px['endpoints'], center_px['endpoints']) - if pt: - points_2d.append(pt) - points_3d.append([across_3d_x[1], center_service_y, 0]) - - # Build list of matched 3D court lines for visualization - matched_lines_3d = [] - # Baseline - matched_lines_3d.append({ - 'name': 'baseline', - 'from': [across_3d_x[0], along_3d_y[0], 0], - 'to': [across_3d_x[0], along_3d_y[1], 0], - }) - # Kitchen - matched_lines_3d.append({ - 'name': 'kitchen', - 'from': [across_3d_x[1], along_3d_y[0], 0], - 'to': [across_3d_x[1], along_3d_y[1], 0], - }) - # Left sideline - matched_lines_3d.append({ - 'name': 'sideline_near', - 'from': [across_3d_x[0], along_3d_y[0], 0], - 'to': [across_3d_x[1], along_3d_y[0], 0], - }) - # Right sideline - matched_lines_3d.append({ - 'name': 'sideline_far', - 'from': [across_3d_x[0], along_3d_y[1], 0], - 'to': [across_3d_x[1], along_3d_y[1], 0], - }) - # Center service (if detected) - if len(along_sorted) >= 3: - matched_lines_3d.append({ - 'name': 'center_service', - 'from': [across_3d_x[0], center_service_y, 0], - 'to': [across_3d_x[1], center_service_y, 0], - }) - - if len(points_2d) < 4: - return { - 'points_2d': None, 'points_3d': None, - 'matched_lines_3d': matched_lines_3d, - 'error': f'Only {len(points_2d)} intersection points found (need >= 4)', - } - - return { - 'points_2d': points_2d, 'points_3d': points_3d, - 'matched_lines_3d': matched_lines_3d, 'error': None, - } + return [line[0].tolist() for line in lines]