blitzid 0.1.0__py3-none-any.whl

blitzid/__init__.py

@@ -0,0 +1,49 @@
+"""BlitzID — DNN-based face detection.
+
+Public API
+----------
+- ``FaceDetectorDNN`` — the SCRFD-based face detector.
+- ``Face`` — detection record (bbox, confidence, landmarks).
+- ``DetectionMetrics`` — metrics dataclass returned by ``detect_face_with_metrics``.
+- ``RapidOCRReader`` — the RapidOCR-based text reader (``ocr`` extra).
+- ``OCRText`` — recognized text line (bbox, text, confidence).
+- ``MRZReader`` — the ICAO 9303 machine-readable zone reader (``ocr`` extra).
+- ``MRZRecord`` — parsed MRZ fields.
+- ``DocumentCropper`` — document localization, perspective crop, and QC.
+- ``QualityReport`` — document crop quality-check verdict.
+- ``BlitzIDError``, ``ModelError``, ``ImageError``, ``MRZError`` — exception hierarchy.
+"""
+
+from .exceptions import BlitzIDError, ImageError, ModelError, MRZError
+from .face._face import DetectionMetrics, Face
+from .face.detector import FaceDetectorDNN
+from .reading.document import DocumentCropper, QualityReport
+from .reading.mrz import MRZReader, MRZRecord
+from .reading.ocr import OCRText, RapidOCRReader
+
+FaceDetectorError = BlitzIDError
+ModelDownloadError = ModelError
+ImageLoadError = ImageError
+ImageProcessingError = ImageError
+InvalidParameterError = BlitzIDError
+
+__all__ = [
+    "BlitzIDError",
+    "DetectionMetrics",
+    "DocumentCropper",
+    "Face",
+    "FaceDetectorDNN",
+    "FaceDetectorError",
+    "ImageError",
+    "ImageLoadError",
+    "ImageProcessingError",
+    "InvalidParameterError",
+    "MRZError",
+    "MRZReader",
+    "MRZRecord",
+    "ModelDownloadError",
+    "ModelError",
+    "OCRText",
+    "QualityReport",
+    "RapidOCRReader",
+]

blitzid/_image.py

@@ -0,0 +1,180 @@
+"""Image loading and validation for face detection.
+
+Handles loading from file paths, numpy arrays, and PIL Images.
+Normalizes all inputs to 3-channel BGR ``np.ndarray``.
+"""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, TypeGuard, Union
+
+import cv2
+import numpy as np
+
+from .exceptions import ImageError
+
+if TYPE_CHECKING:
+    from PIL.Image import Image as PILImageType
+
+ImageInput = Union[Path, str, np.ndarray, "PILImageType"]
+
+MIN_DIMENSION = 10
+MAX_DIMENSION = 10000
+
+
+def _is_pil_image(value: object) -> TypeGuard[PILImageType]:
+    """Return whether *value* is a PIL Image (PIL is an optional dependency)."""
+    try:
+        from PIL.Image import Image as PILImage
+    except ImportError:  # pragma: no cover
+        return False
+    return isinstance(value, PILImage)
+
+
+def load_image(
+    image_input: ImageInput,
+    logger: logging.Logger,
+) -> np.ndarray:
+    """Load and validate image from any source.
+
+    Args:
+        image_input: Path, numpy array, or PIL Image.
+        logger: Logger instance for debug messages.
+
+    Returns:
+        Image as numpy array in BGR format.
+
+    Raises:
+        ImageError: If image cannot be loaded or is invalid.
+    """
+    try:
+        if isinstance(image_input, np.ndarray):
+            img = _load_from_array(image_input, logger)
+        elif _is_pil_image(image_input):
+            img = _load_from_pil(image_input)
+        elif isinstance(image_input, (Path, str)):
+            img = _load_from_path(Path(image_input))
+        else:
+            raise ImageError(f"Unsupported image_input type: {type(image_input)}")
+
+        _validate_image(img)
+        img = _normalize_channels(img)
+        return img
+
+    except (OSError, ValueError, TypeError, cv2.error) as e:
+        raise ImageError(f"Unexpected error loading image: {e}") from e
+
+
+def _load_from_path(path: Path) -> np.ndarray:
+    """Load image from file path."""
+    if not path.exists():
+        raise ImageError(f"Image not found: {path}")
+    if not path.is_file():
+        raise ImageError(f"Path is not a file: {path}")
+    if path.stat().st_size == 0:
+        raise ImageError(f"Image file is empty: {path}")
+
+    img = cv2.imread(str(path))
+    if img is None:
+        raise ImageError(
+            f"Could not read image: {path}. "
+            "File may be corrupted or in an unsupported format."
+        )
+    return img
+
+
+def _load_from_array(array: np.ndarray, logger: logging.Logger) -> np.ndarray:
+    """Validate numpy array."""
+    if array.size == 0:
+        raise ImageError("Image is empty (size = 0)")
+
+    if array.dtype == np.float64:
+        logger.debug("Converting float64 image array to float32")
+        array = array.astype(np.float32)
+    elif array.dtype not in (np.uint8, np.float32):
+        raise ImageError(
+            f"Unsupported image dtype: {array.dtype}. "
+            "Expected uint8 or float32/float64."
+        )
+
+    return np.ascontiguousarray(array)
+
+
+def _load_from_pil(pil_image: PILImageType) -> np.ndarray:
+    """Convert PIL Image to numpy array (BGR)."""
+    img_array = np.array(pil_image)
+    if img_array.size == 0:
+        raise ImageError("PIL Image is empty")
+
+    if len(img_array.shape) == 3 and img_array.shape[2] == 3:
+        img = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
+    elif len(img_array.shape) == 2:
+        img = cv2.cvtColor(img_array, cv2.COLOR_GRAY2BGR)
+    else:
+        img = img_array
+
+    return img
+
+
+def _validate_image(img: np.ndarray) -> None:
+    """Validate image dimensions and properties."""
+    if img.size == 0:
+        raise ImageError("Image is empty (size = 0)")
+    if len(img.shape) < 2:
+        raise ImageError(
+            f"Image must have at least 2 dimensions, got shape {img.shape}"
+        )
+
+    h, w = img.shape[:2]
+    if h < MIN_DIMENSION or w < MIN_DIMENSION:
+        raise ImageError(
+            f"Image too small: {w}x{h}. Minimum size is {MIN_DIMENSION}x{MIN_DIMENSION}"
+        )
+    if h > MAX_DIMENSION or w > MAX_DIMENSION:
+        raise ImageError(
+            f"Image too large: {w}x{h}. Maximum size is {MAX_DIMENSION}x{MAX_DIMENSION}"
+        )
+
+
+def _normalize_channels(img: np.ndarray) -> np.ndarray:
+    """Convert to 3-channel BGR."""
+    if len(img.shape) == 2:
+        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+    elif len(img.shape) == 3:
+        channels = img.shape[2]
+        if channels == 1:
+            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+        elif channels == 4:
+            img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGR)
+        elif channels != 3:
+            raise ImageError(
+                f"Unsupported number of channels: {channels}. "
+                "Expected 1, 3, or 4 channels."
+            )
+    return img
+
+
+def crop_with_padding(
+    img: np.ndarray,
+    bbox: tuple[int, int, int, int],
+    padding: float,
+) -> np.ndarray:
+    """Crop an ``(x, y, w, h)`` box with relative padding, clipped to the image.
+
+    Args:
+        img: BGR image array.
+        bbox: The box as ``(x, y, w, h)`` in image pixels.
+        padding: Relative padding on each side (fraction of w/h).
+
+    Returns:
+        The cropped image region.
+    """
+    height, width = img.shape[:2]
+    x, y, box_w, box_h = bbox
+    pad_w, pad_h = int(box_w * padding), int(box_h * padding)
+    x1, y1 = max(0, x - pad_w), max(0, y - pad_h)
+    x2 = min(width, x + box_w + pad_w)
+    y2 = min(height, y + box_h + pad_h)
+    return img[y1:y2, x1:x2]

blitzid/_models.py

@@ -0,0 +1,142 @@
+"""Model management for the SCRFD face detection network.
+
+Downloads the SCRFD-2.5G ONNX model (InsightFace ``buffalo_m`` detection
+weights) and creates a CPU inference session. Weights live in the
+default models dir: ``BLITZID_MODELS_DIR`` when set, else the
+``platformdirs`` user cache.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import urllib.error
+import urllib.request
+from pathlib import Path
+
+import onnxruntime as ort  # type: ignore[import-untyped]
+from platformdirs import user_cache_dir
+
+from .exceptions import ModelError
+
+_DOWNLOAD_CHUNK_BYTES = 64 * 1024
+
+MODELS_DIR_ENV = "BLITZID_MODELS_DIR"
+
+
+def default_model_dir(subdir: str | None = None) -> Path:
+    """Return the default directory for storing model weights.
+
+    Honors ``BLITZID_MODELS_DIR``; otherwise uses
+    ``platformdirs.user_cache_dir("blitzid")``.
+
+    Args:
+        subdir: Optional subdirectory.
+
+    Returns:
+        Resolved :class:`~pathlib.Path` that is guaranteed to exist.
+    """
+    env_dir = os.environ.get(MODELS_DIR_ENV)
+    base = Path(env_dir) if env_dir else Path(user_cache_dir("blitzid")) / "models"
+
+    if subdir:
+        base = base / subdir
+
+    base.mkdir(parents=True, exist_ok=True)
+    return base
+
+
+class ModelManager:
+    """Manages the SCRFD ONNX model file and inference session."""
+
+    MODEL_URL = (
+        "https://huggingface.co/immich-app/buffalo_m/resolve/main/detection/model.onnx"
+    )
+    MODEL_FILENAME = "scrfd_2.5g.onnx"
+
+    def __init__(
+        self,
+        model_dir: Path,
+        logger: logging.Logger,
+        allow_downloads: bool = True,
+        download_timeout: float = 30.0,
+    ):
+        self.model_dir = Path(model_dir)
+        self.model_dir.mkdir(exist_ok=True, parents=True)
+        self.logger = logger
+        self.allow_downloads = allow_downloads
+        self.download_timeout = download_timeout
+
+        self.model_path = self.model_dir / self.MODEL_FILENAME
+
+    def ensure_model_exists(self) -> None:
+        """Ensure the model file exists, downloading if allowed.
+
+        Raises:
+            ModelError: If the model is missing and downloads are disabled.
+        """
+        if self.model_path.exists():
+            return
+
+        if not self.allow_downloads:
+            raise ModelError(
+                f"Missing model file: {self.model_path.name}. "
+                "Downloads are disabled; set allow_downloads=True."
+            )
+
+        self.logger.info("SCRFD model not found. Downloading...")
+        self._download_model()
+
+    def load_session(self) -> ort.InferenceSession:
+        """Load the SCRFD model as a CPU inference session.
+
+        Raises:
+            ModelError: If the model cannot be loaded.
+        """
+        self.ensure_model_exists()
+
+        self.logger.info("Loading SCRFD face detection model...")
+        try:
+            return ort.InferenceSession(
+                str(self.model_path), providers=["CPUExecutionProvider"]
+            )
+        except Exception as e:
+            raise ModelError(f"Failed to load model {self.model_path.name}: {e}") from e
+
+    def _download_model(self) -> None:
+        """Download the model file.
+
+        Streams to a ``.tmp`` file and renames into place on success, so
+        a partial download never masquerades as a valid model.
+
+        Raises:
+            ModelError: If the download fails.
+        """
+        path = self.model_path
+        url = self.MODEL_URL
+        headers = {"User-Agent": "Mozilla/5.0 (compatible; blitzid face detector)"}
+
+        self.logger.info("Downloading %s...", path.name)
+        tmp_path = path.with_suffix(path.suffix + ".tmp")
+        try:
+            req = urllib.request.Request(url, headers=headers)
+            with (
+                urllib.request.urlopen(req, timeout=self.download_timeout) as response,
+                open(tmp_path, "wb") as out_file,
+            ):
+                while True:
+                    chunk = response.read(_DOWNLOAD_CHUNK_BYTES)
+                    if not chunk:
+                        break
+                    out_file.write(chunk)
+
+            tmp_path.replace(path)
+            self.logger.info(
+                "%s downloaded (%.1f KB)", path.name, path.stat().st_size / 1024
+            )
+        except (urllib.error.URLError, TimeoutError, OSError, ValueError) as e:
+            tmp_path.unlink(missing_ok=True)
+            raise ModelError(
+                f"Error downloading {path.name} from {url} "
+                f"(timeout={self.download_timeout}s): {e}"
+            ) from e

blitzid/api/__init__.py

@@ -0,0 +1,11 @@
+"""blitzid HTTP API — exposes ``app`` (``uvicorn blitzid.api:app``).
+
+This is the one deliberate exception to the "no re-exports in
+subpackage ``__init__``" rule: the ASGI entry point must be importable
+as ``blitzid.api:app``. Requires the ``api`` extra
+(``pip install blitzid[api]``).
+"""
+
+from ._app import app
+
+__all__ = ["app"]

blitzid/api/_analyze.py

@@ -0,0 +1,129 @@
+"""Analysis dispatch for /analyze jobs: decode, run, build sections."""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Callable, Mapping
+from dataclasses import asdict, dataclass, field
+from threading import Lock
+from types import MappingProxyType
+from typing import Any
+
+import numpy as np
+
+from .. import Face, FaceDetectorDNN, MRZReader, RapidOCRReader
+from .._image import crop_with_padding
+from ..exceptions import BlitzIDError, ModelError
+from ._upload import decode_image_bytes, dims_within_bounds, encode_jpeg_b64
+
+FACE_CROP_PADDING = 0.2
+
+
+@dataclass
+class Engines:
+    """Shared analysis engine singletons with one lock per engine.
+
+    The locks serialize access to engines that are not thread-safe:
+    the cache-less detector and the RapidOCR pipeline shared by the
+    OCR and MRZ sections. Fields are None when an engine is
+    unavailable (missing ``ocr`` extra, missing weights).
+    """
+
+    detector: FaceDetectorDNN | None
+    ocr_reader: RapidOCRReader | None
+    mrz_reader: MRZReader | None
+    detector_lock: Lock = field(default_factory=Lock)
+    ocr_lock: Lock = field(default_factory=Lock)
+
+
+def run_job(image: bytes, types: list[str], engines: Engines) -> dict[str, Any]:
+    """Run the requested analyses and build the result dict.
+
+    Per-section failures (``BlitzIDError`` subclasses) are captured as
+    section errors; other sections still return.
+
+    Returns:
+        ``{"state": "done", "<type>": <section>, ...}`` — duplicates
+        deduped, one section per requested type.
+    """
+    img = decode_image_bytes(image)
+    if img is None or not dims_within_bounds(img):
+        return {"state": "failed", "error": "stored image bytes failed validation"}
+    sections: dict[str, Any] = {"state": "done"}
+    for analysis_type in dict.fromkeys(types):
+        sections[analysis_type] = _run_section(analysis_type, img, engines)
+    return sections
+
+
+def _run_section(
+    analysis_type: str, img: np.ndarray, engines: Engines
+) -> dict[str, Any]:
+    """Run one analysis, capturing BlitzIDError as a section error."""
+    start = time.perf_counter()
+    try:
+        body = _SECTION_RUNNERS[analysis_type](img, engines)
+    except BlitzIDError as e:
+        body = {"error": str(e)}
+    body["processing_time_ms"] = round((time.perf_counter() - start) * 1000, 1)
+    return body
+
+
+def _face_item(img: np.ndarray, face: Face) -> dict[str, Any]:
+    """Build one face result item (bbox, confidence, landmarks, crop)."""
+    crop = crop_with_padding(img, face.bbox, FACE_CROP_PADDING)
+    return {
+        "bbox": list(face.bbox),
+        "confidence": face.confidence,
+        "landmarks": [[x, y] for x, y in face.landmarks],
+        "crop_base64": encode_jpeg_b64(crop),
+    }
+
+
+def _run_face(img: np.ndarray, engines: Engines) -> dict[str, Any]:
+    """Detect faces and build the face section."""
+    detector = engines.detector
+    if detector is None:
+        raise ModelError("face engine unavailable (SCRFD weights not loaded)")
+    with engines.detector_lock:
+        faces = detector.detect_face_landmarks(img)
+    return {"faces": [_face_item(img, face) for face in faces]}
+
+
+def _run_ocr(img: np.ndarray, engines: Engines) -> dict[str, Any]:
+    """Read text lines and build the ocr section."""
+    reader = engines.ocr_reader
+    if reader is None:
+        raise ModelError("ocr analysis requires the blitzid[ocr] extra")
+    with engines.ocr_lock:
+        texts = reader.read(img)
+    return {
+        "lines": [
+            {
+                "bbox": list(text.bbox),
+                "text": text.text,
+                "confidence": text.confidence,
+            }
+            for text in texts
+        ]
+    }
+
+
+def _run_mrz(img: np.ndarray, engines: Engines) -> dict[str, Any]:
+    """Read the MRZ and build the mrz section (shares the OCR engine)."""
+    reader = engines.mrz_reader
+    if reader is None:
+        raise ModelError("mrz analysis requires the blitzid[ocr] extra")
+    with engines.ocr_lock:
+        record = reader.read(img)
+    return {"record": asdict(record)}
+
+
+_SECTION_RUNNERS: Mapping[str, Callable[[np.ndarray, Engines], dict[str, Any]]] = (
+    MappingProxyType(
+        {
+            "face": _run_face,
+            "ocr": _run_ocr,
+            "mrz": _run_mrz,
+        }
+    )
+)