Multi-bbox predictions (inherited from PR 73)

advanced_topics.md

@@ -7,7 +7,6 @@
 - [Alternative installation variants of the SpeciesNet Python package](#alternative-installation-variants)
 - [SpeciesNet model variants](#speciesnet-model-variants)
 - [Alternative input formats for run_model](#alternative-input-formats-for-run_model)
-- [Alternative output formats for run_model](#alternative-output-formats-for-run_model)
 - [Contributing code](#contributing-code)
 - [Build status](#build-status)
 
@@ -61,7 +60,7 @@ Depending on how you plan to run SpeciesNet, you may want to install additional
 
 There are two variants of the SpeciesNet classifier, which lend themselves to different ensemble strategies:
 
-- [v4.0.3a](model_cards/v4.0.1a.md) (default): Always-crop model, i.e. we run the detector first and crop the image to the top detection bounding box before feeding it to the species classifier.
+- [v4.0.3a](model_cards/v4.0.1a.md) (default): Always-crop model, i.e. we run the detector first, then run the classifier on the pixels within each detection above a confidence threshold.
 - [v4.0.3b](model_cards/v4.0.1b.md): Full-image model, i.e. we run both the detector and the species classifier on the full image, independently.
 
 Both links point to the model cards for the 4.0.1 models; model cards were not updated for the 4.0.3 release, which only included changes to geofencing rules and minor taxonomy updates.
@@ -103,53 +102,6 @@ admin1_region is currently only supported in the US, where valid values for admi
 
 Latitude and longitude are only used to determine admin1_region, so if you are specifying a state code, you don't need to specify latitude and longitude.
 
-## Alternative output formats for run_model
-
-The main README describes the typical output format for run_model, when you are running the full ensemble (detector and classifier).  This section describes the output format when you are running only the classifier or detector.
-
-### Classifier-only inference
-
-```text
-{
-    "predictions": [
-        {
-            "filepath": str  => Image filepath.
-            "failures": list[str] (optional)  => List of internal components that failed during prediction (in this case, only "CLASSIFIER" can be in that list). If absent, the prediction was successful.
-            "classifications": {  => dict (optional)  => Top-5 classifications. Included only if "CLASSIFIER" if not part of the "failures" field.
-                "classes": list[str]  => List of top-5 classes predicted by the classifier, matching the decreasing order of their scores below.
-                "scores": list[float]  => List of scores corresponding to top-5 classes predicted by the classifier, in decreasing order.
-                "target_classes": list[str] (optional)  => List of target classes, only present if target classes are passed as arguments.
-                "target_logits": list[float] (optional)  => Raw confidence scores (logits) of the target classes, only present if target classes are passed as arguments.
-            }
-        },
-        ...  => A response will contain one prediction for each instance in the request.
-    ]
-}
-```
-
-### Detector-only inference
-
-```text
-{
-    "predictions": [
-        {
-            "filepath": str  => Image filepath.
-            "failures": list[str] (optional)  => List of internal components that failed during prediction (in this case, only "DETECTOR" can be in that list). If absent, the prediction was successful.
-            "detections": [  => list (optional)  => List of detections with confidence scores > 0.01, in decreasing order of their scores. Included only if "DETECTOR" if not part of the "failures" field.
-                {
-                    "category": str  => Detection class "1" (= animal), "2" (= human) or "3" (= vehicle) from MegaDetector's raw output.
-                    "label": str  => Detection class "animal", "human" or "vehicle", matching the "category" field above. Added for readability purposes.
-                    "conf": float  => Confidence score of the current detection.
-                    "bbox": list[float]  => Bounding box coordinates, in (xmin, ymin, width, height) format, of the current detection. Coordinates are normalized to the [0.0, 1.0] range, relative to the image dimensions.
-                },
-                ...  => A prediction can contain zero or multiple detections.
-            ]
-        },
-        ...  => A response will contain one prediction for each instance in the request.
-    ]
-}
-```
-
 ## Contributing code
 
 If you're interested in contributing to our repo, rather than installing via pip, we recommend cloning the repo, creating a Python virtual environment for development, and installing the package locally using the following command:

speciesnet/classifier.py

@@ -127,16 +127,18 @@ def preprocess(
         img: Optional[PIL.Image.Image],
         bboxes: Optional[list[BBox]] = None,
         resize: bool = True,
-    ) -> Optional[PreprocessedImage]:
+    ) -> Optional[list[PreprocessedImage]]:
         """Preprocesses an image according to this classifier's needs.
 
         This method prepares an input image for classification. It handles
         image loading, cropping, and resizing to the expected
         input size for the classifier model.
 
         In `always_crop` mode images are cropped according to the bounding boxes
-        provided. In `full_image` mode the top and bottom of the image are cropped
-        to prevent learning correlations between camera brand and species priors.
+        provided, producing one crop per bounding box (and no crops at all if no
+        bounding boxes are provided). In `full_image` mode, bounding boxes are ignored,
+        and the top and bottom of the image are cropped to prevent learning
+        correlations between camera brand and species priors, producing a single crop.
         See the paper for more details.
 
         Args:
@@ -149,130 +151,187 @@ def preprocess(
                 Whether to resize the image to some expected dimensions.
 
         Returns:
-            A preprocessed image, or `None` if no PIL image was provided initially.
+            A list of preprocessed images (one element per bounding box in `always_crop`
+            mode, a single element in `full_image` mode), or `None` if no image was
+            provided.
         """
-
         if img is None:
             return None
 
         img_tensor = F.pil_to_tensor(img)  # HWC to CHW.
         img_tensor = F.convert_image_dtype(img_tensor, torch.float32)
 
+        results = []
         if self.model_info.type_ == "always_crop":
-            # Crop to top bbox if available, otherwise leave image uncropped.
-            if bboxes:
-                img_tensor = F.crop(
+            for bbox in bboxes or []:
+                crop = F.crop(
                     img_tensor,
-                    int(bboxes[0].ymin * img.height),
-                    int(bboxes[0].xmin * img.width),
-                    int(bboxes[0].height * img.height),
-                    int(bboxes[0].width * img.width),
+                    int(bbox.ymin * img.height),
+                    int(bbox.xmin * img.width),
+                    int(bbox.height * img.height),
+                    int(bbox.width * img.width),
                 )
+                if resize:
+                    crop = F.resize(
+                        crop,
+                        [
+                            SpeciesNetClassifier.IMG_SIZE,
+                            SpeciesNetClassifier.IMG_SIZE,
+                        ],
+                        antialias=False,
+                    )
+                crop = F.convert_image_dtype(crop, torch.uint8)
+                crop = crop.permute([1, 2, 0])  # CHW to HWC.
+                results.append(PreprocessedImage(crop.numpy(), img.width, img.height))
+
         elif self.model_info.type_ == "full_image":
             # Crop top and bottom of image.
             target_height = max(
                 int(img.height * (1.0 - SpeciesNetClassifier.MAX_CROP_RATIO)),
                 img.height - SpeciesNetClassifier.MAX_CROP_SIZE,
             )
-            img_tensor = F.center_crop(img_tensor, [target_height, img.width])
+            crop = F.center_crop(img_tensor, [target_height, img.width])
 
-        if resize:
-            img_tensor = F.resize(
-                img_tensor,
-                [SpeciesNetClassifier.IMG_SIZE, SpeciesNetClassifier.IMG_SIZE],
-                antialias=False,
-            )
+            if resize:
+                crop = F.resize(
+                    crop,
+                    [SpeciesNetClassifier.IMG_SIZE, SpeciesNetClassifier.IMG_SIZE],
+                    antialias=False,
+                )
 
-        img_tensor = F.convert_image_dtype(img_tensor, torch.uint8)
-        img_tensor = img_tensor.permute([1, 2, 0])  # CHW to HWC.
-        return PreprocessedImage(img_tensor.numpy(), img.width, img.height)
+            crop = F.convert_image_dtype(crop, torch.uint8)
+            crop = crop.permute([1, 2, 0])  # CHW to HWC.
+            results.append(PreprocessedImage(crop.numpy(), img.width, img.height))
+
+        return results
 
     def predict(
-        self, filepath: str, img: Optional[PreprocessedImage]
+        self,
+        filepath: str,
+        imgs: Optional[list[PreprocessedImage]],
+        num_classes: Optional[int] = 5,
     ) -> dict[str, Any]:
-        """Runs inference on a given preprocessed image.
+        """Runs inference on the preprocessed crops of a given image.
 
         Args:
             filepath:
                 Location of image to run inference on. Used for reporting purposes only,
                 and not for loading the image.
-            img:
-                Preprocessed image to run inference on. If `None`, a failure message is
-                reported back.
+            imgs:
+                List of preprocessed image crops to run inference on. If `None`, a
+                failure message is reported back. If empty, there is nothing to
+                classify and an empty list of classifications is reported back.
+            num_classes:
+                Number of top classes to report for each crop. If `None`, all classes
+                are reported.
 
         Returns:
-            A dict containing either the top-5 classifications for the given image (in
-            decreasing order of confidence scores), or a failure message if no
-            preprocessed image was provided.
+            A dict containing either the top classifications for each of the given
+            image crops (in decreasing order of confidence scores), or a failure message
+            if no preprocessed image was provided.
         """
 
-        return self.batch_predict([filepath], [img])[0]
+        return self.batch_predict([filepath], [imgs], num_classes=num_classes)[0]
 
     def batch_predict(
-        self, filepaths: list[str], imgs: list[Optional[PreprocessedImage]]
+        self,
+        filepaths: list[str],
+        imgs_list: list[Optional[list[PreprocessedImage]]],
+        num_classes: Optional[int] = 5,
     ) -> list[dict[str, Any]]:
         """Runs inference on a batch of preprocessed images.
 
         Args:
             filepaths:
                 List of image locations to run inference on. Used for reporting purposes
                 only, and not for loading the images.
-            imgs:
-                List of preprocessed images to run inference on. If an image is `None`,
-                a corresponding failure message is reported back.
+            imgs_list:
+                List of lists of preprocessed image crops to run inference on, one list
+                per image. If a list is `None`, a corresponding failure message is
+                reported back. If a list is empty, an empty list of classifications is
+                reported back.
+            num_classes:
+                Number of top classes to report for each crop. If `None`, all classes
+                are reported.
 
         Returns:
-            A list of dict results. Each dict result contains either the top-5
-            classifications for the corresponding image (in decreasing order of
-            confidence scores), or a failure message if no preprocessed image was
-            provided.
+            A list of dict results. Each dict result contains `classifications_list`
+            with the top classifications for each corresponding crop image, in the same
+            order as the crops (and in decreasing order of confidence scores for each
+            crop), or a failure message if no preprocessed image was provided.
         """
 
         predictions = {}
 
         inference_filepaths = []
         batch_arr = []
-        for filepath, img in zip(filepaths, imgs):
-            if img is None:
+
+        for filepath, imgs in zip(filepaths, imgs_list):
+            if imgs is None:
                 predictions[filepath] = {
                     "filepath": filepath,
                     "failures": [Failure.CLASSIFIER.name],
                 }
-            else:
+                continue
+
+            predictions[filepath] = {"filepath": filepath, "classifications_list": []}
+
+            for img in imgs:
                 inference_filepaths.append(filepath)
-                batch_arr.append(img.arr / 255)
+                batch_arr.append(img.arr / 255.0)
+
         if not batch_arr:
-            return list(predictions.values())
+            return [predictions[fp] for fp in filepaths if fp in predictions]
+
         batch_arr = np.stack(batch_arr, axis=0, dtype=np.float32)
 
-        # On MPS devices, make sure all MPS operations happen serially.  This
-        # lock is a no-op on non-MPS devices.
-        with mps_inference_lock(self.device):
-            batch_tensor = torch.from_numpy(batch_arr).to(self.device)
-            logits = self.model(batch_tensor).cpu()
+        # Process crops in chunks of at most `max_crop_batch_size` so that a batch of
+        # images with many bounding boxes does not cause excessive memory allocation
+        # (especially on MPS unified memory) or hold the MPS lock for too long.
+        max_crop_batch_size = 8
+        logits_list = []
+        for idx in range(0, len(batch_arr), max_crop_batch_size):
+            chunk_arr = batch_arr[idx : idx + max_crop_batch_size]
+            # On MPS devices, make sure all MPS operations happen serially.  This
+            # lock is a no-op on non-MPS devices.
+            with mps_inference_lock(self.device):
+                chunk_tensor = torch.from_numpy(chunk_arr).to(self.device)
+                logits_list.append(self.model(chunk_tensor).cpu())
+                del chunk_tensor
+        if self.device == "mps":
+            with mps_inference_lock(self.device):
+                torch.mps.empty_cache()
+        logits = torch.cat(logits_list, dim=0)
         scores = torch.softmax(logits, dim=-1)
-        scores, indices = torch.topk(scores, k=5, dim=-1)
-
-        for file_idx, (filepath, scores_arr, indices_arr) in enumerate(
-            zip(inference_filepaths, scores.numpy(), indices.numpy())
+        num_labels = scores.shape[-1]
+        k = num_labels if num_classes is None else min(num_classes, num_labels)
+        scores, indices = torch.topk(scores, k=k, dim=-1)
+
+        for crop_idx, (filepath, scores_arr, indices_arr) in enumerate(
+            zip(
+                inference_filepaths,
+                scores.numpy(),
+                indices.numpy(),
+            )
         ):
 
-            predictions[filepath] = {
-                "filepath": filepath,
-                "classifications": {
-                    "classes": [self.labels[idx] for idx in indices_arr],
-                    "scores": scores_arr.tolist(),
-                },
+            classification = {
+                "classes": [self.labels[idx] for idx in indices_arr],
+                "scores": scores_arr.tolist(),
             }
 
             if hasattr(self, "target_idx"):
-                predictions[filepath]["classifications"].update(
+                classification.update(
                     {
                         "target_classes": self.target_labels,
                         "target_logits": [
-                            float(logits[file_idx][idx]) for idx in self.target_idx
+                            float(logits[crop_idx][idx]) for idx in self.target_idx
                         ],
                     }
                 )
 
-        return [predictions[filepath] for filepath in filepaths]
+            predictions[filepath]["classifications_list"].append(classification)
+
+        return [
+            predictions[filepath] for filepath in filepaths if filepath in predictions
+        ]