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