feedmill-recounter/app.py

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

from __future__ import annotations

import os

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

from src.job import JobQueue
from src.model_registry import scan_models
from src.preview import extract_thumbnail, extract_sample_frames

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.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)


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


@app.route("/")
def index():
    models = scan_models(MODELS_DIR)
    return render_template("index.html", models=models, 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((".mp4", ".avi", ".mkv", ".mov", ".webm")):
        return "Invalid video file type", 400

    selected_models = request.form.getlist("models")
    models = scan_models(MODELS_DIR)
    by_name = {m.filename: m for m in models}

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

    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)

    job = job_queue.add_job(
        video_path=video_path,
        model_configs=model_configs,
        class_filters=class_filters,
    )

    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)


@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
    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
    return send_file(file_path, as_attachment=True)


@app.route("/api/models")
def api_models():
    models = scan_models(MODELS_DIR)
    return jsonify([
        {
            "filename": m.filename,
            "stem": m.stem,
            "known_classes": m.known_classes,
        }
        for m in models
    ])


@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>")
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,
                "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,
    })


@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/jobs/<job_id>/frame")
def api_job_frame(job_id):
    """Return current latest_frame as JPEG (live preview during processing)."""
    job = job_queue.get_job(job_id)
    if job is None:
        return "Job not found", 404
    if job.latest_frame is None:
        return "No frame available", 404
    from flask import Response
    return Response(job.latest_frame, mimetype="image/jpeg")


def main():
    host = os.getenv("WEB_HOST", "0.0.0.0")
    port = int(os.getenv("WEB_PORT", "9000"))
    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()