from __future__ import annotations import json import re from dataclasses import asdict, dataclass from pathlib import Path import numpy as np from PIL import Image from src import config @dataclass class GalleryRecord: person: str gallery_filename: str query_filename: str gallery_path: str query_path: str class DemoGallery: def __init__(self) -> None: gallery_npz = Path(config.LFW_FEATURES_PATH) manifest_path = Path(config.LFW_VISIBLE_MANIFEST) example_manifest_path = Path(config.LFW_EXAMPLES_MANIFEST) data = np.load(gallery_npz, allow_pickle=True) self.base_features = data["features"].astype(np.float32) self.base_people = data["persons"].astype(str) self.base_filenames = data["filenames"].astype(str) self.base_identity_count = int(len(set(self.base_people.tolist()))) self.feature_dim = int(self.base_features.shape[1]) calibration_npz = Path(config.LFW_CALIBRATION_PATH) calibration = np.load(calibration_npz, allow_pickle=True) self.calibration_features = calibration["features"].astype(np.float32) self.calibration_people = calibration["persons"].astype(str) self.calibration_filenames = calibration["filenames"].astype(str) with open(manifest_path) as fh: self.base_records = [GalleryRecord(**row) for row in json.load(fh)] with open(example_manifest_path) as fh: self.examples = json.load(fh) self.features = self.base_features self.people = self.base_people self.filenames = self.base_filenames self._reload_enrollments() def _empty_features(self) -> np.ndarray: return np.empty((0, self.feature_dim), dtype=np.float32) def _person_index(self, person: str) -> int | None: normalized = person.strip().lower() for idx, existing in enumerate(self.enrolled_people.tolist()): if str(existing).strip().lower() == normalized: return idx return None @staticmethod def _slugify(person: str) -> str: slug = re.sub(r"[^a-zA-Z0-9]+", "_", person.strip()).strip("_") return slug or "user" def _reload_enrollments(self) -> None: if config.ENROLLMENT_FEATURES_PATH.exists(): data = np.load(config.ENROLLMENT_FEATURES_PATH, allow_pickle=True) self.enrolled_features = data["features"].astype(np.float32) self.enrolled_people = data["persons"].astype(str) self.enrolled_filenames = data["filenames"].astype(str) else: self.enrolled_features = self._empty_features() self.enrolled_people = np.array([], dtype=str) self.enrolled_filenames = np.array([], dtype=str) if config.ENROLLMENT_MANIFEST_PATH.exists(): with open(config.ENROLLMENT_MANIFEST_PATH) as fh: self.enrolled_records = [GalleryRecord(**row) for row in json.load(fh)] else: self.enrolled_records = [] self.records = self.enrolled_records + self.base_records combined_people = self.base_people.tolist() + self.enrolled_people.tolist() self.identity_count = int(len(set(combined_people))) self.enrolled_identity_count = int(len(set(self.enrolled_people.tolist()))) def _save_enrollments(self) -> None: np.savez( config.ENROLLMENT_FEATURES_PATH, features=self.enrolled_features.astype(np.float32), persons=self.enrolled_people.astype(str), filenames=self.enrolled_filenames.astype(str), ) with open(config.ENROLLMENT_MANIFEST_PATH, "w") as fh: json.dump([asdict(record) for record in self.enrolled_records], fh, indent=2) def has_enrollments(self) -> bool: return len(self.enrolled_people) > 0 def enroll(self, person: str, image: Image.Image, embedding: np.ndarray) -> tuple[GalleryRecord, bool]: person = " ".join(person.split()) if not person: raise ValueError("Person name cannot be empty") filename = f"enrolled_{self._slugify(person)}.jpg" image_path = config.ENROLLMENT_IMAGES_DIR / filename image.convert("RGB").save(image_path, format="JPEG", quality=95) embedding_row = np.asarray(embedding, dtype=np.float32).reshape(1, -1) if embedding_row.shape[1] != self.feature_dim: raise ValueError(f"Embedding dimension mismatch: expected {self.feature_dim}, got {embedding_row.shape[1]}") record = GalleryRecord( person=person, gallery_filename=filename, query_filename=filename, gallery_path=str(image_path), query_path=str(image_path), ) existing_idx = self._person_index(person) replaced = existing_idx is not None if replaced: assert existing_idx is not None self.enrolled_features[existing_idx : existing_idx + 1] = embedding_row self.enrolled_people[existing_idx] = person self.enrolled_filenames[existing_idx] = filename self.enrolled_records[existing_idx] = record else: self.enrolled_features = np.vstack([self.enrolled_features, embedding_row]) self.enrolled_people = np.concatenate([self.enrolled_people, np.array([person], dtype=str)]) self.enrolled_filenames = np.concatenate([self.enrolled_filenames, np.array([filename], dtype=str)]) self.enrolled_records.append(record) self._save_enrollments() self._reload_enrollments() return record, replaced def summary(self) -> dict[str, object]: return { "gallery_size": int(len(self.base_people) + len(self.enrolled_people)), "identity_count": self.identity_count, "visible_gallery_size": int(len(self.records)), "calibration_size": int(len(self.calibration_people)), "people": [str(x.person) for x in self.records], "dimensions": self.feature_dim, "enrolled_gallery_size": int(len(self.enrolled_people)), "enrolled_identity_count": self.enrolled_identity_count, "demo_gallery_size": int(len(self.base_people)), "demo_identity_count": self.base_identity_count, } def example_choices(self) -> list[tuple[str, str]]: return [(item["person"], str(config.EXAMPLES_DIR / item["filename"])) for item in self.examples] def load_example_image(self, filename: str) -> Image.Image: return Image.open(config.EXAMPLES_DIR / filename).convert("RGB") def gallery_image_path(self, gallery_filename: str, person: str | None = None) -> Path: direct = Path(gallery_filename) if direct.is_absolute() and direct.exists(): return direct enrolled_direct = config.ENROLLMENT_IMAGES_DIR / gallery_filename if enrolled_direct.exists(): return enrolled_direct base_direct = config.GALLERY_DIR / gallery_filename if base_direct.exists(): return base_direct def resolve_record_path(record: GalleryRecord) -> Path | None: for candidate in ( Path(record.gallery_path), config.GALLERY_DIR / record.gallery_path, config.EXAMPLES_DIR / record.gallery_path, config.GALLERY_DIR / record.gallery_filename, config.EXAMPLES_DIR / record.gallery_filename, ): if candidate.exists(): return candidate return None if person is not None: for record in self.enrolled_records + self.base_records: if record.person == person: candidate = resolve_record_path(record) if candidate is not None: return candidate for record in self.enrolled_records + self.base_records: if record.gallery_filename == gallery_filename: candidate = resolve_record_path(record) if candidate is not None: return candidate raise FileNotFoundError(f"No visible gallery image found for {gallery_filename}") def gallery_items(self, page: int = 1, per_page: int = 24) -> list[tuple[str, str]]: page = max(1, int(page or 1)) start = (page - 1) * per_page end = start + per_page items: list[tuple[str, str]] = [] for record in self.records[start:end]: items.append((str(self.gallery_image_path(record.gallery_filename, record.person)), record.person)) return items def total_pages(self, per_page: int = 24) -> int: return max(1, (len(self.records) + per_page - 1) // per_page) def database_info(self, page: int = 1, per_page: int = 24) -> str: summary = self.summary() total_pages = self.total_pages(per_page) page = min(max(1, int(page or 1)), total_pages) start = (page - 1) * per_page + 1 end = min(page * per_page, len(self.records)) return ( f"{summary['identity_count']} identities | {summary['gallery_size']} images\n" f"Enrolled: {summary['enrolled_identity_count']} identities / {summary['enrolled_gallery_size']} images\n" f"Visible: {start}-{end} of {summary['visible_gallery_size']} | Page {page}/{total_pages}" )