Compare commits
6
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
8f2834775b | ||
|
|
75f38ec7ed | ||
|
|
5fa67a733c | ||
|
|
46548881cb | ||
|
|
4782dc5bfa | ||
|
|
1611485a66 |
+1
-1
@@ -1,5 +1,5 @@
|
||||
# VideoDetect — Docker Compose environment variables
|
||||
# Copy this file to .env and fill in real values before running docker-compose up.
|
||||
|
||||
API_HOST=https://10.0.1.7:3000
|
||||
API_HOST=https://10.0.1.34:3000
|
||||
|
||||
|
||||
@@ -9,7 +9,7 @@ services:
|
||||
container_name: videodetect-aiworker
|
||||
restart: "no"
|
||||
environment:
|
||||
- API_BASE_URL=${API_BASE_URL:-http://10.0.1.7:3000}
|
||||
- API_BASE_URL=${API_BASE_URL:-http://10.0.1.34:3000}
|
||||
- OLLAMA_URL=${OLLAMA_URL:-http://10.0.1.7:11434}
|
||||
- SECONDS_PER_FRAME=${SECONDS_PER_FRAME:-30}
|
||||
- TASK_COUNT=${TASK_COUNT:-50}
|
||||
|
||||
+78
-15
@@ -3,6 +3,7 @@ use strict;
|
||||
use Dancer2;
|
||||
use Dancer2::Plugin::Database;
|
||||
use Data::Dumper;
|
||||
use JSON ();
|
||||
use SQL::SplitStatement;
|
||||
use HTTP::Lite;
|
||||
|
||||
@@ -83,12 +84,42 @@ post '/api/v1/video' => sub {
|
||||
|
||||
get '/api/v1/tasks' => sub {
|
||||
my $status = query_parameters->get("status") || 'PENDING';
|
||||
my $sth = database->prepare("SELECT id, video_id, status FROM tasks WHERE status=? ORDER BY id");
|
||||
$sth->execute($status);
|
||||
my $ref = $sth->fetchall_hashref('id');
|
||||
my $type = query_parameters->get("type");
|
||||
|
||||
unless ($type) {
|
||||
send_error("Missing 'type' query parameter", 400);
|
||||
}
|
||||
|
||||
# Optional pagination
|
||||
my $page = query_parameters->get("page") || 1;
|
||||
my $per_page = query_parameters->get("per_page") || 20;
|
||||
my $offset = ($page - 1) * $per_page;
|
||||
|
||||
my $sth = database->prepare(
|
||||
"SELECT id, video_id, task_type AS type, status FROM tasks WHERE status=? AND task_type=? ORDER BY id LIMIT ? OFFSET ?"
|
||||
);
|
||||
$sth->execute($status, $type, $per_page, $offset);
|
||||
my $ref = [];
|
||||
while(my $row = $sth->fetchrow_hashref()) {
|
||||
push @$ref, $row;
|
||||
}
|
||||
$sth->finish();
|
||||
|
||||
return $ref;
|
||||
# Total count for pagination metadata
|
||||
my $count_sth = database->prepare(
|
||||
"SELECT COUNT(*) AS cnt FROM tasks WHERE status=? AND task_type=?"
|
||||
);
|
||||
$count_sth->execute($status, $type);
|
||||
my $cnt = ($count_sth->fetchrow_hashref() || { cnt => 0 })->{cnt};
|
||||
$count_sth->finish();
|
||||
|
||||
return {
|
||||
tasks => $ref,
|
||||
page => $page,
|
||||
per_page => $per_page,
|
||||
total => $cnt,
|
||||
total_pages => max(1, int(($cnt + $per_page - 1) / $per_page)),
|
||||
};
|
||||
};
|
||||
|
||||
get '/api/v1/task/:task' => sub {
|
||||
@@ -129,13 +160,13 @@ get '/api/v1/nexttask/:type' => sub {
|
||||
};
|
||||
|
||||
post '/api/v1/task/:task/complete' => sub {
|
||||
my $task = route_parameters->get("task");
|
||||
my $task_id = route_parameters->get("task");
|
||||
|
||||
# Accept both JSON body and form-encoded data for assign_key/results.
|
||||
my ($assign_key, $results);
|
||||
my $ct = request_header('Content-Type') || '';
|
||||
if ($ct eq 'application/json') {
|
||||
my $body = decode_json(request->body());
|
||||
if ($ct =~ m{application/json}i) {
|
||||
my $body = JSON::decode_json(request->body());
|
||||
$assign_key = $body->{assign_key};
|
||||
$results = $body->{results};
|
||||
} else {
|
||||
@@ -143,24 +174,31 @@ post '/api/v1/task/:task/complete' => sub {
|
||||
$results = body_parameters->get("results");
|
||||
}
|
||||
|
||||
unless ($assign_key) {
|
||||
send_error("Missing assign_key", 400);
|
||||
}
|
||||
|
||||
# Verify the task is assigned to the worker
|
||||
# Optional worker verification: if assign_key is provided, require IN_PROGRESS + matching key.
|
||||
# When omitted (e.g., UI-driven workflows), skip worker checks and just complete the task.
|
||||
if ($assign_key) {
|
||||
my $sth = database->prepare("SELECT * FROM tasks WHERE id=? AND assign_key=? AND status='IN_PROGRESS'");
|
||||
$sth->execute($task, $assign_key);
|
||||
$sth->execute($task_id, $assign_key);
|
||||
my $ref = $sth->fetchrow_hashref();
|
||||
$sth->finish();
|
||||
|
||||
unless(ref $ref) {
|
||||
send_error("Task not assigned to this worker or not in progress", 403);
|
||||
}
|
||||
} else {
|
||||
# No assign_key provided — caller may be an unauthenticated UI, etc.
|
||||
my $sth = database->prepare("SELECT id FROM tasks WHERE id=?");
|
||||
$sth->execute($task_id);
|
||||
unless ($sth->fetchrow_hashref()) {
|
||||
send_error("Task not found", 404);
|
||||
}
|
||||
$sth->finish();
|
||||
}
|
||||
|
||||
# Update the task as completed (store results as JSON text for MariaDB)
|
||||
my $results_json = defined $results ? encode_json($results) : 'null';
|
||||
my $results_json = defined $results ? JSON::encode_json($results) : 'null';
|
||||
my $update_sth = database->prepare("UPDATE tasks SET status='COMPLETED', results=?, updated_at=NOW() WHERE id=?");
|
||||
$update_sth->execute($results_json, $task);
|
||||
$update_sth->execute($results_json, $task_id);
|
||||
if($update_sth->rows == 0) {
|
||||
send_error("Failed to complete task", 500);
|
||||
}
|
||||
@@ -183,4 +221,29 @@ post '/api/v1/task' => sub {
|
||||
return { id => $task_id };
|
||||
};
|
||||
|
||||
post '/api/v1/task/:task/unassign' => sub {
|
||||
my $task_id = route_parameters->get("task");
|
||||
|
||||
my $update_sth = database->prepare(
|
||||
"UPDATE tasks SET status='PENDING', assign_key=NULL, assigned_at=NULL, updated_at=NOW() WHERE id=?"
|
||||
);
|
||||
$update_sth->execute($task_id);
|
||||
|
||||
if ($update_sth->rows == 0) {
|
||||
send_error("Task not found or already in PENDING state", 404);
|
||||
}
|
||||
|
||||
|
||||
return { message => 'Task unassigned' };
|
||||
};
|
||||
|
||||
start();
|
||||
|
||||
|
||||
sub max {
|
||||
return (sort { $b <=> $a } @_)[0];
|
||||
}
|
||||
|
||||
sub min {
|
||||
return (sort { $a <=> $b } @_)[0];
|
||||
}
|
||||
|
||||
+1
-1
@@ -18,7 +18,7 @@ plugins:
|
||||
dbi_params:
|
||||
RaiseError: 1
|
||||
AutoCommit: 1
|
||||
on_connect_do: ["SET NAMES 'utf8'", "SET CHARACTER SET 'utf8'" ]
|
||||
on_connect_do: ["SET NAMES 'utf8mb4'", "SET CHARACTER SET 'utf8mb4'" ]
|
||||
log_queries: 1
|
||||
|
||||
server:
|
||||
|
||||
+4
-10
@@ -7,7 +7,7 @@ services:
|
||||
MYSQL_ROOT_PASSWORD: ${DB_ROOT_PASSWORD:-rootpass}
|
||||
MYSQL_DATABASE: ${DB_NAME:-videodetect}
|
||||
MYSQL_USER: ${DB_USER:-videodetect}
|
||||
MYSQL_PASSWORD: ${DB_PASSWORD:-changeme_videodetect}
|
||||
MYSQL_PASSWORD: ${DB_PASSWORD:-videodetect}
|
||||
ports:
|
||||
- "3306:3306"
|
||||
volumes:
|
||||
@@ -29,20 +29,14 @@ services:
|
||||
build:
|
||||
context: ./ui
|
||||
dockerfile: Dockerfile
|
||||
container_name: videodetect-ui
|
||||
restart: unless-stopped
|
||||
ports:
|
||||
- "5000:5000"
|
||||
environment:
|
||||
- DB_HOST=${DB_HOST:-mariadb}
|
||||
- DB_PORT=${DB_PORT:-3306}
|
||||
- DB_NAME=${DB_NAME:-videodetect}
|
||||
- DB_USER=${DB_USER:-videodetect}
|
||||
- DB_PASSWORD=${DB_PASSWORD:-changeme_videodetect}
|
||||
- API_BASE=http://10.0.1.34:3000/api/v1
|
||||
volumes:
|
||||
- ./ui:/app
|
||||
- nas_input:/data/input:ro
|
||||
- ${NAS_OUTPUT_PATH:-/mnt/nas/output}:/data/output
|
||||
- nas_input:/data:ro
|
||||
networks:
|
||||
- videodetect-network
|
||||
depends_on:
|
||||
@@ -66,7 +60,7 @@ services:
|
||||
- DB_PORT=${DB_PORT:-3306}
|
||||
- DB_NAME=${DB_NAME:-videodetect}
|
||||
- DB_USER=${DB_USER:-videodetect}
|
||||
- DB_PASSWORD=${DB_PASSWORD:-changeme_videodetect}
|
||||
- DB_PASSWORD=${DB_PASSWORD:-videodetect}
|
||||
volumes:
|
||||
- ./api:/app
|
||||
networks:
|
||||
|
||||
Executable
+5
@@ -0,0 +1,5 @@
|
||||
#!/bin/bash
|
||||
curl http://10.0.1.7:11434/api/generate -d '{
|
||||
"model": "qwen3.6:35b-a3b",
|
||||
"keep_alive": "60m"
|
||||
}'
|
||||
+2
-2
@@ -8,7 +8,7 @@ DB_ROOT_PASSWORD=videodetect
|
||||
DB_PASSWORD=videodetect
|
||||
DB_USER=videodetect
|
||||
DB_NAME=videodetect
|
||||
DB_HOST=10.0.1.7
|
||||
DB_HOST=10.0.1.34
|
||||
|
||||
API_HOST=https://10.0.1.7:3000
|
||||
API_HOST=https://10.0.1.34:3000
|
||||
|
||||
|
||||
+1
-1
@@ -18,7 +18,7 @@ my $do_delete = (defined $ENV{'DO_DELETE'}) ? $ENV{'DO_DELETE'} : 1;
|
||||
$SIG{INT} = sub { die "Interrupted\n"; };
|
||||
|
||||
my $dbh = DBI->connect(
|
||||
"DBI:mysql:database=$db_name;host=$db_host;mysql_enable_utf8=1",
|
||||
"DBI:mysql:database=$db_name;host=$db_host;mysql_enable_utf8mb4=1",
|
||||
$db_user,
|
||||
$db_pass,
|
||||
{ mysql_auto_reconnect => 1 }
|
||||
|
||||
@@ -1 +0,0 @@
|
||||
"""Active learning pipeline: label ingestion, training, validation, and deployment."""
|
||||
@@ -1,217 +0,0 @@
|
||||
"""
|
||||
Label ingestion for active learning.
|
||||
|
||||
Queries the review_queue for annotated labels, extracts face crops,
|
||||
and builds a versioned training dataset with stratified train/val split.
|
||||
"""
|
||||
|
||||
import csv
|
||||
import json
|
||||
import logging
|
||||
import random
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Train split fraction (val = 1 - TRAIN_FRAC)
|
||||
_TRAIN_FRAC = 0.80
|
||||
|
||||
|
||||
class LabelIngestor:
|
||||
"""Ingest annotated review data and build a versioned dataset."""
|
||||
|
||||
def __init__(self, db_connector, config):
|
||||
self._db = db_connector
|
||||
self._cfg = config
|
||||
self._training_path = Path(config.get("storage.training_path", "/data/training"))
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Public API
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def count_annotated(self) -> int:
|
||||
"""Return number of annotated rows ready for training."""
|
||||
rows = self._db.fetchall(
|
||||
"SELECT COUNT(*) AS cnt FROM review_queue "
|
||||
"WHERE annotated = TRUE AND ground_truth IS NOT NULL"
|
||||
)
|
||||
return int((rows or [{"cnt": 0}])[0]["cnt"])
|
||||
|
||||
def ingest(self, version: str, seed: int = 42) -> Optional[str]:
|
||||
"""
|
||||
Build a versioned dataset directory from the review queue.
|
||||
|
||||
Returns the dataset path on success, or None if no data is available.
|
||||
"""
|
||||
rows = self._fetch_annotated_rows()
|
||||
if not rows:
|
||||
logger.warning("No annotated samples found; skipping dataset build")
|
||||
return None
|
||||
|
||||
dataset_path = self._training_path / version
|
||||
crops_pos = dataset_path / "crops" / "class_1"
|
||||
crops_neg = dataset_path / "crops" / "class_0"
|
||||
for d in (crops_pos, crops_neg):
|
||||
d.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
records = self._extract_crops(rows, crops_pos, crops_neg)
|
||||
if not records:
|
||||
logger.error("No crops could be extracted from annotated rows")
|
||||
return None
|
||||
|
||||
train_records, val_records = _stratified_split(records, _TRAIN_FRAC, seed)
|
||||
|
||||
self._write_labels_csv(dataset_path / "labels.csv", records)
|
||||
self._write_metadata(dataset_path / "metadata.json", version, records, train_records, val_records)
|
||||
|
||||
logger.info(
|
||||
"Dataset v%s built: %d total (%d train / %d val) at %s",
|
||||
version, len(records), len(train_records), len(val_records), dataset_path,
|
||||
)
|
||||
return str(dataset_path)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Internal helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _fetch_annotated_rows(self) -> List[Dict]:
|
||||
rows = self._db.fetchall(
|
||||
"""SELECT rq.video_id, rq.ground_truth,
|
||||
pl.confidence_scores,
|
||||
v.file_path
|
||||
FROM review_queue rq
|
||||
JOIN videos v ON v.id = rq.video_id
|
||||
LEFT JOIN (
|
||||
SELECT video_id, confidence_scores,
|
||||
ROW_NUMBER() OVER (PARTITION BY video_id ORDER BY processed_at DESC) AS rn
|
||||
FROM processing_logs
|
||||
) pl ON pl.video_id = rq.video_id AND pl.rn = 1
|
||||
WHERE rq.annotated = TRUE AND rq.ground_truth IS NOT NULL"""
|
||||
)
|
||||
return rows or []
|
||||
|
||||
def _extract_crops(
|
||||
self,
|
||||
rows: List[Dict],
|
||||
crops_pos: Path,
|
||||
crops_neg: Path,
|
||||
) -> List[Dict]:
|
||||
"""Copy or re-extract face crops into the dataset directory."""
|
||||
records: List[Dict] = []
|
||||
scratch_path = Path(self._cfg.get("storage.scratch_path", "/scratch"))
|
||||
|
||||
for row in rows:
|
||||
video_id = row["video_id"]
|
||||
label = int(bool(row["ground_truth"]))
|
||||
dest_dir = crops_pos if label == 1 else crops_neg
|
||||
|
||||
try:
|
||||
contributing = _parse_json_field(row.get("confidence_scores"))
|
||||
crop_paths = self._resolve_crop_paths(video_id, contributing, scratch_path)
|
||||
|
||||
for crop_path in crop_paths:
|
||||
src = Path(crop_path)
|
||||
if not src.exists():
|
||||
logger.debug("Crop not found, skipping: %s", src)
|
||||
continue
|
||||
dest = dest_dir / f"video{video_id}_{src.name}"
|
||||
shutil.copy2(src, dest)
|
||||
records.append({"crop_path": str(dest), "label": label, "video_id": video_id})
|
||||
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to extract crops for video %s: %s", video_id, exc)
|
||||
|
||||
return records
|
||||
|
||||
def _resolve_crop_paths(
|
||||
self,
|
||||
video_id: int,
|
||||
contributing_frames: List,
|
||||
scratch_path: Path,
|
||||
) -> List[str]:
|
||||
"""Return paths of existing crop files for this video."""
|
||||
paths: List[str] = []
|
||||
|
||||
# contributing_frames may be a list of frame scores or frame dicts
|
||||
if contributing_frames:
|
||||
for item in contributing_frames:
|
||||
if isinstance(item, dict) and "crop_path" in item:
|
||||
paths.append(item["crop_path"])
|
||||
elif isinstance(item, str):
|
||||
paths.append(item)
|
||||
|
||||
# Fall back: glob any crop files already written to scratch
|
||||
if not paths:
|
||||
pattern = f"*video{video_id}*.jpg"
|
||||
paths = [str(p) for p in scratch_path.rglob(pattern)]
|
||||
|
||||
return paths
|
||||
|
||||
@staticmethod
|
||||
def _write_labels_csv(csv_path: Path, records: List[Dict]):
|
||||
with open(csv_path, "w", newline="", encoding="utf-8") as f:
|
||||
writer = csv.DictWriter(f, fieldnames=["crop_path", "label", "video_id"])
|
||||
writer.writeheader()
|
||||
writer.writerows(records)
|
||||
|
||||
@staticmethod
|
||||
def _write_metadata(
|
||||
meta_path: Path,
|
||||
version: str,
|
||||
all_records: List[Dict],
|
||||
train_records: List[Dict],
|
||||
val_records: List[Dict],
|
||||
):
|
||||
pos = sum(1 for r in all_records if r["label"] == 1)
|
||||
neg = len(all_records) - pos
|
||||
meta = {
|
||||
"version": version,
|
||||
"total_samples": len(all_records),
|
||||
"train_samples": len(train_records),
|
||||
"val_samples": len(val_records),
|
||||
"class_counts": {"class_0": neg, "class_1": pos},
|
||||
}
|
||||
meta_path.write_text(json.dumps(meta, indent=2), encoding="utf-8")
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Module-level helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _stratified_split(
|
||||
records: List[Dict],
|
||||
train_frac: float,
|
||||
seed: int,
|
||||
) -> Tuple[List[Dict], List[Dict]]:
|
||||
"""Return (train, val) with stratification by label."""
|
||||
rng = random.Random(seed)
|
||||
|
||||
by_label: Dict[int, List[Dict]] = {}
|
||||
for r in records:
|
||||
by_label.setdefault(r["label"], []).append(r)
|
||||
|
||||
train, val = [], []
|
||||
for label_records in by_label.values():
|
||||
shuffled = list(label_records)
|
||||
rng.shuffle(shuffled)
|
||||
split = max(1, int(len(shuffled) * train_frac))
|
||||
train.extend(shuffled[:split])
|
||||
val.extend(shuffled[split:])
|
||||
|
||||
rng.shuffle(train)
|
||||
rng.shuffle(val)
|
||||
return train, val
|
||||
|
||||
|
||||
def _parse_json_field(value) -> List:
|
||||
if not value:
|
||||
return []
|
||||
if isinstance(value, list):
|
||||
return value
|
||||
try:
|
||||
parsed = json.loads(value)
|
||||
return parsed if isinstance(parsed, list) else []
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return []
|
||||
@@ -1,82 +0,0 @@
|
||||
"""
|
||||
Active learning pipeline orchestrator.
|
||||
|
||||
Entry point that coordinates:
|
||||
1. Check annotated sample count against minimum threshold
|
||||
2. Ingest labels and build versioned dataset
|
||||
3. Fine-tune classification head
|
||||
4. Validate candidate model against quality gates
|
||||
5. Deploy if gates pass; skip otherwise
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from .label_ingestor import LabelIngestor
|
||||
from .registry import ModelRegistry
|
||||
from .trainer import Trainer
|
||||
from .validator import Validator
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ActiveLearningPipeline:
|
||||
"""Orchestrate the full active learning cycle."""
|
||||
|
||||
def __init__(self, db_connector, config):
|
||||
self._db = db_connector
|
||||
self._cfg = config
|
||||
self._ingestor = LabelIngestor(db_connector, config)
|
||||
self._registry = ModelRegistry(db_connector, config)
|
||||
|
||||
def run(self, version: str) -> bool:
|
||||
"""
|
||||
Execute a full active learning cycle for the given model version.
|
||||
|
||||
Returns True if a new model was deployed, False otherwise.
|
||||
"""
|
||||
al = self._cfg.get_section("active_learning")
|
||||
min_samples: int = int(al.get("min_annotated_samples", 100))
|
||||
|
||||
# 1. Check threshold
|
||||
n_annotated = self._ingestor.count_annotated()
|
||||
logger.info("Annotated samples available: %d (minimum: %d)", n_annotated, min_samples)
|
||||
if n_annotated < min_samples:
|
||||
logger.info("Insufficient annotated samples; skipping active learning cycle")
|
||||
return False
|
||||
|
||||
# 2. Ingest labels
|
||||
seed = int(self._cfg.get("active_learning.seed", 42))
|
||||
dataset_path = self._ingestor.ingest(version, seed=seed)
|
||||
if not dataset_path:
|
||||
logger.error("Label ingestion produced no dataset; aborting")
|
||||
return False
|
||||
|
||||
# 3. Train
|
||||
trainer = Trainer(self._cfg)
|
||||
checkpoint_path = trainer.train(dataset_path, version)
|
||||
if not checkpoint_path:
|
||||
logger.error("Training failed; aborting active learning cycle")
|
||||
return False
|
||||
|
||||
# 4. Validate
|
||||
current_f1 = self._registry.get_active_f1()
|
||||
validator = Validator(self._cfg, current_f1=current_f1)
|
||||
metrics = validator.validate(checkpoint_path, dataset_path)
|
||||
|
||||
self._registry.register_candidate(
|
||||
version, checkpoint_path,
|
||||
f1=metrics.get("f1", 0.0),
|
||||
ece=metrics.get("ece", 1.0),
|
||||
)
|
||||
|
||||
if not metrics.get("gates_passed", False):
|
||||
logger.warning(
|
||||
"Quality gates failed for %s; skipping deployment. Details: %s",
|
||||
version, metrics.get("gate_details"),
|
||||
)
|
||||
return False
|
||||
|
||||
# 5. Deploy
|
||||
deployed = self._registry.deploy(version, checkpoint_path)
|
||||
return deployed
|
||||
@@ -1,223 +0,0 @@
|
||||
"""
|
||||
Model registry and deployment for active learning.
|
||||
|
||||
Handles:
|
||||
- Promoting a candidate model to ACTIVE in the DB
|
||||
- Converting the PyTorch checkpoint to ONNX then TensorRT
|
||||
- Hot-reloading the classifier in worker processes via SIGUSER1
|
||||
- Rollback to the previous ACTIVE model on failure
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import signal
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Signal used to notify the worker process to reload its model
|
||||
_RELOAD_SIGNAL = signal.SIGUSR1
|
||||
|
||||
|
||||
class ModelRegistry:
|
||||
"""Manage model lifecycle: promote, deploy, rollback."""
|
||||
|
||||
def __init__(self, db_connector, config):
|
||||
self._db = db_connector
|
||||
self._cfg = config
|
||||
self._models_path = Path(config.get("storage.models_path", "/models"))
|
||||
al = config.get_section("active_learning")
|
||||
dep = al.get("deployment", {})
|
||||
self._auto_deploy: bool = bool(dep.get("auto_deploy", True))
|
||||
self._hot_reload: bool = bool(dep.get("hot_reload", True))
|
||||
self._rollback_enabled: bool = bool(dep.get("rollback_enabled", True))
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Public API
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_active_version(self) -> Optional[str]:
|
||||
"""Return the version string of the currently ACTIVE model."""
|
||||
rows = self._db.fetchall(
|
||||
"SELECT version FROM models WHERE status = 'ACTIVE' ORDER BY deployed_at DESC LIMIT 1"
|
||||
)
|
||||
if rows:
|
||||
return rows[0]["version"]
|
||||
return None
|
||||
|
||||
def get_active_f1(self) -> float:
|
||||
"""Return the F1 score of the currently ACTIVE model (0 if unknown)."""
|
||||
rows = self._db.fetchall(
|
||||
"SELECT f1_score FROM models WHERE status = 'ACTIVE' ORDER BY deployed_at DESC LIMIT 1"
|
||||
)
|
||||
if rows and rows[0]["f1_score"] is not None:
|
||||
return float(rows[0]["f1_score"])
|
||||
return 0.0
|
||||
|
||||
def register_candidate(self, version: str, checkpoint_path: str, f1: float, ece: float):
|
||||
"""Insert or update the candidate model record."""
|
||||
self._db.execute(
|
||||
"""INSERT INTO models (version, status, path, f1_score, ece_score, created_at)
|
||||
VALUES (%s, 'CANDIDATE', %s, %s, %s, NOW())
|
||||
ON DUPLICATE KEY UPDATE
|
||||
status = 'CANDIDATE', path = VALUES(path),
|
||||
f1_score = VALUES(f1_score), ece_score = VALUES(ece_score)""",
|
||||
(version, checkpoint_path, f1, ece),
|
||||
)
|
||||
logger.info("Registered candidate model %s (f1=%.4f, ece=%.4f)", version, f1, ece)
|
||||
|
||||
def deploy(self, version: str, checkpoint_path: str) -> bool:
|
||||
"""
|
||||
Convert, promote and hot-reload the candidate model.
|
||||
|
||||
Returns True if deployment succeeded, False otherwise.
|
||||
"""
|
||||
if not self._auto_deploy:
|
||||
logger.info("Auto-deploy disabled; skipping deployment of %s", version)
|
||||
return False
|
||||
|
||||
try:
|
||||
engine_path = self._build_engine(version, checkpoint_path)
|
||||
self._promote(version, engine_path)
|
||||
if self._hot_reload:
|
||||
self._reload_worker()
|
||||
logger.info("Model %s deployed successfully", version)
|
||||
return True
|
||||
except Exception as exc:
|
||||
logger.error("Deployment of %s failed: %s", version, exc, exc_info=True)
|
||||
if self._rollback_enabled:
|
||||
self.rollback(version)
|
||||
return False
|
||||
|
||||
def rollback(self, failed_version: str):
|
||||
"""Restore the most recent ARCHIVED model to ACTIVE."""
|
||||
rows = self._db.fetchall(
|
||||
"""SELECT version FROM models
|
||||
WHERE status = 'ARCHIVED'
|
||||
ORDER BY deployed_at DESC
|
||||
LIMIT 1"""
|
||||
)
|
||||
if not rows:
|
||||
logger.error("No archived model available for rollback")
|
||||
return
|
||||
|
||||
prev_version = rows[0]["version"]
|
||||
self._db.execute(
|
||||
"UPDATE models SET status = 'ARCHIVED' WHERE version = %s",
|
||||
(failed_version,),
|
||||
)
|
||||
self._db.execute(
|
||||
"UPDATE models SET status = 'ACTIVE' WHERE version = %s",
|
||||
(prev_version,),
|
||||
)
|
||||
logger.warning("Rolled back to model %s (failed: %s)", prev_version, failed_version)
|
||||
if self._hot_reload:
|
||||
self._reload_worker()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Internal helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _build_engine(self, version: str, checkpoint_path: str) -> str:
|
||||
"""Export checkpoint to ONNX then build a TensorRT FP32 engine."""
|
||||
candidate_dir = self._models_path / "candidate"
|
||||
onnx_path = candidate_dir / f"{version}.onnx"
|
||||
engine_path = candidate_dir / f"{version}.trt"
|
||||
|
||||
self._export_onnx(checkpoint_path, onnx_path)
|
||||
self._build_trt(onnx_path, engine_path)
|
||||
return str(engine_path)
|
||||
|
||||
@staticmethod
|
||||
def _export_onnx(checkpoint_path: str, onnx_path: Path):
|
||||
"""Convert PyTorch checkpoint to ONNX."""
|
||||
try:
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torchvision.models as models
|
||||
|
||||
device = torch.device("cpu")
|
||||
model = models.mobilenet_v3_small(weights=None)
|
||||
model.classifier = nn.Sequential(
|
||||
nn.Linear(576, 256),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(256, 2),
|
||||
)
|
||||
state = torch.load(checkpoint_path, map_location=device)
|
||||
model.load_state_dict(state)
|
||||
model.eval()
|
||||
|
||||
dummy = torch.zeros(1, 3, 224, 224)
|
||||
torch.onnx.export(
|
||||
model,
|
||||
dummy,
|
||||
str(onnx_path),
|
||||
input_names=["input"],
|
||||
output_names=["output"],
|
||||
dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
|
||||
opset_version=11,
|
||||
)
|
||||
logger.info("Exported ONNX model to %s", onnx_path)
|
||||
except ImportError as exc:
|
||||
logger.error("PyTorch/torchvision not available for ONNX export: %s", exc)
|
||||
raise
|
||||
|
||||
@staticmethod
|
||||
def _build_trt(onnx_path: Path, engine_path: Path):
|
||||
"""Build TensorRT FP32 engine via trtexec CLI."""
|
||||
cmd = [
|
||||
"trtexec",
|
||||
f"--onnx={onnx_path}",
|
||||
f"--saveEngine={engine_path}",
|
||||
"--fp32",
|
||||
"--maxBatch=32",
|
||||
"--workspace=2048",
|
||||
]
|
||||
logger.info("Building TensorRT engine: %s", " ".join(cmd))
|
||||
result = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(
|
||||
f"trtexec failed (rc={result.returncode}): {result.stderr[:500]}"
|
||||
)
|
||||
logger.info("TensorRT engine saved to %s", engine_path)
|
||||
|
||||
def _promote(self, version: str, engine_path: str):
|
||||
"""Atomically promote candidate to ACTIVE, archive current ACTIVE."""
|
||||
self._db.execute(
|
||||
"""UPDATE models
|
||||
SET status = 'ARCHIVED'
|
||||
WHERE status = 'ACTIVE'""",
|
||||
)
|
||||
self._db.execute(
|
||||
"""UPDATE models
|
||||
SET status = 'ACTIVE', path = %s, deployed_at = NOW()
|
||||
WHERE version = %s""",
|
||||
(engine_path, version),
|
||||
)
|
||||
logger.info("Promoted model %s to ACTIVE (engine=%s)", version, engine_path)
|
||||
|
||||
@staticmethod
|
||||
def _reload_worker():
|
||||
"""Send SIGUSR1 to the worker process to trigger a model hot-reload."""
|
||||
pid_file = Path("/tmp/videodetect_worker.pid")
|
||||
if pid_file.exists():
|
||||
try:
|
||||
pid = int(pid_file.read_text().strip())
|
||||
os.kill(pid, _RELOAD_SIGNAL)
|
||||
logger.info("Sent SIGUSR1 to worker PID %d for hot-reload", pid)
|
||||
except (ValueError, ProcessLookupError, PermissionError) as exc:
|
||||
logger.warning("Could not signal worker process: %s", exc)
|
||||
else:
|
||||
# Fallback: restart via docker-compose
|
||||
logger.warning("No PID file found; attempting docker-compose restart")
|
||||
try:
|
||||
subprocess.run(
|
||||
["docker-compose", "restart", "worker"],
|
||||
capture_output=True,
|
||||
timeout=60,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error("docker-compose restart failed: %s", exc)
|
||||
@@ -1,344 +0,0 @@
|
||||
"""
|
||||
Training pipeline for active learning head-only fine-tuning.
|
||||
|
||||
Loads a pre-trained MobileNetV3 backbone (frozen), replaces the classification
|
||||
head, and fine-tunes on versioned crop datasets using AdamW + early stopping.
|
||||
All training is FP32 for Tesla P40 compatibility (CC 5.2, no Tensor Cores).
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Trainer:
|
||||
"""Head-only fine-tuning pipeline for MobileNetV3 classifier."""
|
||||
|
||||
def __init__(self, config):
|
||||
self._cfg = config
|
||||
self._models_path = Path(config.get("storage.models_path", "/models"))
|
||||
self._candidate_path = self._models_path / "candidate"
|
||||
self._candidate_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
al = config.get_section("active_learning")
|
||||
tr = al.get("training", {})
|
||||
self._epochs: int = int(tr.get("epochs", 20))
|
||||
self._batch_size: int = int(tr.get("batch_size", 32))
|
||||
self._lr: float = float(tr.get("learning_rate", 1e-3))
|
||||
self._weight_decay: float = float(tr.get("weight_decay", 1e-2))
|
||||
self._patience: int = int(tr.get("early_stopping_patience", 5))
|
||||
self._lr_factor: float = float(tr.get("lr_factor", 0.5))
|
||||
self._lr_patience: int = int(tr.get("lr_patience", 3))
|
||||
|
||||
aug = al.get("augmentation", {})
|
||||
self._aug_flip: bool = bool(aug.get("horizontal_flip", True))
|
||||
self._aug_jitter: bool = bool(aug.get("color_jitter", True))
|
||||
self._aug_affine: bool = bool(aug.get("affine", True))
|
||||
self._aug_degrees: int = int(aug.get("affine_degrees", 10))
|
||||
self._aug_scale: float = float(aug.get("affine_scale", 0.1))
|
||||
|
||||
seed = int(config.get("active_learning.seed", 42))
|
||||
self._seed = seed
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Public API
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def train(self, dataset_path: str, version: str) -> Optional[str]:
|
||||
"""
|
||||
Fine-tune the classification head on the versioned dataset.
|
||||
|
||||
Returns the path to the best checkpoint or None on failure.
|
||||
"""
|
||||
try:
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
self._seed_everything(self._seed)
|
||||
self._assert_cuda_compatibility(torch)
|
||||
|
||||
train_loader, val_loader = self._build_dataloaders(dataset_path, torch)
|
||||
model = self._build_model(nn)
|
||||
best_path = self._run_training(model, train_loader, val_loader, version, torch, nn)
|
||||
return best_path
|
||||
except ImportError as exc:
|
||||
logger.error("PyTorch not available: %s", exc)
|
||||
return None
|
||||
except Exception as exc:
|
||||
logger.error("Training failed: %s", exc, exc_info=True)
|
||||
return None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Model construction
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _build_model(self, nn):
|
||||
"""Load MobileNetV3-Small backbone, freeze it, attach trainable head."""
|
||||
import torchvision.models as models
|
||||
|
||||
model = models.mobilenet_v3_small(weights=models.MobileNet_V3_Small_Weights.IMAGENET1K_V1)
|
||||
|
||||
# Freeze backbone
|
||||
for param in model.features.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
# Replace classification head
|
||||
model.classifier = nn.Sequential(
|
||||
nn.Linear(576, 256),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(256, 2),
|
||||
)
|
||||
return model
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Data loading
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _build_dataloaders(self, dataset_path: str, torch) -> Tuple:
|
||||
from torch.utils.data import DataLoader
|
||||
import torchvision.transforms as T
|
||||
|
||||
train_transform = self._build_train_transform(T)
|
||||
val_transform = self._build_val_transform(T)
|
||||
|
||||
train_dataset = _CropDataset(
|
||||
Path(dataset_path) / "crops", split="train",
|
||||
transform=train_transform, seed=self._seed,
|
||||
)
|
||||
val_dataset = _CropDataset(
|
||||
Path(dataset_path) / "crops", split="val",
|
||||
transform=val_transform, seed=self._seed,
|
||||
)
|
||||
|
||||
g = torch.Generator()
|
||||
g.manual_seed(self._seed)
|
||||
|
||||
train_loader = DataLoader(
|
||||
train_dataset,
|
||||
batch_size=self._batch_size,
|
||||
shuffle=True,
|
||||
num_workers=2,
|
||||
pin_memory=True,
|
||||
generator=g,
|
||||
)
|
||||
val_loader = DataLoader(
|
||||
val_dataset,
|
||||
batch_size=self._batch_size,
|
||||
shuffle=False,
|
||||
num_workers=2,
|
||||
pin_memory=True,
|
||||
)
|
||||
logger.info(
|
||||
"Dataset: %d train / %d val samples",
|
||||
len(train_dataset), len(val_dataset),
|
||||
)
|
||||
return train_loader, val_loader
|
||||
|
||||
def _build_train_transform(self, T):
|
||||
ops = [T.Resize((224, 224))]
|
||||
if self._aug_flip:
|
||||
ops.append(T.RandomHorizontalFlip())
|
||||
if self._aug_jitter:
|
||||
ops.append(T.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1))
|
||||
if self._aug_affine:
|
||||
ops.append(T.RandomAffine(
|
||||
degrees=self._aug_degrees,
|
||||
scale=(1.0 - self._aug_scale, 1.0 + self._aug_scale),
|
||||
))
|
||||
ops += [T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]
|
||||
return T.Compose(ops)
|
||||
|
||||
def _build_val_transform(self, T):
|
||||
return T.Compose([
|
||||
T.Resize((224, 224)),
|
||||
T.ToTensor(),
|
||||
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
||||
])
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Training loop
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _run_training(self, model, train_loader, val_loader, version: str, torch, nn) -> Optional[str]:
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
model = model.to(device)
|
||||
logger.info("Training on device: %s", device)
|
||||
|
||||
class_weights = self._compute_class_weights(train_loader, torch, device)
|
||||
criterion = nn.CrossEntropyLoss(weight=class_weights)
|
||||
|
||||
optimizer = torch.optim.AdamW(
|
||||
filter(lambda p: p.requires_grad, model.parameters()),
|
||||
lr=self._lr,
|
||||
weight_decay=self._weight_decay,
|
||||
betas=(0.9, 0.999),
|
||||
)
|
||||
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
|
||||
optimizer, mode="max", factor=self._lr_factor, patience=self._lr_patience,
|
||||
)
|
||||
|
||||
best_f1 = -1.0
|
||||
best_path: Optional[str] = None
|
||||
no_improve = 0
|
||||
|
||||
for epoch in range(1, self._epochs + 1):
|
||||
train_loss = self._train_epoch(model, train_loader, criterion, optimizer, device, torch)
|
||||
val_f1 = self._val_epoch(model, val_loader, device, torch)
|
||||
scheduler.step(val_f1)
|
||||
|
||||
ckpt_path = self._candidate_path / f"{version}_epoch_{epoch:02d}.pt"
|
||||
torch.save(model.state_dict(), ckpt_path)
|
||||
logger.info(
|
||||
"Epoch %d/%d train_loss=%.4f val_f1=%.4f",
|
||||
epoch, self._epochs, train_loss, val_f1,
|
||||
)
|
||||
|
||||
if val_f1 > best_f1:
|
||||
best_f1 = val_f1
|
||||
best_path = str(self._candidate_path / f"{version}_best.pt")
|
||||
torch.save(model.state_dict(), best_path)
|
||||
no_improve = 0
|
||||
else:
|
||||
no_improve += 1
|
||||
if no_improve >= self._patience:
|
||||
logger.info("Early stopping at epoch %d (patience=%d)", epoch, self._patience)
|
||||
break
|
||||
|
||||
logger.info("Training complete. Best val_f1=%.4f checkpoint=%s", best_f1, best_path)
|
||||
return best_path
|
||||
|
||||
@staticmethod
|
||||
def _train_epoch(model, loader, criterion, optimizer, device, torch) -> float:
|
||||
model.train()
|
||||
total_loss = 0.0
|
||||
for images, labels in loader:
|
||||
images, labels = images.to(device), labels.to(device)
|
||||
optimizer.zero_grad()
|
||||
outputs = model(images)
|
||||
loss = criterion(outputs, labels)
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
total_loss += loss.item() * len(images)
|
||||
return total_loss / max(len(loader.dataset), 1)
|
||||
|
||||
@staticmethod
|
||||
def _val_epoch(model, loader, device, torch) -> float:
|
||||
from sklearn.metrics import f1_score
|
||||
|
||||
model.eval()
|
||||
all_preds, all_labels = [], []
|
||||
with torch.no_grad():
|
||||
for images, labels in loader:
|
||||
images = images.to(device)
|
||||
outputs = model(images)
|
||||
preds = outputs.argmax(dim=1).cpu().numpy()
|
||||
all_preds.extend(preds)
|
||||
all_labels.extend(labels.numpy())
|
||||
|
||||
return float(f1_score(all_labels, all_preds, average="macro", zero_division=0))
|
||||
|
||||
@staticmethod
|
||||
def _compute_class_weights(train_loader, torch, device):
|
||||
"""Compute inverse-frequency class weights to handle imbalance."""
|
||||
counts = np.zeros(2, dtype=np.int64)
|
||||
for _, labels in train_loader:
|
||||
for lbl in labels.numpy():
|
||||
counts[int(lbl)] += 1
|
||||
total = counts.sum()
|
||||
if total == 0 or counts.min() == 0:
|
||||
return None
|
||||
weights = total / (2.0 * counts.astype(np.float32))
|
||||
return torch.tensor(weights, dtype=torch.float32, device=device)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Misc
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _assert_cuda_compatibility(torch):
|
||||
if torch.cuda.is_available():
|
||||
cuda_ver = torch.version.cuda or ""
|
||||
torch_ver = torch.__version__
|
||||
logger.info("CUDA version: %s PyTorch: %s", cuda_ver, torch_ver)
|
||||
major = int(cuda_ver.split(".")[0]) if cuda_ver else 0
|
||||
if major > 11:
|
||||
logger.warning(
|
||||
"CUDA %s detected; system requires CUDA ≤ 11.8 for Tesla P40", cuda_ver
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _seed_everything(seed: int):
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
try:
|
||||
import torch
|
||||
torch.manual_seed(seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
torch.backends.cudnn.deterministic = True
|
||||
torch.backends.cudnn.benchmark = False
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Dataset helper
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
class _CropDataset:
|
||||
"""Reads face crop images from class_0/ and class_1/ subdirectories."""
|
||||
|
||||
_EXTENSIONS = {".jpg", ".jpeg", ".png"}
|
||||
|
||||
def __init__(self, crops_root: Path, split: str, transform, seed: int):
|
||||
all_samples = self._discover(crops_root)
|
||||
train_samples, val_samples = _split_samples(all_samples, 0.80, seed)
|
||||
self._samples = train_samples if split == "train" else val_samples
|
||||
self._transform = transform
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._samples)
|
||||
|
||||
def __getitem__(self, idx: int):
|
||||
from PIL import Image
|
||||
path, label = self._samples[idx]
|
||||
image = Image.open(path).convert("RGB")
|
||||
if self._transform:
|
||||
image = self._transform(image)
|
||||
return image, label
|
||||
|
||||
def _discover(self, crops_root: Path):
|
||||
samples = []
|
||||
for label, subdir in ((0, "class_0"), (1, "class_1")):
|
||||
d = crops_root / subdir
|
||||
if d.is_dir():
|
||||
for p in sorted(d.iterdir()):
|
||||
if p.suffix.lower() in self._EXTENSIONS:
|
||||
samples.append((p, label))
|
||||
return samples
|
||||
|
||||
|
||||
def _split_samples(samples, train_frac: float, seed: int):
|
||||
"""Stratified split mirroring LabelIngestor logic."""
|
||||
rng = random.Random(seed)
|
||||
by_label: Dict[int, list] = {}
|
||||
for path, label in samples:
|
||||
by_label.setdefault(label, []).append((path, label))
|
||||
|
||||
train, val = [], []
|
||||
for group in by_label.values():
|
||||
shuffled = list(group)
|
||||
rng.shuffle(shuffled)
|
||||
split = max(1, int(len(shuffled) * train_frac))
|
||||
train.extend(shuffled[:split])
|
||||
val.extend(shuffled[split:])
|
||||
|
||||
return train, val
|
||||
@@ -1,204 +0,0 @@
|
||||
"""
|
||||
Candidate model validation for active learning.
|
||||
|
||||
Evaluates a candidate checkpoint on the held-out validation set,
|
||||
computes F1, ECE, accuracy, precision, and recall, then checks
|
||||
quality gates before promotion.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Validator:
|
||||
"""Evaluate a candidate model and enforce quality gates."""
|
||||
|
||||
def __init__(self, config, current_f1: float = 0.0):
|
||||
al = config.get_section("active_learning")
|
||||
val_cfg = al.get("validation", {})
|
||||
self._min_f1_delta: float = float(val_cfg.get("min_f1_improvement", 0.02))
|
||||
self._max_ece: float = float(val_cfg.get("max_ece", 0.08))
|
||||
self._current_f1 = current_f1
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Public API
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def validate(self, checkpoint_path: str, dataset_path: str) -> Dict:
|
||||
"""
|
||||
Run validation and return a results dict including gate pass/fail.
|
||||
|
||||
Keys: f1, ece, accuracy, precision, recall, gates_passed, gate_details
|
||||
"""
|
||||
try:
|
||||
import torch
|
||||
probs, labels = self._run_inference(checkpoint_path, dataset_path, torch)
|
||||
return self._compute_metrics(probs, labels)
|
||||
except ImportError as exc:
|
||||
logger.error("PyTorch not available: %s", exc)
|
||||
return self._empty_result()
|
||||
except Exception as exc:
|
||||
logger.error("Validation failed: %s", exc, exc_info=True)
|
||||
return self._empty_result()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Inference
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _run_inference(self, checkpoint_path: str, dataset_path: str, torch) -> Tuple[np.ndarray, np.ndarray]:
|
||||
import torch.nn as nn
|
||||
import torchvision.models as models
|
||||
import torchvision.transforms as T
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
model = self._load_checkpoint(checkpoint_path, nn, models, torch, device)
|
||||
model.eval()
|
||||
|
||||
transform = T.Compose([
|
||||
T.Resize((224, 224)),
|
||||
T.ToTensor(),
|
||||
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
||||
])
|
||||
|
||||
from .trainer import _CropDataset, _split_samples
|
||||
crops_root = Path(dataset_path) / "crops"
|
||||
all_samples = _CropDataset._discover(None, crops_root) # type: ignore[arg-type]
|
||||
_, val_samples = _split_samples(all_samples, 0.80, seed=42)
|
||||
|
||||
if not val_samples:
|
||||
raise ValueError("Validation set is empty")
|
||||
|
||||
val_dataset = _ValDataset(val_samples, transform)
|
||||
loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=2, pin_memory=True)
|
||||
|
||||
all_probs, all_labels = [], []
|
||||
with torch.no_grad():
|
||||
for images, batch_labels in loader:
|
||||
images = images.to(device)
|
||||
outputs = model(images)
|
||||
probs_batch = torch.softmax(outputs, dim=1)[:, 1].cpu().numpy()
|
||||
all_probs.extend(probs_batch)
|
||||
all_labels.extend(batch_labels.numpy())
|
||||
|
||||
return np.array(all_probs, dtype=np.float32), np.array(all_labels, dtype=np.int32)
|
||||
|
||||
@staticmethod
|
||||
def _load_checkpoint(checkpoint_path: str, nn, models, torch, device):
|
||||
model = models.mobilenet_v3_small(weights=None)
|
||||
model.classifier = nn.Sequential(
|
||||
nn.Linear(576, 256),
|
||||
nn.ReLU(),
|
||||
nn.Dropout(0.3),
|
||||
nn.Linear(256, 2),
|
||||
)
|
||||
state = torch.load(checkpoint_path, map_location=device)
|
||||
model.load_state_dict(state)
|
||||
return model.to(device)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Metrics
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _compute_metrics(self, probs: np.ndarray, labels: np.ndarray) -> Dict:
|
||||
from sklearn.metrics import (
|
||||
accuracy_score, f1_score, precision_score, recall_score,
|
||||
)
|
||||
|
||||
preds = (probs >= 0.5).astype(np.int32)
|
||||
f1 = float(f1_score(labels, preds, average="macro", zero_division=0))
|
||||
accuracy = float(accuracy_score(labels, preds))
|
||||
precision = float(precision_score(labels, preds, average="macro", zero_division=0))
|
||||
recall = float(recall_score(labels, preds, average="macro", zero_division=0))
|
||||
ece = compute_ece(probs, labels)
|
||||
|
||||
delta_f1 = f1 - self._current_f1
|
||||
gate_f1 = delta_f1 > self._min_f1_delta
|
||||
gate_ece = ece < self._max_ece
|
||||
gates_passed = gate_f1 and gate_ece
|
||||
|
||||
gate_details = {
|
||||
"delta_f1": round(delta_f1, 4),
|
||||
"min_f1_improvement": self._min_f1_delta,
|
||||
"gate_f1_passed": gate_f1,
|
||||
"ece": round(ece, 4),
|
||||
"max_ece": self._max_ece,
|
||||
"gate_ece_passed": gate_ece,
|
||||
}
|
||||
|
||||
logger.info(
|
||||
"Validation results: f1=%.4f delta_f1=%.4f ece=%.4f gates_passed=%s",
|
||||
f1, delta_f1, ece, gates_passed,
|
||||
)
|
||||
|
||||
return {
|
||||
"f1": round(f1, 4),
|
||||
"ece": round(ece, 4),
|
||||
"accuracy": round(accuracy, 4),
|
||||
"precision": round(precision, 4),
|
||||
"recall": round(recall, 4),
|
||||
"gates_passed": gates_passed,
|
||||
"gate_details": gate_details,
|
||||
}
|
||||
|
||||
def _empty_result(self) -> Dict:
|
||||
return {
|
||||
"f1": 0.0, "ece": 1.0, "accuracy": 0.0,
|
||||
"precision": 0.0, "recall": 0.0,
|
||||
"gates_passed": False,
|
||||
"gate_details": {},
|
||||
}
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# ECE calculation (module-level for reuse in tests)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def compute_ece(predictions: np.ndarray, labels: np.ndarray, n_bins: int = 15) -> float:
|
||||
"""
|
||||
Expected Calibration Error.
|
||||
|
||||
Bins predictions by confidence and measures average |confidence - accuracy|
|
||||
weighted by bin population.
|
||||
"""
|
||||
bin_boundaries = np.linspace(0, 1, n_bins + 1)
|
||||
ece = 0.0
|
||||
n = len(predictions)
|
||||
if n == 0:
|
||||
return 0.0
|
||||
|
||||
for i in range(n_bins):
|
||||
mask = (predictions >= bin_boundaries[i]) & (predictions < bin_boundaries[i + 1])
|
||||
if mask.sum() > 0:
|
||||
bin_confidence = float(predictions[mask].mean())
|
||||
bin_accuracy = float(labels[mask].mean())
|
||||
ece += (mask.sum() / n) * abs(bin_confidence - bin_accuracy)
|
||||
|
||||
return float(ece)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Minimal dataset wrapper for validation
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
class _ValDataset:
|
||||
def __init__(self, samples, transform):
|
||||
self._samples = samples
|
||||
self._transform = transform
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self._samples)
|
||||
|
||||
def __getitem__(self, idx: int):
|
||||
from PIL import Image
|
||||
path, label = self._samples[idx]
|
||||
image = Image.open(path).convert("RGB")
|
||||
if self._transform:
|
||||
image = self._transform(image)
|
||||
return image, label
|
||||
@@ -1,39 +0,0 @@
|
||||
"""Confidence aggregation strategies for video-level scoring."""
|
||||
|
||||
import logging
|
||||
from typing import List
|
||||
|
||||
import numpy as np
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def aggregate(
|
||||
confidences: List[float],
|
||||
strategy: str = "max",
|
||||
alpha: float = 1.0,
|
||||
beta: float = 0.1,
|
||||
top_k: int = 3,
|
||||
) -> float:
|
||||
"""Aggregate per-crop confidences into a single video-level score C ∈ [0,1]."""
|
||||
if not confidences:
|
||||
return 0.0
|
||||
|
||||
arr = np.array(confidences, dtype=np.float64)
|
||||
|
||||
if strategy == "max":
|
||||
return float(arr.max())
|
||||
|
||||
if strategy == "weighted_mean":
|
||||
mean = float(arr.mean())
|
||||
var = float(arr.var())
|
||||
raw = alpha * mean + beta * var
|
||||
# Clamp to [0,1]; the formula is a linear combination, not a softmax
|
||||
return float(np.clip(raw, 0.0, 1.0))
|
||||
|
||||
if strategy == "top_k_mean":
|
||||
k = max(1, min(top_k, len(arr)))
|
||||
return float(np.partition(arr, -k)[-k:].mean())
|
||||
|
||||
logger.warning("Unknown aggregation strategy '%s', falling back to max", strategy)
|
||||
return float(arr.max())
|
||||
@@ -1,136 +0,0 @@
|
||||
"""HTTP client for VideoDetect's Perl REST API (api/app.pl).
|
||||
|
||||
Wraps the three worker-facing endpoints:
|
||||
GET /api/v1/nexttask/:type — claim next pending task
|
||||
GET /api/v1/video/:id — fetch video metadata
|
||||
POST /api/v1/task/:task/complete — submit results
|
||||
|
||||
Uses the ``requests`` library; raises :class:`ApiError` on unexpected HTTP status codes.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Optional
|
||||
|
||||
import requests
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Exception
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class ApiError(Exception):
|
||||
"""Raised when the API returns an unexpected error status."""
|
||||
|
||||
def __init__(self, status_code: int, message: str):
|
||||
self.status_code = status_code
|
||||
self.message = message
|
||||
super().__init__(f"API error {status_code}: {message}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Client
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class ApiClient:
|
||||
"""Thin HTTP client for the Dancer2 REST API."""
|
||||
|
||||
def __init__(self, base_url: str = "http://localhost:3000"):
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.session = requests.Session()
|
||||
logger.info("API client configured: base_url=%s", self.base_url)
|
||||
|
||||
# -- helpers -------------------------------------------------------------
|
||||
|
||||
def _get(self, path: str, **kwargs: Any) -> dict:
|
||||
url = f"{self.base_url}{path}"
|
||||
resp = self.session.get(url, **kwargs)
|
||||
if resp.status_code == 404:
|
||||
logger.debug("GET %s → 404", url)
|
||||
return {} # caller distinguishes "not found" from real data
|
||||
resp.raise_for_status()
|
||||
return resp.json()
|
||||
|
||||
def _post(self, path: str, json_body: dict[str, Any]) -> dict:
|
||||
"""POST a JSON body (Content-Type: application/json)."""
|
||||
url = f"{self.base_url}{path}"
|
||||
logger.debug("POST %s → %s", url, json.dumps(json_body))
|
||||
resp = self.session.post(url, json=json_body)
|
||||
resp.raise_for_status()
|
||||
return resp.json()
|
||||
|
||||
# -- public API ----------------------------------------------------------
|
||||
|
||||
def get_next_task(self, task_type: str) -> Optional[dict[str, Any]]:
|
||||
"""Claim the next pending task.
|
||||
|
||||
Returns ``{"task": {...}, "assign_key": "worker_NNN"}`` on success,
|
||||
or ``None`` when no PENDING tasks remain (HTTP 404).
|
||||
"""
|
||||
path = f"/api/v1/nexttask/{task_type}"
|
||||
result = self._get(path)
|
||||
if not result:
|
||||
return None
|
||||
assert "task" in result and "assign_key" in result, \
|
||||
f"Unexpected response shape: {result}"
|
||||
return result
|
||||
|
||||
def get_video(self, video_id: int) -> dict[str, Any]:
|
||||
"""Fetch full video metadata by integer ID.
|
||||
|
||||
Raises :class:`ApiError` if the video is not found (404).
|
||||
"""
|
||||
path = f"/api/v1/video/{video_id}"
|
||||
result = self._get(path)
|
||||
if not result:
|
||||
raise ApiError(404, f"Video {video_id} not found")
|
||||
assert "id" in result, f"Unexpected video response shape: {result}"
|
||||
return result
|
||||
|
||||
def submit_results(
|
||||
self,
|
||||
task_id: int,
|
||||
assign_key: str,
|
||||
results: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
"""Submit processing results for a claimed task.
|
||||
|
||||
Args:
|
||||
task_id: The database ID returned by ``get_next_task``.
|
||||
assign_key: The worker token returned by ``get_next_task``.
|
||||
results: A JSON-serializable dict (status, confidence, routing_decision, etc.).
|
||||
|
||||
Returns the server response dict on success.
|
||||
Raises :class:`ApiError` on 4xx/5xx.
|
||||
"""
|
||||
path = f"/api/v1/task/{task_id}/complete"
|
||||
return self._post(path, {
|
||||
"assign_key": assign_key,
|
||||
"results": results,
|
||||
})
|
||||
|
||||
def create_task(self, task_type: str, video_id: int) -> dict[str, Any]:
|
||||
"""Create a new pending task.
|
||||
|
||||
Args:
|
||||
task_type: Task type string (e.g. 'AISCAN', 'REVIEW').
|
||||
video_id: The video to associate the task with.
|
||||
|
||||
Returns the server response dict on success.
|
||||
Raises :class:`ApiError` on 4xx/5xx.
|
||||
"""
|
||||
path = "/api/v1/tasks"
|
||||
return self._post(path, {
|
||||
"task_type": task_type,
|
||||
"video_id": video_id,
|
||||
"status": "PENDING",
|
||||
})
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close the underlying HTTP session."""
|
||||
self.session.close()
|
||||
@@ -1,138 +0,0 @@
|
||||
"""
|
||||
Face crop classifier.
|
||||
|
||||
Loads a MobileNetV3 model via ONNX/TensorRT FP32, applies temperature-scaled
|
||||
softmax, and returns per-crop probabilities for the target class.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Any, List, Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ImageNet normalisation constants
|
||||
_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
|
||||
_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
|
||||
|
||||
|
||||
def calibrated_softmax(logits: np.ndarray, temperature: float = 1.0) -> np.ndarray:
|
||||
"""Softmax with temperature scaling; returns class probabilities."""
|
||||
scaled = logits / max(temperature, 1e-8)
|
||||
shifted = scaled - scaled.max(axis=-1, keepdims=True) # numerical stability
|
||||
exp = np.exp(shifted)
|
||||
return exp / exp.sum(axis=-1, keepdims=True)
|
||||
|
||||
|
||||
class FaceClassifier:
|
||||
"""Classify face crops and return calibrated target-class probabilities."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
engine_path: str,
|
||||
temperature: float = 1.0,
|
||||
input_size: int = 224,
|
||||
device: str = "cuda",
|
||||
):
|
||||
self.engine_path = engine_path
|
||||
self.temperature = temperature
|
||||
self.input_size = input_size
|
||||
self.device = device
|
||||
self._session: Optional[Any] = None
|
||||
self._load_model()
|
||||
|
||||
def _load_model(self):
|
||||
path = Path(self.engine_path)
|
||||
if not path.exists():
|
||||
logger.warning("Classifier engine not found at %s; using placeholder", self.engine_path)
|
||||
return
|
||||
|
||||
suffix = path.suffix.lower()
|
||||
if suffix == ".onnx":
|
||||
try:
|
||||
import onnxruntime as ort
|
||||
providers = (["CUDAExecutionProvider"] if self.device.startswith("cuda")
|
||||
else ["CPUExecutionProvider"])
|
||||
self._session = ort.InferenceSession(str(path), providers=providers)
|
||||
logger.info("Loaded ONNX classifier from %s", self.engine_path)
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to load ONNX classifier: %s", exc)
|
||||
elif suffix in (".trt", ".engine", ".plan"):
|
||||
try:
|
||||
import tensorrt as trt
|
||||
with trt.Logger() as trt_logger, open(path, "rb") as f:
|
||||
runtime = trt.Runtime(trt_logger)
|
||||
self._session = runtime.deserialize_cuda_engine(f.read())
|
||||
logger.info("Loaded TensorRT classifier from %s", self.engine_path)
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to load TensorRT classifier: %s", exc)
|
||||
else:
|
||||
logger.warning("Unsupported classifier format: %s", suffix)
|
||||
|
||||
def classify(self, crop_paths: List[str], batch_size: int = 16) -> List[float]:
|
||||
"""Return target-class probabilities (p ∈ [0,1]) for each crop path."""
|
||||
results: List[float] = []
|
||||
for i in range(0, len(crop_paths), batch_size):
|
||||
batch = crop_paths[i : i + batch_size]
|
||||
results.extend(self._classify_batch(batch))
|
||||
return results
|
||||
|
||||
def _classify_batch(self, crop_paths: List[str]) -> List[float]:
|
||||
preprocessed = []
|
||||
for path in crop_paths:
|
||||
try:
|
||||
preprocessed.append(self._preprocess(path))
|
||||
except Exception as exc:
|
||||
logger.warning("Preprocessing failed for %s: %s", path, exc)
|
||||
preprocessed.append(np.zeros((3, self.input_size, self.input_size), dtype=np.float32))
|
||||
|
||||
batch = np.stack(preprocessed, axis=0) # (N, 3, H, W)
|
||||
logits = self._infer(batch) # (N, 2)
|
||||
probs = calibrated_softmax(logits, self.temperature)
|
||||
return probs[:, 1].tolist() # target-class column
|
||||
|
||||
def _preprocess(self, crop_path: str) -> np.ndarray:
|
||||
from PIL import Image
|
||||
image = Image.open(crop_path).convert("RGB")
|
||||
if image.size != (self.input_size, self.input_size):
|
||||
image = image.resize((self.input_size, self.input_size), Image.Resampling.BILINEAR)
|
||||
arr = np.array(image, dtype=np.float32) / 255.0
|
||||
arr = (arr - _MEAN) / _STD
|
||||
image.close()
|
||||
return np.transpose(arr, (2, 0, 1)) # HWC → CHW
|
||||
|
||||
def _infer(self, batch: np.ndarray) -> np.ndarray:
|
||||
if self._session is None:
|
||||
return self._placeholder_logits(batch.shape[0])
|
||||
|
||||
try:
|
||||
if hasattr(self._session, "run"):
|
||||
input_name = self._session.get_inputs()[0].name
|
||||
output = self._session.run(None, {input_name: batch})
|
||||
return np.array(output[0])
|
||||
|
||||
# TensorRT path
|
||||
import pycuda.autoinit # noqa: F401
|
||||
import pycuda.driver as cuda
|
||||
|
||||
context = self._session.create_execution_context()
|
||||
out_shape = (batch.shape[0], 2)
|
||||
d_in = cuda.mem_alloc(batch.nbytes)
|
||||
d_out = cuda.mem_alloc(np.prod(out_shape) * np.dtype(np.float32).itemsize)
|
||||
stream = cuda.Stream()
|
||||
cuda.memcpy_htod_async(d_in, batch, stream)
|
||||
context.execute_async_v2(bindings=[int(d_in), int(d_out)], stream_handle=stream.handle)
|
||||
output = np.empty(out_shape, dtype=np.float32)
|
||||
cuda.memcpy_dtoh_async(output, d_out, stream)
|
||||
stream.synchronize()
|
||||
return output
|
||||
except Exception as exc:
|
||||
logger.warning("Classifier inference failed: %s", exc)
|
||||
return self._placeholder_logits(batch.shape[0])
|
||||
|
||||
@staticmethod
|
||||
def _placeholder_logits(n: int) -> np.ndarray:
|
||||
"""Return neutral logits when no model is loaded."""
|
||||
return np.zeros((n, 2), dtype=np.float32)
|
||||
@@ -1,107 +0,0 @@
|
||||
"""
|
||||
Codec Validator
|
||||
|
||||
Validates video codecs against a whitelist/blacklist.
|
||||
Flags unsupported codecs as UNSCANNABLE.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Set
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Default supported codecs
|
||||
DEFAULT_WHITELIST: Set[str] = {
|
||||
"avc1", # H.264
|
||||
"hevc", # H.265
|
||||
"vp8", # VP8
|
||||
"vp9", # VP9
|
||||
"av01", # AV1
|
||||
"mjpeg", # MJPEG
|
||||
"mp4v", # MPEG-4
|
||||
"h264", # H.264 (alternative name)
|
||||
"h265", # H.265 (alternative name)
|
||||
}
|
||||
|
||||
# Known unsupported codecs
|
||||
DEFAULT_BLACKLIST: Set[str] = {
|
||||
"theora", # Theora
|
||||
"divx", # DivX
|
||||
"xvid", # Xvid
|
||||
"prores", # ProRes (requires special handling)
|
||||
"avc", # Old AVC naming
|
||||
"hev1", # HEVC alternative
|
||||
}
|
||||
|
||||
|
||||
class CodecValidator:
|
||||
"""Validate video codecs against whitelist/blacklist."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
whitelist: Set[str] = None,
|
||||
blacklist: Set[str] = None,
|
||||
default_status_on_error: str = "UNSCANNABLE",
|
||||
):
|
||||
self.whitelist = whitelist or DEFAULT_WHITELIST
|
||||
self.blacklist = blacklist or DEFAULT_BLACKLIST
|
||||
self.default_status_on_error = default_status_on_error
|
||||
|
||||
def is_supported(self, codec: str) -> bool:
|
||||
"""Check if a codec is supported."""
|
||||
codec_lower = codec.lower().strip()
|
||||
|
||||
# Check blacklist first
|
||||
if codec_lower in self.blacklist:
|
||||
logger.info("Codec '%s' is blacklisted", codec)
|
||||
return False
|
||||
|
||||
# Check whitelist
|
||||
if codec_lower in self.whitelist:
|
||||
return True
|
||||
|
||||
# Not in either list - assume unsupported
|
||||
logger.warning("Codec '%s' not in whitelist or blacklist", codec)
|
||||
return False
|
||||
|
||||
def validate(self, codec: str) -> tuple:
|
||||
"""
|
||||
Validate a codec and return (is_supported, reason).
|
||||
|
||||
Returns:
|
||||
(True, "supported") if codec is in whitelist
|
||||
(False, "blacklisted") if codec is in blacklist
|
||||
(False, "not_in_whitelist") if codec is unknown
|
||||
"""
|
||||
codec_lower = codec.lower().strip()
|
||||
|
||||
if codec_lower in self.blacklist:
|
||||
return False, "blacklisted"
|
||||
|
||||
if codec_lower in self.whitelist:
|
||||
return True, "supported"
|
||||
|
||||
return False, "not_in_whitelist"
|
||||
|
||||
def add_to_whitelist(self, codec: str):
|
||||
"""Add a codec to the whitelist."""
|
||||
self.whitelist.add(codec.lower().strip())
|
||||
logger.info("Added codec '%s' to whitelist", codec)
|
||||
|
||||
def add_to_blacklist(self, codec: str):
|
||||
"""Add a codec to the blacklist."""
|
||||
self.blacklist.add(codec.lower().strip())
|
||||
logger.info("Added codec '%s' to blacklist", codec)
|
||||
|
||||
def get_status(self, codec: str) -> str:
|
||||
"""Get the status for a codec (for DB update)."""
|
||||
is_supported, reason = self.validate(codec)
|
||||
if is_supported:
|
||||
return "PENDING"
|
||||
return self.default_status_on_error
|
||||
|
||||
def __repr__(self):
|
||||
return (
|
||||
f"CodecValidator(whitelist={len(self.whitelist)}, "
|
||||
f"blacklist={len(self.blacklist)})"
|
||||
)
|
||||
@@ -1,126 +0,0 @@
|
||||
"""
|
||||
Configuration loader for VideoDetect.
|
||||
|
||||
Loads config.yaml with environment variable overrides.
|
||||
Provides typed accessors and validates all required fields.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import yaml
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Default config path
|
||||
DEFAULT_CONFIG_PATH = Path(__file__).parent.parent / "config.yaml"
|
||||
|
||||
|
||||
class Config:
|
||||
"""Typed configuration manager with environment variable overrides."""
|
||||
|
||||
def __init__(self, config_path: Optional[str] = None):
|
||||
self._path = Path(config_path) if config_path else DEFAULT_CONFIG_PATH
|
||||
self._data: Dict[str, Any] = {}
|
||||
self._load()
|
||||
|
||||
def _load(self):
|
||||
"""Load config from YAML file."""
|
||||
if not self._path.exists():
|
||||
raise FileNotFoundError(f"Config file not found: {self._path}")
|
||||
|
||||
with open(self._path, "r") as f:
|
||||
self._data = yaml.safe_load(f) or {}
|
||||
|
||||
# Apply environment variable overrides
|
||||
self._apply_env_overrides()
|
||||
logger.info("Config loaded from %s", self._path)
|
||||
|
||||
def _apply_env_overrides(self):
|
||||
"""Override config values with environment variables."""
|
||||
overrides = {
|
||||
"DB_HOST": ("database.host", None),
|
||||
"DB_PORT": ("database.port", int),
|
||||
"DB_NAME": ("database.name", None),
|
||||
"DB_USER": ("database.user", None),
|
||||
"DB_PASSWORD": ("database.password", None),
|
||||
"DB_POOL_SIZE": ("database.pool_size", int),
|
||||
"DB_POOL_MIN": ("database.pool_min", int),
|
||||
"GPU_MAX_MEMORY_GB": ("gpu.max_memory_gb", float),
|
||||
"SAMPLING_INTERVAL_SECONDS": ("sampling.interval_seconds", int),
|
||||
"T_HIGH": ("thresholds.T_high", float),
|
||||
"T_LOW": ("thresholds.T_low", float),
|
||||
"CONFIG_PATH": (None, None), # handled separately
|
||||
}
|
||||
|
||||
for env_key, (config_path, type_fn) in overrides.items():
|
||||
env_val = os.environ.get(env_key)
|
||||
if env_val is not None:
|
||||
if config_path is None:
|
||||
continue
|
||||
if type_fn is not None:
|
||||
env_val = type_fn(env_val)
|
||||
self._set_nested(self._data, config_path, env_val)
|
||||
logger.debug("Config override: %s=%s (from %s)", config_path, env_val, env_key)
|
||||
|
||||
@staticmethod
|
||||
def _set_nested(data: Dict, path: str, value):
|
||||
"""Set a value in a nested dict using dot notation."""
|
||||
keys = path.split(".")
|
||||
d = data
|
||||
for key in keys[:-1]:
|
||||
d = d.setdefault(key, {})
|
||||
d[keys[-1]] = value
|
||||
|
||||
def get(self, path: str, default=None):
|
||||
"""Get a config value using dot notation."""
|
||||
keys = path.split(".")
|
||||
d = self._data
|
||||
for key in keys:
|
||||
if isinstance(d, dict):
|
||||
d = d.get(key, default)
|
||||
else:
|
||||
return default
|
||||
return d if d is not None else default
|
||||
|
||||
def get_section(self, section: str) -> Dict:
|
||||
"""Get an entire config section as a dict."""
|
||||
return self._data.get(section, {})
|
||||
|
||||
def validate(self, required_keys: list) -> list:
|
||||
"""Validate that required config keys exist. Returns list of missing keys."""
|
||||
missing = []
|
||||
for key in required_keys:
|
||||
if self.get(key) is None:
|
||||
missing.append(key)
|
||||
if missing:
|
||||
logger.error("Missing required config keys: %s", missing)
|
||||
return missing
|
||||
|
||||
@property
|
||||
def data(self) -> Dict[str, Any]:
|
||||
"""Access raw config data."""
|
||||
return self._data
|
||||
|
||||
def __repr__(self):
|
||||
return f"Config(path={self._path}, sections={list(self._data.keys())})"
|
||||
|
||||
|
||||
# Module-level singleton
|
||||
_config: Optional[Config] = None
|
||||
|
||||
|
||||
def get_config(config_path: Optional[str] = None) -> Config:
|
||||
"""Get or create the global config singleton."""
|
||||
global _config
|
||||
if _config is None:
|
||||
_config = Config(config_path)
|
||||
return _config
|
||||
|
||||
|
||||
def reset_config():
|
||||
"""Reset the config singleton (useful for testing)."""
|
||||
global _config
|
||||
_config = None
|
||||
@@ -1,137 +0,0 @@
|
||||
"""
|
||||
Crash recovery and idempotent processing guarantees.
|
||||
|
||||
On startup: re-queues any video stuck in PROCESSING past the lock timeout.
|
||||
During processing: saves per-video checkpoint JSON to scratch for resume.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_CHECKPOINT_DIR = "checkpoints"
|
||||
|
||||
|
||||
class CrashRecovery:
|
||||
"""Handle crash recovery, checkpointing, and idempotent re-queueing."""
|
||||
|
||||
def __init__(self, db_connector, config):
|
||||
self._db = db_connector
|
||||
scratch = config.get("storage.scratch_path", "/scratch")
|
||||
self._checkpoint_dir = Path(scratch) / _CHECKPOINT_DIR
|
||||
mon = config.get_section("monitoring")
|
||||
cr = mon.get("crash_recovery", {})
|
||||
self._lock_timeout_minutes: int = int(cr.get("lock_timeout_minutes", 5))
|
||||
self._auto_requeue: bool = bool(cr.get("auto_requeue", True))
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Startup recovery
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def recover_on_startup(self) -> int:
|
||||
"""
|
||||
Re-queue videos stuck in PROCESSING past the lock timeout.
|
||||
|
||||
Returns the count of videos re-queued.
|
||||
"""
|
||||
if not self._auto_requeue:
|
||||
logger.info("auto_requeue disabled; skipping crash recovery")
|
||||
return 0
|
||||
|
||||
try:
|
||||
result = self._db.execute(
|
||||
"""UPDATE videos
|
||||
SET status = 'PENDING', updated_at = NOW()
|
||||
WHERE status = 'PROCESSING'
|
||||
AND updated_at < NOW() - INTERVAL %s MINUTE""",
|
||||
(self._lock_timeout_minutes,),
|
||||
)
|
||||
count = result if isinstance(result, int) else 0
|
||||
if count:
|
||||
logger.warning("Crash recovery: re-queued %d stuck PROCESSING videos", count)
|
||||
else:
|
||||
logger.info("Crash recovery: no stuck videos found")
|
||||
return count
|
||||
except Exception as exc:
|
||||
logger.error("Crash recovery query failed: %s", exc)
|
||||
return 0
|
||||
|
||||
def list_stuck_videos(self) -> List[Dict]:
|
||||
"""Return videos currently stuck past the lock timeout (read-only)."""
|
||||
try:
|
||||
rows = self._db.fetchall(
|
||||
"""SELECT id, file_path, updated_at
|
||||
FROM videos
|
||||
WHERE status = 'PROCESSING'
|
||||
AND updated_at < NOW() - INTERVAL %s MINUTE""",
|
||||
(self._lock_timeout_minutes,),
|
||||
)
|
||||
return rows or []
|
||||
except Exception as exc:
|
||||
logger.error("Failed to list stuck videos: %s", exc)
|
||||
return []
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Per-video checkpointing
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def save_checkpoint(self, video_id: int, state: str, progress: Dict) -> None:
|
||||
"""Persist a checkpoint JSON for the given video."""
|
||||
self._checkpoint_dir.mkdir(parents=True, exist_ok=True)
|
||||
checkpoint = {
|
||||
"video_id": video_id,
|
||||
"state": state,
|
||||
"progress": progress,
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
}
|
||||
path = self._checkpoint_dir / f"{video_id}.json"
|
||||
try:
|
||||
path.write_text(json.dumps(checkpoint, indent=2), encoding="utf-8")
|
||||
logger.debug("Saved checkpoint for video %d (state=%s)", video_id, state)
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to save checkpoint for video %d: %s", video_id, exc)
|
||||
|
||||
def load_checkpoint(self, video_id: int) -> Optional[Dict]:
|
||||
"""Return the stored checkpoint for video_id, or None if absent."""
|
||||
path = self._checkpoint_dir / f"{video_id}.json"
|
||||
if not path.exists():
|
||||
return None
|
||||
try:
|
||||
return json.loads(path.read_text(encoding="utf-8"))
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to read checkpoint for video %d: %s", video_id, exc)
|
||||
return None
|
||||
|
||||
def delete_checkpoint(self, video_id: int) -> None:
|
||||
"""Remove the checkpoint file after successful processing."""
|
||||
path = self._checkpoint_dir / f"{video_id}.json"
|
||||
if path.exists():
|
||||
try:
|
||||
path.unlink()
|
||||
logger.debug("Deleted checkpoint for video %d", video_id)
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to delete checkpoint for video %d: %s", video_id, exc)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Idempotency guard
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def is_already_completed(self, video_id: int) -> bool:
|
||||
"""
|
||||
Return True if this video already has a COMPLETED status.
|
||||
|
||||
Prevents duplicate processing if a job is accidentally re-submitted.
|
||||
"""
|
||||
try:
|
||||
rows = self._db.fetchall(
|
||||
"SELECT status FROM videos WHERE id = %s",
|
||||
(video_id,),
|
||||
)
|
||||
return bool(rows and rows[0].get("status") == "COMPLETED")
|
||||
except Exception as exc:
|
||||
logger.warning("Idempotency check failed for video %d: %s", video_id, exc)
|
||||
return False
|
||||
@@ -1,227 +0,0 @@
|
||||
"""
|
||||
Database connection layer with connection pooling for VideoDetect.
|
||||
|
||||
Provides:
|
||||
- Connection pooling via DBUtils + PyMySQL
|
||||
- Automatic reconnection on disconnect (ping + retry on transient errors)
|
||||
- Context manager support
|
||||
- Prepared statements for all queries
|
||||
- Transaction support for atomic state transitions
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import logging
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
from dbutils.pooled_db import PooledDB
|
||||
import pymysql
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DBConnector:
|
||||
"""Thread-safe database connection pool manager."""
|
||||
|
||||
# Transient errors that indicate a dropped/stale connection and are safe
|
||||
# to retry with a fresh connection from the pool.
|
||||
_RETRYABLE_ERRORS = (pymysql.err.InterfaceError, pymysql.err.OperationalError)
|
||||
_MAX_RETRIES = 3
|
||||
_RETRY_BACKOFF_SECONDS = 0.5
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host: str = "mariadb",
|
||||
port: int = 3306,
|
||||
database: str = "videodetect",
|
||||
user: str = "videodetect",
|
||||
password: str = "videodetect123",
|
||||
pool_size: int = 20,
|
||||
pool_min: int = 5,
|
||||
pool_recycle: int = 3600,
|
||||
):
|
||||
# PooledDB parameters
|
||||
pool_config = {
|
||||
"creator": pymysql,
|
||||
"maxconnections": pool_size,
|
||||
"mincached": pool_min,
|
||||
"maxcached": pool_size,
|
||||
"maxusage": 200,
|
||||
"blocking": True,
|
||||
# ping(0) = check if connection is alive BEFORE returning it from
|
||||
# the pool. This is the correct DBUtils mechanism for detecting
|
||||
# dead/stale connections (idle timeout, network blip, DB restart).
|
||||
"ping": 0,
|
||||
}
|
||||
|
||||
# PyMySQL connection parameters
|
||||
connection_params = {
|
||||
"host": host,
|
||||
"port": port,
|
||||
"database": database,
|
||||
"user": user,
|
||||
"password": password,
|
||||
"charset": "utf8mb4",
|
||||
"cursorclass": pymysql.cursors.DictCursor,
|
||||
"read_timeout": 30,
|
||||
"write_timeout": 30,
|
||||
}
|
||||
|
||||
self._pool = PooledDB(**pool_config, **connection_params)
|
||||
self._pool_recycle = pool_recycle
|
||||
logger.info(
|
||||
"DB pool initialized: host=%s db=%s pool_size=%d min=%d ping=0",
|
||||
host, database, pool_size, pool_min,
|
||||
)
|
||||
|
||||
def get_connection(self):
|
||||
"""Get a connection from the pool."""
|
||||
return self._pool.connection()
|
||||
|
||||
def _run_query(self, query: str, params, transaction: bool, fetch: str):
|
||||
"""Run a single query on a fresh connection, with retry on transient errors.
|
||||
|
||||
A pooled connection can be dropped by the server (idle timeout, network
|
||||
blip, DB restart). When that happens the query fails with an
|
||||
InterfaceError/OperationalError. We discard the dead connection and
|
||||
retry with a fresh one from the pool.
|
||||
|
||||
Args:
|
||||
query: SQL statement.
|
||||
params: Bound parameters (or None).
|
||||
transaction: If True, wrap the query in a transaction.
|
||||
fetch: One of "none", "one", or "all" controlling the return value.
|
||||
|
||||
Returns:
|
||||
- "none": the affected row count.
|
||||
- "one": a single row (dict) or None.
|
||||
- "all": a list of rows (dicts).
|
||||
"""
|
||||
last_exc = None
|
||||
for attempt in range(1, self._MAX_RETRIES + 1):
|
||||
conn = self.get_connection()
|
||||
try:
|
||||
cursor = conn.cursor()
|
||||
if transaction:
|
||||
conn.begin()
|
||||
cursor.execute(query, params or ())
|
||||
if fetch == "one":
|
||||
result = cursor.fetchone()
|
||||
elif fetch == "all":
|
||||
result = cursor.fetchall()
|
||||
else:
|
||||
result = cursor.rowcount
|
||||
if transaction:
|
||||
conn.commit()
|
||||
cursor.close()
|
||||
conn.close()
|
||||
return result
|
||||
except self._RETRYABLE_ERRORS as e:
|
||||
last_exc = e
|
||||
logger.warning(
|
||||
"Transient DB error (attempt %d/%d): %s",
|
||||
attempt, self._MAX_RETRIES, e,
|
||||
)
|
||||
# Discard the dead connection; a fresh one is taken on retry.
|
||||
self._safe_close(conn)
|
||||
if attempt < self._MAX_RETRIES:
|
||||
time.sleep(self._RETRY_BACKOFF_SECONDS * attempt)
|
||||
except Exception:
|
||||
# Non-transient error: roll back if in a transaction and re-raise.
|
||||
if transaction:
|
||||
self._safe_rollback(conn)
|
||||
self._safe_close(conn)
|
||||
raise
|
||||
raise last_exc
|
||||
|
||||
@staticmethod
|
||||
def _safe_close(conn):
|
||||
try:
|
||||
conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _safe_rollback(conn):
|
||||
try:
|
||||
conn.rollback()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@contextlib.contextmanager
|
||||
def get_cursor(self, transaction: bool = False):
|
||||
"""Context manager for getting a cursor with optional transaction support.
|
||||
|
||||
Note: this does not retry on transient errors; use execute/fetchone/
|
||||
fetchall for retry-safe single-statement queries.
|
||||
"""
|
||||
conn = self.get_connection()
|
||||
cursor = None
|
||||
try:
|
||||
cursor = conn.cursor()
|
||||
if transaction:
|
||||
conn.begin()
|
||||
yield cursor
|
||||
if transaction:
|
||||
conn.commit()
|
||||
except Exception:
|
||||
if transaction:
|
||||
self._safe_rollback(conn)
|
||||
raise
|
||||
finally:
|
||||
if cursor is not None:
|
||||
try:
|
||||
cursor.close()
|
||||
except Exception:
|
||||
pass
|
||||
self._safe_close(conn)
|
||||
|
||||
@contextlib.contextmanager
|
||||
def transaction(self):
|
||||
"""Context manager for a full transaction."""
|
||||
conn = self.get_connection()
|
||||
try:
|
||||
conn.begin()
|
||||
yield conn
|
||||
conn.commit()
|
||||
except Exception:
|
||||
self._safe_rollback(conn)
|
||||
raise
|
||||
finally:
|
||||
self._safe_close(conn)
|
||||
|
||||
def execute(self, query: str, params=None, transaction: bool = False):
|
||||
"""Execute a query and return affected rows. Retries on transient errors."""
|
||||
return self._run_query(query, params, transaction, fetch="none")
|
||||
|
||||
def fetchone(self, query: str, params=None):
|
||||
"""Execute a query and return one row. Retries on transient errors."""
|
||||
return self._run_query(query, params, transaction=False, fetch="one")
|
||||
|
||||
def fetchall(self, query: str, params=None):
|
||||
"""Execute a query and return all rows. Retries on transient errors."""
|
||||
return self._run_query(query, params, transaction=False, fetch="all")
|
||||
|
||||
def initialize_schema(self, schema_path: str = "db/schema.sql"):
|
||||
"""Initialize the database schema from SQL file."""
|
||||
with open(schema_path, "r") as f:
|
||||
sql = f.read()
|
||||
# Split on semicolons and execute each statement
|
||||
statements = [s.strip() for s in sql.split(";") if s.strip()]
|
||||
for stmt in statements:
|
||||
if stmt.startswith("--"):
|
||||
continue
|
||||
try:
|
||||
self.execute(stmt)
|
||||
except Exception as e:
|
||||
logger.debug("Schema statement skipped (may already exist): %s", e)
|
||||
logger.info("Schema initialized from %s", schema_path)
|
||||
|
||||
def health_check(self) -> bool:
|
||||
"""Check if the database is reachable."""
|
||||
try:
|
||||
result = self.fetchone("SELECT 1")
|
||||
return result is not None
|
||||
except Exception as e:
|
||||
logger.error("Health check failed: %s", e)
|
||||
return False
|
||||
@@ -1,219 +0,0 @@
|
||||
"""
|
||||
Confidence distribution drift detection.
|
||||
|
||||
Compares recent video confidence scores against a stored baseline and fires
|
||||
alerts when the distribution shifts beyond configured thresholds.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Callable, Dict, List, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DriftDetector:
|
||||
"""Detect and alert on confidence distribution drift."""
|
||||
|
||||
def __init__(self, db_connector, config, alert_fn: Optional[Callable[[str], None]] = None):
|
||||
self._db = db_connector
|
||||
mon = config.get_section("monitoring")
|
||||
alerts = mon.get("alerts", {})
|
||||
self._drift_threshold: float = float(alerts.get("confidence_drift_threshold", 0.10))
|
||||
self._review_max_size: int = int(alerts.get("review_queue_max_size", 1000))
|
||||
self._review_max_age_h: int = int(alerts.get("review_queue_max_age_hours", 24))
|
||||
self._throughput_min: float = float(alerts.get("throughput_min_videos_per_hour", 20))
|
||||
self._throughput_window_h: int = int(alerts.get("throughput_min_duration_hours", 1))
|
||||
self._error_rate_threshold: float = float(alerts.get("error_rate_threshold", 0.05))
|
||||
self._error_rate_window_h: int = int(alerts.get("error_rate_window_hours", 1))
|
||||
self._alert_fn = alert_fn or _log_alert
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Public API
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def run_all_checks(self) -> Dict[str, bool]:
|
||||
"""
|
||||
Run all alert checks and return a dict of {check_name: triggered}.
|
||||
|
||||
Fires the alert function for each triggered condition.
|
||||
"""
|
||||
results = {}
|
||||
results["confidence_drift"] = self.check_confidence_drift()
|
||||
results["review_queue_growth"] = self.check_review_queue_growth()
|
||||
results["low_throughput"] = self.check_low_throughput()
|
||||
results["high_error_rate"] = self.check_error_rate()
|
||||
return results
|
||||
|
||||
def check_confidence_drift(self) -> bool:
|
||||
"""
|
||||
Compare p(C > 0.5) in recent processing logs to the stored baseline.
|
||||
|
||||
Returns True if drift exceeds the threshold.
|
||||
"""
|
||||
current = self._fetch_recent_confidences(hours=24 * 7)
|
||||
baseline = self._fetch_baseline_confidences()
|
||||
if not current or not baseline:
|
||||
logger.debug("Not enough data for drift check")
|
||||
return False
|
||||
|
||||
triggered = detect_drift(current, baseline, self._drift_threshold)
|
||||
if triggered:
|
||||
current_high = _high_conf_fraction(current)
|
||||
baseline_high = _high_conf_fraction(baseline)
|
||||
shift = abs(current_high - baseline_high)
|
||||
self._alert_fn(
|
||||
f"Confidence drift detected: {shift:.1%} shift in p(C>0.5) "
|
||||
f"(current={current_high:.1%}, baseline={baseline_high:.1%}, "
|
||||
f"threshold={self._drift_threshold:.1%})"
|
||||
)
|
||||
return triggered
|
||||
|
||||
def check_review_queue_growth(self) -> bool:
|
||||
"""Alert if the unannotated review queue has been above max_size for max_age hours."""
|
||||
try:
|
||||
cutoff = datetime.now(timezone.utc) - timedelta(hours=self._review_max_age_h)
|
||||
rows = self._db.fetchall(
|
||||
"""SELECT COUNT(*) AS cnt
|
||||
FROM review_queue
|
||||
WHERE annotated = FALSE
|
||||
AND created_at < %s""",
|
||||
(cutoff.replace(tzinfo=None),),
|
||||
)
|
||||
count = int((rows or [{"cnt": 0}])[0]["cnt"])
|
||||
if count > self._review_max_size:
|
||||
self._alert_fn(
|
||||
f"Review queue growth alert: {count} unannotated items older than "
|
||||
f"{self._review_max_age_h}h (threshold={self._review_max_size})"
|
||||
)
|
||||
return True
|
||||
return False
|
||||
except Exception as exc:
|
||||
logger.warning("Review queue growth check failed: %s", exc)
|
||||
return False
|
||||
|
||||
def check_low_throughput(self) -> bool:
|
||||
"""Alert if throughput dropped below the minimum for the last window hours."""
|
||||
try:
|
||||
cutoff = datetime.now(timezone.utc) - timedelta(hours=self._throughput_window_h)
|
||||
rows = self._db.fetchall(
|
||||
"""SELECT COUNT(*) AS cnt
|
||||
FROM videos
|
||||
WHERE status = 'COMPLETED'
|
||||
AND last_processed_time >= %s""",
|
||||
(cutoff.replace(tzinfo=None),),
|
||||
)
|
||||
count = int((rows or [{"cnt": 0}])[0]["cnt"])
|
||||
videos_per_hour = count / max(self._throughput_window_h, 1)
|
||||
if videos_per_hour < self._throughput_min:
|
||||
self._alert_fn(
|
||||
f"Low throughput alert: {videos_per_hour:.1f} videos/hour "
|
||||
f"over last {self._throughput_window_h}h "
|
||||
f"(minimum={self._throughput_min})"
|
||||
)
|
||||
return True
|
||||
return False
|
||||
except Exception as exc:
|
||||
logger.warning("Throughput check failed: %s", exc)
|
||||
return False
|
||||
|
||||
def check_error_rate(self) -> bool:
|
||||
"""Alert if the fraction of ERROR/UNSCANNABLE videos exceeds the threshold."""
|
||||
try:
|
||||
cutoff = datetime.now(timezone.utc) - timedelta(hours=self._error_rate_window_h)
|
||||
rows = self._db.fetchall(
|
||||
"""SELECT status, COUNT(*) AS cnt
|
||||
FROM videos
|
||||
WHERE last_processed_time >= %s
|
||||
AND status IN ('COMPLETED', 'ERROR', 'UNSCANNABLE')
|
||||
GROUP BY status""",
|
||||
(cutoff.replace(tzinfo=None),),
|
||||
)
|
||||
counts: Dict[str, int] = {r["status"]: int(r["cnt"]) for r in (rows or [])}
|
||||
total = sum(counts.values())
|
||||
if total == 0:
|
||||
return False
|
||||
errors = counts.get("ERROR", 0) + counts.get("UNSCANNABLE", 0)
|
||||
rate = errors / total
|
||||
if rate > self._error_rate_threshold:
|
||||
self._alert_fn(
|
||||
f"High error rate alert: {rate:.1%} over last {self._error_rate_window_h}h "
|
||||
f"({errors}/{total} videos, threshold={self._error_rate_threshold:.1%})"
|
||||
)
|
||||
return True
|
||||
return False
|
||||
except Exception as exc:
|
||||
logger.warning("Error rate check failed: %s", exc)
|
||||
return False
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# DB helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _fetch_recent_confidences(self, hours: int = 168) -> List[float]:
|
||||
"""Return confidence scores from the last `hours` hours."""
|
||||
try:
|
||||
cutoff = datetime.now(timezone.utc) - timedelta(hours=hours)
|
||||
rows = self._db.fetchall(
|
||||
"""SELECT confidence_score
|
||||
FROM processing_logs
|
||||
WHERE processed_at >= %s
|
||||
AND confidence_score IS NOT NULL""",
|
||||
(cutoff.replace(tzinfo=None),),
|
||||
)
|
||||
return [float(r["confidence_score"]) for r in (rows or [])]
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to fetch recent confidences: %s", exc)
|
||||
return []
|
||||
|
||||
def _fetch_baseline_confidences(self) -> List[float]:
|
||||
"""Return the baseline confidence distribution stored in the DB."""
|
||||
try:
|
||||
rows = self._db.fetchall(
|
||||
"""SELECT confidence_score
|
||||
FROM processing_logs
|
||||
WHERE confidence_score IS NOT NULL
|
||||
ORDER BY processed_at ASC
|
||||
LIMIT 10000"""
|
||||
)
|
||||
return [float(r["confidence_score"]) for r in (rows or [])]
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to fetch baseline confidences: %s", exc)
|
||||
return []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Module-level pure functions (testable without DB)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def detect_drift(
|
||||
current_confidences: List[float],
|
||||
baseline_confidences: List[float],
|
||||
threshold: float = 0.10,
|
||||
) -> bool:
|
||||
"""
|
||||
Return True if |p(C>0.5)_current - p(C>0.5)_baseline| > threshold.
|
||||
|
||||
Both lists must be non-empty.
|
||||
"""
|
||||
if not current_confidences or not baseline_confidences:
|
||||
return False
|
||||
current_high = _high_conf_fraction(current_confidences)
|
||||
baseline_high = _high_conf_fraction(baseline_confidences)
|
||||
shift = abs(current_high - baseline_high)
|
||||
if shift > threshold:
|
||||
logger.warning(
|
||||
"Drift detected: shift=%.2f%% (current=%.2f%% baseline=%.2f%%)",
|
||||
shift * 100, current_high * 100, baseline_high * 100,
|
||||
)
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _high_conf_fraction(confidences: List[float]) -> float:
|
||||
"""Return the fraction of scores > 0.5."""
|
||||
return sum(1 for c in confidences if c > 0.5) / len(confidences)
|
||||
|
||||
|
||||
def _log_alert(message: str):
|
||||
logger.warning("ALERT: %s", message)
|
||||
@@ -1,104 +0,0 @@
|
||||
"""Export a YOLOv8n face-detection model to ONNX and build a TensorRT engine."""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def export_onnx(output_dir: Path, input_size: int = 640) -> Path:
|
||||
"""Export YOLOv8n to ONNX with the required input shape."""
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
onnx_path = output_dir / "face_detector.onnx"
|
||||
metadata_path = output_dir / "model.json"
|
||||
|
||||
try:
|
||||
from ultralytics import YOLO
|
||||
|
||||
model = YOLO("yolov8n.pt")
|
||||
model.export(
|
||||
format="onnx",
|
||||
imgsz=input_size,
|
||||
half=False,
|
||||
simplify=True,
|
||||
dynamic=False,
|
||||
)
|
||||
exported = Path("yolov8n.onnx")
|
||||
if exported.exists():
|
||||
exported.rename(onnx_path)
|
||||
logger.info("Exported ONNX model to %s", onnx_path)
|
||||
except Exception as exc:
|
||||
logger.error("Failed to export ONNX model: %s", exc)
|
||||
raise
|
||||
|
||||
metadata = {
|
||||
"input_shape": [1, 3, input_size, input_size],
|
||||
"mean": [0.485, 0.456, 0.406],
|
||||
"std": [0.229, 0.224, 0.225],
|
||||
"confidence_threshold": 0.25,
|
||||
"iou_threshold": 0.45,
|
||||
}
|
||||
metadata_path.write_text(json.dumps(metadata, indent=2))
|
||||
return onnx_path
|
||||
|
||||
|
||||
def build_tensorrt_engine(onnx_path: Path, output_dir: Path, max_batch_size: int = 32) -> Path:
|
||||
"""Build and serialize a TensorRT FP32 engine from an ONNX file."""
|
||||
engine_path = output_dir / "face_detector.trt"
|
||||
|
||||
try:
|
||||
import tensorrt as trt
|
||||
|
||||
logger = trt.Logger(trt.Logger.WARNING)
|
||||
builder = trt.Builder(logger)
|
||||
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
|
||||
parser = trt.OnnxParser(network, logger)
|
||||
|
||||
onnx_data = onnx_path.read_bytes()
|
||||
if not parser.parse(onnx_data):
|
||||
for error in range(parser.num_errors):
|
||||
logger.error(parser.get_error(error))
|
||||
raise RuntimeError("ONNX parsing failed")
|
||||
|
||||
config = builder.create_builder_config()
|
||||
config.max_workspace_size = 4 << 30 # 4GB
|
||||
config.set_flag(trt.BuilderFlag.FP32)
|
||||
|
||||
profile = builder.create_optimization_profile()
|
||||
input_name = network.get_input(0).name
|
||||
profile.set_shape(
|
||||
input_name,
|
||||
(1, 3, 640, 640),
|
||||
(max_batch_size // 2, 3, 640, 640),
|
||||
(max_batch_size, 3, 640, 640),
|
||||
)
|
||||
config.add_optimization_profile(profile)
|
||||
|
||||
engine = builder.build_engine(network, config)
|
||||
if engine is None:
|
||||
raise RuntimeError("TensorRT engine build failed")
|
||||
|
||||
engine_path.write_bytes(engine.serialize())
|
||||
logger.info("Serialized TensorRT engine to %s", engine_path)
|
||||
return engine_path
|
||||
except Exception as exc:
|
||||
logger.error("Failed to build TensorRT engine: %s", exc)
|
||||
raise
|
||||
|
||||
|
||||
def main():
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
parser = argparse.ArgumentParser(description="Export YOLOv8n face detector to ONNX/TensorRT")
|
||||
parser.add_argument("--output-dir", type=Path, default=Path("models/face_detector"))
|
||||
parser.add_argument("--input-size", type=int, default=640)
|
||||
parser.add_argument("--max-batch-size", type=int, default=32)
|
||||
args = parser.parse_args()
|
||||
|
||||
onnx_path = export_onnx(args.output_dir, args.input_size)
|
||||
build_tensorrt_engine(onnx_path, args.output_dir, args.max_batch_size)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,295 +0,0 @@
|
||||
"""
|
||||
Face detection runner.
|
||||
|
||||
Wraps a YOLOv8n model exported to ONNX/TensorRT, runs batched inference on
|
||||
sampled frames, applies NMS, and produces 224×224 face crops for the classifier.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Detection:
|
||||
"""A single face detection result."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
frame_path: str,
|
||||
x1: float,
|
||||
y1: float,
|
||||
x2: float,
|
||||
y2: float,
|
||||
confidence: float,
|
||||
crop_path: Optional[str] = None,
|
||||
):
|
||||
self.frame_path = frame_path
|
||||
self.x1 = x1
|
||||
self.y1 = y1
|
||||
self.x2 = x2
|
||||
self.y2 = y2
|
||||
self.confidence = confidence
|
||||
self.crop_path = crop_path
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"frame_path": self.frame_path,
|
||||
"x1": self.x1,
|
||||
"y1": self.y1,
|
||||
"x2": self.x2,
|
||||
"y2": self.y2,
|
||||
"confidence": self.confidence,
|
||||
"crop_path": self.crop_path,
|
||||
}
|
||||
|
||||
|
||||
class FaceDetector:
|
||||
"""Run batched face detection on video frames."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
engine_path: str,
|
||||
input_size: int = 640,
|
||||
confidence_threshold: float = 0.25,
|
||||
iou_threshold: float = 0.45,
|
||||
max_faces_per_frame: int = 10,
|
||||
max_faces_per_video: int = 100,
|
||||
device: str = "cuda",
|
||||
):
|
||||
self.engine_path = engine_path
|
||||
self.input_size = input_size
|
||||
self.confidence_threshold = confidence_threshold
|
||||
self.iou_threshold = iou_threshold
|
||||
self.max_faces_per_frame = max_faces_per_frame
|
||||
self.max_faces_per_video = max_faces_per_video
|
||||
self.device = device
|
||||
self._session: Optional[Any] = None
|
||||
self._load_model()
|
||||
|
||||
def _load_model(self):
|
||||
"""Load the inference backend."""
|
||||
path = Path(self.engine_path)
|
||||
if not path.exists():
|
||||
logger.warning("Face detector engine not found at %s; using placeholder", self.engine_path)
|
||||
return
|
||||
|
||||
suffix = path.suffix.lower()
|
||||
if suffix == ".onnx":
|
||||
try:
|
||||
import onnxruntime as ort
|
||||
|
||||
providers = ["CUDAExecutionProvider"] if self.device.startswith("cuda") else ["CPUExecutionProvider"]
|
||||
self._session = ort.InferenceSession(str(path), providers=providers)
|
||||
logger.info("Loaded ONNX face detector from %s", self.engine_path)
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to load ONNX face detector: %s", exc)
|
||||
elif suffix in (".trt", ".engine", ".plan"):
|
||||
try:
|
||||
import tensorrt as trt
|
||||
|
||||
with trt.Logger() as trt_logger, open(path, "rb") as f:
|
||||
runtime = trt.Runtime(trt_logger)
|
||||
self._session = runtime.deserialize_cuda_engine(f.read())
|
||||
logger.info("Loaded TensorRT face detector from %s", self.engine_path)
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to load TensorRT face detector: %s", exc)
|
||||
else:
|
||||
logger.warning("Unsupported face detector format: %s", suffix)
|
||||
|
||||
def detect_faces(
|
||||
self,
|
||||
frame_paths: List[str],
|
||||
batch_size: int = 16,
|
||||
) -> List[List[Detection]]:
|
||||
"""Detect faces in a list of frame image paths."""
|
||||
all_results: List[List[Detection]] = []
|
||||
for i in range(0, len(frame_paths), batch_size):
|
||||
batch = frame_paths[i : i + batch_size]
|
||||
batch_results = self._detect_batch(batch)
|
||||
all_results.extend(batch_results)
|
||||
return all_results
|
||||
|
||||
def _detect_batch(self, frame_paths: List[str]) -> List[List[Detection]]:
|
||||
"""Run detection on one batch and return detections per frame."""
|
||||
preprocessed = []
|
||||
for path in frame_paths:
|
||||
try:
|
||||
preprocessed.append(self._preprocess(path))
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to preprocess %s: %s", path, exc)
|
||||
preprocessed.append(np.zeros((3, self.input_size, self.input_size), dtype=np.float32))
|
||||
|
||||
batch_input = np.stack(preprocessed, axis=0)
|
||||
|
||||
outputs = self._infer(batch_input)
|
||||
|
||||
results = []
|
||||
for idx, path in enumerate(frame_paths):
|
||||
try:
|
||||
detections = self._parse_outputs(outputs, idx, path)
|
||||
detections = self._nms(detections)
|
||||
detections = detections[: self.max_faces_per_frame]
|
||||
results.append(detections)
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to parse detections for %s: %s", path, exc)
|
||||
results.append([])
|
||||
return results
|
||||
|
||||
def _preprocess(self, frame_path: str) -> np.ndarray:
|
||||
"""Load and normalize a frame for the detector."""
|
||||
from PIL import Image
|
||||
|
||||
image = Image.open(frame_path).convert("RGB")
|
||||
image = image.resize((self.input_size, self.input_size), Image.Resampling.BILINEAR)
|
||||
arr = np.array(image, dtype=np.float32) / 255.0
|
||||
mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
|
||||
std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
|
||||
arr = (arr - mean) / std
|
||||
return np.transpose(arr, (2, 0, 1)) # HWC -> CHW
|
||||
|
||||
def _infer(self, batch_input: np.ndarray) -> Any:
|
||||
"""Run inference on the preprocessed batch."""
|
||||
if self._session is None:
|
||||
return self._placeholder_outputs(batch_input.shape[0])
|
||||
|
||||
try:
|
||||
if hasattr(self._session, "run"):
|
||||
input_name = self._session.get_inputs()[0].name
|
||||
return self._session.run(None, {input_name: batch_input})
|
||||
|
||||
# TensorRT execution
|
||||
import pycuda.driver as cuda
|
||||
import pycuda.autoinit # noqa: F401
|
||||
|
||||
context = self._session.create_execution_context()
|
||||
output_shape = (batch_input.shape[0], 84, 8400) # YOLOv8n default shape
|
||||
d_input = cuda.mem_alloc(batch_input.nbytes)
|
||||
d_output = cuda.mem_alloc(np.prod(output_shape) * np.dtype(np.float32).itemsize)
|
||||
stream = cuda.Stream()
|
||||
cuda.memcpy_htod_async(d_input, batch_input, stream)
|
||||
context.execute_async_v2(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)
|
||||
output = np.empty(output_shape, dtype=np.float32)
|
||||
cuda.memcpy_dtoh_async(output, d_output, stream)
|
||||
stream.synchronize()
|
||||
return output
|
||||
except Exception as exc:
|
||||
logger.warning("Face detector inference failed: %s", exc)
|
||||
return self._placeholder_outputs(batch_input.shape[0])
|
||||
|
||||
def _placeholder_outputs(self, batch_size: int) -> np.ndarray:
|
||||
"""Return an empty output tensor when no model is loaded."""
|
||||
return np.zeros((batch_size, 84, 8400), dtype=np.float32)
|
||||
|
||||
def _parse_outputs(self, outputs: Any, batch_index: int, frame_path: str) -> List[Detection]:
|
||||
"""Parse raw inference outputs into Detection objects."""
|
||||
if isinstance(outputs, list):
|
||||
raw = outputs[0][batch_index] # (84, 8400)
|
||||
else:
|
||||
raw = outputs[batch_index]
|
||||
|
||||
# YOLOv8 output layout: (cx, cy, w, h, cls scores...)
|
||||
raw = raw.T # (8400, 84)
|
||||
scores = raw[:, 4:].max(axis=1)
|
||||
mask = scores >= self.confidence_threshold
|
||||
candidates = raw[mask]
|
||||
scores = scores[mask]
|
||||
|
||||
detections = []
|
||||
for row, score in zip(candidates, scores):
|
||||
cx, cy, w, h = row[:4]
|
||||
x1 = cx - w / 2
|
||||
y1 = cy - h / 2
|
||||
x2 = cx + w / 2
|
||||
y2 = cy + h / 2
|
||||
detections.append(
|
||||
Detection(
|
||||
frame_path=frame_path,
|
||||
x1=float(x1),
|
||||
y1=float(y1),
|
||||
x2=float(x2),
|
||||
y2=float(y2),
|
||||
confidence=float(score),
|
||||
)
|
||||
)
|
||||
return detections
|
||||
|
||||
def _nms(self, detections: List[Detection]) -> List[Detection]:
|
||||
"""Apply greedy Non-Maximum Suppression."""
|
||||
if not detections:
|
||||
return []
|
||||
|
||||
sorted_dets = sorted(detections, key=lambda d: d.confidence, reverse=True)
|
||||
kept: List[Detection] = []
|
||||
|
||||
while sorted_dets:
|
||||
current = sorted_dets.pop(0)
|
||||
kept.append(current)
|
||||
sorted_dets = [
|
||||
det
|
||||
for det in sorted_dets
|
||||
if self._iou(current, det) <= self.iou_threshold
|
||||
]
|
||||
if len(kept) >= self.max_faces_per_frame:
|
||||
break
|
||||
|
||||
return kept
|
||||
|
||||
@staticmethod
|
||||
def _iou(a: Detection, b: Detection) -> float:
|
||||
"""Compute IoU between two detections."""
|
||||
x1 = max(a.x1, b.x1)
|
||||
y1 = max(a.y1, b.y1)
|
||||
x2 = min(a.x2, b.x2)
|
||||
y2 = min(a.y2, b.y2)
|
||||
|
||||
inter_area = max(0, x2 - x1) * max(0, y2 - y1)
|
||||
area_a = (a.x2 - a.x1) * (a.y2 - a.y1)
|
||||
area_b = (b.x2 - b.x1) * (b.y2 - b.y1)
|
||||
union_area = area_a + area_b - inter_area
|
||||
if union_area == 0:
|
||||
return 0.0
|
||||
return inter_area / union_area
|
||||
|
||||
def extract_crops(
|
||||
self,
|
||||
detections: List[Detection],
|
||||
output_dir: str,
|
||||
crop_size: Tuple[int, int] = (224, 224),
|
||||
) -> List[Detection]:
|
||||
"""Extract resized face crops from original frames and update detections."""
|
||||
from PIL import Image
|
||||
|
||||
output_dir_path = Path(output_dir)
|
||||
output_dir_path.mkdir(parents=True, exist_ok=True)
|
||||
cropped: List[Detection] = []
|
||||
|
||||
for idx, det in enumerate(detections):
|
||||
try:
|
||||
image = Image.open(det.frame_path).convert("RGB")
|
||||
width, height = image.size
|
||||
|
||||
x1 = int(max(0, det.x1 * width / self.input_size))
|
||||
y1 = int(max(0, det.y1 * height / self.input_size))
|
||||
x2 = int(min(width, det.x2 * width / self.input_size))
|
||||
y2 = int(min(height, det.y2 * height / self.input_size))
|
||||
|
||||
if x2 <= x1 or y2 <= y1:
|
||||
image.close()
|
||||
continue
|
||||
|
||||
crop = image.crop((x1, y1, x2, y2)).resize(crop_size, Image.Resampling.BILINEAR)
|
||||
image.close()
|
||||
crop_path = output_dir_path / f"crop_{int(time.time() * 1000)}_{idx}.jpg"
|
||||
crop.save(crop_path, "JPEG", quality=95)
|
||||
crop.close()
|
||||
det.crop_path = str(crop_path)
|
||||
cropped.append(det)
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to extract crop for %s: %s", det.frame_path, exc)
|
||||
|
||||
return cropped
|
||||
@@ -1,163 +0,0 @@
|
||||
"""
|
||||
Frame sampling module for VideoDetect.
|
||||
|
||||
Extracts uniform temporal frames from videos using FFmpeg, handles resolution
|
||||
constraints, and writes JPEG frames to a scratch directory.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from typing import Iterable, List, Optional, Sequence, Tuple
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def calculate_timestamps(duration: Optional[float], interval: float) -> List[float]:
|
||||
"""Return uniformly spaced timestamps in seconds for a video duration."""
|
||||
if duration is None or duration <= 0:
|
||||
return [0.0]
|
||||
if interval <= 0:
|
||||
interval = 30.0
|
||||
|
||||
count = max(1, int(duration / interval))
|
||||
step = duration / count
|
||||
return [round(i * step, 3) for i in range(count)]
|
||||
|
||||
|
||||
class FrameSampler:
|
||||
"""Extract sampled frames from videos using FFmpeg."""
|
||||
|
||||
def __init__(self, interval_seconds: int = 30, quality: int = 2, output_format: str = "jpeg"):
|
||||
self.interval_seconds = interval_seconds
|
||||
self.quality = quality
|
||||
self.output_format = output_format.lower()
|
||||
|
||||
def extract_frame(
|
||||
self,
|
||||
video_path: str,
|
||||
output_path: str,
|
||||
timestamp: float,
|
||||
resolution: Optional[Tuple[int, int]] = None,
|
||||
timeout_seconds: int = 30,
|
||||
) -> Optional[str]:
|
||||
"""Extract a single frame at the provided timestamp."""
|
||||
output = Path(output_path)
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
cmd = self._build_ffmpeg_command(
|
||||
video_path=video_path,
|
||||
output_path=str(output),
|
||||
timestamp=timestamp,
|
||||
resolution=resolution,
|
||||
timeout_seconds=timeout_seconds,
|
||||
)
|
||||
|
||||
try:
|
||||
completed = subprocess.run(
|
||||
cmd,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=timeout_seconds + 5,
|
||||
check=False,
|
||||
)
|
||||
except subprocess.TimeoutExpired as exc:
|
||||
logger.warning("Frame extraction timed out for %s at %.3fs", video_path, timestamp)
|
||||
return None
|
||||
|
||||
if completed.returncode != 0:
|
||||
logger.warning(
|
||||
"Frame extraction failed for %s at %.3fs: %s",
|
||||
video_path,
|
||||
timestamp,
|
||||
completed.stderr.strip(),
|
||||
)
|
||||
return None
|
||||
|
||||
if output.exists():
|
||||
return str(output)
|
||||
return None
|
||||
|
||||
def extract_frames(
|
||||
self,
|
||||
video_path: str,
|
||||
output_dir: str,
|
||||
duration: Optional[float] = None,
|
||||
interval_seconds: Optional[int] = None,
|
||||
resolution: Optional[Tuple[int, int]] = None,
|
||||
timeout_seconds: int = 30,
|
||||
) -> List[str]:
|
||||
"""Extract uniform frames for the video and return output paths."""
|
||||
resolved_interval = interval_seconds or self.interval_seconds
|
||||
timestamps = calculate_timestamps(duration=duration, interval=float(resolved_interval))
|
||||
|
||||
extracted: List[str] = []
|
||||
for ts in timestamps:
|
||||
frame_name = self._frame_name_from_timestamp(ts, video_path)
|
||||
output_path = str(Path(output_dir) / frame_name)
|
||||
result = self.extract_frame(
|
||||
video_path=video_path,
|
||||
output_path=output_path,
|
||||
timestamp=ts,
|
||||
resolution=resolution,
|
||||
timeout_seconds=timeout_seconds,
|
||||
)
|
||||
if result:
|
||||
extracted.append(result)
|
||||
|
||||
return extracted
|
||||
|
||||
def _build_ffmpeg_command(
|
||||
self,
|
||||
video_path: str,
|
||||
output_path: str,
|
||||
timestamp: float,
|
||||
resolution: Optional[Tuple[int, int]] = None,
|
||||
timeout_seconds: int = 30,
|
||||
) -> List[str]:
|
||||
"""Build an FFmpeg command for extracting a single frame."""
|
||||
cmd = ["ffmpeg", "-y", "-ss", str(timestamp), "-i", video_path]
|
||||
|
||||
vf_filters = []
|
||||
if resolution is not None:
|
||||
target_width, target_height = self._resolve_target_resolution(resolution)
|
||||
vf_filters.append(f"scale={target_width}:{target_height}:force_original_aspect_ratio=decrease")
|
||||
|
||||
if vf_filters:
|
||||
cmd.extend(["-vf", ",".join(vf_filters)])
|
||||
|
||||
cmd.extend([
|
||||
"-vframes",
|
||||
"1",
|
||||
"-q:v",
|
||||
str(self.quality),
|
||||
"-f",
|
||||
self.output_format,
|
||||
output_path,
|
||||
])
|
||||
return cmd
|
||||
|
||||
@staticmethod
|
||||
def _resolve_target_resolution(resolution: Tuple[int, int]) -> Tuple[int, int]:
|
||||
"""Downscale frames larger than 1080p to fit VRAM constraints."""
|
||||
width, height = resolution
|
||||
if max(width, height) > 1080:
|
||||
if width >= 3840 or height >= 2160:
|
||||
return 1920, 1080
|
||||
if width >= 2560 or height >= 1440:
|
||||
return 1280, 720
|
||||
return width, height
|
||||
|
||||
@staticmethod
|
||||
def _frame_name_from_timestamp(timestamp: float, video_path: str) -> str:
|
||||
video_id = Path(video_path).stem
|
||||
timestamp_ms = int(round(timestamp * 1000))
|
||||
return f"{video_id}_{timestamp_ms}.jpg"
|
||||
|
||||
|
||||
def get_target_resolution(resolution: Optional[Tuple[int, int]]) -> Tuple[int, int]:
|
||||
"""Compatibility helper for resolution handling."""
|
||||
if resolution is None:
|
||||
return (0, 0)
|
||||
return FrameSampler._resolve_target_resolution(resolution)
|
||||
@@ -1,101 +0,0 @@
|
||||
"""GPU memory management and batch-size auto-tuning helpers."""
|
||||
|
||||
import logging
|
||||
from typing import Optional, Tuple
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class GPUMemoryManager:
|
||||
"""Track GPU memory and recommend safe batch sizes."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
device: int = 0,
|
||||
max_memory_gb: float = 18.0,
|
||||
reduce_threshold_gb: float = 16.0,
|
||||
increase_threshold_gb: float = 10.0,
|
||||
initial_batch_size: int = 16,
|
||||
min_batch_size: int = 1,
|
||||
max_batch_size: int = 32,
|
||||
):
|
||||
self.device = device
|
||||
self.max_memory_gb = max_memory_gb
|
||||
self.reduce_threshold_gb = reduce_threshold_gb
|
||||
self.increase_threshold_gb = increase_threshold_gb
|
||||
self.batch_size = initial_batch_size
|
||||
self.min_batch_size = min_batch_size
|
||||
self.max_batch_size = max_batch_size
|
||||
|
||||
def get_memory_stats(self) -> Tuple[float, float]:
|
||||
"""Return (allocated_gb, reserved_gb) for the managed device."""
|
||||
try:
|
||||
import torch
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
return 0.0, 0.0
|
||||
|
||||
allocated = torch.cuda.memory_allocated(self.device) / (1024**3)
|
||||
reserved = torch.cuda.memory_reserved(self.device) / (1024**3)
|
||||
return allocated, reserved
|
||||
except Exception as exc:
|
||||
logger.debug("Could not query GPU memory: %s", exc)
|
||||
return 0.0, 0.0
|
||||
|
||||
def adjust_batch_size(self) -> int:
|
||||
"""Adjust the current batch size based on memory pressure."""
|
||||
allocated, _ = self.get_memory_stats()
|
||||
|
||||
if allocated >= self.reduce_threshold_gb:
|
||||
new_batch_size = max(self.min_batch_size, int(self.batch_size * 0.75))
|
||||
if new_batch_size < self.batch_size:
|
||||
logger.info(
|
||||
"GPU memory high (%.2f GB), reducing batch size %d -> %d",
|
||||
allocated,
|
||||
self.batch_size,
|
||||
new_batch_size,
|
||||
)
|
||||
self.batch_size = new_batch_size
|
||||
elif allocated <= self.increase_threshold_gb:
|
||||
new_batch_size = min(self.max_batch_size, int(self.batch_size * 1.25))
|
||||
if new_batch_size > self.batch_size:
|
||||
logger.info(
|
||||
"GPU memory low (%.2f GB), increasing batch size %d -> %d",
|
||||
allocated,
|
||||
self.batch_size,
|
||||
new_batch_size,
|
||||
)
|
||||
self.batch_size = new_batch_size
|
||||
|
||||
return self.batch_size
|
||||
|
||||
def is_memory_critical(self) -> bool:
|
||||
"""Return True if allocated memory is close to the hard limit."""
|
||||
allocated, _ = self.get_memory_stats()
|
||||
return allocated >= self.max_memory_gb
|
||||
|
||||
def empty_cache(self):
|
||||
"""Try to release unused cached GPU memory."""
|
||||
try:
|
||||
import torch
|
||||
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
except Exception as exc:
|
||||
logger.debug("Could not empty GPU cache: %s", exc)
|
||||
|
||||
@property
|
||||
def current_batch_size(self) -> int:
|
||||
return self.batch_size
|
||||
|
||||
|
||||
def get_available_vram_gb(device: int = 0) -> float:
|
||||
"""Return total available VRAM in GB, or 0.0 if CUDA is unavailable."""
|
||||
try:
|
||||
import torch
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
return 0.0
|
||||
return torch.cuda.get_device_properties(device).total_memory / (1024**3)
|
||||
except Exception:
|
||||
return 0.0
|
||||
@@ -1,134 +0,0 @@
|
||||
"""
|
||||
Worker health check HTTP endpoint.
|
||||
|
||||
Exposes GET /health returning a JSON status document.
|
||||
Designed to be started in a background thread alongside the main worker.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
from http.server import BaseHTTPRequestHandler, HTTPServer
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_START_TIME = time.monotonic()
|
||||
|
||||
|
||||
class HealthStatus:
|
||||
"""Mutable singleton updated by the worker as it processes videos."""
|
||||
|
||||
def __init__(self):
|
||||
self._lock = threading.Lock()
|
||||
self._data: Dict[str, Any] = {
|
||||
"status": "starting",
|
||||
"gpu_available": False,
|
||||
"gpu_memory_used_gb": 0.0,
|
||||
"queue_depth": 0,
|
||||
"videos_processed_today": 0,
|
||||
"last_error": None,
|
||||
}
|
||||
|
||||
def update(self, **kwargs):
|
||||
with self._lock:
|
||||
self._data.update(kwargs)
|
||||
|
||||
def snapshot(self) -> Dict[str, Any]:
|
||||
with self._lock:
|
||||
data = dict(self._data)
|
||||
data["uptime_seconds"] = round(time.monotonic() - _START_TIME, 1)
|
||||
return data
|
||||
|
||||
|
||||
# Module-level singleton shared between the HTTP handler and the worker
|
||||
health = HealthStatus()
|
||||
|
||||
|
||||
class _HealthHandler(BaseHTTPRequestHandler):
|
||||
"""Minimal HTTP handler — only responds to GET /health."""
|
||||
|
||||
def do_GET(self): # noqa: N802
|
||||
if self.path.rstrip("/") == "/health":
|
||||
body = json.dumps(health.snapshot(), default=str).encode()
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
else:
|
||||
self.send_response(404)
|
||||
self.end_headers()
|
||||
|
||||
def log_message(self, fmt, *args): # silence default access log
|
||||
logger.debug("health_check: " + fmt, *args)
|
||||
|
||||
|
||||
def start_health_server(port: int = 8080) -> threading.Thread:
|
||||
"""
|
||||
Start the health HTTP server in a daemon thread.
|
||||
|
||||
Returns the thread so callers can join if needed.
|
||||
"""
|
||||
server = HTTPServer(("0.0.0.0", port), _HealthHandler)
|
||||
|
||||
def _serve():
|
||||
logger.info("Health check server listening on port %d", port)
|
||||
server.serve_forever()
|
||||
|
||||
t = threading.Thread(target=_serve, daemon=True, name="health-server")
|
||||
t.start()
|
||||
return t
|
||||
|
||||
|
||||
def refresh_health_from_system(db_connector=None, scratch_path: str = "/scratch"):
|
||||
"""
|
||||
Update the health singleton with current GPU, queue, and disk state.
|
||||
|
||||
Call this periodically from the worker main loop.
|
||||
"""
|
||||
updates: Dict[str, Any] = {"status": "healthy"}
|
||||
|
||||
# GPU state
|
||||
try:
|
||||
import pynvml
|
||||
pynvml.nvmlInit()
|
||||
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
|
||||
mem = pynvml.nvmlDeviceGetMemoryInfo(handle)
|
||||
updates["gpu_available"] = True
|
||||
updates["gpu_memory_used_gb"] = round(mem.used / 1024**3, 2)
|
||||
pynvml.nvmlShutdown()
|
||||
except Exception:
|
||||
updates["gpu_available"] = _torch_gpu_available()
|
||||
|
||||
# Queue depth from DB
|
||||
if db_connector is not None:
|
||||
try:
|
||||
rows = db_connector.fetchall(
|
||||
"SELECT COUNT(*) AS cnt FROM videos WHERE status IN ('PENDING','PROCESSING')"
|
||||
)
|
||||
updates["queue_depth"] = int((rows or [{"cnt": 0}])[0]["cnt"])
|
||||
|
||||
today = datetime.now(timezone.utc).date().isoformat()
|
||||
rows2 = db_connector.fetchall(
|
||||
"SELECT COUNT(*) AS cnt FROM videos "
|
||||
"WHERE status = 'COMPLETED' AND DATE(last_processed_time) = %s",
|
||||
(today,),
|
||||
)
|
||||
updates["videos_processed_today"] = int((rows2 or [{"cnt": 0}])[0]["cnt"])
|
||||
except Exception as exc:
|
||||
logger.debug("Health DB query failed: %s", exc)
|
||||
|
||||
health.update(**updates)
|
||||
|
||||
|
||||
def _torch_gpu_available() -> bool:
|
||||
try:
|
||||
import torch
|
||||
return torch.cuda.is_available()
|
||||
except ImportError:
|
||||
return False
|
||||
@@ -1,87 +0,0 @@
|
||||
"""
|
||||
Logging configuration for VideoDetect.
|
||||
|
||||
Sets up JSON structured logging via python-json-logger with:
|
||||
- Configurable log levels
|
||||
- Standardized field names
|
||||
- Log rotation
|
||||
- All logs to stdout for Docker capture
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import logging.handlers
|
||||
import sys
|
||||
from datetime import datetime, timezone
|
||||
|
||||
|
||||
def setup_logging(
|
||||
level: str = "INFO",
|
||||
log_format: str = "json",
|
||||
rotation_max_bytes: int = 104857600, # 100MB
|
||||
rotation_backup_count: int = 10,
|
||||
):
|
||||
"""Configure structured logging for the application."""
|
||||
log_level = getattr(logging, level.upper(), logging.INFO)
|
||||
|
||||
# Root logger
|
||||
root_logger = logging.getLogger()
|
||||
root_logger.setLevel(log_level)
|
||||
|
||||
# Remove existing handlers
|
||||
root_logger.handlers.clear()
|
||||
|
||||
# Handler: stdout (for Docker capture)
|
||||
handler = logging.StreamHandler(sys.stdout)
|
||||
handler.setLevel(log_level)
|
||||
|
||||
if log_format == "json":
|
||||
formatter = JsonFormatter()
|
||||
else:
|
||||
formatter = logging.Formatter(
|
||||
"%(asctime)s %(levelname)s %(name)s %(message)s",
|
||||
datefmt="%Y-%m-%dT%H:%M:%S%z",
|
||||
)
|
||||
|
||||
handler.setFormatter(formatter)
|
||||
root_logger.addHandler(handler)
|
||||
|
||||
# Handler: file rotation (for persistence)
|
||||
file_handler = logging.handlers.RotatingFileHandler(
|
||||
filename="/logs/videodetect.log",
|
||||
maxBytes=rotation_max_bytes,
|
||||
backupCount=rotation_backup_count,
|
||||
encoding="utf-8",
|
||||
)
|
||||
file_handler.setLevel(log_level)
|
||||
file_handler.setFormatter(formatter)
|
||||
root_logger.addHandler(file_handler)
|
||||
|
||||
logging.info("Logging configured: level=%s format=%s", level, log_format)
|
||||
|
||||
|
||||
class JsonFormatter(logging.Formatter):
|
||||
"""JSON log formatter with standardized fields."""
|
||||
|
||||
def format(self, record: logging.LogRecord) -> str:
|
||||
log_data = {
|
||||
"timestamp": datetime.fromtimestamp(record.created, tz=timezone.utc).isoformat(),
|
||||
"level": record.levelname,
|
||||
"logger": record.name,
|
||||
"message": record.getMessage(),
|
||||
"module": record.module,
|
||||
"function": record.funcName,
|
||||
"line": record.lineno,
|
||||
}
|
||||
|
||||
# Add extra fields
|
||||
if hasattr(record, "video_id"):
|
||||
log_data["video_id"] = record.video_id
|
||||
if hasattr(record, "metadata"):
|
||||
log_data["metadata"] = record.metadata
|
||||
|
||||
# Add exception info if present
|
||||
if record.exc_info and record.exc_info[0] is not None:
|
||||
log_data["exception"] = self.formatException(record.exc_info)
|
||||
|
||||
return json.dumps(log_data, default=str)
|
||||
-261
@@ -1,261 +0,0 @@
|
||||
"""
|
||||
Prometheus metrics definitions for VideoDetect.
|
||||
|
||||
All metric objects are module-level singletons — import and call them directly.
|
||||
Call start_metrics_server() once at worker startup to expose /metrics.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import threading
|
||||
from typing import Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Populated on first _ensure_prometheus() call; None if library unavailable.
|
||||
_prometheus_available: Optional[bool] = None
|
||||
_server_started = False
|
||||
_server_lock = threading.Lock()
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Lazy metric definitions — avoid import errors when prometheus_client absent
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _ensure_prometheus():
|
||||
global _prometheus_available
|
||||
if _prometheus_available is not None:
|
||||
return _prometheus_available
|
||||
try:
|
||||
import prometheus_client # noqa: F401
|
||||
_prometheus_available = True
|
||||
except ImportError:
|
||||
logger.warning("prometheus_client not installed; metrics disabled")
|
||||
_prometheus_available = False
|
||||
return _prometheus_available
|
||||
|
||||
|
||||
class _NoOpMetric:
|
||||
"""Drop-in replacement used when prometheus_client is unavailable."""
|
||||
def labels(self, **_):
|
||||
return self
|
||||
def inc(self, *_, **__):
|
||||
pass
|
||||
def set(self, *_, **__):
|
||||
pass
|
||||
def observe(self, *_, **__):
|
||||
pass
|
||||
|
||||
|
||||
def _counter(name, doc, labelnames=()):
|
||||
if not _ensure_prometheus():
|
||||
return _NoOpMetric()
|
||||
from prometheus_client import Counter
|
||||
return Counter(name, doc, list(labelnames))
|
||||
|
||||
|
||||
def _gauge(name, doc, labelnames=()):
|
||||
if not _ensure_prometheus():
|
||||
return _NoOpMetric()
|
||||
from prometheus_client import Gauge
|
||||
return Gauge(name, doc, list(labelnames))
|
||||
|
||||
|
||||
def _histogram(name, doc, labelnames=(), buckets=None):
|
||||
if not _ensure_prometheus():
|
||||
return _NoOpMetric()
|
||||
from prometheus_client import Histogram
|
||||
kwargs = {"labelnames": list(labelnames)}
|
||||
if buckets is not None:
|
||||
kwargs["buckets"] = buckets
|
||||
return Histogram(name, doc, **kwargs)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Counters
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
videos_processed_total = _counter(
|
||||
"videos_processed_total",
|
||||
"Total videos processed",
|
||||
["routing_decision"],
|
||||
)
|
||||
|
||||
videos_processed_by_model_total = _counter(
|
||||
"videos_processed_by_model_total",
|
||||
"Total videos processed per model version",
|
||||
["model_version"],
|
||||
)
|
||||
|
||||
frames_extracted_total = _counter(
|
||||
"frames_extracted_total",
|
||||
"Total frames extracted across all videos",
|
||||
)
|
||||
|
||||
faces_detected_total = _counter(
|
||||
"faces_detected_total",
|
||||
"Total face crops detected",
|
||||
)
|
||||
|
||||
inference_errors_total = _counter(
|
||||
"inference_errors_total",
|
||||
"Total inference errors by type",
|
||||
["error_type"],
|
||||
)
|
||||
|
||||
retry_attempts_total = _counter(
|
||||
"retry_attempts_total",
|
||||
"Total retry attempts by pipeline step",
|
||||
["step"],
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Gauges
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
gpu_utilization_percent = _gauge(
|
||||
"gpu_utilization_percent",
|
||||
"GPU utilization as a percentage",
|
||||
["gpu_id"],
|
||||
)
|
||||
|
||||
gpu_memory_used_bytes = _gauge(
|
||||
"gpu_memory_used_bytes",
|
||||
"GPU memory currently allocated in bytes",
|
||||
["gpu_id"],
|
||||
)
|
||||
|
||||
gpu_memory_free_bytes = _gauge(
|
||||
"gpu_memory_free_bytes",
|
||||
"GPU memory currently free in bytes",
|
||||
["gpu_id"],
|
||||
)
|
||||
|
||||
queue_depth_pending = _gauge(
|
||||
"queue_depth_pending",
|
||||
"Number of videos in PENDING state",
|
||||
)
|
||||
|
||||
queue_depth_processing = _gauge(
|
||||
"queue_depth_processing",
|
||||
"Number of videos currently in PROCESSING state",
|
||||
)
|
||||
|
||||
queue_depth_review = _gauge(
|
||||
"queue_depth_review",
|
||||
"Number of unannotated videos in the review queue",
|
||||
)
|
||||
|
||||
scratch_usage_bytes = _gauge(
|
||||
"scratch_usage_bytes",
|
||||
"Scratch space used in bytes",
|
||||
)
|
||||
|
||||
scratch_usage_percent = _gauge(
|
||||
"scratch_usage_percent",
|
||||
"Scratch space used as a percentage of total",
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Histograms
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_CONFIDENCE_BUCKETS = [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
|
||||
|
||||
video_processing_duration_seconds = _histogram(
|
||||
"video_processing_duration_seconds",
|
||||
"End-to-end processing time per video in seconds",
|
||||
["routing_decision"],
|
||||
)
|
||||
|
||||
confidence_score_distribution = _histogram(
|
||||
"confidence_score_distribution",
|
||||
"Distribution of video-level confidence scores",
|
||||
buckets=_CONFIDENCE_BUCKETS,
|
||||
)
|
||||
|
||||
frame_count_per_video = _histogram(
|
||||
"frame_count_per_video",
|
||||
"Number of frames sampled per video",
|
||||
buckets=[1, 2, 5, 10, 20, 50, 100, 200, 500],
|
||||
)
|
||||
|
||||
face_count_per_video = _histogram(
|
||||
"face_count_per_video",
|
||||
"Number of face crops detected per video",
|
||||
buckets=[0, 1, 2, 5, 10, 20, 50, 100],
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Server startup
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def start_metrics_server(port: int = 9090):
|
||||
"""Start the Prometheus HTTP metrics endpoint (idempotent)."""
|
||||
global _server_started
|
||||
with _server_lock:
|
||||
if _server_started:
|
||||
return
|
||||
if not _ensure_prometheus():
|
||||
logger.warning("Cannot start metrics server: prometheus_client unavailable")
|
||||
return
|
||||
try:
|
||||
from prometheus_client import start_http_server
|
||||
start_http_server(port)
|
||||
_server_started = True
|
||||
logger.info("Prometheus metrics server started on port %d", port)
|
||||
except Exception as exc:
|
||||
logger.error("Failed to start metrics server: %s", exc)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Convenience update functions
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def update_queue_depths(db_connector):
|
||||
"""Query the DB and refresh all queue-depth gauges."""
|
||||
try:
|
||||
rows = db_connector.fetchall(
|
||||
"""SELECT status, COUNT(*) AS cnt
|
||||
FROM videos
|
||||
WHERE status IN ('PENDING', 'PROCESSING')
|
||||
GROUP BY status"""
|
||||
)
|
||||
counts = {r["status"]: int(r["cnt"]) for r in (rows or [])}
|
||||
queue_depth_pending.set(counts.get("PENDING", 0))
|
||||
queue_depth_processing.set(counts.get("PROCESSING", 0))
|
||||
|
||||
review_rows = db_connector.fetchall(
|
||||
"SELECT COUNT(*) AS cnt FROM review_queue WHERE annotated = FALSE"
|
||||
)
|
||||
queue_depth_review.set(int((review_rows or [{"cnt": 0}])[0]["cnt"]))
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to update queue depth metrics: %s", exc)
|
||||
|
||||
|
||||
def update_gpu_metrics():
|
||||
"""Poll nvidia-smi (via pynvml) and refresh GPU gauges."""
|
||||
try:
|
||||
import pynvml
|
||||
pynvml.nvmlInit()
|
||||
count = pynvml.nvmlDeviceGetCount()
|
||||
for i in range(count):
|
||||
handle = pynvml.nvmlDeviceGetHandleByIndex(i)
|
||||
util = pynvml.nvmlDeviceGetUtilizationRates(handle)
|
||||
mem = pynvml.nvmlDeviceGetMemoryInfo(handle)
|
||||
gpu_id = str(i)
|
||||
gpu_utilization_percent.labels(gpu_id=gpu_id).set(util.gpu)
|
||||
gpu_memory_used_bytes.labels(gpu_id=gpu_id).set(mem.used)
|
||||
gpu_memory_free_bytes.labels(gpu_id=gpu_id).set(mem.free)
|
||||
pynvml.nvmlShutdown()
|
||||
except Exception as exc:
|
||||
logger.debug("GPU metrics unavailable: %s", exc)
|
||||
|
||||
|
||||
def update_scratch_metrics(scratch_path: str = "/scratch"):
|
||||
"""Measure scratch disk usage and refresh gauges."""
|
||||
try:
|
||||
import shutil
|
||||
total, used, free = shutil.disk_usage(scratch_path)
|
||||
scratch_usage_bytes.set(used)
|
||||
scratch_usage_percent.set(round(used / max(total, 1) * 100, 2))
|
||||
except Exception as exc:
|
||||
logger.debug("Scratch metrics unavailable: %s", exc)
|
||||
-182
@@ -1,182 +0,0 @@
|
||||
"""
|
||||
Video Probing Module
|
||||
|
||||
Extracts video metadata (codec, resolution, duration, frame rate) using ffprobe.
|
||||
Handles errors gracefully for corrupt or unsupported files.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class VideoMetadata:
|
||||
"""Metadata extracted from a video file."""
|
||||
file_path: str
|
||||
codec: Optional[str] = None
|
||||
resolution_w: Optional[int] = None
|
||||
resolution_h: Optional[int] = None
|
||||
duration: Optional[float] = None
|
||||
frame_rate: Optional[float] = None
|
||||
file_size: Optional[int] = None
|
||||
status: str = "PENDING"
|
||||
error_message: Optional[str] = None
|
||||
probe_time: float = field(default_factory=time.time)
|
||||
|
||||
@property
|
||||
def is_valid(self) -> bool:
|
||||
return self.status == "PENDING" and self.codec is not None
|
||||
|
||||
@property
|
||||
def is_unscannable(self) -> bool:
|
||||
return self.status == "UNSCANNABLE"
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"file_path": self.file_path,
|
||||
"codec": self.codec,
|
||||
"resolution_w": self.resolution_w,
|
||||
"resolution_h": self.resolution_h,
|
||||
"duration": self.duration,
|
||||
"frame_rate": self.frame_rate,
|
||||
"file_size": self.file_size,
|
||||
"status": self.status,
|
||||
"error_message": self.error_message,
|
||||
}
|
||||
|
||||
|
||||
class VideoProber:
|
||||
"""Probe video files using ffprobe to extract metadata."""
|
||||
|
||||
def __init__(self, timeout: int = 10):
|
||||
self.timeout = timeout
|
||||
|
||||
def probe(self, file_path: str) -> VideoMetadata:
|
||||
"""Probe a video file and return metadata."""
|
||||
metadata = VideoMetadata(file_path=file_path)
|
||||
|
||||
# Get file size
|
||||
try:
|
||||
metadata.file_size = self._get_file_size(file_path)
|
||||
except OSError as e:
|
||||
metadata.status = "UNSCANNABLE"
|
||||
metadata.error_message = f"Cannot read file: {e}"
|
||||
logger.warning("File access error for %s: %s", file_path, e)
|
||||
return metadata
|
||||
|
||||
# Run ffprobe
|
||||
try:
|
||||
result = self._run_ffprobe(file_path)
|
||||
if result is None:
|
||||
metadata.status = "UNSCANNABLE"
|
||||
metadata.error_message = "ffprobe timed out or failed"
|
||||
logger.warning("ffprobe failed for %s", file_path)
|
||||
return metadata
|
||||
|
||||
metadata = self._parse_ffprobe_output(result, file_path)
|
||||
except subprocess.TimeoutExpired:
|
||||
metadata.status = "UNSCANNABLE"
|
||||
metadata.error_message = f"ffprobe timed out after {self.timeout}s"
|
||||
logger.warning("ffprobe timeout for %s", file_path)
|
||||
except Exception as e:
|
||||
metadata.status = "UNSCANNABLE"
|
||||
metadata.error_message = f"ffprobe error: {e}"
|
||||
logger.warning("ffprobe error for %s: %s", file_path, e)
|
||||
|
||||
return metadata
|
||||
|
||||
def _run_ffprobe(self, file_path: str) -> Optional[str]:
|
||||
"""Run ffprobe and return JSON output."""
|
||||
cmd = [
|
||||
"ffprobe",
|
||||
"-v", "error",
|
||||
"-select_streams", "v:0",
|
||||
"-show_entries", "stream=codec_name,width,height,r_frame_rate,duration",
|
||||
"-show_entries", "format=duration,size",
|
||||
"-of", "json",
|
||||
file_path,
|
||||
]
|
||||
|
||||
try:
|
||||
output = subprocess.run(
|
||||
cmd,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=self.timeout,
|
||||
)
|
||||
|
||||
if output.returncode != 0:
|
||||
logger.debug("ffprobe stderr for %s: %s", file_path, output.stderr)
|
||||
return None
|
||||
|
||||
return output.stdout
|
||||
except subprocess.TimeoutExpired:
|
||||
return None
|
||||
|
||||
def _parse_ffprobe_output(self, output: str, file_path: str) -> VideoMetadata:
|
||||
"""Parse ffprobe JSON output into VideoMetadata."""
|
||||
metadata = VideoMetadata(file_path=file_path)
|
||||
|
||||
try:
|
||||
data = json.loads(output)
|
||||
except json.JSONDecodeError as e:
|
||||
metadata.status = "UNSCANNABLE"
|
||||
metadata.error_message = f"Invalid ffprobe output: {e}"
|
||||
return metadata
|
||||
|
||||
# Extract stream info
|
||||
streams = data.get("streams", [])
|
||||
if not streams:
|
||||
metadata.status = "UNSCANNABLE"
|
||||
metadata.error_message = "No video streams found"
|
||||
return metadata
|
||||
|
||||
stream = streams[0]
|
||||
metadata.codec = stream.get("codec_name")
|
||||
metadata.resolution_w = stream.get("width")
|
||||
metadata.resolution_h = stream.get("height")
|
||||
|
||||
# Parse frame rate
|
||||
r_frame_rate = stream.get("r_frame_rate")
|
||||
if r_frame_rate and "/" in r_frame_rate:
|
||||
try:
|
||||
num, den = r_frame_rate.split("/")
|
||||
metadata.frame_rate = float(num) / float(den) if float(den) != 0 else None
|
||||
except (ValueError, ZeroDivisionError):
|
||||
pass
|
||||
|
||||
# Extract format info
|
||||
format_info = data.get("format", {})
|
||||
if format_info.get("duration"):
|
||||
try:
|
||||
metadata.duration = float(format_info["duration"])
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
|
||||
if format_info.get("size"):
|
||||
try:
|
||||
metadata.file_size = int(format_info["size"])
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
|
||||
# Validate we got essential data
|
||||
if not metadata.codec:
|
||||
metadata.status = "UNSCANNABLE"
|
||||
metadata.error_message = "No codec information available"
|
||||
elif not metadata.duration:
|
||||
metadata.status = "UNSCANNABLE"
|
||||
metadata.error_message = "No duration information available"
|
||||
|
||||
return metadata
|
||||
|
||||
@staticmethod
|
||||
def _get_file_size(file_path: str) -> int:
|
||||
"""Get file size in bytes."""
|
||||
return os.path.getsize(file_path)
|
||||
@@ -1,26 +0,0 @@
|
||||
"""Threshold-based routing of video-level confidence to MATCH / REVIEW / SKIP."""
|
||||
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
MATCH = "MATCH"
|
||||
REVIEW = "REVIEW"
|
||||
SKIP = "SKIP"
|
||||
|
||||
|
||||
def route(
|
||||
confidence: float,
|
||||
t_high: float = 0.75,
|
||||
t_low: float = 0.45,
|
||||
) -> str:
|
||||
"""Return the routing decision for a video-level confidence score."""
|
||||
if confidence >= t_high:
|
||||
decision = MATCH
|
||||
elif confidence >= t_low:
|
||||
decision = REVIEW
|
||||
else:
|
||||
decision = SKIP
|
||||
|
||||
logger.debug("route: C=%.4f t_high=%.2f t_low=%.2f → %s", confidence, t_high, t_low, decision)
|
||||
return decision
|
||||
@@ -1,42 +0,0 @@
|
||||
"""Utilities for managing scratch space used by frame sampling."""
|
||||
|
||||
import logging
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ScratchManager:
|
||||
"""Create and clean per-video frame scratch directories."""
|
||||
|
||||
def __init__(self, base_path: str = "/scratch", video_id: Optional[str] = None, auto_cleanup: bool = True):
|
||||
self.base_path = Path(base_path)
|
||||
self.video_id = video_id or "unknown"
|
||||
self.auto_cleanup = auto_cleanup
|
||||
self.frame_dir = self.base_path / self.video_id / "frames"
|
||||
|
||||
def ensure_frame_dir(self) -> Path:
|
||||
"""Create the per-video scratch directory if it does not exist."""
|
||||
self.frame_dir.mkdir(parents=True, exist_ok=True)
|
||||
return self.frame_dir
|
||||
|
||||
def cleanup(self) -> None:
|
||||
"""Remove the frame scratch directory only."""
|
||||
if self.auto_cleanup and self.frame_dir.exists():
|
||||
shutil.rmtree(self.frame_dir)
|
||||
logger.info("Cleaned scratch frames for %s", self.video_id)
|
||||
|
||||
def cleanup_all(self) -> None:
|
||||
"""Remove the entire per-video scratch directory (frames + crops)."""
|
||||
video_root = self.base_path / self.video_id
|
||||
if self.auto_cleanup and video_root.exists():
|
||||
shutil.rmtree(video_root)
|
||||
logger.info("Cleaned all scratch data for %s", self.video_id)
|
||||
|
||||
def usage_bytes(self) -> int:
|
||||
"""Return the total size of the scratch directory in bytes."""
|
||||
if not self.frame_dir.exists():
|
||||
return 0
|
||||
return sum(path.stat().st_size for path in self.frame_dir.rglob("*") if path.is_file())
|
||||
@@ -1,297 +0,0 @@
|
||||
"""Standalone AI-processing pipeline for the task-worker architecture.
|
||||
|
||||
``process_video(video_dict, config)`` runs the full scan pipeline on a single
|
||||
video file and returns a results dict that can be submitted via the REST API.
|
||||
|
||||
It reuses existing classes from the codebase:
|
||||
* ``VideoProber`` – ffprobe metadata extraction
|
||||
* ``ScratchManager`` – per-video temp dirs under /scratch
|
||||
* ``FrameSampler`` – uniform frame extraction via ffmpeg
|
||||
* ``FaceDetector`` – TensorRT face detection on sampled frames
|
||||
* ``FaceClassifier`` – MobileNetV3 classification of face crops
|
||||
* ``aggregator.aggregate`` – per-crop → video-level confidence
|
||||
* ``router.route`` – routing thresholds (MATCH / REVIEW / SKIP)
|
||||
|
||||
GPU model loading is lazily initialised and cached at module level so that
|
||||
sequential invocations within the same Python process reuse the same loaded
|
||||
engines. This is important because each short-lived worker may be asked to
|
||||
process several tasks in a row (``--tasks N``).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Lazy imports — avoid heavy imports until actually needed
|
||||
# ---------------------------------------------------------------------------
|
||||
_face_detector_instance: Optional[Any] = None
|
||||
_classifier_instance: Optional[Any] = None
|
||||
|
||||
|
||||
def _get_face_detector(config: dict) -> Any:
|
||||
"""Return a singleton ``FaceDetector`` (initialised once per process)."""
|
||||
global _face_detector_instance
|
||||
if _face_detector_instance is not None:
|
||||
return _face_detector_instance
|
||||
|
||||
from face_detector import FaceDetector # noqa: local import, heavy
|
||||
|
||||
cfg = (config or {}).get("face_detection", {})
|
||||
engine_path = cfg.get(
|
||||
"model_path", "/models/face_detector/face_detector.trt"
|
||||
)
|
||||
_face_detector_instance = FaceDetector(
|
||||
engine_path=engine_path,
|
||||
input_size=int(cfg.get("input_size", 640)),
|
||||
confidence_threshold=float(cfg.get("confidence_threshold", 0.25)),
|
||||
iou_threshold=float(cfg.get("iou_threshold", 0.45)),
|
||||
max_faces_per_frame=int(cfg.get("max_faces_per_frame", 10)),
|
||||
max_faces_per_video=int(cfg.get("max_faces_per_video", 100)),
|
||||
)
|
||||
return _face_detector_instance
|
||||
|
||||
|
||||
def _get_classifier(config: dict) -> Any:
|
||||
"""Return a singleton ``FaceClassifier`` (initialised once per process)."""
|
||||
global _classifier_instance
|
||||
if _classifier_instance is not None:
|
||||
return _classifier_instance
|
||||
|
||||
from classifier import FaceClassifier # noqa: local import, heavy
|
||||
|
||||
cfg = (config or {}).get("classifier", {})
|
||||
engine_path = cfg.get(
|
||||
"model_path", "/models/classifier/classifier.trt"
|
||||
)
|
||||
_classifier_instance = FaceClassifier(
|
||||
engine_path=engine_path,
|
||||
temperature=float(cfg.get("temperature", 1.0)),
|
||||
input_size=int(cfg.get("input_size", 224)),
|
||||
)
|
||||
return _classifier_instance
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Public API
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def get_model_version() -> str:
|
||||
"""Return the currently active model version from the DB (best-effort)."""
|
||||
try:
|
||||
from db_connector import DatabaseConnector # noqa: local import, heavy
|
||||
db = DatabaseConnector("default")
|
||||
row = db.fetchone(
|
||||
"SELECT version FROM models WHERE status = 'ACTIVE' LIMIT 1"
|
||||
)
|
||||
if row:
|
||||
return row["version"]
|
||||
except Exception as exc:
|
||||
logger.debug("Could not fetch active model version: %s", exc)
|
||||
return "v0.0.0-placeholder"
|
||||
|
||||
|
||||
def process_video(video: Dict[str, Any], config: Optional[Dict] = None) -> Dict[str, Any]:
|
||||
"""Run the full AI scan pipeline on a single video file.
|
||||
|
||||
Args:
|
||||
video: Metadata dict (as returned by ``ApiClient.get_video``).
|
||||
Must contain at minimum ``"id"`` and ``"file_path"``.
|
||||
config: The loaded ``config.yaml`` dict (or None for defaults).
|
||||
|
||||
Returns:
|
||||
A results dict suitable for submission via the REST API.
|
||||
"""
|
||||
config = config or {}
|
||||
start_time = time.time()
|
||||
video_id = video["id"]
|
||||
file_path = video["file_path"]
|
||||
|
||||
logger.info("Processing video %d: %s", video_id, file_path)
|
||||
|
||||
sampling_cfg = (config or {}).get("sampling", {})
|
||||
storage_cfg = (config or {}).get("storage", {})
|
||||
gpu_cfg = (config or {}).get("gpu", {})
|
||||
batching_cfg = (config or {}).get("batching", {})
|
||||
agg_cfg = (config or {}).get("aggregation", {})
|
||||
routing_cfg = (config or {}).get("routing", {})
|
||||
|
||||
try:
|
||||
# 1. Probe metadata (use DB metadata if already populated, re-probe)
|
||||
prober = VideoProber(timeout=10)
|
||||
metadata = prober.probe(file_path)
|
||||
if metadata.is_unscannable:
|
||||
raise RuntimeError(metadata.error_message or "Video metadata could not be determined")
|
||||
|
||||
# Use video from API if available, fall back to probe results
|
||||
resolution_w = video.get("resolution_w") or metadata.resolution_w or 640
|
||||
resolution_h = video.get("resolution_h") or metadata.resolution_h or 480
|
||||
duration = video.get("duration") or metadata.duration or 10.0
|
||||
|
||||
logger.info(
|
||||
"Video %d: codec=%s resolution=%dx%d duration=%.1fs",
|
||||
video_id, metadata.codec or "unknown",
|
||||
resolution_w, resolution_h, duration,
|
||||
)
|
||||
|
||||
# 2. Set up scratch space
|
||||
scratch = ScratchManager(
|
||||
base_path=storage_cfg.get("scratch_path", "/scratch"),
|
||||
video_id=str(video_id),
|
||||
auto_cleanup=True,
|
||||
)
|
||||
frame_dir = scratch.ensure_frame_dir()
|
||||
|
||||
# 3. Sample frames
|
||||
interval = int(sampling_cfg.get("interval_seconds", 30))
|
||||
quality = int(sampling_cfg.get("quality", 2))
|
||||
sampler = FrameSampler(interval_seconds=interval, quality=quality)
|
||||
|
||||
resolution = (resolution_w, resolution_h)
|
||||
extracted_frames = sampler.extract_frames(
|
||||
video_path=file_path,
|
||||
output_dir=str(frame_dir),
|
||||
duration=duration,
|
||||
interval_seconds=interval,
|
||||
resolution=resolution,
|
||||
timeout_seconds=30,
|
||||
)
|
||||
|
||||
if not extracted_frames:
|
||||
scratch.cleanup_all()
|
||||
raise RuntimeError("No frames were extracted")
|
||||
|
||||
logger.info(
|
||||
"Video %d: extracted %d frame(s)", video_id, len(extracted_frames)
|
||||
)
|
||||
|
||||
# 4. Face detection
|
||||
detector = _get_face_detector(config)
|
||||
batch_size = _get_batch_size(gpu_cfg, batching_cfg)
|
||||
|
||||
detections_per_frame = detector.detect_faces(
|
||||
extracted_frames, batch_size=batch_size
|
||||
)
|
||||
|
||||
all_detections = sorted(
|
||||
[d for dets in detections_per_frame for d in dets],
|
||||
key=lambda d: d.confidence,
|
||||
reverse=True,
|
||||
)
|
||||
total_faces = len(all_detections)
|
||||
|
||||
# 5. Classify (or SKIP if no faces)
|
||||
if total_faces == 0:
|
||||
logger.info(
|
||||
"Video %d: no faces detected → routing to SKIP", video_id
|
||||
)
|
||||
scratch.cleanup_all()
|
||||
return _failed_result(video_id, start_time, 0.0, "SKIP")
|
||||
|
||||
crop_dir = Path(scratch.frame_dir.parent) / "crops"
|
||||
cropped_detections = detector.extract_crops(
|
||||
all_detections, output_dir=str(crop_dir), crop_size=(224, 224)
|
||||
)
|
||||
crop_paths = [d.crop_path for d in cropped_detections if d.crop_path]
|
||||
|
||||
classifier = _get_classifier(config)
|
||||
frame_confidences = classifier.classify(
|
||||
crop_paths, batch_size=batch_size
|
||||
)
|
||||
|
||||
# 6. Aggregate confidence
|
||||
from aggregator import aggregate as agg_func # noqa: local import
|
||||
|
||||
video_confidence = agg_func(
|
||||
frame_confidences,
|
||||
strategy=agg_cfg.get("strategy", "max"),
|
||||
alpha=float(agg_cfg.get("alpha", 1.0)),
|
||||
beta=float(agg_cfg.get("beta", 0.1)),
|
||||
top_k=int(agg_cfg.get("top_k", 3)),
|
||||
)
|
||||
|
||||
# 7. Route
|
||||
from router import route as route_fn # noqa: local import
|
||||
|
||||
routing = route_fn(
|
||||
video_confidence,
|
||||
t_high=float(routing_cfg.get("T_high", 0.75)),
|
||||
t_low=float(routing_cfg.get("T_low", 0.45)),
|
||||
)
|
||||
|
||||
processing_time = time.time() - start_time
|
||||
logger.info(
|
||||
"Video %d complete: C=%.4f routing=%s faces=%d frames=%d time=%.1fs",
|
||||
video_id, video_confidence, routing, total_faces,
|
||||
len(extracted_frames), processing_time,
|
||||
)
|
||||
|
||||
scratch.cleanup_all()
|
||||
|
||||
return {
|
||||
"status": "COMPLETED",
|
||||
"confidence": round(video_confidence, 4),
|
||||
"routing_decision": routing,
|
||||
"face_count": total_faces,
|
||||
"frame_count": len(extracted_frames),
|
||||
"model_version": get_model_version(),
|
||||
"processing_time_seconds": round(processing_time, 2),
|
||||
"error": None,
|
||||
}
|
||||
|
||||
except Exception as exc:
|
||||
logger.error("Processing video %d failed: %s", video_id, exc, exc_info=True)
|
||||
try:
|
||||
scratch.cleanup_all() # noqa: undefined-name guard below
|
||||
except NameError:
|
||||
pass # scratch was never created (probe/probe error)
|
||||
|
||||
return {
|
||||
"status": "FAILED",
|
||||
"confidence": 0.0,
|
||||
"routing_decision": "REVIEW",
|
||||
"face_count": 0,
|
||||
"frame_count": 0,
|
||||
"model_version": get_model_version(),
|
||||
"processing_time_seconds": round(time.time() - start_time, 2),
|
||||
"error": str(exc),
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _get_batch_size(gpu_cfg: dict, batching_cfg: dict) -> int:
|
||||
"""Compute an appropriate batch size based on GPU config."""
|
||||
from gpu_manager import GPUMemoryManager # noqa: local import, heavy
|
||||
|
||||
mgr = GPUMemoryManager(
|
||||
max_memory_gb=gpu_cfg.get("max_memory_gb", 18.0),
|
||||
reduce_threshold_gb=batching_cfg.get("vram_reduce_threshold_gb", 16.0),
|
||||
increase_threshold_gb=batching_cfg.get("vram_increase_threshold_gb", 10.0),
|
||||
initial_batch_size=batching_cfg.get("max_batch_size", 16),
|
||||
)
|
||||
return mgr.current_batch_size
|
||||
|
||||
|
||||
def _failed_result(
|
||||
video_id: int, start_time: float, confidence: float, routing: str
|
||||
) -> Dict[str, Any]:
|
||||
"""Build a SKIP/early-exit result dict (no faces case)."""
|
||||
return {
|
||||
"status": "COMPLETED",
|
||||
"confidence": round(confidence, 4),
|
||||
"routing_decision": routing,
|
||||
"face_count": 0,
|
||||
"frame_count": 0,
|
||||
"model_version": get_model_version(),
|
||||
"processing_time_seconds": round(time.time() - start_time, 2),
|
||||
"error": None,
|
||||
}
|
||||
@@ -1,186 +0,0 @@
|
||||
"""VideoDetect task worker — short-lived process for AI scanning.
|
||||
|
||||
Usage::
|
||||
|
||||
python3 -m src.task_worker # process 1 task and exit
|
||||
python3 -m src.task_worker --tasks 50 # process up to 50 tasks
|
||||
python3 src/task_worker.py --tasks 10 # direct invocation
|
||||
|
||||
Environment variables:
|
||||
API_BASE_URL (default http://localhost:3000)
|
||||
Base URL of the Dancer2 REST API.
|
||||
|
||||
TASK_COUNT (default 1)
|
||||
Maximum number of tasks to process before exiting.
|
||||
|
||||
On success exits with code 0; on unrecoverable errors exits non-zero.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import signal
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# Ensure the src/ directory is on sys.path so our submodules import cleanly
|
||||
_src = Path(__file__).resolve().parent
|
||||
if str(_src) not in sys.path:
|
||||
sys.path.insert(0, str(_src))
|
||||
|
||||
from api_client import ApiClient, ApiError
|
||||
from task_processor import process_video, get_model_version
|
||||
from config_loader import get_config
|
||||
from logging_config import setup_logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Graceful shutdown
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_shutdown_requested = False
|
||||
|
||||
|
||||
def _signal_handler(signum: int, frame) -> None:
|
||||
logger.info("Received signal %d — finishing current task then exiting", signum)
|
||||
global _shutdown_requested
|
||||
_shutdown_requested = True
|
||||
|
||||
|
||||
signal.signal(signal.SIGTERM, _signal_handler)
|
||||
signal.signal(signal.SIGINT, _signal_handler)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI entry point
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="VideoDetect task worker — process AI scan tasks",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tasks",
|
||||
type=int,
|
||||
default=int(os.environ.get("TASK_COUNT", "1")),
|
||||
help="Number of tasks to process before exiting (default: 1 or $TASK_COUNT)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--api-url",
|
||||
default=os.environ.get("API_BASE_URL", "http://localhost:3000"),
|
||||
help="Base URL of the Dancer2 REST API (default: http://localhost:3000)",
|
||||
)
|
||||
return parser.parse_args(argv)
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
"""Run the task worker loop.
|
||||
|
||||
Returns exit code (0 = success).
|
||||
"""
|
||||
args = parse_args(argv)
|
||||
|
||||
# Load config for the processing pipeline
|
||||
config = get_config()
|
||||
|
||||
# Setup JSON logging
|
||||
log_cfg = (config or {}).get("logging", {})
|
||||
setup_logging(
|
||||
level=log_cfg.get("level", "INFO"),
|
||||
log_format=log_cfg.get("format", "json"),
|
||||
rotation_max_bytes=log_cfg.get("rotation_max_bytes", 104857600),
|
||||
rotation_backup_count=log_cfg.get("rotation_backup_count", 10),
|
||||
)
|
||||
|
||||
model_version = get_model_version()
|
||||
logger.info(
|
||||
"Task worker starting. Will process up to %d task(s). Model version: %s",
|
||||
args.tasks,
|
||||
model_version,
|
||||
)
|
||||
|
||||
api = ApiClient(base_url=args.api_url)
|
||||
total_attempted = 0
|
||||
total_succeeded = 0
|
||||
total_failed = 0
|
||||
|
||||
try:
|
||||
for i in range(1, args.tasks + 1):
|
||||
if _shutdown_requested:
|
||||
logger.info("Shutdown requested after %d/%d tasks", i - 1, args.tasks)
|
||||
break
|
||||
|
||||
logger.info("=== Processing task %d/%d ===", i, args.tasks)
|
||||
|
||||
# 1. Claim a task
|
||||
response = api.get_next_task("AISCAN")
|
||||
if response is None:
|
||||
logger.info("No more tasks available (API returned no task)")
|
||||
break
|
||||
|
||||
task_id = response["task"]["id"]
|
||||
assign_key = response["assign_key"]
|
||||
video_id = response["task"]["video_id"]
|
||||
total_attempted += 1
|
||||
|
||||
logger.info(
|
||||
"Claiming task %d (video_id=%d) for AISCAN", task_id, video_id
|
||||
)
|
||||
|
||||
# 2. Fetch video metadata
|
||||
try:
|
||||
video = api.get_video(video_id)
|
||||
except ApiError as exc:
|
||||
logger.error("Failed to fetch video %d via API: %s", video_id, exc)
|
||||
total_failed += 1
|
||||
continue
|
||||
|
||||
# 3. Process the video (AI pipeline)
|
||||
results = process_video(video, config=config)
|
||||
|
||||
if results["status"] == "COMPLETED":
|
||||
logger.info(
|
||||
"Video %d complete: C=%.4f routing=%s faces=%d frames=%d time=%.1fs",
|
||||
video_id,
|
||||
results.get("confidence", 0),
|
||||
results.get("routing_decision", "?"),
|
||||
results.get("face_count", 0),
|
||||
results.get("frame_count", 0),
|
||||
results.get("processing_time_seconds", 0),
|
||||
)
|
||||
total_succeeded += 1
|
||||
else:
|
||||
logger.warning(
|
||||
"Video %d failed: %s",
|
||||
video_id,
|
||||
results.get("error", "unknown error"),
|
||||
)
|
||||
total_failed += 1
|
||||
|
||||
# 4. Submit results via API
|
||||
try:
|
||||
api.submit_results(task_id, assign_key, results)
|
||||
logger.info("Submitted results for task %d via API", task_id)
|
||||
except ApiError as exc:
|
||||
logger.error(
|
||||
"Failed to submit results for task %d: %s", task_id, exc
|
||||
)
|
||||
|
||||
finally:
|
||||
api.close()
|
||||
|
||||
logger.info(
|
||||
"Task worker finished: %d attempted, %d succeeded, %d failed",
|
||||
total_attempted,
|
||||
total_succeeded,
|
||||
total_failed,
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
+1
-1
@@ -1,4 +1,4 @@
|
||||
Flask==3.0.0
|
||||
gunicorn==21.2.0
|
||||
PyMySQL==1.1.0
|
||||
Jinja2==3.1.2
|
||||
requests>=2.28.0
|
||||
|
||||
@@ -1,68 +0,0 @@
|
||||
# Base image: CUDA 11.8 runtime on Ubuntu 22.04
|
||||
FROM nvidia/cuda:11.8.0-runtime-ubuntu22.04
|
||||
|
||||
# Avoid interactive prompts during build
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
ENV TZ=Etc/UTC
|
||||
|
||||
# System dependencies
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
# Python
|
||||
python3.10 \
|
||||
python3.10-dev \
|
||||
python3.10-venv \
|
||||
python3-pip \
|
||||
# FFmpeg
|
||||
ffmpeg \
|
||||
# OpenCV dependencies
|
||||
libgl1-mesa-glx \
|
||||
libglib2.0-0 \
|
||||
libsm6 \
|
||||
libxext6 \
|
||||
libxrender-dev \
|
||||
# Utilities
|
||||
curl \
|
||||
wget \
|
||||
git \
|
||||
ca-certificates \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Create non-root user
|
||||
RUN groupadd -g 1000 appuser && \
|
||||
useradd -m -u 1000 -g appuser -s /bin/bash appuser
|
||||
|
||||
# Set working directory
|
||||
WORKDIR /app
|
||||
|
||||
# Create necessary directories
|
||||
RUN mkdir -p /scratch /models /data/training /data/output /logs && \
|
||||
chown -R appuser:appuser /app /scratch /models /data /logs
|
||||
|
||||
# Copy requirements first for better caching
|
||||
COPY worker/requirements.txt /app/requirements.txt
|
||||
|
||||
# Create and activate virtual environment
|
||||
RUN python3.10 -m venv /opt/venv
|
||||
ENV PATH="/opt/venv/bin:$PATH"
|
||||
|
||||
# Install Python dependencies
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
|
||||
# Copy application code
|
||||
COPY src/ /app/src/
|
||||
COPY config.yaml /app/config.yaml
|
||||
|
||||
# Set environment variables
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
ENV PATH="/opt/venv/bin:$PATH"
|
||||
ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
|
||||
|
||||
# Health check
|
||||
HEALTHCHECK --interval=30s --timeout=10s --retries=3 \
|
||||
CMD python3 -c "import torch; assert torch.cuda.is_available()" || exit 1
|
||||
|
||||
# Switch to non-root user
|
||||
USER appuser
|
||||
|
||||
# Default command: process tasks from the API until drained
|
||||
CMD ["python3", "-m", "src.task_worker"]
|
||||
@@ -1,40 +0,0 @@
|
||||
# VideoDetect — Worker service (standalone compose file)
|
||||
# Run: docker compose -f worker/docker-compose.yml up -d
|
||||
|
||||
services:
|
||||
worker:
|
||||
build:
|
||||
context: ./
|
||||
dockerfile: Dockerfile
|
||||
container_name: videodetect-worker
|
||||
restart: "no"
|
||||
environment:
|
||||
- API_BASE_URL=${API_BASE_URL:-http://videodetect-api:3000}
|
||||
- TASK_COUNT=${TASK_COUNT:-50}
|
||||
- REVIEW_CONFIDENCE=${REVIEW_CONFIDENCE:-0.75}
|
||||
- LOG_FILE=/logs/videodetect.log
|
||||
volumes:
|
||||
- nas_input:/data/input:ro
|
||||
- scratch_data:/scratch
|
||||
- $PWD/worker-logs:/logs
|
||||
networks:
|
||||
- videodetect_videodetect-network
|
||||
deploy:
|
||||
resources:
|
||||
limits:
|
||||
memory: 4G
|
||||
nvidia-gpus: "1"
|
||||
|
||||
volumes:
|
||||
nas_input:
|
||||
driver: local
|
||||
driver_opts:
|
||||
type: nfs
|
||||
o: addr=10.0.0.2,ro,nfsvers=4,hard,intr
|
||||
device: ":/mnt/Bulk/Homes/ryan/Prawns"
|
||||
scratch_data:
|
||||
driver: local
|
||||
|
||||
networks:
|
||||
videodetect_videodetect-network:
|
||||
external: true
|
||||
@@ -1,40 +0,0 @@
|
||||
# Core ML Frameworks
|
||||
torch==2.1.0
|
||||
--extra-index-url https://download.pytorch.org/whl/cu118
|
||||
|
||||
# TensorRT (compatible with CUDA 11.8)
|
||||
# Note: Install from NVIDIA repo or wheel for your specific platform
|
||||
# tensorrt==8.6.1 # Uncomment after verifying compatibility
|
||||
|
||||
# ONNX Runtime with GPU support
|
||||
onnxruntime-gpu==1.16.3
|
||||
|
||||
# Computer Vision
|
||||
opencv-python-headless==4.8.1.78
|
||||
Pillow==10.1.0
|
||||
|
||||
# Database
|
||||
PyMySQL==1.1.0
|
||||
DBUtils==3.1.0
|
||||
|
||||
# Data Processing
|
||||
numpy==1.24.4
|
||||
pandas==2.1.4
|
||||
pyarrow==14.0.0
|
||||
|
||||
# Configuration
|
||||
pyyaml==6.0.1
|
||||
|
||||
# Logging
|
||||
python-json-logger==2.0.7
|
||||
|
||||
# Monitoring
|
||||
prometheus-client==0.19.0
|
||||
|
||||
# Utilities
|
||||
tqdm==4.66.1
|
||||
filelock==3.13.1
|
||||
|
||||
# Testing (dev dependencies - install with [dev] extra)
|
||||
# pytest==7.4.3
|
||||
# pytest-cov==4.1.0
|
||||
Reference in New Issue
Block a user