feat: add ood score

Author: Yuyan-C

trapdata/api/api.py

@@ -18,6 +18,7 @@
     MothClassifierGlobal,
     MothClassifierPanama,
     MothClassifierPanama2024,
+    MothClassifierPanamaPlus2025,
     MothClassifierQuebecVermont,
     MothClassifierTuringAnguilla,
     MothClassifierTuringCostaRica,
@@ -39,6 +40,7 @@
 
 
 CLASSIFIER_CHOICES = {
+    "panama_plus_moths_2025": MothClassifierPanamaPlus2025,
     "panama_moths_2023": MothClassifierPanama,
     "panama_moths_2024": MothClassifierPanama2024,
     "quebec_vermont_moths_2023": MothClassifierQuebecVermont,

trapdata/api/models/classification.py

@@ -15,6 +15,7 @@
     TuringAnguillaSpeciesClassifier,
     TuringCostaRicaSpeciesClassifier,
     UKDenmarkMothSpeciesClassifier2024,
+    PanamaPlusWithOODClassifier2025,
 )
 
 from ..datasets import ClassificationImageDataset
@@ -25,6 +26,7 @@
     SourceImage,
 )
 from .base import APIInferenceBaseClass
+from trapdata.ml.models.base import ClassifierResult
 
 
 class APIMothClassifier(
@@ -188,3 +190,47 @@ class MothClassifierTuringAnguilla(APIMothClassifier, TuringAnguillaSpeciesClass
 
 class MothClassifierGlobal(APIMothClassifier, GlobalMothSpeciesClassifier):
     pass
+
+
+class MothClassifierPanamaPlus2025(APIMothClassifier, PanamaPlusWithOODClassifier2025):
+    def post_process_batch(self, logits: torch.Tensor):
+        """
+        Return the labels, softmax/calibrated scores, and the original logits for
+        each image in the batch.
+
+        Almost like the base class method, but we need to return the logits as well.
+        """
+        predictions = torch.nn.functional.softmax(logits, dim=1)
+        predictions = predictions.cpu().numpy()
+
+        ood_scores = None
+        if self.class_prior:
+            _, ood_scores = torch.max(predictions - self.class_prior, dim=-1)
+        else:
+            _, ood_scores = torch.max(predictions, dim=-1)
+
+        batch_results = []
+        for softmax_scores in predictions:
+            # Get all class indices and their corresponding scores
+            class_indices = np.arange(len(softmax_scores))
+            labels = [self.category_map[i] for i in class_indices]
+
+            print("labels type", type(labels))
+            print("logits type", type(logits))
+            print("label type", type(softmax_scores))
+            print("label type", type(ood_scores))
+
+            exit()
+
+            # TODO: Change batch_results
+            result = ClassifierResult(
+                labels=labels,
+                logits=logits,
+                softmax_scores=softmax_scores,
+                ood_scores=ood_scores,
+            )
+            batch_results.append(result)
+
+        logger.debug(f"Post-processing result batch: {batch_results}")
+
+        return batch_results

trapdata/common/logs.py

@@ -2,12 +2,16 @@
 
 import structlog
 
+# structlog.configure(
+#     wrapper_class=structlog.make_filtering_bound_logger(logging.INFO),
+# )
+
 structlog.configure(
-    wrapper_class=structlog.make_filtering_bound_logger(logging.INFO),
+    wrapper_class=structlog.make_filtering_bound_logger(logging.CRITICAL),
 )
 
-
 logger = structlog.get_logger()
+logging.disable(logging.CRITICAL)
 
 # import logging
 # from rich.logging import RichHandler

trapdata/db/models/detections.py

@@ -27,6 +27,7 @@ class DetectionListItem(BaseModel):
     model_name: Optional[str]
     in_queue: bool
     notes: Optional[str]
+    ood_score: Optional[str]
 
     # PyDantic complains because we have an attribute called `model_name`
     model_config = ConfigDict(protected_namespaces=[])  # type:ignore
@@ -43,6 +44,7 @@ class DetectionDetail(DetectionListItem):
     timestamp: Optional[str]
     bbox_center: Optional[tuple[int, int]]
     area_pixels: Optional[int]
+    ood_score: Optional[float]
 
 
 class DetectedObject(db.Base):
@@ -76,6 +78,7 @@ class DetectedObject(db.Base):
     sequence_previous_id = sa.Column(sa.Integer)
     sequence_previous_cost = sa.Column(sa.Float)
     cnn_features = sa.Column(sa.JSON)
+    ood_score = sa.Column(sa.Float)
 
     # @TODO add updated & created timestamps to all db models
 
@@ -288,6 +291,7 @@ def report_data(self) -> DetectionDetail:
             last_detected=self.last_detected,
             notes=self.notes,
             in_queue=self.in_queue,
+            ood_score=self.ood_score
         )
 
     def report_data_simple(self):
@@ -510,7 +514,9 @@ def get_species_for_image(db_path, image_id):
 def num_species_for_event(
     db_path, monitoring_session, classification_threshold: float = 0.6
 ) -> int:
-    query = sa.select(sa.func.count(DetectedObject.specific_label.distinct()),).where(
+    query = sa.select(
+        sa.func.count(DetectedObject.specific_label.distinct()),
+    ).where(
         (DetectedObject.specific_label_score >= classification_threshold)
         & (DetectedObject.monitoring_session == monitoring_session)
     )
@@ -522,7 +528,9 @@ def num_species_for_event(
 def num_occurrences_for_event(
     db_path, monitoring_session, classification_threshold: float = 0.6
 ) -> int:
-    query = sa.select(sa.func.count(DetectedObject.sequence_id.distinct()),).where(
+    query = sa.select(
+        sa.func.count(DetectedObject.sequence_id.distinct()),
+    ).where(
         (DetectedObject.specific_label_score >= classification_threshold)
         & (DetectedObject.monitoring_session == monitoring_session)
     )

trapdata/ml/models/__init__.py

@@ -29,7 +29,6 @@ def get_default_model(choices: EnumMeta) -> str:
 )
 DEFAULT_OBJECT_DETECTOR = get_default_model(ObjectDetectorChoice)
 
-
 binary_classifiers = {Model.name: Model for Model in BinaryClassifier.__subclasses__()}
 BinaryClassifierChoice = ModelChoiceEnum(
     "BinaryClassifierChoice",

trapdata/ml/models/base.py

@@ -1,4 +1,5 @@
 import json
+import pandas as pd
 from typing import Union
 
 import numpy as np
@@ -14,6 +15,8 @@
 from trapdata.db.models.queue import QueueManager
 from trapdata.ml.utils import StopWatch, get_device, get_or_download_file
 
+from dataclasses import dataclass
+
 
 class BatchEmptyException(Exception):
     pass
@@ -88,6 +91,7 @@ class InferenceBaseClass:
     queue: QueueManager
     dataset: torch.utils.data.Dataset
     dataloader: torch.utils.data.DataLoader
+    training_csv_path: str | None = None
 
     def __init__(
         self,
@@ -105,6 +109,7 @@ def __init__(
 
         self.device = self.device or get_device()
         self.category_map = self.get_labels(self.labels_path)
+        self.class_prior = self.get_class_prior(self.training_csv_path)
         self.num_classes = self.num_classes or len(self.category_map)
         self.weights = self.get_weights(self.weights_path)
         self.transforms = self.get_transforms()
@@ -183,6 +188,25 @@ def fetch_gbif_ids(labels):
         else:
             return {}
 
+    def get_class_prior(self, training_csv_path):
+        if training_csv_path:
+            local_path = get_or_download_file(
+                training_csv_path,
+                self.user_data_path or torch.hub.get_dir(),
+                prefix="models",
+            )
+            df_train = pd.read_csv(local_path)
+            categories = sorted(list(df_train["speciesKey"].unique()))
+            categories_map = {categ: id for id, categ in enumerate(categories)}
+            df_train["label"] = df_train["speciesKey"].map(categories_map)
+            cls_idx = df_train["label"].astype(int).values
+            num_classes = df_train["label"].nunique()
+            cls_num = np.bincount(cls_idx, minlength=num_classes)
+            targets = cls_num / cls_num.sum()
+            return targets
+        else:
+            return None
+
     def get_model(self) -> torch.nn.Module:
         """
         This method must be implemented by a subclass.
@@ -330,3 +354,12 @@ def run(self):
             logger.info(f"{self.name} Batch -- Done")
 
         logger.info(f"{self.name} -- Done")
+
+
+@dataclass
+class ClassifierResult:
+    # TODO: add types
+    labels = None
+    logits = None
+    softmax_scores = None
+    ood_scores = None

trapdata/ml/models/classification.py

@@ -1,14 +1,21 @@
+from typing import Union
+from sqlalchemy.engine.url import URL as URL
 import timm
 import torch
 import torch.utils.data
 import torchvision
 
 from trapdata import constants, logger
+from trapdata.common.schemas import FilePath
 from trapdata.db.models.detections import save_classified_objects
 from trapdata.db.models.queue import DetectedObjectQueue, UnclassifiedObjectQueue
 
 from .base import InferenceBaseClass, imagenet_normalization
 
+import numpy as np
+import os
+from trapdata.ml.utils import get_or_download_file
+
 
 class ClassificationIterableDatabaseDataset(torch.utils.data.IterableDataset):
     def __init__(self, queue, image_transforms, batch_size=4):
@@ -318,6 +325,24 @@ def save_results(self, object_ids, batch_output, *args, **kwargs):
         save_classified_objects(self.db_path, object_ids, classified_objects_data)
 
 
+
+# class SpeciesClassifierWithOOD(SpeciesClassifier):
+#     def save_results(self, object_ids, batch_output, *args, **kwargs):
+#         # Here we are saving the specific taxon labels
+#         classified_objects_data = [
+#             {
+#                 "specific_label": label,
+#                 "specific_label_score": score,
+#                 "model_name": self.name,
+#                 "in_queue": True,  # Put back in queue for the feature extractor & tracking
+#             }
+#             for label, score in batch_output
+#         ]
+#         save_classified_objects(self.db_path, object_ids, classified_objects_data)
+
+
+
+
 class QuebecVermontMothSpeciesClassifierMixedResolution(
     SpeciesClassifier, Resnet50ClassifierLowRes
 ):
@@ -456,7 +481,7 @@ class QuebecVermontMothSpeciesClassifier2024(SpeciesClassifier, Resnet50TimmClas
     )
     weights_path = (
         "https://object-arbutus.cloud.computecanada.ca/ami-models/moths/classification/"
-        "quebec-vermont_resnet50_baseline_20240417_950de764.pth"
+        "=-vermont_resnet50_baseline_20240417_950de764.pth"
     )
     labels_path = (
         "https://object-arbutus.cloud.computecanada.ca/ami-models/moths/classification/"
@@ -505,12 +530,12 @@ class PanamaMothSpeciesClassifier2024(SpeciesClassifier, Resnet50TimmClassifier)
     )
 
 
-class PanamaMothSpeciesClassifier2025(SpeciesClassifier, Resnet50TimmClassifier):
+class PanamaPlusWithOODClassifier2025(SpeciesClassifier, Resnet50TimmClassifier):
     input_size = 128
     normalization = imagenet_normalization
     lookup_gbif_names = False
 
-    name = "Panama Species Classifier - Mar 2025"
+    name = "Panama Plus Species Classifier with OOD detection - Mar 2025"
     description = (
         "Trained on March 13th, 2025 for 2360 species. "
         "https://wandb.ai/moth-ai/panama_classifier/runs/81f5ssv9/overview"
@@ -523,5 +548,21 @@ class PanamaMothSpeciesClassifier2025(SpeciesClassifier, Resnet50TimmClassifier)
 
     labels_path = (
         "https://object-arbutus.cloud.computecanada.ca/ami-models/moths/classification/"
-        "panama_plus_category_map-with_names.json"  
+        "panama_plus_category_map-with_names.json"
     )
+
+    training_csv_path = "https://object-arbutus.cloud.computecanada.ca/ami-models/moths/classification/panama_plus_train.csv"
+
+    def save_results(self, object_ids, batch_output, *args, **kwargs):
+        # Here we are saving the specific taxon labels
+        classified_objects_data = [
+            {
+                "specific_label": label,
+                "specific_label_score": score,
+                "model_name": self.name,
+                "in_queue": True,  # Put back in queue for the feature extractor & tracking
+            }
+            for label, score in batch_output
+        ]
+        save_classified_objects(self.db_path, object_ids, classified_objects_data)
+

trapdata/ml/models/localization.py

@@ -214,6 +214,14 @@ def get_model(self):
         state_dict = checkpoint.get("model_state_dict") or checkpoint
         model.load_state_dict(state_dict)
         model = model.to(self.device)
+
+        # Get the state dictionary
+        state_dict = model.state_dict()
+
+        # Print the shape of each tensor in the state_dict
+        for name, param in state_dict.items():
+            print(f"{name}: {param.shape}")
+
         model.eval()
         self.model = model
         return self.model
@@ -267,6 +275,14 @@ def get_model(self):
         checkpoint = torch.load(self.weights, map_location=self.device)
         state_dict = checkpoint.get("model_state_dict") or checkpoint
         model.load_state_dict(state_dict)
+
+        # Get the state dictionary
+        state_dict = model.state_dict()
+
+        # Print the shape of each tensor in the state_dict
+        for name, param in state_dict.items():
+            print(f"{name}: {param.shape}")
+
         model = model.to(self.device)
         model.eval()
         self.model = model

trapdata/ml/models/tracking.py

@@ -23,6 +23,7 @@
     PanamaMothSpeciesClassifierMixedResolution2023,
     QuebecVermontMothSpeciesClassifierMixedResolution,
     UKDenmarkMothSpeciesClassifierMixedResolution,
+    PanamaPlusWithOODClassifier2025
 )
 from trapdata.ml.utils import get_device
 
@@ -501,6 +502,12 @@ class PanamaFeatureExtractor(
 ):
     name = "Features from Panama species model"
 
+class PanamaPlusFeatureExtractor(
+    FeatureExtractor, 
+    PanamaPlusWithOODClassifier2025
+):
+    name = "Features from Panama Plus species model"
+
 
 def clear_sequences(monitoring_session: MonitoringSession, session: orm.Session):
     logger.info(f"Clearing existing sequences for {monitoring_session.day}")

trapdata/ml/pipeline.py

@@ -23,6 +23,7 @@ def start_pipeline(
         num_workers=settings.num_workers,
         single=single,
     )
+
     if object_detector.queue.queue_count() > 0:
         object_detector.run()
         logger.info("Localization complete")