457 lines
15 KiB
Python
457 lines
15 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,
|
|
)
|
|
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
|
|
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.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)
|
|
|
|
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,
|
|
)
|
|
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
|
|
|
|
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,
|
|
)
|
|
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/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,
|
|
"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", "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()
|