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