pickle_vision/src/web/app.py

"""
Flask web application for Pickle Vision referee system.
Three tabs: Detection, Court, Trajectory.
"""

import os
import cv2
import time
import json
import base64
import threading
import numpy as np
from pathlib import Path
from flask import Flask, Response, render_template, jsonify, request

app = Flask(__name__, template_folder=os.path.join(os.path.dirname(__file__), 'templates'))

# Global state — set by main.py before starting
state = {
    'cameras': {},           # sensor_id -> {frame: bytes, lock, fps, detections}
    'trajectory': None,      # TrajectoryModel instance
    'event_detector': None,
    'calibrators': {},       # sensor_id -> CameraCalibrator
    'calibration_dir': '',   # path to save calibration files
    'events': [],            # recent VAR events
    'last_var': None,        # last VAR event with snapshot
    'calibrate_fn': None,    # callback for calibration trigger
}


@app.route('/')
@app.route('/<tab>')
def index(tab='camera'):
    return render_template('index.html', active_tab=tab)


@app.route('/frame/<int:sensor_id>')
def frame(sensor_id):
    cam = state['cameras'].get(sensor_id)
    if not cam:
        return "Camera not found", 404
    with cam['lock']:
        jpg = cam.get('frame')
    if jpg is None:
        return "No frame yet", 503
    return Response(jpg, mimetype='image/jpeg',
                    headers={'Cache-Control': 'no-cache, no-store'})


@app.route('/api/stats')
def api_stats():
    result = {}
    for k, v in state['cameras'].items():
        result[str(k)] = {
            'fps': v.get('fps', 0),
            'detections': v.get('detections', 0),
        }
    return jsonify(result)


@app.route('/api/trajectory')
def api_trajectory():
    traj = state.get('trajectory')
    if traj is None:
        return jsonify({'points': [], 'speed': None, 'landing': None})

    points = traj.get_recent(120)
    speed = traj.get_speed()
    landing = traj.predict_landing()

    return jsonify({
        'points': points,
        'speed': speed,
        'landing': {'x': landing[0], 'y': landing[1], 't': landing[2]} if landing else None,
    })


@app.route('/api/events')
def api_events():
    return jsonify(state.get('events', [])[-20:])


@app.route('/api/var/last')
def api_var_last():
    """Get last VAR event with snapshot image."""
    last = state.get('last_var')
    if not last:
        return jsonify(None)

    result = {
        'event': last['event'],
        'ago_seconds': time.time() - last['event']['timestamp'],
        'snapshot_b64': last.get('snapshot_b64'),
    }
    return jsonify(result)


@app.route('/api/calibration/status')
def api_calibration_status():
    cals = {}
    for sid, cal in state.get('calibrators', {}).items():
        cals[str(sid)] = cal.calibrated
    # System is ready when at least one camera is calibrated
    any_calibrated = any(cals.values()) if cals else False
    return jsonify({**cals, 'system_ready': any_calibrated})


@app.route('/api/calibration/trigger', methods=['POST'])
def api_calibration_trigger():
    """Trigger one-click court calibration from current camera frames."""
    fn = state.get('calibrate_fn')
    if fn is None:
        return jsonify({'ok': False, 'error': 'Calibration not available'}), 500

    import traceback
    try:
        result = fn()
        # Check if any camera failed
        any_ok = any(r.get('ok') for r in result.values())
        return jsonify({'ok': any_ok, 'result': result})
    except Exception as e:
        tb = traceback.format_exc()
        print(f"[CALIBRATION ERROR]\n{tb}")
        return jsonify({'ok': False, 'error': f'{type(e).__name__}: {e}'}), 500


@app.route('/api/calibration/data')
def api_calibration_data():
    """Return calibration data for 3D scene reconstruction."""
    cals = state.get('calibrators', {})
    result = {}
    for sid, cal in cals.items():
        if not cal.calibrated:
            continue
        # Camera position in world coordinates
        cam_pos = (-cal.rotation_matrix.T @ cal.translation_vec).flatten()
        # Camera look direction (Z axis of camera in world coords)
        look_dir = cal.rotation_matrix.T @ np.array([0, 0, 1.0])

        result[str(sid)] = {
            'position': cam_pos.tolist(),
            'look_direction': look_dir.tolist(),
            'focal_length': float(cal.camera_matrix[0, 0]),
            'image_size': [int(cal.camera_matrix[0, 2] * 2),
                           int(cal.camera_matrix[1, 2] * 2)],
        }
    return jsonify(result)


def run(host='0.0.0.0', port=8080):
    app.run(host=host, port=port, threaded=True)