| """ |
| Inference script for CAM-AI4G-v1 (Colombian Amazon, AI for Good Lab) |
| |
| Model: AI4GAmazonRainforest |
| Input: 224x224 RGB, ImageNet-normalised (applied by PytorchWildlife) |
| Framework: PyTorch (PytorchWildlife, ResNet) |
| Classes: 36 Amazon genera |
| Developer: AI for Good Lab, Microsoft |
| |
| Ported from AddaxAI's legacy classify_detections.py (pywildlife), with |
| the double-transform bug fixed. See below. |
| |
| Author: Peter van Lunteren |
| """ |
|
|
| from __future__ import annotations |
|
|
| from pathlib import Path |
|
|
| import numpy as np |
| import supervision as sv |
| import torch |
| from PIL import Image, ImageFile |
| from PytorchWildlife.models import classification as pw_classification |
|
|
| |
| ImageFile.LOAD_TRUNCATED_IMAGES = True |
|
|
| |
| |
| |
|
|
|
|
| class ModelInference: |
| """Colombian Amazon classifier (PytorchWildlife AI4GAmazonRainforest).""" |
|
|
| def __init__(self, model_dir: Path, model_path: Path) -> None: |
| self.model_dir = Path(model_dir) |
| self.model_path = Path(model_path) |
| self.model = None |
|
|
| |
| |
| |
|
|
| def check_gpu(self) -> bool: |
| try: |
| if torch.backends.mps.is_built() and torch.backends.mps.is_available(): |
| return True |
| except Exception: |
| pass |
| return torch.cuda.is_available() |
|
|
| def load_model(self) -> None: |
| self.model = pw_classification.AI4GAmazonRainforest( |
| weights=str(self.model_path) |
| ) |
|
|
| def get_crop( |
| self, image: Image.Image, bbox: tuple[float, float, float, float] |
| ) -> Image.Image: |
| """ |
| Plain crop, no squaring and no padding. |
| |
| Kept on supervision's crop_image so it stays identical to the |
| legacy adapter, down to how it rounds and how numpy slicing |
| handles a box that runs past an edge. |
| """ |
| img_width, img_height = image.size |
| left = int(round(bbox[0] * img_width)) |
| top = int(round(bbox[1] * img_height)) |
| right = int(round(bbox[2] * img_width)) + left |
| bottom = int(round(bbox[3] * img_height)) + top |
| return Image.fromarray( |
| sv.crop_image( |
| np.array(image.convert("RGB")), xyxy=[left, top, right, bottom] |
| ) |
| ) |
|
|
| def get_classification(self, crop: Image.Image) -> list[list]: |
| """ |
| Per-crop inference. Returns [[name, prob], ...] for all classes. |
| |
| PytorchWildlife's single_image_classification does the whole |
| preprocessing itself: Image.fromarray -> resize 224 -> ToTensor |
| -> ImageNet Normalize -> forward. So it takes the plain RGB crop. |
| |
| The legacy adapter ran Classification_Inference_Transform first |
| and then cast the result back to uint8 before handing it over: |
| |
| preprocessed = trans_clf(PIL_crop) # float, [-2.1, 2.5] |
| preprocessed = preprocessed.permute(1,2,0).numpy().astype(np.uint8) |
| model.single_image_classification(preprocessed) # transforms it AGAIN |
| |
| That cast is destructive. The normalised floats sit in about |
| [-2.1, 2.5], so uint8 truncates them to {0, 1, 2} and wraps the |
| negatives: the 224x224x3 array reaching the model held three |
| distinct values, i.e. a near-black image. Measured on a real |
| crop, legacy answered `Mazama` at 0.362 with `Unknown` at 0.357 |
| behind it (a model with no signal, guessing), where the same |
| weights on the same crop answer `Bos` at 0.917. |
| |
| So this is not a faithful-port question: the legacy path never |
| classified the image at all. Passing the crop through once, the |
| way PytorchWildlife intends, is the fix. |
| """ |
| assert self.model is not None |
| return self.model.single_image_classification( |
| np.array(crop.convert("RGB")) |
| )["all_confidences"] |
|
|
| def get_class_names(self) -> dict[str, str]: |
| """1-indexed mapping {id: class_name} for the output JSON.""" |
| assert self.model is not None |
| names = self.model.CLASS_NAMES |
| return {str(i + 1): names[i] for i in range(len(names))} |
|
|
