jetson
f9e411f407
- flask-socketio + threading mode (no eventlet monkey-patching) - preview_queue in pipeline: downscale to 640p, JPEG encode, bounded queue - encoder thread per job reads queue → SocketIO broadcast to room - Frontend: SocketIO client replaces HTTP polling, handles preview_frame event - Configurable: preview_every_n, preview_max_dim, preview_jpeg_quality - 62/62 tests pass
374 lines
11 KiB
Python
374 lines
11 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 flask_socketio import SocketIO, join_room, leave_room
|
|
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
|
|
|
|
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
|
|
|
|
socketio = SocketIO(app, cors_allowed_origins="*", async_mode='threading')
|
|
|
|
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')
|
|
|
|
|
|
@socketio.on('connect')
|
|
def on_connect():
|
|
pass
|
|
|
|
|
|
@socketio.on('disconnect')
|
|
def on_disconnect():
|
|
pass
|
|
|
|
|
|
@socketio.on('join_job')
|
|
def on_join_job(data):
|
|
job_id = data.get('job_id')
|
|
if job_id:
|
|
join_room(job_id)
|
|
|
|
|
|
@socketio.on('leave_job')
|
|
def on_leave_job(data):
|
|
job_id = data.get('job_id')
|
|
if job_id:
|
|
leave_room(job_id)
|
|
|
|
|
|
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)
|
|
|
|
job = job_queue.add_job(
|
|
video_path=video_path,
|
|
model_configs=model_configs,
|
|
class_filters=class_filters,
|
|
)
|
|
job.start_preview_broadcaster(job.job_id, socketio)
|
|
|
|
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
|
|
|
|
job = job_queue.add_job(
|
|
video_path=video_path,
|
|
model_configs=model_configs,
|
|
class_filters=class_filters,
|
|
)
|
|
job.start_preview_broadcaster(job.job_id, socketio)
|
|
|
|
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/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/<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}")
|
|
socketio.run(app, host=host, port=port, debug=debug, allow_unsafe_werkzeug=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|