""" 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 []