#!/usr/bin/env python3 """ Pickle Vision - Referee System main entry point for Jetson. Dual CSI cameras, real-time YOLO detection, trajectory tracking, VAR triggers. """ import sys import os import cv2 import time import base64 import argparse import threading import numpy as np # Add project root to path sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from ultralytics import YOLO from src.streaming.camera_reader import CameraReader from src.streaming.ring_buffer import FrameRingBuffer from src.physics.trajectory import TrajectoryModel from src.physics.event_detector import EventDetector from src.calibration.camera_calibrator import ( CameraCalibrator, get_half_court_3d_points, COURT_LENGTH, COURT_WIDTH, HALF_COURT_LENGTH ) from src.web.app import app, state BALL_CLASS_ID = 32 # sports ball in COCO # Global references set in main() _cam_readers = {} _args = None def auto_calibrate(): """One-click calibration: detect court rectangle from current frames, compute camera 3D position via solvePnP using known court dimensions. Returns debug images showing detected court quad + computed camera position. """ results = {} for sensor_id, reader in _cam_readers.items(): frame = reader.grab() if frame is None: results[str(sensor_id)] = {'ok': False, 'error': 'No frame available'} continue h, w = frame.shape[:2] side = 'left' if sensor_id == 0 else 'right' debug_frame = frame.copy() # Detect court rectangle — find 4 corners of the court detection = _detect_court_corners(frame, side) corners_pixel = detection.get('corners') if detection else None if corners_pixel is None: error_detail = detection.get('error', 'Unknown') if detection else 'No court detected' # Draw error on debug frame cv2.putText(debug_frame, f"FAILED: {error_detail}", (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}: {error_detail}', 'debug_image': base64.b64encode(jpeg.tobytes()).decode('ascii'), } continue # Draw detected court quadrilateral on debug frame pts = corners_pixel.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) # Get known 3D coordinates for this half-court corners_3d = get_half_court_3d_points(side) # Calibrate — errors propagate cal = CameraCalibrator() cal.calibrate( np.array(corners_pixel, dtype=np.float32), corners_3d, w, h ) # Camera position in world coordinates cam_pos = (-cal.rotation_matrix.T @ cal.translation_vec).flatten() # Draw camera position on debug frame cv2.putText(debug_frame, f"Camera: X={cam_pos[0]:.2f}m Y={cam_pos[1]:.2f}m Z={cam_pos[2]:.2f}m", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) cv2.putText(debug_frame, f"Court: {side} half", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1) _, jpeg = cv2.imencode('.jpg', debug_frame, [cv2.IMWRITE_JPEG_QUALITY, 85]) debug_b64 = base64.b64encode(jpeg.tobytes()).decode('ascii') # Save calibration 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 results[str(sensor_id)] = { 'ok': True, 'camera_position': cam_pos.tolist(), 'debug_image': debug_b64, } 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_corners(frame, side): """Detect 4 corners of the court rectangle from camera frame. Uses edge detection + Hough lines to find the court boundaries, then finds 4 intersection points forming the court quadrilateral. Returns dict with: corners: 4x2 numpy array (TL, TR, BR, BL) or None error: description string if detection failed """ gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) blur = cv2.GaussianBlur(gray, (5, 5), 0) edges = cv2.Canny(blur, 50, 150) lines = cv2.HoughLinesP(edges, 1, np.pi / 180, threshold=80, minLineLength=100, maxLineGap=20) if lines is None or len(lines) < 4: n = 0 if lines is None else len(lines) return {'corners': None, 'error': f'Found {n} lines, need at least 4'} # Classify into horizontal / vertical horizontals = [] verticals = [] for line in lines: x1, y1, x2, y2 = line[0] angle = abs(np.arctan2(y2 - y1, x2 - x1) * 180 / np.pi) if angle < 30 or angle > 150: horizontals.append(line[0]) elif 60 < angle < 120: verticals.append(line[0]) if len(horizontals) < 2 or len(verticals) < 2: return {'corners': None, 'error': f'{len(horizontals)} horiz + {len(verticals)} vert lines, need 2+ each'} # Take outermost lines as court boundaries h_sorted = sorted(horizontals, key=lambda l: (l[1] + l[3]) / 2) v_sorted = sorted(verticals, key=lambda l: (l[0] + l[2]) / 2) top, bottom = h_sorted[0], h_sorted[-1] left, right = v_sorted[0], v_sorted[-1] def intersect(l1, l2): x1, y1, x2, y2 = l1 x3, y3, x4, y4 = l2 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)] corners = [ intersect(top, left), intersect(top, right), intersect(bottom, right), intersect(bottom, left), ] if any(c is None for c in corners): return {'corners': None, 'error': 'Court lines are parallel, cannot find intersections'} return {'corners': np.array(corners, dtype=np.float32), 'error': None} def _capture_var_snapshot(frame, event): """Create a snapshot for VAR event and store it in state.""" # Draw VAR overlay on frame snapshot = frame.copy() cv2.putText(snapshot, "VAR DETECT", (10, 40), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (0, 0, 255), 3) cv2.putText(snapshot, f"Line: {event['line']} Dist: {event['distance_m']*100:.0f}cm", (10, 80), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2) _, jpeg = cv2.imencode('.jpg', snapshot, [cv2.IMWRITE_JPEG_QUALITY, 85]) b64 = base64.b64encode(jpeg.tobytes()).decode('ascii') state['last_var'] = { 'event': event, 'snapshot_b64': b64, } def detection_loop(cam_readers, model, conf_threshold, ring_buffer): """Main detection loop: alternate cameras, run YOLO, update state.""" frame_counts = {sid: 0 for sid in cam_readers} start_times = {sid: time.time() for sid in cam_readers} trajectory = state['trajectory'] event_detector = state['event_detector'] calibrators = state['calibrators'] while True: for sensor_id, reader in cam_readers.items(): cam = state['cameras'][sensor_id] frame = reader.grab() if frame is None: continue now = time.time() # Store raw frame in ring buffer for VAR ring_buffer.push(frame, now, sensor_id) # Only run detection if camera is calibrated cal = calibrators.get(sensor_id) is_calibrated = cal and cal.calibrated det_count = 0 best_detection = None best_conf = 0 if is_calibrated: # YOLO detection — only after calibration results = model(frame, verbose=False, classes=[BALL_CLASS_ID], conf=conf_threshold) for r in results: for box in r.boxes: x1, y1, x2, y2 = map(int, box.xyxy[0]) conf = float(box.conf[0]) cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 3) label = f"Ball {conf:.0%}" (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2) cv2.rectangle(frame, (x1, y1 - th - 10), (x1 + tw, y1), (0, 255, 0), -1) cv2.putText(frame, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 2) det_count += 1 if conf > best_conf: best_conf = conf cx = (x1 + x2) / 2 cy = (y1 + y2) / 2 diameter = max(x2 - x1, y2 - y1) best_detection = (cx, cy, diameter, conf) # Update trajectory and check VAR if best_detection: px, py, diam, conf = best_detection pos_3d = cal.pixel_to_3d(px, py, diam) if pos_3d: trajectory.add_observation( pos_3d[0], pos_3d[1], pos_3d[2], now, frame_counts[sensor_id], sensor_id, conf ) # Check for close calls event = event_detector.check(trajectory) if event: print(f"[VAR] Close call! Line: {event['line']}, " f"Distance: {event['distance_m']*100:.0f}cm") state['events'].append(event) _capture_var_snapshot(frame, event) # FPS tracking frame_counts[sensor_id] += 1 elapsed = time.time() - start_times[sensor_id] fps_actual = frame_counts[sensor_id] / elapsed if elapsed > 0 else 0 # HUD overlay cv2.putText(frame, f"CAM {sensor_id} | FPS: {fps_actual:.1f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 255), 2) if det_count > 0: cv2.putText(frame, f"Balls: {det_count}", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2) # Encode JPEG and update shared state _, jpeg = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 80]) with cam['lock']: cam['frame'] = jpeg.tobytes() cam['fps'] = fps_actual cam['detections'] = det_count if frame_counts[sensor_id] % 150 == 0: speed = trajectory.get_speed() speed_str = f"{speed:.1f} m/s" if speed else "N/A" print(f"[CAM {sensor_id}] Frame {frame_counts[sensor_id]}, " f"FPS: {fps_actual:.1f}, Det: {det_count}, Speed: {speed_str}") def main(): global _cam_readers, _args parser = argparse.ArgumentParser(description='Pickle Vision Referee System') parser.add_argument('--width', type=int, default=1280) parser.add_argument('--height', type=int, default=720) parser.add_argument('--fps', type=int, default=30) parser.add_argument('--model', type=str, default='yolov8n.pt') parser.add_argument('--conf', type=float, default=0.25) parser.add_argument('--port', type=int, default=8080) parser.add_argument('--buffer-seconds', type=int, default=10, help='Ring buffer size in seconds for VAR clips') parser.add_argument('--calibration-dir', type=str, default=os.path.join(os.path.dirname(__file__), 'config'), help='Directory with calibration JSON files') args = parser.parse_args() _args = args # Load YOLO model print(f"Loading YOLO model: {args.model}") model = YOLO(args.model) try: model.to("cuda") print("Inference on CUDA") except Exception: print("CUDA unavailable, using CPU") # Initialize shared state state['trajectory'] = TrajectoryModel(fps=args.fps) state['event_detector'] = EventDetector(trigger_distance_m=0.3) state['calibration_dir'] = args.calibration_dir state['calibrate_fn'] = auto_calibrate ring_buffer = FrameRingBuffer(max_seconds=args.buffer_seconds, fps=args.fps) # Start with empty calibrators — user must calibrate via UI os.makedirs(args.calibration_dir, exist_ok=True) for sensor_id in [0, 1]: state['calibrators'][sensor_id] = CameraCalibrator() # Start camera readers cam_readers = {} for sensor_id in [0, 1]: state['cameras'][sensor_id] = { 'frame': None, 'lock': threading.Lock(), 'fps': 0, 'detections': 0 } cam_readers[sensor_id] = CameraReader(sensor_id, args.width, args.height, args.fps) _cam_readers = cam_readers # Wait for at least one camera print("Waiting for cameras...") for _ in range(100): if any(r.grab() is not None for r in cam_readers.values()): break time.sleep(0.1) # Start detection loop det_thread = threading.Thread( target=detection_loop, args=(cam_readers, model, args.conf, ring_buffer), daemon=True ) det_thread.start() time.sleep(2) print(f"\n{'=' * 50}") print(f" Pickle Vision Referee System") print(f" Open in browser: http://192.168.1.253:{args.port}") print(f"{'=' * 50}\n") app.run(host='0.0.0.0', port=args.port, threaded=True) if __name__ == '__main__': main()