Add FiftyOne dataset inference example + multi-GPU inference

README.md

@@ -54,7 +54,7 @@ This makes visdet significantly easier to install and deploy compared to the ori
 
 Without access full training logs (loss plots etc.), it can be impossible to know if you have your own implementation wrong or not. Ideally, eventually we integrate the docs, and the experiment results into the same one living documentation. We run hyperparameter search, you get the new best hyperparameters.
 
-Goals of the repo: 
+Goals of the repo:
 - Even more open than your typical open-source project, logs available, roadmap available.
 - Emphasis on all of DevEx, educating new users, production deployments and research.2
 

docs/user-guide/inference.md

@@ -54,6 +54,84 @@ for image_file in image_files:
     # Process result...
 ```
 
+### Multi-GPU Inference
+
+To run inference on multiple GPUs in a single process, pass multiple CUDA devices to `init_detector` and then infer a batch (a list) of images:
+
+```python
+from visdet.apis import inference_detector, init_detector
+
+model = init_detector(config_file, checkpoint_file, device="cuda:0,1")
+results = inference_detector(model, image_files)  # list[DetDataSample]
+```
+
+## FiftyOne (Voxel51) Dataset
+
+If you use [FiftyOne](https://voxel51.com/fiftyone/) for dataset management and visualization, you can run `visdet` inference over a `fiftyone.Dataset` and attach the predictions back onto each sample.
+
+```python
+import fiftyone as fo
+
+from visdet.apis import inference_detector, init_detector
+from visdet.utils import detections_to_fiftyone
+
+# 1) Load your model
+config_file = "configs/faster_rcnn/faster_rcnn_r50_fpn_1x_coco.py"
+checkpoint_file = "checkpoints/faster_rcnn_r50_fpn_1x_coco.pth"
+model = init_detector(config_file, checkpoint_file, device="cuda:0")
+
+# 2) Create (or load) a FiftyOne dataset
+#    You can also do: dataset = fo.load_dataset("my-dataset")
+dataset = fo.Dataset.from_images_dir(
+    "path/to/images",
+    name="my-images",
+    overwrite=True,
+)
+dataset.compute_metadata()  # populates sample.metadata.width/height
+
+classes = model.dataset_meta.get("classes", [])
+score_thr = 0.3
+
+# 3) Run inference and attach detections
+for sample in dataset.iter_samples(progress=True):
+    data_sample = inference_detector(model, sample.filepath)
+
+    pred = data_sample.pred_instances
+    pred = pred[pred.scores > score_thr]
+
+    width = sample.metadata.width
+    height = sample.metadata.height
+
+    dets = []
+    for bbox, label_id, score in zip(
+        pred.bboxes.cpu().numpy(),
+        pred.labels.cpu().numpy(),
+        pred.scores.cpu().numpy(),
+        strict=False,
+    ):
+        x1, y1, x2, y2 = bbox.tolist()
+
+        dets.append(
+            {
+                "label": classes[int(label_id)] if classes else str(int(label_id)),
+                "bounding_box": [
+                    x1 / width,
+                    y1 / height,
+                    (x2 - x1) / width,
+                    (y2 - y1) / height,
+                ],
+                "confidence": float(score),
+            }
+        )
+
+    sample["predictions"] = detections_to_fiftyone(dets)
+    sample.save()
+
+# 4) Visualize in the FiftyOne App
+session = fo.launch_app(dataset)
+session.wait()
+```
+
 ## Test Dataset
 
 Evaluate model on a test dataset:

tests/test_runtime/test_inference_device_parsing.py

@@ -0,0 +1,19 @@
+import pytest
+
+from visdet.apis import inference as inference_api
+
+
+def test_parse_inference_devices_single_device():
+    assert inference_api._parse_inference_devices("cpu") == ("cpu", None)
+    assert inference_api._parse_inference_devices("cuda:0") == ("cuda:0", None)
+
+
+def test_parse_inference_devices_multi_device():
+    assert inference_api._parse_inference_devices("cuda:0,1") == ("cuda:0", [0, 1])
+    assert inference_api._parse_inference_devices(["cuda:0", "cuda:1"]) == ("cuda:0", [0, 1])
+    assert inference_api._parse_inference_devices([0, 2, 3]) == ("cuda:0", [0, 2, 3])
+
+
+def test_parse_inference_devices_rejects_non_cuda_sequences():
+    with pytest.raises(ValueError):
+        inference_api._parse_inference_devices(["cpu", "cuda:0"])  # type: ignore[arg-type]

visdet/apis/inference.py

@@ -16,7 +16,7 @@
 # from visdet.cv.ops import RoIPool  # Removed - eliminating C++ ops
 from visdet.cv.transforms import Compose
 from visdet.engine.config import Config
-from visdet.engine.dataset import default_collate
+from visdet.engine.dataset import default_collate, pseudo_collate
 from visdet.engine.model.utils import revert_sync_batchnorm
 from visdet.engine.registry import init_default_scope
 from visdet.engine.runner import load_checkpoint
@@ -27,11 +27,42 @@
 from visdet.utils import ConfigType, get_test_pipeline_cfg
 
 
+def _parse_inference_devices(device: str | Sequence[str] | Sequence[int]) -> tuple[str, list[int] | None]:
+    if isinstance(device, str):
+        if device.startswith("cuda") and "," in device:
+            parts = [p.strip() for p in device.split(",") if p.strip()]
+            device_ids = [_device_to_id(p) for p in parts]
+            return f"cuda:{device_ids[0]}", device_ids
+        return device, None
+
+    device_ids = [_device_to_id(d) for d in device]
+    return f"cuda:{device_ids[0]}", device_ids
+
+
+def _device_to_id(device: str | int) -> int:
+    if isinstance(device, int):
+        return device
+
+    if device.isdigit():
+        return int(device)
+
+    if device == "cuda":
+        return 0
+
+    if device.startswith("cuda:"):
+        return int(device.split(":", 1)[1])
+
+    raise ValueError(
+        "Multi-device inference only supports CUDA devices. "
+        f"Got device={device!r}; expected e.g. 'cuda:0' or an int device id"
+    )
+
+
 def init_detector(
     config: str | Path | Config,
     checkpoint: str | None = None,
     palette: str = "none",
-    device: str = "cuda:0",
+    device: str | Sequence[str] | Sequence[int] = "cuda:0",
     cfg_options: dict | None = None,
 ) -> nn.Module:
     """Initialize a detector from config file.
@@ -45,8 +76,15 @@ def init_detector(
             is stored in checkpoint, use checkpoint's palette first, otherwise
             use externally passed palette. Currently, supports 'coco', 'voc',
             'citys' and 'random'. Defaults to none.
-        device (str): The device where the anchors will be put on.
-            Defaults to cuda:0.
+        device (str | Sequence[str] | Sequence[int]):
+            The device(s) to run inference on.
+
+            - Single device examples: ``"cpu"``, ``"cuda:0"``
+            - Multi-GPU single-process example: ``"cuda:0,1"`` or
+              ``["cuda:0", "cuda:1"]``
+
+            When multiple CUDA devices are provided, the model is wrapped in
+            :class:`visdet.engine.model.wrappers.MMDataParallel`.
         cfg_options (dict, optional): Options to override some settings in
             the used config.
 
@@ -113,7 +151,22 @@ def init_detector(
                 model.dataset_meta["palette"] = "random"
 
     model.cfg = config  # save the config in the model for convenience
-    model.to(device)
+
+    primary_device, device_ids = _parse_inference_devices(device)
+    model.to(primary_device)
+
+    if device_ids is not None and len(device_ids) > 1:
+        if not torch.cuda.is_available():
+            raise RuntimeError(f"CUDA is not available, cannot use multi-GPU device={device!r}")
+        if max(device_ids) >= torch.cuda.device_count():
+            raise RuntimeError(
+                f"Requested device_ids={device_ids} but only {torch.cuda.device_count()} CUDA device(s) are available"
+            )
+
+        from visdet.engine.model import MMDataParallel
+
+        model = MMDataParallel(model, device_ids=device_ids, output_device=device_ids[0])
+
     model.eval()
     return model
 
@@ -162,8 +215,8 @@ def inference_detector(
 
     # RoIPool check removed - eliminating C++ ops
 
-    result_list = []
-    for _i, img in enumerate(imgs):
+    data_list = []
+    for img in imgs:
         # prepare data
         if isinstance(img, np.ndarray):
             # TODO: remove img_id.
@@ -176,22 +229,18 @@ def inference_detector(
             data_["text"] = text_prompt
             data_["custom_entities"] = custom_entities
 
-        # build the data pipeline
-        data_ = test_pipeline(data_)
-
-        data_["inputs"] = [data_["inputs"]]
-        data_["data_samples"] = [data_["data_samples"]]
+        data_list.append(test_pipeline(data_))
 
-        # forward the model
-        with torch.inference_mode():
-            results = model.test_step(data_)[0]
+    batch = pseudo_collate(data_list)
 
-        result_list.append(results)
+    # forward the model
+    with torch.inference_mode():
+        results = model.test_step(batch)
 
     if not is_batch:
-        return result_list[0]
+        return results[0]
     else:
-        return result_list
+        return results
 
 
 # TODO: Awaiting refactoring

visdet/engine/model/__init__.py

@@ -34,6 +34,7 @@
     xavier_init,
 )
 from visdet.engine.model.wrappers import (
+    MMDataParallel,
     MMDistributedDataParallel,
     MMSeparateDistributedDataParallel,
     is_model_wrapper,
@@ -49,6 +50,7 @@
     # "ExponentialMovingAverage",  # Not imported - EMA is in hooks module
     "ImgDataPreprocessor",
     "KaimingInit",
+    "MMDataParallel",
     "MMDistributedDataParallel",
     "MMSeparateDistributedDataParallel",
     "ModuleDict",

visdet/engine/model/wrappers/__init__.py

@@ -3,11 +3,12 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from visdet.engine.utils.dl_utils import TORCH_VERSION
 from visdet.engine.utils.version_utils import digit_version
-from visdet.engine.model.wrappers.distributed import MMDistributedDataParallel
+from visdet.engine.model.wrappers.distributed import MMDataParallel, MMDistributedDataParallel
 from visdet.engine.model.wrappers.seperate_distributed import MMSeparateDistributedDataParallel
 from visdet.engine.model.wrappers.utils import is_model_wrapper
 
 __all__ = [
+    "MMDataParallel",
     "MMDistributedDataParallel",
     "MMSeparateDistributedDataParallel",
     "is_model_wrapper",

visdet/engine/model/wrappers/distributed.py

@@ -18,6 +18,29 @@
 MODEL_WRAPPERS.register_module(module=DataParallel)
 
 
+@MODEL_WRAPPERS.register_module(force=True)
+class MMDataParallel(DataParallel):
+    """A data parallel model wrapper for inference.
+
+    Compared with :class:`torch.nn.parallel.DataParallel`, this wrapper adds
+    :meth:`val_step` and :meth:`test_step` so it can be used with visdet's
+    higher-level APIs that expect these methods.
+
+    Note:
+        This wrapper parallelizes the model forward pass across multiple GPUs
+        within a single process. For best performance, pass a batch of images
+        (e.g. a list of image paths/arrays) to :func:`visdet.apis.inference_detector`.
+    """
+
+    def val_step(self, data: dict | tuple | list) -> list:
+        data = self.module.data_preprocessor(data, False)
+        return self(**data, mode="predict")
+
+    def test_step(self, data: dict | tuple | list) -> list:
+        data = self.module.data_preprocessor(data, False)
+        return self(**data, mode="predict")
+
+
 @MODEL_WRAPPERS.register_module(force=True)
 class MMDistributedDataParallel(DistributedDataParallel):
     """A distributed model wrapper used for training,testing and validation in