Compare commits
6
Commits
639c29c2c4
...
8f2834775b
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
8f2834775b | ||
|
|
75f38ec7ed | ||
|
|
5fa67a733c | ||
|
|
46548881cb | ||
|
|
4782dc5bfa | ||
|
|
1611485a66 |
+1
-1
@@ -1,5 +1,5 @@
|
|||||||
# VideoDetect — Docker Compose environment variables
|
# VideoDetect — Docker Compose environment variables
|
||||||
# Copy this file to .env and fill in real values before running docker-compose up.
|
# 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
|
container_name: videodetect-aiworker
|
||||||
restart: "no"
|
restart: "no"
|
||||||
environment:
|
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}
|
- OLLAMA_URL=${OLLAMA_URL:-http://10.0.1.7:11434}
|
||||||
- SECONDS_PER_FRAME=${SECONDS_PER_FRAME:-30}
|
- SECONDS_PER_FRAME=${SECONDS_PER_FRAME:-30}
|
||||||
- TASK_COUNT=${TASK_COUNT:-50}
|
- TASK_COUNT=${TASK_COUNT:-50}
|
||||||
|
|||||||
+84
-21
@@ -3,6 +3,7 @@ use strict;
|
|||||||
use Dancer2;
|
use Dancer2;
|
||||||
use Dancer2::Plugin::Database;
|
use Dancer2::Plugin::Database;
|
||||||
use Data::Dumper;
|
use Data::Dumper;
|
||||||
|
use JSON ();
|
||||||
use SQL::SplitStatement;
|
use SQL::SplitStatement;
|
||||||
use HTTP::Lite;
|
use HTTP::Lite;
|
||||||
|
|
||||||
@@ -83,12 +84,42 @@ post '/api/v1/video' => sub {
|
|||||||
|
|
||||||
get '/api/v1/tasks' => sub {
|
get '/api/v1/tasks' => sub {
|
||||||
my $status = query_parameters->get("status") || 'PENDING';
|
my $status = query_parameters->get("status") || 'PENDING';
|
||||||
my $sth = database->prepare("SELECT id, video_id, status FROM tasks WHERE status=? ORDER BY id");
|
my $type = query_parameters->get("type");
|
||||||
$sth->execute($status);
|
|
||||||
my $ref = $sth->fetchall_hashref('id');
|
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();
|
$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 {
|
get '/api/v1/task/:task' => sub {
|
||||||
@@ -129,13 +160,13 @@ get '/api/v1/nexttask/:type' => sub {
|
|||||||
};
|
};
|
||||||
|
|
||||||
post '/api/v1/task/:task/complete' => 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.
|
# Accept both JSON body and form-encoded data for assign_key/results.
|
||||||
my ($assign_key, $results);
|
my ($assign_key, $results);
|
||||||
my $ct = request_header('Content-Type') || '';
|
my $ct = request_header('Content-Type') || '';
|
||||||
if ($ct eq 'application/json') {
|
if ($ct =~ m{application/json}i) {
|
||||||
my $body = decode_json(request->body());
|
my $body = JSON::decode_json(request->body());
|
||||||
$assign_key = $body->{assign_key};
|
$assign_key = $body->{assign_key};
|
||||||
$results = $body->{results};
|
$results = $body->{results};
|
||||||
} else {
|
} else {
|
||||||
@@ -143,24 +174,31 @@ post '/api/v1/task/:task/complete' => sub {
|
|||||||
$results = body_parameters->get("results");
|
$results = body_parameters->get("results");
|
||||||
}
|
}
|
||||||
|
|
||||||
unless ($assign_key) {
|
# Optional worker verification: if assign_key is provided, require IN_PROGRESS + matching key.
|
||||||
send_error("Missing assign_key", 400);
|
# 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_id, $assign_key);
|
||||||
|
my $ref = $sth->fetchrow_hashref();
|
||||||
|
$sth->finish();
|
||||||
|
|
||||||
# Verify the task is assigned to the worker
|
unless(ref $ref) {
|
||||||
my $sth = database->prepare("SELECT * FROM tasks WHERE id=? AND assign_key=? AND status='IN_PROGRESS'");
|
send_error("Task not assigned to this worker or not in progress", 403);
|
||||||
$sth->execute($task, $assign_key);
|
}
|
||||||
my $ref = $sth->fetchrow_hashref();
|
} else {
|
||||||
$sth->finish();
|
# No assign_key provided — caller may be an unauthenticated UI, etc.
|
||||||
|
my $sth = database->prepare("SELECT id FROM tasks WHERE id=?");
|
||||||
unless(ref $ref) {
|
$sth->execute($task_id);
|
||||||
send_error("Task not assigned to this worker or not in progress", 403);
|
unless ($sth->fetchrow_hashref()) {
|
||||||
|
send_error("Task not found", 404);
|
||||||
|
}
|
||||||
|
$sth->finish();
|
||||||
}
|
}
|
||||||
|
|
||||||
# Update the task as completed (store results as JSON text for MariaDB)
|
# 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=?");
|
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) {
|
if($update_sth->rows == 0) {
|
||||||
send_error("Failed to complete task", 500);
|
send_error("Failed to complete task", 500);
|
||||||
}
|
}
|
||||||
@@ -183,4 +221,29 @@ post '/api/v1/task' => sub {
|
|||||||
return { id => $task_id };
|
return { id => $task_id };
|
||||||
};
|
};
|
||||||
|
|
||||||
start();
|
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:
|
dbi_params:
|
||||||
RaiseError: 1
|
RaiseError: 1
|
||||||
AutoCommit: 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
|
log_queries: 1
|
||||||
|
|
||||||
server:
|
server:
|
||||||
|
|||||||
+4
-10
@@ -7,7 +7,7 @@ services:
|
|||||||
MYSQL_ROOT_PASSWORD: ${DB_ROOT_PASSWORD:-rootpass}
|
MYSQL_ROOT_PASSWORD: ${DB_ROOT_PASSWORD:-rootpass}
|
||||||
MYSQL_DATABASE: ${DB_NAME:-videodetect}
|
MYSQL_DATABASE: ${DB_NAME:-videodetect}
|
||||||
MYSQL_USER: ${DB_USER:-videodetect}
|
MYSQL_USER: ${DB_USER:-videodetect}
|
||||||
MYSQL_PASSWORD: ${DB_PASSWORD:-changeme_videodetect}
|
MYSQL_PASSWORD: ${DB_PASSWORD:-videodetect}
|
||||||
ports:
|
ports:
|
||||||
- "3306:3306"
|
- "3306:3306"
|
||||||
volumes:
|
volumes:
|
||||||
@@ -29,20 +29,14 @@ services:
|
|||||||
build:
|
build:
|
||||||
context: ./ui
|
context: ./ui
|
||||||
dockerfile: Dockerfile
|
dockerfile: Dockerfile
|
||||||
container_name: videodetect-ui
|
|
||||||
restart: unless-stopped
|
restart: unless-stopped
|
||||||
ports:
|
ports:
|
||||||
- "5000:5000"
|
- "5000:5000"
|
||||||
environment:
|
environment:
|
||||||
- DB_HOST=${DB_HOST:-mariadb}
|
- API_BASE=http://10.0.1.34:3000/api/v1
|
||||||
- DB_PORT=${DB_PORT:-3306}
|
|
||||||
- DB_NAME=${DB_NAME:-videodetect}
|
|
||||||
- DB_USER=${DB_USER:-videodetect}
|
|
||||||
- DB_PASSWORD=${DB_PASSWORD:-changeme_videodetect}
|
|
||||||
volumes:
|
volumes:
|
||||||
- ./ui:/app
|
- ./ui:/app
|
||||||
- nas_input:/data/input:ro
|
- nas_input:/data:ro
|
||||||
- ${NAS_OUTPUT_PATH:-/mnt/nas/output}:/data/output
|
|
||||||
networks:
|
networks:
|
||||||
- videodetect-network
|
- videodetect-network
|
||||||
depends_on:
|
depends_on:
|
||||||
@@ -66,7 +60,7 @@ services:
|
|||||||
- DB_PORT=${DB_PORT:-3306}
|
- DB_PORT=${DB_PORT:-3306}
|
||||||
- DB_NAME=${DB_NAME:-videodetect}
|
- DB_NAME=${DB_NAME:-videodetect}
|
||||||
- DB_USER=${DB_USER:-videodetect}
|
- DB_USER=${DB_USER:-videodetect}
|
||||||
- DB_PASSWORD=${DB_PASSWORD:-changeme_videodetect}
|
- DB_PASSWORD=${DB_PASSWORD:-videodetect}
|
||||||
volumes:
|
volumes:
|
||||||
- ./api:/app
|
- ./api:/app
|
||||||
networks:
|
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_PASSWORD=videodetect
|
||||||
DB_USER=videodetect
|
DB_USER=videodetect
|
||||||
DB_NAME=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"; };
|
$SIG{INT} = sub { die "Interrupted\n"; };
|
||||||
|
|
||||||
my $dbh = DBI->connect(
|
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_user,
|
||||||
$db_pass,
|
$db_pass,
|
||||||
{ mysql_auto_reconnect => 1 }
|
{ 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
|
Flask==3.0.0
|
||||||
gunicorn==21.2.0
|
gunicorn==21.2.0
|
||||||
PyMySQL==1.1.0
|
|
||||||
Jinja2==3.1.2
|
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