""" 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()