237 lines
7.1 KiB
Python
237 lines
7.1 KiB
Python
# 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": r.model_name,
|
|
"loading": r.loading_count,
|
|
"unloading": r.unloading_count,
|
|
"net": r.net_count,
|
|
"output": 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
|
|
# latest_frame is not set yet (Task 5 adds it) — return 404 for now
|
|
return "No frame available", 404
|
|
|
|
|
|
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()
|