bytetrack-counter-dashboard/recounting_dashboard.py

#!/usr/bin/env python3
"""
Recounting dashboard — consumes live counter + recount APIs,
lists OUTPUT_DIR MP4 files, streams MP4 via ffmpeg RTSP → go2rtc for preview.
"""
import json
import logging
import os
import re
import subprocess
import threading
import time
import requests
from datetime import datetime
from pathlib import Path

from flask import Flask, render_template, jsonify, request, Response, send_file
from werkzeug.serving import WSGIRequestHandler

from dotenv import load_dotenv
load_dotenv()

log = logging.getLogger('werkzeug')
log.setLevel(logging.ERROR)

app = Flask(__name__, template_folder="templates")
app.config["SECRET_KEY"] = os.getenv("SECRET_KEY", "change-me-in-production")

OUTPUT_DIR = os.getenv("OUTPUT_DIR", "/opt/bytetrack-counter")
LIVE_API_URL = os.getenv("LIVE_API_URL", "http://localhost:5000")
RECOUNT_API_URL = os.getenv("RECOUNT_API_URL", "http://localhost:5001")
GO2RTC_URL = os.getenv("GO2RTC_URL", "http://localhost:1984")
RTSP_PORT = os.getenv("RTSP_PORT", "8554")
STREAM_NAME = os.getenv("RECOUNT_STREAM_NAME", "stream_1")

SITE_NAME = os.getenv("SITE_NAME", "RECOUNT")

DASHBOARD_PORT = int(os.getenv("RECOUNTING_DASHBOARD_PORT", "5002"))
DASHBOARD_HOST = os.getenv("DASHBOARD_HOST", "0.0.0.0")
FLASK_DEBUG = os.getenv("FLASK_DEBUG", "false").lower() == "true"

_http_session = requests.Session()
_http_session.timeout = 3

_FILENAME_RE = re.compile(r"batch[_-]?(\d+)", re.IGNORECASE)

_streaming_mp4 = None
_streaming_lock = threading.Lock()
_streaming_proc = None
_current_file = None


def _parse_filename(filename):
    m = _FILENAME_RE.search(filename)
    if m:
        return int(m.group(1)), None
    return None, None


def _api_get(base_url, path, default=None):
    try:
        resp = _http_session.get(f"{base_url}{path}")
        if resp.status_code == 200:
            return resp.json()
    except Exception:
        pass
    return default


def _folder_date(filepath):
    output = Path(OUTPUT_DIR)
    rel = Path(filepath).relative_to(output)
    parts = rel.parts
    if len(parts) >= 2:
        return parts[0]
    return None


def _stop_ffmpeg_locked():
    global _streaming_proc, _streaming_mp4
    if _streaming_proc:
        try:
            _streaming_proc.kill()
            _streaming_proc.wait(timeout=3)
        except Exception:
            pass
        _streaming_proc = None
    _streaming_mp4 = None


@app.route("/")
def index():
    return render_template("recounting.html", site_name=SITE_NAME, output_dir=OUTPUT_DIR)


@app.route("/api/live-progress")
def api_live_progress():
    data = _api_get(LIVE_API_URL, "/api/current-batch")
    if data and data.get("success"):
        return jsonify(data)
    return jsonify({"success": False, "count": 0, "batch_number": None, "error": "Live unreachable"}), 200


@app.route("/api/recount-progress")
def api_recount_progress():
    data = _api_get(RECOUNT_API_URL, "/api/current-batch")
    if data and data.get("success"):
        return jsonify(data)
    return jsonify({"success": False, "count": 0, "batch_number": None, "error": "Recount unreachable"}), 200


@app.route("/api/mp4-files")
def api_mp4_files():
    folders = {}
    output = Path(OUTPUT_DIR)
    if not output.exists():
        return jsonify([])

    for d in sorted(output.iterdir(), reverse=True):
        if not d.is_dir():
            continue
        folder_name = d.name
        mp4s = sorted(d.glob("*.mp4"), key=lambda p: _parse_filename(p.name)[0] or 0)
        if not mp4s:
            continue
        files = []
        for f in mp4s:
            st = f.stat()
            files.append({
                "name": f.name,
                "path": str(f),
                "size": st.st_size,
                "mtime": datetime.fromtimestamp(st.st_mtime).isoformat(),
            })
        folders[folder_name] = files

    result = []
    for folder, files in sorted(folders.items(), reverse=True):
        result.append({"folder": folder, "files": files})
    return jsonify(result)


@app.route("/api/start-recount", methods=["POST"])
def start_recount():
    global _streaming_mp4, _streaming_proc, _current_file
    data = request.get_json(force=True) or {}
    mp4_path = data.get("path", "")
    if not mp4_path:
        return jsonify({"success": False, "error": "Missing 'path'"}), 400
    if not os.path.isfile(mp4_path):
        return jsonify({"success": False, "error": f"File not found: {mp4_path}"}), 404

    with _streaming_lock:
        _stop_ffmpeg_locked()

    try:
        _http_session.post(f"{RECOUNT_API_URL}/api/reset", timeout=5)
        time.sleep(1)
    except Exception:
        pass

    with _streaming_lock:
        _current_file = mp4_path
        cmd = [
            "ffmpeg", "-re", "-i", mp4_path,
            "-vf", "scale=640:480",
            "-c:v", "libx264", "-preset", "veryfast",
            "-c:a", "aac",
            "-f", "rtsp", f"rtsp://localhost:{RTSP_PORT}/{STREAM_NAME}",
        ]
        _streaming_proc = subprocess.Popen(cmd, stderr=subprocess.DEVNULL)
        _streaming_mp4 = mp4_path

    return jsonify({
        "success": True,
        "stream_url": "/api/proxy-stream",
        "file": Path(mp4_path).name,
        "path": mp4_path,
    })


@app.route("/api/proxy-stream")
def api_proxy_stream():
    try:
        resp = _http_session.get(f"{RECOUNT_API_URL}/api/live-video", stream=True, timeout=5)
        if resp.status_code != 200:
            return Response("stream unavailable", status=502)

        def generate():
            for chunk in resp.iter_content(chunk_size=8192):
                if chunk:
                    yield chunk
            resp.close()

        return Response(
            generate(),
            mimetype=resp.headers.get("Content-Type", "multipart/x-mixed-replace; boundary=frame"),
        )
    except Exception:
        return Response("stream unavailable", status=502)


@app.route("/api/stop-recount", methods=["POST"])
def stop_recount():
    global _current_file
    with _streaming_lock:
        _stop_ffmpeg_locked()
        _current_file = None
    return jsonify({"success": True})


@app.route("/api/state")
def api_state():
    with _streaming_lock:
        streaming = _streaming_proc is not None and _streaming_proc.poll() is None
        mp4 = _streaming_mp4
        current_file = _current_file
    result = {
        "streaming": streaming,
    }
    if current_file:
        result["file"] = current_file
        result["file_name"] = Path(current_file).name
        result["stream_url"] = "/api/proxy-stream"
        batch_num, date_str = _parse_filename(Path(current_file).name)
        if batch_num is not None:
            result["batch_number"] = batch_num
            result["batch_date"] = date_str
    return jsonify(result)


@app.route("/api/download-mp4")
def api_download_mp4():
    fpath = request.args.get("file", "")
    if not fpath or not os.path.isfile(fpath):
        return jsonify({"success": False, "error": "File not found"}), 404
    return send_file(fpath, as_attachment=True, download_name=Path(fpath).name)


@app.route("/api/batch-result")
def api_batch_result():
    fpath = request.args.get("file", "")
    filename = Path(fpath).name
    batch_num, _ = _parse_filename(filename)
    if batch_num is None:
        return jsonify({"success": False, "error": "Cannot parse batch number from filename"}), 200

    date_str = _folder_date(fpath)
    if not date_str:
        return jsonify({"success": False, "error": "Cannot determine date from folder"}), 200

    def _match_in_batches(batches):
        for b in batches:
            if b.get("batch_number") == batch_num:
                return b
        return None

    data = _api_get(LIVE_API_URL, f"/api/day-detail/{date_str}")
    if data:
        b = _match_in_batches(data.get("batches", []))
        if b:
            return jsonify({
                "success": True,
                "batch_number": batch_num,
                "date": date_str,
                "count": b["count"],
                "start_time": b.get("start_time"),
                "end_time": b.get("end_time"),
            })

    recent = _api_get(LIVE_API_URL, "/api/recent-batches?limit=100")
    if recent:
        for entry in recent if isinstance(recent, list) else recent.get("batches", []):
            if isinstance(entry, dict) and entry.get("batch_number") == batch_num:
                return jsonify({
                    "success": True,
                    "batch_number": batch_num,
                    "date": entry.get("date", date_str),
                    "count": entry["count"],
                    "start_time": entry.get("start_time"),
                    "end_time": entry.get("end_time"),
                })

    return jsonify({"success": False, "error": f"Batch #{batch_num} not found"}), 200


if __name__ == "__main__":
    WSGIRequestHandler.protocol_version = "HTTP/1.1"
    print(f"Recounting dashboard at http://{DASHBOARD_HOST}:{DASHBOARD_PORT}")
    print(f"Live API: {LIVE_API_URL}")
    print(f"Recount API: {RECOUNT_API_URL}")
    print(f"go2rtc: {GO2RTC_URL}")
    print(f"RTSP: rtsp://localhost:{RTSP_PORT}/{STREAM_NAME}")
    print(f"OUTPUT_DIR: {OUTPUT_DIR}")
    app.run(host=DASHBOARD_HOST, port=DASHBOARD_PORT, debug=FLASK_DEBUG)