feedmill-recounter/app.py

# app.py
"""Flask web UI for feedmill_recounter — port 9050."""

from __future__ import annotations

import os
import subprocess
import threading

from dotenv import load_dotenv
from flask import (
    Flask, render_template, request, redirect,
    url_for, send_file, jsonify,
)
import time as _time
from flask import Response
from werkzeug.utils import secure_filename

from pathlib import Path

from src.job import JobQueue
from src.model_registry import scan_models, scan_model_groups, ModelConfig
from src.preview import extract_thumbnail, extract_sample_frames, probe_video
from src.zone_config import load_zone_config, list_zone_presets, build_static_roi, get_active_zone

load_dotenv()

app = Flask(__name__, template_folder="templates", static_folder="static")
app.config["SECRET_KEY"] = os.getenv("SECRET_KEY", "change-me")
app.config["MAX_CONTENT_LENGTH"] = 2 * 1024 * 1024 * 1024  # 2GB

MODELS_DIR = os.getenv("MODELS_DIR", "./models")
UPLOAD_DIR = os.getenv("UPLOAD_DIR", "./uploads")
OUTPUT_DIR = os.path.abspath(os.getenv("OUTPUT_DIR", "./output"))

os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(OUTPUT_DIR, exist_ok=True)

job_queue = JobQueue(output_dir=OUTPUT_DIR)

VIDEO_EXTENSIONS = ('.mp4', '.avi', '.mkv', '.mov', '.webm')
ZONES_PATH = os.getenv("ZONES_PATH", "./zones.json")


def _parse_model_configs(form):
    """Parse model selection from request form.

    Returns:
        (model_configs, class_filters) tuple.
    """
    selected_stems = form.getlist("model_stems")
    selected_models = form.getlist("models")
    model_configs = []
    class_filters = {}

    if selected_stems:
        groups = scan_model_groups(MODELS_DIR)
        groups_by_stem = {g.stem: g for g in groups}

        for stem in selected_stems:
            if stem not in groups_by_stem:
                continue
            group = groups_by_stem[stem]
            fmt = form.get(f"format_{stem}", group.default_format)
            if fmt not in group.format_paths:
                continue
            model_configs.append(ModelConfig(
                filename=os.path.basename(group.format_paths[fmt]),
                path=group.format_paths[fmt],
                stem=stem,
                known_classes=list(group.known_classes),
            ))
            filter_val = form.get(f"filter_{stem}", "")
            if not filter_val or filter_val in ("default",):
                pass
            elif filter_val == "all":
                class_filters[stem] = None
            else:
                class_filters[stem] = filter_val.split(",")
    elif selected_models:
        models = scan_models(MODELS_DIR)
        by_name = {m.filename: m for m in models}

        for name in selected_models:
            if name in by_name:
                model_configs.append(by_name[name])
                filter_val = form.get(f"filter_{name}", "")
                if not filter_val or filter_val in ("default",):
                    pass
                elif filter_val == "all":
                    class_filters[name] = None
                else:
                    class_filters[name] = filter_val.split(",")

    return model_configs, class_filters


@app.template_filter("basename")
def basename_filter(path):
    """Extract filename from path for templates."""
    return os.path.basename(path)


@app.route("/")
def index():
    model_groups = scan_model_groups(MODELS_DIR)
    return render_template("index.html", model_groups=model_groups, models_dir=MODELS_DIR)


@app.route("/upload", methods=["POST"])
def upload():
    video = request.files.get("video")
    if not video or not video.filename:
        return "No video uploaded", 400

    safe_name = secure_filename(video.filename)
    if not safe_name or not safe_name.lower().endswith(VIDEO_EXTENSIONS):
        return "Invalid video file type", 400

    model_configs, class_filters = _parse_model_configs(request.form)
    if not model_configs:
        return "No models selected", 400

    video_path = os.path.join(UPLOAD_DIR, safe_name)
    base, ext = os.path.splitext(video_path)
    n = 1
    while os.path.exists(video_path):
        video_path = f"{base}_{n}{ext}"
        n += 1
    video.save(video_path)

    output_height = request.form.get("output_height", "")
    zone_mode = request.form.get("zone_mode", "dynamic")
    fixed_zone = None
    zone_preset = ""
    if zone_mode == "fixed" and os.path.isfile(ZONES_PATH):
        preset_name = request.form.get("zone_preset", "")
        zone = get_active_zone(ZONES_PATH, preset_name or None)
        if zone:
            import cv2
            cap = cv2.VideoCapture(video_path)
            vw = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            vh = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            cap.release()
            fixed_zone = build_static_roi(zone, vw, vh)
            zone_preset = preset_name or "default"

    job = job_queue.add_job(
        video_path=video_path,
        model_configs=model_configs,
        class_filters=class_filters,
        zone_mode=zone_mode,
        fixed_zone=fixed_zone,
        zone_preset=zone_preset,
        output_max_height=(int(output_height) if output_height else None),
    )
    return redirect(url_for("status", job_id=job.job_id))


@app.route("/upload/reuse", methods=["POST"])
def upload_reuse():
    """Re-analyze an existing uploaded video."""
    video_path = request.form.get("video_path", "")
    if not video_path or not os.path.isfile(video_path):
        return "Video not found", 400

    model_configs, class_filters = _parse_model_configs(request.form)
    if not model_configs:
        return "No models selected", 400

    output_height = request.form.get("output_height", "")
    zone_mode = request.form.get("zone_mode", "dynamic")
    fixed_zone = None
    zone_preset = ""
    if zone_mode == "fixed" and os.path.isfile(ZONES_PATH):
        preset_name = request.form.get("zone_preset", "")
        zone = get_active_zone(ZONES_PATH, preset_name or None)
        if zone:
            import cv2
            cap = cv2.VideoCapture(video_path)
            vw = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            vh = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            cap.release()
            fixed_zone = build_static_roi(zone, vw, vh)
            zone_preset = preset_name or "default"

    job = job_queue.add_job(
        video_path=video_path,
        model_configs=model_configs,
        class_filters=class_filters,
        zone_mode=zone_mode,
        fixed_zone=fixed_zone,
        zone_preset=zone_preset,
        output_max_height=(int(output_height) if output_height else None),
    )
    return redirect(url_for("status", job_id=job.job_id))


@app.route("/status/<job_id>")
def status(job_id):
    job = job_queue.get_job(job_id)
    if job is None:
        return "Job not found", 404
    return render_template("status.html", job=job)


@app.route("/jobs")
def jobs_list():
    jobs = job_queue.list_jobs()
    return render_template("jobs.html", jobs=jobs)


# Serializes compression encodes: encodes are rare, one global lock is enough.
_COMPRESS_LOCK = threading.Lock()
_AUDIO_BITRATE = 64_000
_MIN_VIDEO_BITRATE = 100_000


def _remove_quiet(path: str) -> None:
    try:
        os.remove(path)
    except OSError:
        pass


def _compressed_cache_valid(out_path: str, file_path: str) -> bool:
    """Cache hit only if the compressed file exists and is newer than source."""
    try:
        return (
            os.path.isfile(out_path)
            and os.path.getmtime(out_path) >= os.path.getmtime(file_path)
        )
    except OSError:
        return False


def _compress_for_download(file_path: str, max_bytes: int = 250 * 1024 * 1024) -> str | None:
    """Compress file_path with ffmpeg if larger than max_bytes.

    Returns path to `<name>_compressed.mp4` next to the input, or None to
    fall back to the original file (small file, ffmpeg missing/failed, or
    result still over the limit).
    """
    if not os.path.isfile(file_path):
        return None
    if os.path.getsize(file_path) <= max_bytes:
        return None
    out_path = os.path.splitext(file_path)[0] + "_compressed.mp4"
    if _compressed_cache_valid(out_path, file_path):
        return out_path
    with _COMPRESS_LOCK:
        if _compressed_cache_valid(out_path, file_path):
            return out_path
        try:
            duration = float(probe_video(file_path).get("duration") or 0)
        except Exception:
            return None
        if duration <= 0:
            return None
        video_bitrate = max(
            int((max_bytes * 8 * 0.92) / duration) - _AUDIO_BITRATE,
            _MIN_VIDEO_BITRATE,
        )
        tmp_path = out_path + ".tmp"
        try:
            subprocess.run(
                ["ffmpeg", "-y", "-i", file_path, "-c:v", "libx264",
                 "-preset", "veryfast", "-b:v", str(video_bitrate),
                 "-c:a", "aac", "-b:a", str(_AUDIO_BITRATE), "-f", "mp4",
                 tmp_path],
                capture_output=True, timeout=3600, check=True,
                stdin=subprocess.DEVNULL,
            )
        except (OSError, subprocess.SubprocessError):
            _remove_quiet(tmp_path)
            return None
        if not os.path.isfile(tmp_path) or os.path.getsize(tmp_path) > max_bytes:
            _remove_quiet(tmp_path)
            return None
        os.replace(tmp_path, out_path)
    return out_path


@app.route("/download/<job_id>/<filename>")
def download(job_id, filename):
    job = job_queue.get_job(job_id)
    if job is None:
        return "Job not found", 404
    if job.status.name in ("PENDING", "RUNNING"):
        return "Job still processing", 409
    safe_filename = secure_filename(filename)
    output_dir_abs = os.path.abspath(job.output_dir)
    file_path = os.path.abspath(os.path.join(job.output_dir, safe_filename))
    if os.path.commonpath([output_dir_abs, file_path]) != output_dir_abs:
        return "File not found", 404
    if not os.path.isfile(file_path):
        return "File not found", 404
    compressed = _compress_for_download(file_path)
    return send_file(compressed or file_path, as_attachment=True)


@app.route("/api/videos")
def api_videos():
    """List previously uploaded videos."""
    videos = []
    for f in sorted(Path(UPLOAD_DIR).iterdir(), key=lambda x: x.stat().st_mtime, reverse=True):
        if f.is_file() and f.suffix.lower() in VIDEO_EXTENSIONS:
            videos.append({
                "filename": f.name,
                "size": f.stat().st_size,
                "mtime": f.stat().st_mtime,
                "path": str(f),
            })
    return jsonify(videos)


@app.route("/api/videos/<filename>", methods=["DELETE"])
def api_delete_video(filename):
    """Delete an uploaded video."""
    safe = secure_filename(filename)
    if not safe or not safe.lower().endswith(VIDEO_EXTENSIONS):
        return jsonify({"error": "invalid filename"}), 400
    file_path = os.path.join(UPLOAD_DIR, safe)
    if not os.path.isfile(file_path):
        return jsonify({"error": "not found"}), 404
    os.remove(file_path)
    return jsonify({"status": "deleted"})


@app.route("/api/models")
def api_models():
    groups = scan_model_groups(MODELS_DIR)
    return jsonify([
        {
            "stem": g.stem,
            "formats": g.formats,
            "format_paths": g.format_paths,
            "known_classes": g.known_classes,
            "default_format": g.default_format,
        }
        for g in groups
    ])


@app.route("/api/zones")
def api_zones():
    """Return zone presets from zones.json."""
    if not os.path.isfile(ZONES_PATH):
        return jsonify([])
    try:
        presets = list_zone_presets(ZONES_PATH)
        return jsonify(presets)
    except Exception:
        return jsonify([])


@app.route("/api/jobs")
def api_jobs():
    return jsonify([{
        "job_id": j.job_id,
        "status": j.status.name,
        "progress": j.progress,
        "video_path": os.path.basename(j.video_path),
        "results": [
            {
                "model": r.model_name,
                "loading": r.loading_count,
                "unloading": r.unloading_count,
                "net": r.net_count,
            }
            for r in j.results
        ],
    } for j in job_queue.list_jobs()])


@app.route("/api/jobs/<job_id>/cancel", methods=["POST"])
def api_cancel_job(job_id):
    """Cancel a running or pending job."""
    success = job_queue.cancel_job(job_id)
    if success:
        return jsonify({"status": "cancelled"})
    job = job_queue.get_job(job_id)
    if job is None:
        return jsonify({"error": "not found"}), 404
    return jsonify({"error": "cannot cancel job in state " + job.status.name}), 400


@app.route("/api/jobs/active")
def api_active_jobs():
    """Return count of running/pending jobs."""
    active = [j for j in job_queue.list_jobs() if j.status.name in ('RUNNING', 'PENDING')]
    return jsonify({
        "count": len(active),
        "jobs": [{"job_id": j.job_id, "status": j.status.name, "video": os.path.basename(j.video_path)} for j in active],
    })


@app.route("/api/jobs/<job_id>", methods=["DELETE"])
def api_delete_job(job_id):
    """Delete a completed/failed/cancelled job."""
    success = job_queue.delete_job(job_id)
    if success:
        return jsonify({"status": "deleted"})
    job = job_queue.get_job(job_id)
    if job is None:
        return jsonify({"error": "not found"}), 404
    return jsonify({"error": "cannot delete job in state " + job.status.name}), 400


@app.route("/api/jobs/<job_id>/preview", methods=["PUT"])
def api_toggle_preview(job_id):
    """Toggle live preview on/off for a running job."""
    job = job_queue.get_job(job_id)
    if job is None:
        return jsonify({"error": "not found"}), 404
    data = request.get_json(silent=True) or {}
    enabled = data.get("enabled", True)
    with job_queue._lock:
        job.preview_enabled = enabled
    return jsonify({"status": "ok", "preview_enabled": enabled})


@app.route("/api/jobs/<job_id>")
def api_job_detail(job_id):
    job = job_queue.get_job(job_id)
    if job is None:
        return jsonify({"error": "not found"}), 404
    return jsonify({
        "job_id": job.job_id,
        "status": job.status.name,
        "progress": job.progress,
        "current_model": job.current_model,
        "results": [
            {
                "model_name": r.model_name,
                "loading_count": r.loading_count,
                "unloading_count": r.unloading_count,
                "net_count": r.net_count,
                "box_loading_count": r.box_loading_count,
                "box_unloading_count": r.box_unloading_count,
                "box_net_count": r.box_net_count,
                "batch_count": r.batch_count,
                "frame_count": r.frame_count,
                "duration_seconds": r.duration_seconds,
                "output_path": os.path.basename(r.output_path) if r.output_path else None,
            }
            for r in job.results
        ],
        "error": job.error,
        "live_fps": job.live_fps,
        "live_loading": job.live_loading,
        "live_unloading": job.live_unloading,
        "live_frame": job.live_frame,
    })


@app.route("/preview/thumb/<job_id>")
def preview_thumb(job_id):
    """Serve a thumbnail JPEG for a job's video (cached in job output_dir)."""
    job = job_queue.get_job(job_id)
    if job is None:
        return "Job not found", 404
    thumb_path = os.path.join(job.output_dir, "thumbnail.jpg")
    if not os.path.isfile(thumb_path):
        try:
            extract_thumbnail(job.video_path, thumb_path)
        except (FileNotFoundError, RuntimeError):
            return "Could not generate thumbnail", 404
    return send_file(thumb_path, mimetype="image/jpeg")


@app.route("/api/jobs/<job_id>/samples")
def api_job_samples(job_id):
    """Return JSON array of sample frame paths (extract if not cached)."""
    job = job_queue.get_job(job_id)
    if job is None:
        return jsonify({"error": "not found"}), 404
    samples_dir = os.path.join(job.output_dir, "samples")
    cache_file = os.path.join(samples_dir, ".cached")
    if not os.path.isfile(cache_file):
        try:
            paths = extract_sample_frames(job.video_path, samples_dir)
            with open(cache_file, "w") as f:
                f.write("\n".join(paths))
        except (FileNotFoundError, RuntimeError):
            return jsonify({"error": "could not generate samples"}), 404
    else:
        with open(cache_file) as f:
            paths = [line.strip() for line in f if line.strip()]
    rel_paths = [os.path.relpath(p, job.output_dir) for p in paths]
    return jsonify(rel_paths)


@app.route("/api/preview/<job_id>")
def api_preview(job_id):
    job = job_queue.get_job(job_id)
    if job is None:
        return jsonify({"error": "not found"}), 404
    preview_path = job.preview_path
    if not preview_path:
        return jsonify({"error": "no preview path"}), 404

    def generate():
        consecutive_fails = 0
        MAX_FAILS = 30
        while True:
            try:
                with open(preview_path, "rb") as f:
                    jpeg = f.read()
                consecutive_fails = 0
                yield (b"--frame\r\n"
                       b"Content-Type: image/jpeg\r\n\r\n" + jpeg + b"\r\n")
            except FileNotFoundError:
                consecutive_fails += 1
                if consecutive_fails >= MAX_FAILS:
                    return
                _time.sleep(1.0)
                continue
            except Exception:
                consecutive_fails += 1
                if consecutive_fails >= MAX_FAILS:
                    return
                _time.sleep(0.5)
                continue
            _time.sleep(0.05)
    return Response(generate(), mimetype="multipart/x-mixed-replace; boundary=frame")


def main():
    host = os.getenv("WEB_HOST", "0.0.0.0")
    port = int(os.getenv("WEB_PORT", "9050"))
    debug = os.getenv("FLASK_DEBUG", "false").lower() == "true"

    print(f"Feedmill Recounter web UI: http://{host}:{port}")
    app.run(host=host, port=port, debug=debug)


if __name__ == "__main__":
    main()