inference.py

"""
Inference script for AFR-DFV-v1 (African tropical forests, DeepForestVision)

Model: DeepForestVision v1
Input: 224x224 RGB, ImageNet-normalised
Framework: PyTorch (HuggingFace transformers, DINOv2-large)
Classes: 34 African tropical forest species
Developer: MNHN-OFVI

Ported from AddaxAI's legacy classify_detections.py (dfv-v1), with one
deliberate change: the architecture is built from a local config.json
rather than fetched from the Hub. See below.

Files expected in the model directory:
    - DFV.pt        the fine-tuned weights
    - config.json   facebook/dinov2-large's config, copied verbatim

Author: Peter van Lunteren
"""

from __future__ import annotations

from collections import OrderedDict
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
from PIL import Image, ImageFile
from torch import tensor
from torchvision.transforms import InterpolationMode, transforms
from transformers import AutoConfig, AutoModelForImageClassification

# Don't freak out over truncated images
ImageFile.LOAD_TRUNCATED_IMAGES = True

CROP_SIZE = 224
RESIZE_SIZE = 256
BACKBONE = "dinov2_large"

# Class order is the model's output order and must not be reordered.
CLASS_NAMES = [
    'aardvark', 'baboon', 'honey badger', 'bird',
    'black-and-white colobus', 'blue duiker', 'blue monkey',
    'african buffalo', 'bushbuck', 'bushpig', 'chimpanzee', 'civet_genet',
    'elephant', 'galago_potto', 'african golden cat', 'gorilla',
    'guineafowl', 'hyrax', 'side-striped jackal', 'leopard',
    "l'hoest's monkey", 'mandrill', 'mongoose', 'monkey', 'pangolin',
    'porcupine', 'red colobus_red-capped mangabey', 'red duiker', 'rodent',
    'serval', 'spotted hyena', 'squirrel', 'water chevrotain',
    'yellow-backed duiker'
]


class _Model(nn.Module):
    """
    DINOv2-large with a 34-way head.

    The legacy adapter builds this with
    `AutoModelForImageClassification.from_pretrained('facebook/dinov2-large')`,
    which downloads 1.2GB of pretrained weights from the Hub and then
    throws every one of them away: `load_weights` below does a strict
    `load_state_dict`, and DFV.pt carries all 441 keys the model has, so
    nothing of the download survives.

    Building from a local config instead is therefore identical in
    result, needs no network at inference time, and does not make an
    offline machine fail. Verified: from_config + DFV.pt loads strict with
    missing=0, unexpected=0.
    """

    def __init__(self, model_dir: Path) -> None:
        super().__init__()
        config = AutoConfig.from_pretrained(str(model_dir))
        self.base_model = AutoModelForImageClassification.from_config(config)
        self.base_model.classifier = nn.Linear(
            self.base_model.classifier.in_features, len(CLASS_NAMES)
        )
        self.backbone = BACKBONE
        self.nbclasses = len(CLASS_NAMES)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.base_model(x)

    def load_weights(self, path: Path, map_location) -> None:
        weights = torch.load(path, map_location=map_location)
        # The checkpoint was saved from a module that held the backbone
        # directly, so its keys need the base_model prefix.
        renamed = OrderedDict(
            (
                key.replace("dinov2", "base_model.dinov2").replace(
                    "classifier", "base_model.classifier"
                ),
                value,
            )
            for key, value in weights.items()
        )
        self.load_state_dict(renamed)


class ModelInference:
    """DeepForestVision African tropical forest classifier."""

    def __init__(self, model_dir: Path, model_path: Path) -> None:
        self.model_dir = Path(model_dir)
        self.model_path = Path(model_path)
        self.model: _Model | None = None
        self.device: torch.device | None = None
        # Match DeepForestVision's own inference (DFV.py), which preprocesses
        # with AutoImageProcessor.from_pretrained('facebook/dinov2-large'):
        # resize the shortest edge to 256 (bicubic), centre-crop 224,
        # rescale to [0,1], ImageNet-normalize. Replicated with torchvision
        # so no processor config is fetched at inference. The earlier port
        # stretched straight to 224x224 (DeepFaune boilerplate), distorting
        # the aspect ratio DINOv2 is sensitive to, and cropped a
        # square-by-expand box rather than the plain box the processor
        # expects.
        self.preprocess = transforms.Compose([
            transforms.Resize(
                size=RESIZE_SIZE,
                interpolation=InterpolationMode.BICUBIC,
                antialias=True,
            ),
            transforms.CenterCrop(CROP_SIZE),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=tensor([0.485, 0.456, 0.406]),
                std=tensor([0.229, 0.224, 0.225]),
            ),
        ])

    # ------------------------------------------------------------------
    # Required interface
    # ------------------------------------------------------------------

    def check_gpu(self) -> bool:
        if torch.cuda.is_available():
            return True
        try:
            return bool(torch.backends.mps.is_built() and torch.backends.mps.is_available())
        except AttributeError:
            return False

    def load_model(self) -> None:
        # Matches the legacy adapter's device order: CUDA first, then MPS.
        if torch.cuda.is_available():
            self.device = torch.device("cuda")
        else:
            try:
                mps = torch.backends.mps.is_built() and torch.backends.mps.is_available()
            except AttributeError:
                mps = False
            self.device = torch.device("mps" if mps else "cpu")

        model = _Model(self.model_dir)
        model.load_weights(self.model_path, self.device)
        self.model = model.to(self.device).eval()

    def get_crop(
        self, image: Image.Image, bbox: tuple[float, float, float, float]
    ) -> Image.Image:
        """
        Plain box crop, matching DeepForestVision's functions.py, which
        crops with supervision.crop_image (a plain xyxy box: no squaring,
        no pad). The AutoImageProcessor pipeline above does its own
        resize/centre-crop, so the crop handed to it must be the raw box.
        The earlier port squared the box by expanding the shorter side,
        which fed a different region and aspect than upstream.
        """
        width, height = image.size
        left = max(0, int(round(bbox[0] * width)))
        top = max(0, int(round(bbox[1] * height)))
        right = min(width, int(round((bbox[0] + bbox[2]) * width)))
        bottom = min(height, int(round((bbox[1] + bbox[3]) * height)))
        if right <= left or bottom <= top:
            raise ValueError(f"Invalid crop dimensions: ({left},{top}) to ({right},{bottom})")
        return image.crop((left, top, right, bottom))

    def get_classification(self, crop: Image.Image) -> list[list]:
        """Per-crop inference. Returns [[name, prob], ...] for all classes."""
        probs = self._forward(np.stack([self.get_tensor(crop)]))[0]
        return [[CLASS_NAMES[i], float(probs[i])] for i in range(len(probs))]

    def get_class_names(self) -> dict[str, str]:
        """1-indexed mapping {id: class_name} for the output JSON."""
        return {str(i + 1): name for i, name in enumerate(CLASS_NAMES)}

    # ------------------------------------------------------------------
    # Optional batch interface
    # ------------------------------------------------------------------

    def get_tensor(self, crop: Image.Image) -> np.ndarray:
        if crop.mode != "RGB":
            crop = crop.convert("RGB")
        return self.preprocess(crop).numpy()

    def classify_batch(self, batch: np.ndarray) -> list[list[list]]:
        probs = self._forward(batch)
        return [
            [[CLASS_NAMES[j], float(p[j])] for j in range(len(p))]
            for p in probs
        ]

    # ------------------------------------------------------------------
    # Internals
    # ------------------------------------------------------------------

    def _forward(self, batch: np.ndarray) -> np.ndarray:
        assert self.model is not None
        tensor_in = torch.from_numpy(batch).to(self.device)
        with torch.no_grad():
            return self.model(tensor_in).logits.softmax(dim=1).cpu().numpy()