handler.py

"""Hugging Face custom handler with a strict, privacy-preserving input contract."""

from __future__ import annotations

import base64
import binascii
import io
import math
import numbers
import re
import threading
import warnings
from pathlib import Path
from typing import Any, Mapping, Sequence

import numpy as np
import torch
from PIL import Image, ImageFile, ImageOps, UnidentifiedImageError
from torchvision.transforms import functional as vision_functional

from model import load_model

_URI_PREFIX = re.compile(r"^[A-Za-z][A-Za-z0-9+.-]*:")
_IMAGE_DECODE_LOCK = threading.Lock()


class InputValidationError(ValueError):
    """A client-visible validation failure with a sanitized field reference."""

    def __init__(self, message: str, *, field: str) -> None:
        super().__init__(message)
        self.message = message
        self.field = field


def _error_response(
    *,
    model_id: str,
    model_version: str,
    code: str,
    message: str,
    field: str | None = None,
) -> dict[str, Any]:
    error: dict[str, Any] = {"code": code, "message": message}
    if field is not None:
        error["field"] = field
    return {
        "error": error,
        "model_id": model_id,
        "model_version": model_version,
    }


def _check_image_header(
    image: Image.Image,
    *,
    image_config: Mapping[str, Any],
    field: str,
) -> None:
    if image.format not in set(image_config["allowed_formats"]):
        raise InputValidationError(
            "frame format must be JPEG or PNG",
            field=field,
        )
    if bool(getattr(image, "is_animated", False)) or int(
        getattr(image, "n_frames", 1)
    ) != 1:
        raise InputValidationError(
            "animated frames are not accepted",
            field=field,
        )
    width, height = image.size
    if width <= 0 or height <= 0:
        raise InputValidationError("frame dimensions are invalid", field=field)
    if width > int(image_config["max_width"]):
        raise InputValidationError("frame width exceeds the configured limit", field=field)
    if height > int(image_config["max_height"]):
        raise InputValidationError(
            "frame height exceeds the configured limit",
            field=field,
        )
    if width * height > int(image_config["max_pixels"]):
        raise InputValidationError(
            "frame pixel count exceeds the configured limit",
            field=field,
        )


def decode_image(
    value: Any,
    *,
    image_config: Mapping[str, Any],
    field: str,
) -> Image.Image:
    """Decode one raw base64 JPEG/PNG and return an EXIF-free RGB image."""

    if not isinstance(value, str):
        raise InputValidationError("frame must be a base64 string", field=field)
    if len(value) > int(image_config["max_encoded_bytes"]) + 2:
        raise InputValidationError(
            "encoded frame exceeds the configured byte limit",
            field=field,
        )
    encoded = value.strip()
    if not encoded:
        raise InputValidationError("frame must not be empty", field=field)
    if _URI_PREFIX.match(encoded):
        raise InputValidationError(
            "URLs and data URLs are not accepted; send raw base64 only",
            field=field,
        )
    try:
        encoded_bytes = encoded.encode("ascii")
    except UnicodeEncodeError as exc:
        raise InputValidationError("frame is not valid base64", field=field) from exc
    if len(encoded_bytes) > int(image_config["max_encoded_bytes"]):
        raise InputValidationError(
            "encoded frame exceeds the configured byte limit",
            field=field,
        )
    try:
        content = base64.b64decode(encoded_bytes, validate=True)
    except (binascii.Error, ValueError) as exc:
        raise InputValidationError("frame is not valid base64", field=field) from exc
    if len(content) > int(image_config["max_decoded_bytes"]):
        raise InputValidationError(
            "decoded frame exceeds the configured byte limit",
            field=field,
        )

    try:
        with _IMAGE_DECODE_LOCK:
            previous_truncated_setting = ImageFile.LOAD_TRUNCATED_IMAGES
            ImageFile.LOAD_TRUNCATED_IMAGES = False
            try:
                with warnings.catch_warnings():
                    warnings.simplefilter("error", Image.DecompressionBombWarning)
                    with Image.open(io.BytesIO(content)) as candidate:
                        _check_image_header(
                            candidate,
                            image_config=image_config,
                            field=field,
                        )
                        candidate.verify()
                    with Image.open(io.BytesIO(content)) as opened:
                        _check_image_header(
                            opened,
                            image_config=image_config,
                            field=field,
                        )
                        opened.load()
                        oriented = ImageOps.exif_transpose(opened)
                        _check_dimensions_after_orientation(
                            oriented,
                            image_config=image_config,
                            field=field,
                        )
                        converted = oriented.convert("RGB")
                        clean = Image.new("RGB", converted.size)
                        clean.paste(converted)
                        clean.info.clear()
                        return clean
            finally:
                ImageFile.LOAD_TRUNCATED_IMAGES = previous_truncated_setting
    except InputValidationError:
        raise
    except (
        Image.DecompressionBombError,
        Image.DecompressionBombWarning,
        UnidentifiedImageError,
        OSError,
        SyntaxError,
        ValueError,
    ) as exc:
        raise InputValidationError(
            "frame is not a complete, supported JPEG or PNG",
            field=field,
        ) from exc


def _check_dimensions_after_orientation(
    image: Image.Image,
    *,
    image_config: Mapping[str, Any],
    field: str,
) -> None:
    width, height = image.size
    if (
        width <= 0
        or height <= 0
        or width > int(image_config["max_width"])
        or height > int(image_config["max_height"])
        or width * height > int(image_config["max_pixels"])
    ):
        raise InputValidationError(
            "frame dimensions exceed the configured limits",
            field=field,
        )


def preprocess_image(
    image: Image.Image,
    *,
    preprocessing_config: Mapping[str, Any],
) -> torch.Tensor:
    """Apply the exact deterministic v0 evaluation transform."""

    width, height = (int(value) for value in preprocessing_config["output_size"])
    contained = ImageOps.contain(
        image,
        (width, height),
        method=Image.Resampling.BICUBIC,
    )
    canvas = Image.new(
        "RGB",
        (width, height),
        color=tuple(int(value) for value in preprocessing_config["letterbox_rgb"]),
    )
    canvas.paste(
        contained,
        ((width - contained.width) // 2, (height - contained.height) // 2),
    )
    tensor = vision_functional.pil_to_tensor(canvas).to(dtype=torch.float32)
    tensor.div_(float(preprocessing_config["pixel_scale"]))
    return vision_functional.normalize(
        tensor,
        mean=tuple(float(value) for value in preprocessing_config["image_mean"]),
        std=tuple(float(value) for value in preprocessing_config["image_std"]),
    )


def _require_real_number(value: Any, *, field: str) -> float:
    if isinstance(value, bool) or not isinstance(value, numbers.Real):
        raise InputValidationError("value must be a JSON number", field=field)
    parsed = float(value)
    if not math.isfinite(parsed):
        raise InputValidationError("value must be finite", field=field)
    return parsed


def validate_telemetry(
    payload: Any,
    *,
    telemetry_config: Mapping[str, Any],
) -> np.ndarray:
    """Require and bounds-check the exact twelve-feature telemetry object."""

    if not isinstance(payload, dict):
        raise InputValidationError(
            "telemetry must be an object",
            field="inputs.telemetry",
        )
    feature_order = tuple(telemetry_config["feature_order"])
    expected = set(feature_order)
    actual = set(payload)
    missing = sorted(expected - actual)
    unknown = sorted(actual - expected)
    if missing:
        raise InputValidationError(
            f"missing required telemetry field: {missing[0]}",
            field=f"inputs.telemetry.{missing[0]}",
        )
    if unknown:
        raise InputValidationError(
            f"unknown telemetry field: {unknown[0]}",
            field=f"inputs.telemetry.{unknown[0]}",
        )

    values: dict[str, float] = {}
    for name in feature_order:
        field = f"inputs.telemetry.{name}"
        value = _require_real_number(payload[name], field=field)
        definition = telemetry_config["fields"][name]
        minimum = float(definition["minimum"])
        maximum = float(definition["maximum"])
        minimum_ok = value >= minimum
        maximum_ok = (
            value <= maximum
            if definition.get("maximum_inclusive", True)
            else value < maximum
        )
        if not minimum_ok or not maximum_ok:
            closing = "]" if definition.get("maximum_inclusive", True) else ")"
            raise InputValidationError(
                f"value must be in [{minimum}, {maximum}{closing}",
                field=field,
            )
        values[name] = value

    tolerance = float(telemetry_config["unit_circle_norm_tolerance"])
    for sine_name, cosine_name in telemetry_config["unit_circle_pairs"]:
        norm = math.hypot(values[sine_name], values[cosine_name])
        if abs(norm - 1.0) > tolerance:
            raise InputValidationError(
                f"{sine_name} and {cosine_name} must encode a unit-circle pair",
                field=f"inputs.telemetry.{sine_name}",
            )
    return np.asarray([values[name] for name in feature_order], dtype=np.float64)


def normalize_telemetry(
    values: np.ndarray,
    *,
    telemetry_config: Mapping[str, Any],
) -> torch.Tensor:
    """Normalize in float64 exactly as documented, then cast once to float32."""

    normalization = telemetry_config["normalization"]
    mean = np.asarray(normalization["mean"], dtype=np.float64)
    std = np.asarray(normalization["std"], dtype=np.float64)
    if values.shape != mean.shape or mean.shape != std.shape:
        raise RuntimeError("telemetry normalization shape mismatch")
    normalized = (values.astype(np.float64, copy=False) - mean) / std
    return torch.from_numpy(normalized.astype(np.float32, copy=False))


def _validate_timestamp_sequence(value: Any, *, frame_count: int) -> None:
    if not isinstance(value, list) or len(value) != frame_count:
        raise InputValidationError(
            f"frame_timestamps must contain exactly {frame_count} numbers",
            field="inputs.frame_timestamps",
        )
    timestamps = [
        _require_real_number(item, field=f"inputs.frame_timestamps[{index}]")
        for index, item in enumerate(value)
    ]
    if any(later <= earlier for earlier, later in zip(timestamps, timestamps[1:])):
        raise InputValidationError(
            "frame_timestamps must be strictly increasing",
            field="inputs.frame_timestamps",
        )


class EndpointHandler:
    """Hugging Face Inference Endpoints custom handler."""

    def __init__(self, path: str = "") -> None:
        repository_path = Path(path) if path else Path(__file__).resolve().parent
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model, self.config = load_model(repository_path, device=self.device)
        self.model.eval()
        self.model_id = str(self.config["model_id"])
        self.model_version = str(self.config["model_version"])
        self.labels = tuple(self.config["labels"])
        self.temperature = float(self.config["calibration"]["temperature"])
        self.warnings = tuple(str(item) for item in self.config["warnings"])

    def __call__(self, data: dict[str, Any]) -> dict[str, Any]:
        try:
            return self._predict(data)
        except InputValidationError as exc:
            return _error_response(
                model_id=self.model_id,
                model_version=self.model_version,
                code="invalid_request",
                message=exc.message,
                field=exc.field,
            )
        except Exception:
            return _error_response(
                model_id=self.model_id,
                model_version=self.model_version,
                code="internal_error",
                message="inference could not be completed",
            )

    def _predict(self, data: Any) -> dict[str, Any]:
        if not isinstance(data, dict):
            raise InputValidationError(
                "request body must be an object",
                field="request",
            )
        if set(data) != {"inputs"}:
            raise InputValidationError(
                "request body must contain only the inputs object",
                field="request",
            )
        inputs = data["inputs"]
        if not isinstance(inputs, dict):
            raise InputValidationError("inputs must be an object", field="inputs")
        allowed_input_keys = {
            "frames",
            "frame_timestamps",
            "telemetry",
            "abstention_threshold",
        }
        unknown_input_keys = sorted(set(inputs) - allowed_input_keys)
        if unknown_input_keys:
            raise InputValidationError(
                f"unknown input field: {unknown_input_keys[0]}",
                field=f"inputs.{unknown_input_keys[0]}",
            )

        frame_count = int(self.config["input"]["frame_count"])
        encoded_frames = inputs.get("frames")
        if not isinstance(encoded_frames, list) or len(encoded_frames) != frame_count:
            raise InputValidationError(
                f"frames must contain exactly {frame_count} base64 images",
                field="inputs.frames",
            )
        if "frame_timestamps" in inputs:
            _validate_timestamp_sequence(
                inputs["frame_timestamps"],
                frame_count=frame_count,
            )

        threshold_value = inputs.get(
            "abstention_threshold",
            self.config["abstention"]["default_threshold"],
        )
        threshold = _require_real_number(
            threshold_value,
            field="inputs.abstention_threshold",
        )
        if threshold < 0.0 or threshold > 1.0:
            raise InputValidationError(
                "abstention_threshold must be in [0, 1]",
                field="inputs.abstention_threshold",
            )

        images = [
            decode_image(
                encoded,
                image_config=self.config["input"]["image"],
                field=f"inputs.frames[{index}]",
            )
            for index, encoded in enumerate(encoded_frames)
        ]
        image_batch = torch.stack(
            [
                preprocess_image(
                    image,
                    preprocessing_config=self.config["preprocessing"],
                )
                for image in images
            ],
            dim=0,
        ).unsqueeze(0)
        telemetry_values = validate_telemetry(
            inputs.get("telemetry"),
            telemetry_config=self.config["telemetry"],
        )
        telemetry = normalize_telemetry(
            telemetry_values,
            telemetry_config=self.config["telemetry"],
        ).unsqueeze(0)

        image_batch = image_batch.to(self.device)
        telemetry = telemetry.to(self.device)
        self.model.eval()
        with torch.inference_mode():
            with torch.autocast(
                device_type=self.device.type,
                dtype=torch.float16,
                enabled=self.device.type == "cuda",
            ):
                logits = self.model(image_batch, telemetry)
            probabilities_tensor = torch.softmax(
                logits.float() / self.temperature,
                dim=-1,
            )[0].cpu()

        probabilities_array = probabilities_tensor.numpy()
        if (
            probabilities_array.shape != (len(self.labels),)
            or not np.isfinite(probabilities_array).all()
            or not math.isclose(
                float(probabilities_array.sum()),
                1.0,
                rel_tol=0.0,
                abs_tol=1e-5,
            )
        ):
            raise RuntimeError("model returned invalid probabilities")
        best_index = int(np.argmax(probabilities_array))
        confidence = float(probabilities_array[best_index])
        abstained = confidence < threshold
        return {
            "model_id": self.model_id,
            "model_version": self.model_version,
            "predicted_label": None if abstained else self.labels[best_index],
            "probabilities": {
                label: float(probabilities_array[index])
                for index, label in enumerate(self.labels)
            },
            "confidence": confidence,
            "abstained": abstained,
            "abstention_threshold": threshold,
            "input_frame_count": frame_count,
            "warnings": list(self.warnings),
        }


__all__ = [
    "EndpointHandler",
    "InputValidationError",
    "decode_image",
    "normalize_telemetry",
    "preprocess_image",
    "validate_telemetry",
]