Upload 2 files

inference.py

@@ -0,0 +1,237 @@
+"""
+ModelInference for the addax-sppnet model family.
+
+Architecture: SpeciesNet GraphModule backbone (frozen) + a thin nn.Linear
+head fine-tuned per region. Originally written for AddaxAI's legacy
+classify_detections.py (Peter van Lunteren, 13 May 2025); ported here to
+the WebUI's class-based ModelInference interface.
+
+Files expected in the model directory:
+    - <model_fname>.pt         fine-tuned head checkpoint, e.g. final-20260317.pt
+    - <backbone>.pt            frozen SpeciesNet backbone, one of:
+                                 - always_crop_99710272_22x8_v12_epoch_00148.pt
+                                 - full_image_88545560_22x8_v12_epoch_00153.pt
+"""
+
+from __future__ import annotations
+
+# Allow loading checkpoints saved on a Windows runner on a POSIX machine.
+import pathlib
+import platform
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image
+from torchvision import transforms
+
+if platform.system() != "Windows":
+    pathlib.WindowsPath = pathlib.PosixPath  # type: ignore[assignment]
+
+# Don't fail on truncated images during inference.
+from PIL import ImageFile
+
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+
+_BACKBONE_FILENAMES = (
+    "always_crop_99710272_22x8_v12_epoch_00148.pt",
+    "full_image_88545560_22x8_v12_epoch_00153.pt",
+)
+
+
+def _load_fx_checkpoint(weights_path: Path, map_location: str = "cpu") -> nn.Module:
+    """Load a SpeciesNet onnx2torch GraphModule.
+
+    The backbone is shipped as a torch.fx GraphModule. PyTorch 2.4+
+    requires `reduce_graph_module` to be in the safe-globals allowlist
+    when loading with `weights_only=True`; older versions don't have
+    this concept. Try both paths.
+    """
+    try:
+        from torch.fx.graph_module import reduce_graph_module
+        from torch.serialization import add_safe_globals
+        add_safe_globals([reduce_graph_module])
+    except Exception:
+        pass
+
+    try:
+        obj = torch.load(weights_path, map_location=map_location, weights_only=True)
+    except Exception:
+        obj = torch.load(weights_path, map_location=map_location, weights_only=False)
+
+    if hasattr(obj, "state_dict") and hasattr(obj, "forward"):
+        return obj
+    raise ValueError(f"{weights_path} is not a torch.nn.Module GraphModule")
+
+
+class _FXClassifier(nn.Module):
+    """SpeciesNet backbone (frozen) + linear head."""
+
+    def __init__(
+        self,
+        backbone: nn.Module,
+        num_classes: int,
+        img_size: int = 480,
+        input_layout: str = "nhwc",
+    ) -> None:
+        super().__init__()
+        self.backbone = backbone
+        self.input_layout = input_layout.lower()
+
+        for p in self.backbone.parameters():
+            p.requires_grad = False
+        self.backbone.eval()
+
+        # Probe the backbone to discover output feature size at this
+        # img_size + layout combo, so the head matches exactly.
+        with torch.no_grad():
+            x = torch.zeros(1, 3, img_size, img_size)
+            if self.input_layout == "nhwc":
+                x = x.permute(0, 2, 3, 1).contiguous()
+            z = self.backbone(x)
+            z = self._pool(z)
+            in_features = z.shape[1]
+
+        self.head = nn.Linear(in_features, num_classes)
+
+    @staticmethod
+    def _pool(z: torch.Tensor) -> torch.Tensor:
+        if z.ndim == 4:
+            return F.adaptive_avg_pool2d(z, 1).flatten(1)
+        if z.ndim == 3:
+            return z.mean(dim=1)
+        return z.flatten(1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.input_layout == "nhwc":
+            x = x.permute(0, 2, 3, 1).contiguous()
+        z = self.backbone(x)
+        z = self._pool(z)
+        return self.head(z)
+
+
+class ModelInference:
+    """ModelInference for the addax-sppnet family (SpeciesNet backbone + linear head)."""
+
+    def __init__(self, model_dir: Path, model_path: Path) -> None:
+        self.model_dir = Path(model_dir)
+        self.model_path = Path(model_path)
+        self.model: _FXClassifier | None = None
+        self.device: torch.device | None = None
+        self._class_names: list[str] = []
+        self._preprocess: transforms.Compose | None = None
+
+    # ------------------------------------------------------------------
+    # Required interface
+    # ------------------------------------------------------------------
+
+    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:
+        if self.check_gpu():
+            self.device = torch.device(
+                "mps" if torch.backends.mps.is_available() else "cuda"
+            )
+        else:
+            self.device = torch.device("cpu")
+
+        # Load fine-tuned head checkpoint.
+        try:
+            checkpoint = torch.load(
+                self.model_path, map_location=self.device, weights_only=True
+            )
+        except Exception:
+            checkpoint = torch.load(
+                self.model_path, map_location=self.device, weights_only=False
+            )
+
+        # Resolve backbone path. The fine-tuned model ships alongside
+        # one of two known backbone files, depending on the recipe.
+        backbone_path: Path | None = None
+        for name in _BACKBONE_FILENAMES:
+            candidate = self.model_dir / name
+            if candidate.exists():
+                backbone_path = candidate
+                break
+        if backbone_path is None:
+            raise FileNotFoundError(
+                "Backbone weights not found. Expected one of "
+                f"{_BACKBONE_FILENAMES} in {self.model_dir}."
+            )
+
+        backbone = _load_fx_checkpoint(backbone_path, map_location="cpu")
+        model = _FXClassifier(
+            backbone=backbone,
+            num_classes=checkpoint["num_classes"],
+            img_size=checkpoint["img_size"],
+            input_layout=checkpoint["input_layout"],
+        )
+        model.load_state_dict(checkpoint["model"])
+        self.model = model.to(self.device).eval()
+
+        self._class_names = list(checkpoint["class_names"])
+
+        norm = checkpoint["normalize"]
+        img_size = checkpoint["img_size"]
+        self._preprocess = transforms.Compose([
+            transforms.Resize((img_size, img_size), antialias=True),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=norm["mean"], std=norm["std"]),
+        ])
+
+    def get_crop(
+        self, image: Image.Image, bbox: tuple[float, float, float, float]
+    ) -> Image.Image:
+        """Crop the bbox region. SpeciesNet head was trained on tight crops."""
+        W, H = image.size
+        x, y, w, h = bbox
+        left = max(0, int(round(x * W)))
+        top = max(0, int(round(y * H)))
+        right = min(W, int(round((x + w) * W)))
+        bottom = min(H, int(round((y + h) * H)))
+        if right <= left or bottom <= top:
+            return image
+        return image.crop((left, top, right, bottom))
+
+    def get_classification(self, crop: Image.Image) -> list[list]:
+        """Per-image inference. Returns [[name, prob], ...] for all classes."""
+        assert self.model is not None and self._preprocess is not None
+        if crop.mode != "RGB":
+            crop = crop.convert("RGB")
+        tensor = self._preprocess(crop).unsqueeze(0).to(self.device)
+        with torch.no_grad():
+            probs = F.softmax(self.model(tensor), dim=1).cpu().numpy()[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 (5-15x GPU speedup vs per-crop calls)
+    # ------------------------------------------------------------------
+
+    def get_tensor(self, crop: Image.Image) -> np.ndarray:
+        assert self._preprocess is not None
+        if crop.mode != "RGB":
+            crop = crop.convert("RGB")
+        return self._preprocess(crop).numpy()
+
+    def classify_batch(self, batch: np.ndarray) -> list[list[list]]:
+        assert self.model is not None
+        tensor = torch.from_numpy(batch).to(self.device)
+        with torch.no_grad():
+            probs = F.softmax(self.model(tensor), dim=1).cpu().numpy()
+        return [
+            [[self._class_names[j], float(p[j])] for j in range(len(p))]
+            for p in probs
+        ]

taxonomy.csv

@@ -0,0 +1,142 @@
+model_class,class,order,family,genus,species
+dingo,mammalia,carnivora,canidae,canis,lupus dingo
+bird sp,aves,,,,
+frog sp,amphibia,anura,,,
+owl sp,aves,strigiformes,,,
+rodent sp,mammalia,rodentia,,,
+kingfisher sp,aves,coraciiformes,alcedinidae,,
+gecko sp,reptilia,squamata,gekkonidae,,
+skink sp,reptilia,squamata,scincidae,,
+dragon sp,reptilia,squamata,agamidae,,
+quail spp,aves,galliformes,phasianidae,,
+giant stick insect,insecta,phasmida,phasmatidae,extatosoma,tiaratum
+short-beaked echidna,mammalia,monotremata,tachyglossidae,tachyglossus,aculeatus
+cat,mammalia,carnivora,felidae,felis,catus
+northern brown bandicoot,mammalia,peramelemorphia,peramelidae,isoodon,macrourus
+northern quoll,mammalia,dasyuromorphia,dasyuridae,dasyurus,hallucatus
+fawn antechinus,mammalia,dasyuromorphia,dasyuridae,antechinus,bellus
+sandstone pseudantechinus,mammalia,dasyuromorphia,dasyuridae,pseudantechinus,bilarni
+dunnart sp,mammalia,dasyuromorphia,dasyuridae,sminthopsis,
+red-cheeked dunnart,mammalia,dasyuromorphia,dasyuridae,sminthopsis,virginiae
+planigale sp,mammalia,dasyuromorphia,dasyuridae,planigale,maculata
+brush-tailed rabbit-rat,mammalia,rodentia,muridae,conilurus,penicillatus
+pseudomys sp,mammalia,rodentia,muridae,pseudomys,
+black-footed tree-rat,mammalia,rodentia,muridae,mesembriomys,gouldii
+water-rat,mammalia,rodentia,muridae,hydromys,chrysogaster
+northern hopping mouse,mammalia,rodentia,muridae,notomys,aquilo
+grassland melomys,mammalia,rodentia,muridae,melomys,burtoni
+rattus sp,mammalia,rodentia,muridae,rattus,
+pale field-rat,mammalia,rodentia,muridae,rattus,tunneyi
+dusky rat,mammalia,rodentia,muridae,rattus,colletti
+black rat,mammalia,rodentia,muridae,rattus,rattus
+rock-rat sp,mammalia,rodentia,muridae,zyzomys,
+arnhem rock-rat,mammalia,rodentia,muridae,zyzomys,maini
+common rock-rat,mammalia,rodentia,muridae,zyzomys,argurus
+rock ringtail possum,mammalia,diprotodontia,pseudocheiridae,petropseudes,dahli
+rock-wallaby sp,mammalia,diprotodontia,macropodidae,petrogale,
+northern nailtail wallaby,mammalia,diprotodontia,macropodidae,onychogalea,unguifera
+northern brushtail possum,mammalia,diprotodontia,phalangeridae,trichosurus,vulpecula
+horse,mammalia,perissodactyla,equidae,equus,caballus
+donkey,mammalia,perissodactyla,equidae,equus,asinus
+rusa deer,mammalia,artiodactyla,cervidae,rusa,timorensis
+cow,mammalia,artiodactyla,bovidae,bos,
+banteng,mammalia,artiodactyla,bovidae,bos,javanicus
+northern brown snake,reptilia,squamata,elapidae,pseudonaja,nuchalis
+northern bluetongue,reptilia,squamata,scincidae,tiliqua,scincoides
+frill-neck lizard,reptilia,squamata,agamidae,chlamydosaurus,kingii
+gilbert's dragon,reptilia,squamata,agamidae,lophognathus,gilberti
+kimberley rock goanna,reptilia,squamata,varanidae,varanus,glauerti
+mertens' water goanna,reptilia,squamata,varanidae,varanus,mertensi
+yellow-spotted goanna,reptilia,squamata,varanidae,varanus,panoptes
+black-palmed goanna,reptilia,squamata,varanidae,varanus,glebopalma
+spotted tree goanna,reptilia,squamata,varanidae,varanus,scalaris
+goanna sp,reptilia,squamata,varanidae,varanus,
+brown quail,aves,galliformes,phasianidae,synoicus,ypsilophorus
+chestnut rail,aves,gruiformes,rallidae,eulabeornis,castaneoventris
+red-backed button-quail,aves,charadriiformes,turnicidae,turnix,maculosus
+chestnut-backed button-quail,aves,charadriiformes,turnicidae,turnix,castanotus
+rainbow bee-eater,aves,coraciiformes,meropidae,merops,ornatus
+blue-winged kookaburra,aves,coraciiformes,alcedinidae,dacelo,leachii
+forest kingfisher,aves,coraciiformes,alcedinidae,todiramphus,macleayii
+northern rosella,aves,psittaciformes,psittacidae,platycercus,venustus
+red-winged parrot,aves,psittaciformes,psittacidae,aprosmictus,erythropterus
+little corella,aves,psittaciformes,psittacidae,cacatua,sanguinea
+sulphur-crested cockatoo,aves,psittaciformes,psittacidae,cacatua,galerita
+red-tailed black cockatoo,aves,psittaciformes,psittacidae,calyptorhynchus,banksii
+whistling kite,aves,accipitriformes,accipitridae,haliastur,sphenurus
+black-breasted buzzard,aves,accipitriformes,accipitridae,hamirostra,melanosternon
+collared sparrowhawk,aves,accipitriformes,accipitridae,accipiter,cirrocephalus
+brown goshawk,aves,accipitriformes,accipitridae,accipiter,fasciatus
+black bittern,aves,pelecaniformes,ardeidae,botaurus,
+peregrine falcon,aves,falconiformes,falconidae,falco,peregrinus
+brown falcon,aves,falconiformes,falconidae,falco,berigora
+bush stone-curlew,aves,charadriiformes,burhinidae,burhinus,grallarius
+orange-footed scrubfowl,aves,galliformes,megapodiidae,megapodius,reinwardt
+little shrike-thrush,aves,passeriformes,pachycephalidae,colluricincla,megarhyncha
+sandstone shrike-thrush,aves,passeriformes,pachycephalidae,colluricincla,woodwardi
+grey shrike-thrush,aves,passeriformes,pachycephalidae,colluricincla,harmonica
+torresian crow,aves,passeriformes,corvidae,corvus,orru
+striated pardalote,aves,passeriformes,pardalotidae,pardalotus,striatus
+leaden flycatcher,aves,passeriformes,monarchidae,myiagra,rubecula
+magpie-lark,aves,passeriformes,monarchidae,grallina,cyanoleuca
+blue-faced honeyeater,aves,passeriformes,meliphagidae,entomyzon,cyanotis
+friarbird sp,aves,passeriformes,meliphagidae,philemon,
+little friarbird,aves,passeriformes,meliphagidae,philemon,citreogularis
+silver-crowned friarbird,aves,passeriformes,meliphagidae,philemon,argenticeps
+brown honeyeater,aves,passeriformes,meliphagidae,lichmera,indistincta
+yellow-throated miner,aves,passeriformes,meliphagidae,manorina,flavigula
+white-throated grasswren,aves,passeriformes,maluridae,amytornis,woodwardi
+pied butcherbird,aves,passeriformes,cracticidae,cracticus,nigrogularis
+grey butcherbird,aves,passeriformes,cracticidae,cracticus,torquatus
+rainbow pitta,aves,passeriformes,pittidae,pitta,iris
+grey-crowned babbler,aves,passeriformes,pomatostomidae,pomatostomus,temporalis
+rufous songlark,aves,passeriformes,locustellidae,cincloramphus,mathewsi
+crimson finch,aves,passeriformes,estrildidae,neochmia,phaeton
+great bowerbird,aves,passeriformes,ptilonorhynchidae,chlamydera,nuchalis
+common bronzewing,aves,columbiformes,columbidae,phaps,chalcoptera
+spinifex pigeon,aves,columbiformes,columbidae,geophaps,plumifera
+eastern partridge pigeon,aves,columbiformes,columbidae,geophaps,smithii
+bar-shouldered dove,aves,columbiformes,columbidae,geopelia,humeralis
+diamond dove,aves,columbiformes,columbidae,geopelia,cuneata
+peaceful dove,aves,columbiformes,columbidae,geopelia,placida
+emerald dove,aves,columbiformes,columbidae,chalcophaps,indica
+chestnut-quilled rock-pigeon,aves,columbiformes,columbidae,petrophassa,rufipennis
+australian owlet-nightjar,aves,apodiformes,aegothelidae,aegotheles,cristatus
+tawny frogmouth,aves,caprimulgiformes,podargidae,podargus,strigoides
+australian boobook,aves,strigiformes,strigidae,ninox,novaeseelandiae
+antilopine wallaroo,mammalia,diprotodontia,macropodidae,osphranter,antilopinus
+wallaroo sp,mammalia,diprotodontia,macropodidae,osphranter,
+eastern short-eared rock-wallaby,mammalia,diprotodontia,macropodidae,petrogale,wilkinsi
+australian pipit,aves,passeriformes,motacillidae,anthus,australis
+pallid cuckoo,aves,cuculiformes,cuculidae,heteroscenes,pallidus
+cane toad,amphibia,anura,bufonidae,rhinella,marina
+delicate mouse,mammalia,rodentia,muridae,pseudomys,delicatulus
+western chestnut mouse,mammalia,rodentia,muridae,pseudomys,nanus
+carpet python,reptilia,squamata,pythonidae,morelia,spilota variegata
+wedge-tailed eagle,aves,accipitriformes,accipitridae,aquila,audax
+long-tailed finch,aves,passeriformes,estrildidae,poephila,acuticauda
+willie wagtail,aves,passeriformes,rhipiduridae,rhipidura,leucophrys
+northern fantail,aves,passeriformes,rhipiduridae,rhipidura,rufiventris
+pheasant coucal,aves,cuculiformes,cuculidae,centropus,phasianinus
+masked owl,aves,strigiformes,tytonidae,tyto,novaehollandiae
+northern brush-tailed phascogale,mammalia,dasyuromorphia,dasyuridae,phascogale,pirata
+rufous whistler,aves,passeriformes,pachycephalidae,pachycephala,rufiventris
+grey whistler,aves,passeriformes,pachycephalidae,pachycephala,simplex
+banded honeyeater,aves,passeriformes,meliphagidae,cissomela,pectoralis
+singing honeyeater,aves,passeriformes,meliphagidae,gavicalis,virescens
+buffalo,mammalia,artiodactyla,bovidae,bubalus,bubalis
+pig,mammalia,artiodactyla,suidae,sus,scrofa
+yellow-tinted honeyeater,aves,passeriformes,meliphagidae,ptilotula,flavescens
+savanna glider,mammalia,diprotodontia,petauridae,petaurus,
+sand goanna,reptilia,squamata,varanidae,varanus,gouldii
+black-tailed goanna,reptilia,squamata,varanidae,varanus,tristis
+dusky myzomela,aves,passeriformes,meliphagidae,myzomela,obscura
+radjah shelduck,aves,anseriformes,anatidae,radjah,radjah
+double-barred finch,aves,passeriformes,estrildidae,stizoptera,bichenovii
+mangrove robin,aves,passeriformes,petroicidae,peneoenanthe,pulverulenta
+snake sp,reptilia,squamata,,,
+agile wallaby,mammalia,diprotodontia,macropodidae,notamacropus,agilis
+black wallaroo,mammalia,diprotodontia,macropodidae,osphranter,bernardus
+common wallaroo,mammalia,diprotodontia,macropodidae,osphranter,robustus
+black-spotted ridge-tailed goanna,reptilia,squamata,varanidae,varanus,insulanicus
+false detection,,,,,