pipeline/feature_extraction.py

"""
Feature extraction — uses cores.vision for cropping + cores.face for
box conversions.  No duplicated crop logic.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import List, Optional

import numpy as np
from loguru import logger

from cores.vision import BBox, crop_region
from providers.base import Provider, ProviderResult


@dataclass
class FaceCrop:
    image: np.ndarray
    box: dict
    confidence: float = 1.0
    detector: str = ""


@dataclass
class PipelineOutput:
    """The normalized payload the orchestrator consumes."""
    image: np.ndarray
    image_hash: str
    width: int
    height: int
    source: str
    original_bytes: Optional[bytes] = None
    original_format: Optional[str] = None
    face_crops: List[FaceCrop] = field(default_factory=list)
    primary_detector: str = ""
    gallery: Optional[dict] = None
    scrape_url: Optional[str] = None

    @property
    def num_faces(self) -> int:
        return len(self.face_crops)


class FeatureExtractor:
    """Uses a detection provider to pre-extract face crops."""

    def __init__(self, detector: Optional[Provider] = None) -> None:
        self._detector = detector

    def set_detector(self, provider: Provider) -> None:
        self._detector = provider

    def extract(
        self,
        image: np.ndarray,
        image_hash: str,
        width: int,
        height: int,
        source: str,
        original_bytes: Optional[bytes] = None,
        original_format: Optional[str] = None,
    ) -> PipelineOutput:
        crops: List[FaceCrop] = []
        detector_name = ""

        if self._detector is not None and self._detector.is_available():
            try:
                result: ProviderResult = self._detector.execute(image)
                if result.success and result.normalized.get("boxes"):
                    detector_name = result.provider
                    boxes = result.normalized["boxes"]
                    confs = result.normalized.get("confidences", [1.0] * len(boxes))
                    for box_dict, conf in zip(boxes, confs):
                        bbox = BBox(box_dict["x"], box_dict["y"],
                                    box_dict["w"], box_dict["h"])
                        crop = crop_region(image, bbox, margin=0.2)
                        crops.append(FaceCrop(
                            image=crop, box=box_dict,
                            confidence=float(conf), detector=detector_name,
                        ))
            except Exception as e:
                logger.warning(f"Feature extraction failed: {e}")
        else:
            logger.debug("No detector available; pipeline output will have 0 face crops.")

        return PipelineOutput(
            image=image, image_hash=image_hash,
            width=width, height=height, source=source,
            original_bytes=original_bytes, original_format=original_format,
            face_crops=crops, primary_detector=detector_name,
        )