inference.py

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

Model: DeepForestVision v2
Input: 224x224 RGB, ImageNet-normalised
Framework: PyTorch, DINOv3 ViT-B/16 backbone (vendored, torch.hub source="local")
Classes: 61 African tropical forest species / groups (read from the checkpoint)
Developer: MNHN-OFVI (Hugo Magaldi, One Forest Vision initiative)

Ported from DeepForestVisionV2's own classification.py, adapted to AddaxAI's
ModelInference interface (v1's inference.py was the structural template). The
backbone architecture is vendored next to the weights (hubconf.py + dinov3/)
and loaded via torch.hub.load(source="local"), the same way AddaxAI's DINOv2
embedding models load their architecture.

One deliberate change vs upstream: the backbone is built with pretrained=False
(no network fetch, no Meta pretrained weights). The strict load of the
checkpoint below supplies every weight, so the architecture-only build is
identical in result.

Files expected in the model directory:
    - DeepForestVisionV2.pth   checkpoint {"labels", "model_state_dict"}
    - hubconf.py + dinov3/      vendored DINOv3 architecture source

Author: Peter van Lunteren
"""

from __future__ import annotations

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

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

CROP_SIZE = 224
RESIZE_SIZE = 256
BACKBONE = "dinov3_vitb16"


class DinoV3Head(nn.Module):
    """
    DINOv3 ViT backbone with a linear classification head.

    Head input = concat([CLS], mean(patch_tokens)) -> dim = 2 * embed_dim.
    Copied verbatim from DeepForestVisionV2's classification.py so the
    checkpoint's state_dict loads strict (missing=0, unexpected=0).
    """

    def __init__(self, backbone: nn.Module, num_classes: int) -> None:
        super().__init__()
        self.backbone = backbone
        embed_dim = backbone.embed_dim
        self.classifier = (
            nn.Linear(embed_dim * 2, num_classes) if num_classes > 0 else nn.Identity()
        )

    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
        seq = self.backbone.get_intermediate_layers(pixel_values, n=1)[0]
        cls_token = seq[:, 0]
        patch_tokens = seq[:, 1:]
        pooled_patches = patch_tokens.mean(dim=1)
        x = torch.cat([cls_token, pooled_patches], dim=1)
        return self.classifier(x)


class ModelInference:
    """DeepForestVision v2 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: DinoV3Head | None = None
        self.device: torch.device | None = None
        self.class_names: list[str] = []
        # Matches DeepForestVisionV2's build_val_transform: resize the
        # shortest edge to 256 (bicubic), centre-crop 224, rescale to [0,1],
        # ImageNet-normalize.
        self.preprocess = transforms.Compose([
            transforms.Resize(
                size=RESIZE_SIZE,
                interpolation=InterpolationMode.BICUBIC,
            ),
            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:
        # CUDA first, then MPS, then CPU (matches v1 and the other classifiers).
        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")

        # Build the DINOv3 architecture from the vendored source next to the
        # weights (hubconf.py + dinov3/). pretrained=False: no network, no
        # Meta weights; the strict load below supplies every weight.
        backbone = torch.hub.load(
            str(self.model_dir), BACKBONE, source="local", pretrained=False
        )
        checkpoint = torch.load(self.model_path, map_location=self.device)
        self.class_names = list(checkpoint["labels"])
        model = DinoV3Head(backbone, len(self.class_names))
        model.load_state_dict(checkpoint["model_state_dict"], strict=True)
        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 DeepForestVisionV2 (supervision box crop: no
        squaring, no pad). The transform above does its own resize/centre-crop,
        so the crop handed to it must be the raw detection box.
        """
        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 [[self.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(self.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 [
            [[self.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():
            # DinoV3Head returns raw logits (a plain tensor, not an HF output).
            return self.model(tensor_in).softmax(dim=1).cpu().numpy()