From fa9b21856f43b06e52da766278edd2e605a2c87b Mon Sep 17 00:00:00 2001 From: dtoro Date: Fri, 10 Apr 2026 13:49:02 +0200 Subject: [PATCH] feat: replace face pipeline with InsightFace, add content classifier MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Face detection/recognition: - Replace YuNet + SFace with InsightFace buffalo_l (RetinaFace + ArcFace) - 512-d ArcFace embeddings (was 128-d SFace), migration 0006 resizes column - Remove YOLO person-bbox workaround — RetinaFace is accurate enough - Detection threshold 0.65 cleanly separates real faces (0.72+) from false positives on dogs/paintings (0.56-0.61) Content-type classification: - CLIP zero-shot classifier using native PyTorch text encoder + ONNX image encoder for high-quality text-image similarity - Categories: photograph, screenshot, document, receipt, meme, artwork - Writes Tag(kind=content_type) per photo via photo_tags - Margin-based confidence: top-1 vs top-2 score difference - New ClassifierSettings in config (enabled, min_confidence) - Wired into vision_fanout pipeline Tested: 6 real faces from 4 photos (zero false positives), 11/13 photos classified (8 photograph, 2 artwork, 1 meme). Co-Authored-By: Claude Opus 4.6 (1M context) --- .gitignore | 2 +- .../versions/0006_face_embeddings_512d.py | 39 ++++++ backend/app/config.py | 10 +- backend/app/models/face_embedding.py | 2 +- backend/app/services/vision/base.py | 16 +++ .../app/services/vision/bootstrap_models.py | 17 +-- backend/app/services/vision/classify.py | 106 ++++++++++++++++ .../services/vision/insightface_processor.py | 70 +++++++++++ backend/app/services/vision/onnx_backend.py | 10 +- backend/app/services/vision/registry.py | 9 +- backend/app/tasks/celery.py | 1 + backend/app/tasks/vision.py | 114 ++++++++++-------- backend/requirements.txt | 1 + mulita.yml | 7 +- 14 files changed, 331 insertions(+), 73 deletions(-) create mode 100644 backend/alembic/versions/0006_face_embeddings_512d.py create mode 100644 backend/app/services/vision/classify.py create mode 100644 backend/app/services/vision/insightface_processor.py diff --git a/.gitignore b/.gitignore index 99a8a51..ec68f8e 100644 --- a/.gitignore +++ b/.gitignore @@ -65,4 +65,4 @@ docker-compose.override.yml # Thumbnails /thumbs/ -/trash/ \ No newline at end of file +/trash/backend/yolov8n.pt diff --git a/backend/alembic/versions/0006_face_embeddings_512d.py b/backend/alembic/versions/0006_face_embeddings_512d.py new file mode 100644 index 0000000..b8eaf93 --- /dev/null +++ b/backend/alembic/versions/0006_face_embeddings_512d.py @@ -0,0 +1,39 @@ +"""face_embeddings vector 128 -> 512 + +Revision ID: 0006_face_512d +Revises: 0005_face_embeddings +Create Date: 2026-04-10 + +Resize face_embeddings.vector from Vector(128) to Vector(512) for +ArcFace embeddings (InsightFace). Drops existing data and HNSW index, +recreates both. +""" +from typing import Sequence, Union + +from alembic import op + +revision: str = "0006_face_512d" +down_revision: Union[str, None] = "0005_face_embeddings" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # Drop index, truncate (old 128-d vectors are incompatible), resize + op.execute("DROP INDEX IF EXISTS ix_face_embeddings_vector_hnsw") + op.execute("DELETE FROM face_embeddings") + op.execute("ALTER TABLE face_embeddings ALTER COLUMN vector TYPE vector(512)") + op.execute(""" + CREATE INDEX IF NOT EXISTS ix_face_embeddings_vector_hnsw + ON face_embeddings USING hnsw (vector vector_cosine_ops) + """) + + +def downgrade() -> None: + op.execute("DROP INDEX IF EXISTS ix_face_embeddings_vector_hnsw") + op.execute("DELETE FROM face_embeddings") + op.execute("ALTER TABLE face_embeddings ALTER COLUMN vector TYPE vector(128)") + op.execute(""" + CREATE INDEX IF NOT EXISTS ix_face_embeddings_vector_hnsw + ON face_embeddings USING hnsw (vector vector_cosine_ops) + """) diff --git a/backend/app/config.py b/backend/app/config.py index cc3b9ef..5d5b866 100644 --- a/backend/app/config.py +++ b/backend/app/config.py @@ -50,8 +50,13 @@ class FacesSettings(BaseModel): """YuNet + SFace face detection/recognition settings""" enabled: bool = True min_face_size: int = 40 - recognition_threshold: float = 0.6 - cluster_eps: float = 0.25 + recognition_threshold: float = 0.65 + cluster_eps: float = 0.5 + +class ClassifierSettings(BaseModel): + """CLIP zero-shot content classification settings""" + enabled: bool = True + min_confidence: float = 0.3 class VisionSettings(BaseModel): """AI vision pipeline settings. Disabled when running on SQLite @@ -63,6 +68,7 @@ class VisionSettings(BaseModel): ocr: OCRSettings = OCRSettings() detector: DetectorSettings = DetectorSettings() faces: FacesSettings = FacesSettings() + classifier: ClassifierSettings = ClassifierSettings() worker_concurrency: int = 2 class MulitaConfig(BaseModel): diff --git a/backend/app/models/face_embedding.py b/backend/app/models/face_embedding.py index 138e61f..291d18c 100644 --- a/backend/app/models/face_embedding.py +++ b/backend/app/models/face_embedding.py @@ -18,7 +18,7 @@ class FaceEmbedding(Base): id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4())) photo_id = Column(String, ForeignKey('photos.id', ondelete='CASCADE'), nullable=False, index=True) bbox = Column(JSONB) # [x1, y1, x2, y2] normalized 0-1 - vector = Column(Vector(128)) # SFace → 128-d + vector = Column(Vector(512)) # ArcFace → 512-d cluster_id = Column(String, ForeignKey('tags.id', ondelete='SET NULL'), nullable=True, index=True) quality = Column(Float) created_at = Column(DateTime(timezone=True), server_default=func.now()) diff --git a/backend/app/services/vision/base.py b/backend/app/services/vision/base.py index 391c070..71a7856 100644 --- a/backend/app/services/vision/base.py +++ b/backend/app/services/vision/base.py @@ -36,6 +36,13 @@ class FaceDetection: quality: float +@dataclass +class ClassificationResult: + """A content-type classification.""" + label: str + confidence: float + + class Embedder(ABC): """Generates image and text embeddings (e.g. OpenCLIP ViT-B/32).""" @@ -74,6 +81,15 @@ class ObjectDetector(ABC): ... +class ContentClassifier(ABC): + """Classifies images into content types (screenshot, document, etc.).""" + + @abstractmethod + def classify(self, image: np.ndarray) -> list[ClassificationResult]: + """Return content type classifications for an RGB image.""" + ... + + class FaceProcessor(ABC): """Detects faces and extracts recognition embeddings (e.g. YuNet + SFace).""" diff --git a/backend/app/services/vision/bootstrap_models.py b/backend/app/services/vision/bootstrap_models.py index 5ac7fc2..bff5bf5 100644 --- a/backend/app/services/vision/bootstrap_models.py +++ b/backend/app/services/vision/bootstrap_models.py @@ -21,20 +21,9 @@ logger = logging.getLogger(__name__) # (relative_path, url, description) # Models with url=None must be pre-exported via export_models.py. -DOWNLOADS = [ - # YuNet — face detection (Apache 2.0, opencv_zoo, ~233 KB) - ( - "face/yunet.onnx", - "https://github.com/opencv/opencv_zoo/raw/main/models/face_detection_yunet/face_detection_yunet_2023mar.onnx", - "YuNet face detector", - ), - # SFace — face recognition (Apache 2.0, opencv_zoo, ~37 MB) - ( - "face/sface.onnx", - "https://github.com/opencv/opencv_zoo/raw/main/models/face_recognition_sface/face_recognition_sface_2021dec.onnx", - "SFace face recognizer", - ), -] +# InsightFace (RetinaFace + ArcFace) auto-downloads via the insightface +# package on first use — no manual download entries needed. +DOWNLOADS = [] # Models that need manual export via export_models.py EXPORTS = [ diff --git a/backend/app/services/vision/classify.py b/backend/app/services/vision/classify.py new file mode 100644 index 0000000..6de4424 --- /dev/null +++ b/backend/app/services/vision/classify.py @@ -0,0 +1,106 @@ +""" +CLIP zero-shot content-type classifier. + +Uses the native OpenCLIP PyTorch text encoder for high-quality text +embeddings (the ONNX text encoder has degraded quality due to the +eot_indices workaround). Image embeddings use the ONNX visual encoder +which works well. +""" +import logging + +import numpy as np +import torch + +from app.config import VisionSettings +from app.services.vision.base import ContentClassifier, ClassificationResult + +logger = logging.getLogger(__name__) + +CATEGORY_PROMPTS = { + "screenshot": [ + "a screenshot of a computer screen", + "a screenshot of a phone screen", + "a screen capture of a user interface", + ], + "document": [ + "a scanned document", + "a photo of a document with printed text", + "a photo of a page of text on paper", + ], + "receipt": [ + "a photo of a receipt", + "a photo of a bill or invoice", + ], + "meme": [ + "an internet meme with text overlay", + "a funny image with caption text", + ], + "artwork": [ + "a painting or drawing", + "a sketch or illustration", + "digital art or graphic design", + ], + "photograph": [ + "a photograph taken with a camera", + "a real photo of a real scene or person", + "a candid photograph", + ], +} + + +class CLIPContentClassifier(ContentClassifier): + """Zero-shot content classifier using CLIP text-image similarity. + Uses native PyTorch for text encoding, ONNX for image encoding.""" + + def __init__(self, settings: VisionSettings): + import open_clip + + self._min_confidence = settings.classifier.min_confidence + + # Load native model for text encoding only + logger.info("Loading OpenCLIP text encoder for content classification") + model, _, _ = open_clip.create_model_and_transforms( + "ViT-B-32", pretrained="laion2b_s34b_b79k" + ) + model.eval() + self._model = model + self._tokenizer = open_clip.get_tokenizer("ViT-B-32") + + # Get the ONNX image embedder from the registry + from app.services.vision.registry import registry + self._embedder = registry.get_embedder() + + # Pre-compute text embeddings for each category + self._category_embeddings: dict[str, np.ndarray] = {} + for category, prompts in CATEGORY_PROMPTS.items(): + tokens = self._tokenizer(prompts) + with torch.no_grad(): + text_features = model.encode_text(tokens) + text_features /= text_features.norm(dim=-1, keepdim=True) + avg = text_features.mean(dim=0) + avg /= avg.norm() + self._category_embeddings[category] = avg.numpy().astype(np.float32) + + logger.info("Content classifier ready with %d categories", len(self._category_embeddings)) + + def classify(self, image: np.ndarray) -> list[ClassificationResult]: + img_vec = self._embedder.embed_image(image) + + # Cosine similarity against each category + scores = {} + for category, cat_vec in self._category_embeddings.items(): + scores[category] = float(np.dot(img_vec, cat_vec)) + + # Sort by score descending + ranked = sorted(scores.items(), key=lambda x: -x[1]) + best_cat, best_score = ranked[0] + second_score = ranked[1][1] + + margin = best_score - second_score + # Normalize: 0.01 margin → ~0.5 confidence, 0.03+ → ~1.0 + confidence = min(1.0, margin * 30) + + if confidence >= self._min_confidence: + return [ClassificationResult(label=best_cat, confidence=confidence)] + + return [] diff --git a/backend/app/services/vision/insightface_processor.py b/backend/app/services/vision/insightface_processor.py new file mode 100644 index 0000000..aea8bb0 --- /dev/null +++ b/backend/app/services/vision/insightface_processor.py @@ -0,0 +1,70 @@ +""" +Face detection + recognition using InsightFace (RetinaFace + ArcFace). + +Uses the buffalo_l model pack which auto-downloads on first use (~300MB). +Produces 512-d ArcFace embeddings. Non-commercial research license — +fine for homelab self-hosting. +""" +import logging +from pathlib import Path + +import numpy as np + +from app.config import VisionSettings +from app.services.vision.base import FaceProcessor, FaceDetection + +logger = logging.getLogger(__name__) + + +class InsightFaceProcessor(FaceProcessor): + def __init__(self, settings: VisionSettings): + from insightface.app import FaceAnalysis + + model_root = str(Path(settings.models_dir) / "face" / "insightface") + logger.info("Loading InsightFace buffalo_l from %s", model_root) + + self._app = FaceAnalysis( + name="buffalo_l", + root=model_root, + providers=["CPUExecutionProvider"], + ) + self._app.prepare(ctx_id=-1, det_size=(640, 640)) + self._min_det_score = settings.faces.recognition_threshold + + def process(self, image: np.ndarray) -> list[FaceDetection]: + orig_h, orig_w = image.shape[:2] + + # InsightFace expects BGR + bgr = image[:, :, ::-1].copy() + + faces = self._app.get(bgr) + + if not faces: + return [] + + results = [] + for face in faces: + if face.det_score < self._min_det_score: + continue + + # face.bbox is [x1, y1, x2, y2] in pixel coords + x1, y1, x2, y2 = face.bbox + bbox = [ + max(0, float(x1) / orig_w), + max(0, float(y1) / orig_h), + min(1, float(x2) / orig_w), + min(1, float(y2) / orig_h), + ] + + embedding = face.normed_embedding # already L2-normalized, 512-d + results.append(FaceDetection( + bbox=bbox, + embedding=embedding.astype(np.float32), + quality=float(face.det_score), + )) + + return results + + @property + def embedding_dim(self) -> int: + return 512 diff --git a/backend/app/services/vision/onnx_backend.py b/backend/app/services/vision/onnx_backend.py index c5321b6..125a160 100644 --- a/backend/app/services/vision/onnx_backend.py +++ b/backend/app/services/vision/onnx_backend.py @@ -9,7 +9,7 @@ bootstrap_models.py. import logging from app.config import VisionSettings -from app.services.vision.base import Embedder, OCREngine, ObjectDetector, FaceProcessor +from app.services.vision.base import Embedder, OCREngine, ObjectDetector, FaceProcessor, ContentClassifier logger = logging.getLogger(__name__) @@ -33,5 +33,9 @@ class ONNXBackend: return YOLOv8Detector(self._settings) def create_face_processor(self) -> FaceProcessor: - from app.services.vision.faces import YuNetSFaceProcessor - return YuNetSFaceProcessor(self._settings) + from app.services.vision.insightface_processor import InsightFaceProcessor + return InsightFaceProcessor(self._settings) + + def create_classifier(self) -> ContentClassifier: + from app.services.vision.classify import CLIPContentClassifier + return CLIPContentClassifier(self._settings) diff --git a/backend/app/services/vision/registry.py b/backend/app/services/vision/registry.py index 8ce1853..f5ff51b 100644 --- a/backend/app/services/vision/registry.py +++ b/backend/app/services/vision/registry.py @@ -15,7 +15,7 @@ import logging from functools import lru_cache from app.config import settings -from app.services.vision.base import Embedder, OCREngine, ObjectDetector, FaceProcessor +from app.services.vision.base import Embedder, OCREngine, ObjectDetector, FaceProcessor, ContentClassifier logger = logging.getLogger(__name__) @@ -46,6 +46,11 @@ class ModelRegistry: logger.info("Loading face processor (backend=%s)", self._vision.backend) return self._load_backend().create_face_processor() + @lru_cache(maxsize=1) + def get_classifier(self) -> ContentClassifier: + logger.info("Loading content classifier (backend=%s)", self._vision.backend) + return self._load_backend().create_classifier() + @lru_cache(maxsize=1) def _load_backend(self): """Import and instantiate the configured backend.""" @@ -70,6 +75,8 @@ class ModelRegistry: self.get_detector() if self._vision.faces.enabled: self.get_face_processor() + if self._vision.classifier.enabled: + self.get_classifier() logger.info("Vision model warmup complete") diff --git a/backend/app/tasks/celery.py b/backend/app/tasks/celery.py index 33df67a..742ef65 100644 --- a/backend/app/tasks/celery.py +++ b/backend/app/tasks/celery.py @@ -27,6 +27,7 @@ celery_app.conf.update( 'ocr_photo': {'queue': 'vision'}, 'detect_objects': {'queue': 'vision'}, 'extract_faces': {'queue': 'vision'}, + 'classify_content': {'queue': 'vision'}, 'vision_fanout': {'queue': 'vision'}, }, task_default_queue='default', diff --git a/backend/app/tasks/vision.py b/backend/app/tasks/vision.py index 4c30200..94b1748 100644 --- a/backend/app/tasks/vision.py +++ b/backend/app/tasks/vision.py @@ -93,6 +93,8 @@ def vision_fanout(photo_id: str): detect_objects.delay(photo_id) if settings.vision.faces.enabled: extract_faces.delay(photo_id) + if settings.vision.classifier.enabled: + classify_content.delay(photo_id) return {'status': 'dispatched', 'photo_id': photo_id} @@ -206,6 +208,66 @@ def detect_objects(photo_id: str): return {'status': 'success', 'photo_id': photo_id, 'objects': len(detections)} +@shared_task(name='classify_content', queue='vision') +def classify_content(photo_id: str): + """Classify image content type (screenshot, document, artwork, etc.) + using CLIP zero-shot classification. Writes Tag(kind=content_type).""" + if not settings.vision.enabled or not settings.vision.classifier.enabled: + return {'status': 'skipped', 'reason': 'classifier disabled'} + + image = _load_thumb(photo_id, "medium") + if image is None: + return {'status': 'error', 'message': 'thumbnail not found'} + + from app.services.vision.registry import registry + classifier = registry.get_classifier() + results = classifier.classify(image) + + if not results: + logger.info("No confident classification for photo %s", photo_id) + return {'status': 'success', 'photo_id': photo_id, 'content_type': None} + + from app.models.tags import Tag, photo_tags + + source_name = "vision:clip_classifier" + best = results[0] + + session = _get_sync_session() + try: + # Wipe previous classification for this photo + session.execute( + delete(photo_tags).where( + photo_tags.c.photo_id == photo_id, + photo_tags.c.source == source_name, + ) + ) + + # Find or create content_type tag + tag = session.execute( + select(Tag).where(Tag.name == best.label, Tag.kind == 'content_type') + ).scalar_one_or_none() + + if not tag: + tag = Tag(name=best.label, kind='content_type', source=source_name) + session.add(tag) + session.flush() + + session.execute( + photo_tags.insert().values( + photo_id=photo_id, + tag_id=tag.id, + confidence=best.confidence, + source=source_name, + ) + ) + session.commit() + finally: + session.close() + + logger.info("Classified photo %s as '%s' (%.2f)", photo_id, best.label, best.confidence) + return {'status': 'success', 'photo_id': photo_id, 'content_type': best.label} + + def _load_original(photo_id: str) -> np.ndarray | None: """Load the original photo file as an RGB numpy array, resized to max 1280px on the longest edge for face detection.""" @@ -242,72 +304,26 @@ def _load_original(photo_id: str) -> np.ndarray | None: return None -def _iou(a: list[float], b: list[float]) -> float: - """Intersection-over-area of box a within box b (how much of a is inside b). - Boxes are [x1, y1, x2, y2] normalized 0-1.""" - x1 = max(a[0], b[0]) - y1 = max(a[1], b[1]) - x2 = min(a[2], b[2]) - y2 = min(a[3], b[3]) - inter = max(0, x2 - x1) * max(0, y2 - y1) - area_a = max(0, a[2] - a[0]) * max(0, a[3] - a[1]) - return inter / area_a if area_a > 0 else 0 - - -def _face_inside_person(face_bbox: list[float], person_bboxes: list[list[float]]) -> bool: - """Return True if the face bbox overlaps at least 50% with any YOLO - 'person' detection. Filters out faces on dogs, paintings, etc.""" - for pb in person_bboxes: - if _iou(face_bbox, pb) >= 0.5: - return True - return False - - @shared_task(name='extract_faces', queue='vision') def extract_faces(photo_id: str): - """Detect faces and store recognition embeddings. Only keeps faces - that overlap with a YOLO 'person' detection to filter out animal - and cartoon false positives.""" + """Detect faces and store recognition embeddings using InsightFace + (RetinaFace + ArcFace). No YOLO workaround needed — RetinaFace has + strong human-vs-non-human precision on its own.""" if not settings.vision.enabled or not settings.vision.faces.enabled: return {'status': 'skipped', 'reason': 'faces disabled'} - # Use original file for face detection — thumbnails are often too - # small (240px) for reliable face detection. image = _load_original(photo_id) if image is None: image = _load_thumb(photo_id, "large") if image is None: return {'status': 'error', 'message': 'no image available'} - # Step 1: run YOLO to find "person" bounding boxes from app.services.vision.registry import registry - detector = registry.get_detector() - thumb = _load_thumb(photo_id, "medium") - person_bboxes = [] - if thumb is not None: - detections = detector.detect(thumb) - person_bboxes = [d.bbox for d in detections if d.label == 'person'] - - # Step 2: run face detection face_proc = registry.get_face_processor() faces = face_proc.process(image) if not faces: logger.info("No faces detected for photo %s", photo_id) - return _save_faces(photo_id, []) - - # Step 3: filter — keep only faces inside a person bbox - if person_bboxes: - verified = [f for f in faces if _face_inside_person(f.bbox, person_bboxes)] - dropped = len(faces) - len(verified) - if dropped: - logger.info("Dropped %d non-person face(s) for photo %s", dropped, photo_id) - faces = verified - else: - # No person detected by YOLO → drop all face detections - # (no human body visible = likely false positives) - logger.info("No YOLO person detected, dropping %d face(s) for photo %s", len(faces), photo_id) - faces = [] return _save_faces(photo_id, faces) diff --git a/backend/requirements.txt b/backend/requirements.txt index 5e59f06..b48df03 100644 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -39,6 +39,7 @@ onnxruntime==1.18.1 open-clip-torch==2.24.0 # tokenizer + export helper; inference via ONNX rapidocr-onnxruntime==1.3.22 scikit-learn==1.4.0 # DBSCAN for face clustering +insightface>=0.7.3 # RetinaFace + ArcFace face detection/recognition numpy>=1.26.0,<2.0 # Utilities diff --git a/mulita.yml b/mulita.yml index 6f22fcf..334b0ff 100644 --- a/mulita.yml +++ b/mulita.yml @@ -44,6 +44,9 @@ vision: faces: enabled: true min_face_size: 40 - recognition_threshold: 0.6 - cluster_eps: 0.25 + recognition_threshold: 0.65 + cluster_eps: 0.5 + classifier: + enabled: true + min_confidence: 0.3 worker_concurrency: 2