294 lines
10 KiB
Python
294 lines
10 KiB
Python
# src/job.py
|
|
"""Job queue — manages async video processing jobs."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from dataclasses import dataclass, field
|
|
from enum import Enum, auto
|
|
from pathlib import Path
|
|
|
|
from src.model_registry import ModelConfig, find_truck_detector
|
|
from src.pipeline import run_pipeline, run_merged_pipeline, PipelineResult, MergedPipelineResult
|
|
from src.truck_roi import TruckROI
|
|
|
|
|
|
class JobStatus(Enum):
|
|
PENDING = auto()
|
|
RUNNING = auto()
|
|
COMPLETED = auto()
|
|
FAILED = auto()
|
|
CANCELLED = auto()
|
|
|
|
|
|
@dataclass
|
|
class JobResult:
|
|
"""Result from a single model run within a job."""
|
|
|
|
model_name: str
|
|
output_path: str
|
|
loading_count: int
|
|
unloading_count: int
|
|
net_count: int
|
|
batch_count: int
|
|
frame_count: int
|
|
duration_seconds: float
|
|
error: str | None = None
|
|
|
|
|
|
@dataclass
|
|
class Job:
|
|
"""A processing job that runs one or more model configs on a video."""
|
|
|
|
job_id: str
|
|
video_path: str
|
|
model_configs: list[ModelConfig]
|
|
class_filters: dict[str, list[str] | None] = field(default_factory=dict)
|
|
output_dir: str = ""
|
|
status: JobStatus = JobStatus.PENDING
|
|
progress: float = 0.0
|
|
current_model: str = ""
|
|
results: list[JobResult] = field(default_factory=list)
|
|
error: str | None = None
|
|
created_at: float = field(default_factory=time.time)
|
|
completed_at: float | None = None
|
|
latest_frame: bytes | None = None
|
|
preview_path: str = ""
|
|
live_fps: float = 0.0
|
|
live_loading: int = 0
|
|
live_unloading: int = 0
|
|
live_frame: int = 0
|
|
preview_enabled: bool = True
|
|
zone_mode: str = "dynamic"
|
|
fixed_zone: TruckROI | None = None
|
|
zone_preset: str = ""
|
|
|
|
|
|
class JobQueue:
|
|
"""Thread-safe job queue with background worker."""
|
|
|
|
def __init__(self, output_dir: str = "./output") -> None:
|
|
self._output_dir = Path(output_dir)
|
|
self._output_dir.mkdir(parents=True, exist_ok=True)
|
|
self._jobs: dict[str, Job] = {}
|
|
self._lock = threading.Lock()
|
|
self._threads: list[threading.Thread] = []
|
|
|
|
def add_job(
|
|
self,
|
|
video_path: str,
|
|
model_configs: list[ModelConfig],
|
|
class_filters: dict[str, list[str] | None] | None = None,
|
|
zone_mode: str = "dynamic",
|
|
fixed_zone: TruckROI | None = None,
|
|
zone_preset: str = "",
|
|
) -> Job:
|
|
"""Create a new job and enqueue it. Returns the Job (processing starts immediately)."""
|
|
job_id = f"job-{uuid.uuid4().hex[:8]}"
|
|
job = Job(
|
|
job_id=job_id,
|
|
video_path=video_path,
|
|
model_configs=list(model_configs),
|
|
class_filters=class_filters or {},
|
|
output_dir=str(self._output_dir / job_id),
|
|
zone_mode=zone_mode,
|
|
fixed_zone=fixed_zone,
|
|
zone_preset=zone_preset,
|
|
)
|
|
Path(job.output_dir).mkdir(parents=True, exist_ok=True)
|
|
|
|
job.preview_path = f"/tmp/feedmill_preview_{job.job_id}.jpg"
|
|
|
|
with self._lock:
|
|
self._jobs[job_id] = job
|
|
|
|
t = threading.Thread(target=self._run_job, args=(job_id,), daemon=True)
|
|
self._threads.append(t)
|
|
t.start()
|
|
|
|
return job
|
|
|
|
# NOTE: get_job/list_jobs return the live Job object (Flask renders it
|
|
# directly), so readers must treat its fields as eventually consistent —
|
|
# the worker mutates them under self._lock while readers may observe
|
|
# a slightly stale snapshot.
|
|
def get_job(self, job_id: str) -> Job | None:
|
|
with self._lock:
|
|
return self._jobs.get(job_id)
|
|
|
|
def list_jobs(self) -> list[Job]:
|
|
with self._lock:
|
|
return list(self._jobs.values())
|
|
|
|
def cancel_job(self, job_id: str) -> bool:
|
|
with self._lock:
|
|
job = self._jobs.get(job_id)
|
|
if job is None:
|
|
return False
|
|
if job.status in (JobStatus.PENDING, JobStatus.RUNNING):
|
|
job.status = JobStatus.CANCELLED
|
|
return True
|
|
return False
|
|
|
|
def delete_job(self, job_id: str) -> bool:
|
|
"""Delete a completed/failed/cancelled job and its files."""
|
|
import shutil
|
|
with self._lock:
|
|
job = self._jobs.get(job_id)
|
|
if job is None:
|
|
return False
|
|
if job.status in (JobStatus.PENDING, JobStatus.RUNNING):
|
|
return False # can't delete running jobs
|
|
# Remove from dict
|
|
del self._jobs[job_id]
|
|
# Delete output directory
|
|
try:
|
|
shutil.rmtree(job.output_dir, ignore_errors=True)
|
|
except Exception:
|
|
pass
|
|
# Delete preview file
|
|
try:
|
|
os.remove(job.preview_path)
|
|
except OSError:
|
|
pass
|
|
return True
|
|
|
|
def status_counts(self) -> dict[str, int]:
|
|
"""Return counts by status: {pending: N, running: N, completed: N, ...}."""
|
|
counts = {s.name.lower(): 0 for s in JobStatus}
|
|
with self._lock:
|
|
for job in self._jobs.values():
|
|
counts[job.status.name.lower()] += 1
|
|
return counts
|
|
|
|
def _run_job(self, job_id: str) -> None:
|
|
"""Worker: process each model config sequentially."""
|
|
with self._lock:
|
|
job: Job | None = self._jobs.get(job_id)
|
|
if job is None:
|
|
return
|
|
|
|
with self._lock:
|
|
job.status = JobStatus.RUNNING
|
|
total_models = len(job.model_configs)
|
|
|
|
if total_models == 0:
|
|
with self._lock:
|
|
job.status = JobStatus.COMPLETED
|
|
job.completed_at = time.time()
|
|
return
|
|
|
|
try:
|
|
if total_models == 1:
|
|
# Single model: use existing run_pipeline
|
|
model_cfg = job.model_configs[0]
|
|
class_filter = job.class_filters.get(model_cfg.filename)
|
|
truck_model_config = None
|
|
if "truck" not in (model_cfg.known_classes or []):
|
|
truck_model_config = find_truck_detector("./models")
|
|
|
|
output_path = os.path.join(job.output_dir, f"{model_cfg.stem}_annotated.mp4")
|
|
|
|
def _on_live_stats(stats):
|
|
with self._lock:
|
|
job.live_fps = stats["fps"]
|
|
job.live_loading = stats["loading"]
|
|
job.live_unloading = stats["unloading"]
|
|
job.live_frame = stats["frame"]
|
|
|
|
result = run_pipeline(
|
|
video_path=job.video_path,
|
|
model_config=model_cfg,
|
|
output_path=output_path,
|
|
class_filter=class_filter,
|
|
truck_model_config=truck_model_config,
|
|
fixed_zone=job.fixed_zone,
|
|
preview_path=job.preview_path,
|
|
preview_enabled=job.preview_enabled,
|
|
preview_every_n=2,
|
|
cancel_check=lambda: job.status == JobStatus.CANCELLED,
|
|
live_stats_callback=_on_live_stats,
|
|
)
|
|
|
|
with self._lock:
|
|
job.results.append(JobResult(
|
|
model_name=model_cfg.filename,
|
|
output_path=result.output_path,
|
|
loading_count=result.loading_count,
|
|
unloading_count=result.unloading_count,
|
|
net_count=result.net_count,
|
|
batch_count=result.batch_count,
|
|
frame_count=result.frame_count,
|
|
duration_seconds=result.duration_seconds,
|
|
))
|
|
else:
|
|
# Multiple models: run merged pipeline
|
|
with self._lock:
|
|
job.current_model = f"Processing {total_models} models (merged)..."
|
|
|
|
output_path = os.path.join(job.output_dir, "merged_annotated.mp4")
|
|
|
|
def _on_live_stats_merged(stats):
|
|
with self._lock:
|
|
job.live_fps = stats["fps"]
|
|
job.live_loading = stats["loading"]
|
|
job.live_unloading = stats["unloading"]
|
|
job.live_frame = stats["frame"]
|
|
|
|
# Auto-select truck detector
|
|
truck_model_config = None
|
|
if not any("truck" in (c.known_classes or []) for c in job.model_configs):
|
|
truck_model_config = find_truck_detector("./models")
|
|
|
|
result = run_merged_pipeline(
|
|
video_path=job.video_path,
|
|
model_configs=job.model_configs,
|
|
output_path=output_path,
|
|
class_filters=job.class_filters,
|
|
truck_model_config=truck_model_config,
|
|
fixed_zone=job.fixed_zone,
|
|
preview_path=job.preview_path,
|
|
preview_enabled=job.preview_enabled,
|
|
preview_every_n=2,
|
|
cancel_check=lambda: job.status == JobStatus.CANCELLED,
|
|
live_stats_callback=_on_live_stats_merged,
|
|
)
|
|
|
|
with self._lock:
|
|
job.results.append(JobResult(
|
|
model_name="Merged (" + ", ".join(c.stem for c in job.model_configs) + ")",
|
|
output_path=result.output_path,
|
|
loading_count=result.loading_count,
|
|
unloading_count=result.unloading_count,
|
|
net_count=result.net_count,
|
|
batch_count=result.batch_count,
|
|
frame_count=result.frame_count,
|
|
duration_seconds=result.duration_seconds,
|
|
))
|
|
|
|
with self._lock:
|
|
if job.status != JobStatus.CANCELLED:
|
|
job.status = JobStatus.COMPLETED
|
|
job.progress = 1.0
|
|
job.latest_frame = None
|
|
|
|
except Exception as e:
|
|
with self._lock:
|
|
if job.status != JobStatus.CANCELLED:
|
|
job.status = JobStatus.FAILED
|
|
job.error = str(e)
|
|
job.latest_frame = None
|
|
|
|
finally:
|
|
try:
|
|
os.remove(job.preview_path)
|
|
except OSError:
|
|
pass
|
|
|
|
with self._lock:
|
|
job.completed_at = time.time()
|
|
job.current_model = ""
|