Repository: microsoft/torchgeo

Issue:
Using VHR10 dataset with AugmentationSequential
https://github.com/microsoft/torchgeo/blob/0c7d46be10c47fdb370364fddb7328d1b0661ac2/torchgeo/datasets/nwpu.py#L143

sample (dict) has the key: `'label'`, but should have the key `'mask'`

The reason:

If using torchgeo augmentations
https://github.com/microsoft/torchgeo/blob/main/torchgeo/transforms/transforms.py#L40

The parameter `data_keys` in `kornia.augmentation.AugmentationSequential` accepts only `[“input”, “mask”, “bbox”, “bbox_xyxy”, “bbox_xywh”, “keypoints”]` as stated in [kornia doc](https://kornia.readthedocs.io/en/latest/augmentation.container.html#augmentation-sequential)



Discussion:

Resolving patch:
diff --git a/torchgeo/datasets/nwpu.py b/torchgeo/datasets/nwpu.py
--- a/torchgeo/datasets/nwpu.py
+++ b/torchgeo/datasets/nwpu.py
@@ -4,10 +4,12 @@
 """NWPU VHR-10 dataset."""
 
 import os
-from typing import Any, Callable, Dict, Optional
+from typing import Any, Callable, Dict, List, Optional
 
+import matplotlib.pyplot as plt
 import numpy as np
 import torch
+from matplotlib import patches
 from PIL import Image
 from torch import Tensor
 
@@ -15,6 +17,95 @@
 from .utils import check_integrity, download_and_extract_archive, download_url
 
 
+def convert_coco_poly_to_mask(
+    segmentations: List[int], height: int, width: int
+) -> Tensor:
+    """Convert coco polygons to mask tensor.
+
+    Args:
+        segmentations (List[int]): polygon coordinates
+        height (int): image height
+        width (int): image width
+
+    Returns:
+        Tensor: Mask tensor
+    """
+    from pycocotools import mask as coco_mask  # noqa: F401
+
+    masks = []
+    for polygons in segmentations:
+        rles = coco_mask.frPyObjects(polygons, height, width)
+        mask = coco_mask.decode(rles)
+        mask = torch.as_tensor(mask, dtype=torch.uint8)
+        mask = mask.any(dim=2)
+        masks.append(mask)
+    if masks:
+        masks_tensor = torch.stack(masks, dim=0)
+    else:
+        masks_tensor = torch.zeros((0, height, width), dtype=torch.uint8)
+    return masks_tensor
+
+
+class ConvertCocoAnnotations:
+    """Callable for converting the boxes, masks and labels into tensors.
+
+    This is a modified version of ConvertCocoPolysToMask() from torchvision found in
+    https://github.com/pytorch/vision/blob/v0.14.0/references/detection/coco_utils.py
+    """
+
+    def __call__(self, sample: Dict[str, Any]) -> Dict[str, Any]:
+        """Converts MS COCO fields (boxes, masks & labels) from list of ints to tensors.
+
+        Args:
+            sample: Sample
+
+        Returns:
+            Processed sample
+        """
+        image = sample["image"]
+        _, h, w = image.size()
+        target = sample["label"]
+
+        image_id = target["image_id"]
+        image_id = torch.tensor([image_id])
+
+        anno = target["annotations"]
+
+        anno = [obj for obj in anno if obj["iscrowd"] == 0]
+
+        bboxes = [obj["bbox"] for obj in anno]
+        # guard against no boxes via resizing
+        boxes = torch.as_tensor(bboxes, dtype=torch.float32).reshape(-1, 4)
+        boxes[:, 2:] += boxes[:, :2]
+        boxes[:, 0::2].clamp_(min=0, max=w)
+        boxes[:, 1::2].clamp_(min=0, max=h)
+
+        categories = [obj["category_id"] for obj in anno]
+        classes = torch.tensor(categories, dtype=torch.int64)
+
+        if "segmentation" in anno[0]:
+            segmentations = [obj["segmentation"] for obj in anno]
+        else:
+            segmentations = []
+        masks = convert_coco_poly_to_mask(segmentations, h, w)
+
+        keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0])
+        boxes = boxes[keep]
+        classes = classes[keep]
+
+        target = {"boxes": boxes, "labels": classes, "image_id": image_id}
+        if masks.nelement() > 0:
+            masks = masks[keep]
+            target["masks"] = masks
+
+        # for conversion to coco api
+        area = torch.tensor([obj["area"] for obj in anno])
+        iscrowd = torch.tensor([obj["iscrowd"] for obj in anno])
+        target["area"] = area
+        target["iscrowd"] = iscrowd
+        return {"image": image, "label": target}
+
+
 class VHR10(NonGeoDataset):
     """NWPU VHR-10 dataset.
 
@@ -76,6 +167,20 @@ class VHR10(NonGeoDataset):
         "md5": "7c76ec50c17a61bb0514050d20f22c08",
     }
 
+    categories = [
+        "background",
+        "airplane",
+        "ships",
+        "storage tank",
+        "baseball diamond",
+        "tennis court",
+        "basketball court",
+        "ground track field",
+        "harbor",
+        "bridge",
+        "vehicle",
+    ]
+
     def __init__(
         self,
         root: str = "data",
@@ -96,6 +201,7 @@ def __init__(
 
         Raises:
             AssertionError: if ``split`` argument is invalid
+            ImportError: if ``split="positive"`` and pycocotools is not installed
             RuntimeError: if ``download=False`` and data is not found, or checksums
                 don't match
         """
@@ -130,6 +236,9 @@ def __init__(
                 )
             )
 
+            self.coco_convert = ConvertCocoAnnotations()
+            self.ids = list(sorted(self.coco.imgs.keys()))
+
     def __getitem__(self, index: int) -> Dict[str, Any]:
         """Return an index within the dataset.
 
@@ -140,7 +249,19 @@ def __getitem__(self, index: int) -> Dict[str, Any]:
             data and label at that index
         """
         id_ = index % len(self) + 1
-        sample = {"image": self._load_image(id_), "label": self._load_target(id_)}
+
+        sample: Dict[str, Any] = {
+            "image": self._load_image(id_),
+            "label": self._load_target(id_),
+        }
+
+        if sample["label"]["annotations"]:
+            sample = self.coco_convert(sample)
+            sample["labels"] = sample["label"]["labels"]
+            sample["boxes"] = sample["label"]["boxes"]
+            if "masks" in sample["label"]:
+                sample["masks"] = sample["label"]["masks"]
+            del sample["label"]
 
         if self.transforms is not None:
             sample = self.transforms(sample)
@@ -154,7 +275,7 @@ def __len__(self) -> int:
             length of the dataset
         """
         if self.split == "positive":
-            return 650
+            return len(self.ids)
         else:
             return 150
 
@@ -192,7 +313,7 @@ def _load_target(self, id_: int) -> Dict[str, Any]:
         # Images in the "negative" image set have no annotations
         annot = []
         if self.split == "positive":
-            annot = self.coco.loadAnns(self.coco.getAnnIds(id_))
+            annot = self.coco.loadAnns(self.coco.getAnnIds(id_ - 1))
 
         target = dict(image_id=id_, annotations=annot)
 
@@ -244,3 +365,183 @@ def _download(self) -> None:
                 self.target_meta["filename"],
                 self.target_meta["md5"] if self.checksum else None,
             )
+
+    def plot(
+        self,
+        sample: Dict[str, Tensor],
+        show_titles: bool = True,
+        suptitle: Optional[str] = None,
+        show_feats: Optional[str] = "both",
+        box_alpha: float = 0.7,
+        mask_alpha: float = 0.7,
+    ) -> plt.Figure:
+        """Plot a sample from the dataset.
+
+        Args:
+            sample: a sample returned by :meth:`__getitem__`
+            suptitle: optional string to use as a suptitle
+            show_titles: flag indicating whether to show titles above each panel
+            show_feats: optional string to pick features to be shown: boxes, masks, both
+            box_alpha: alpha value of box
+            mask_alpha: alpha value of mask
+
+        Returns:
+            a matplotlib Figure with the rendered sample
+
+        Raises:
+            AssertionError: if ``show_feats`` argument is invalid
+            ImportError: if plotting masks and scikit-image is not installed
+
+        .. versionadded:: 0.4
+        """
+        assert show_feats in {"boxes", "masks", "both"}
+
+        if self.split == "negative":
+            plt.imshow(sample["image"].permute(1, 2, 0))
+            axs = plt.gca()
+            axs.axis("off")
+
+            if suptitle is not None:
+                plt.suptitle(suptitle)
+            return plt.gcf()
+
+        if show_feats != "boxes":
+            try:
+                from skimage.measure import find_contours  # noqa: F401
+            except ImportError:
+                raise ImportError(
+                    "scikit-image is not installed and is required to plot masks."
+                )
+
+        image = sample["image"].permute(1, 2, 0).numpy()
+        boxes = sample["boxes"].cpu().numpy()
+        labels = sample["labels"].cpu().numpy()
+        if "masks" in sample:
+            masks = [mask.squeeze().cpu().numpy() for mask in sample["masks"]]
+
+        n_gt = len(boxes)
+
+        ncols = 1
+        show_predictions = "prediction_labels" in sample
+
+        if show_predictions:
+            show_pred_boxes = False
+            show_pred_masks = False
+            prediction_labels = sample["prediction_labels"].numpy()
+            prediction_scores = sample["prediction_scores"].numpy()
+            if "prediction_boxes" in sample:
+                prediction_boxes = sample["prediction_boxes"].numpy()
+                show_pred_boxes = True
+            if "prediction_masks" in sample:
+                prediction_masks = sample["prediction_masks"].numpy()
+                show_pred_masks = True
+
+            n_pred = len(prediction_labels)
+            ncols += 1
+
+        # Display image
+        fig, axs = plt.subplots(ncols=ncols, figsize=(ncols * 10, 10))
+        if not isinstance(axs, np.ndarray):
+            axs = [axs]
+        axs[0].imshow(image)
+        axs[0].axis("off")
+
+        cm = plt.get_cmap("gist_rainbow")
+        for i in range(n_gt):
+            class_num = labels[i]
+            color = cm(class_num / len(self.categories))
+
+            # Add bounding boxes
+            x1, y1, x2, y2 = boxes[i]
+            if show_feats in {"boxes", "both"}:
+                p = patches.Rectangle(
+                    (x1, y1),
+                    x2 - x1,
+                    y2 - y1,
+                    linewidth=2,
+                    alpha=box_alpha,
+                    linestyle="dashed",
+                    edgecolor=color,
+                    facecolor="none",
+                )
+                axs[0].add_patch(p)
+
+            # Add labels
+            label = self.categories[class_num]
+            caption = label
+            axs[0].text(
+                x1, y1 - 8, caption, color="white", size=11, backgroundcolor="none"
+            )
+
+            # Add masks
+            if show_feats in {"masks", "both"} and "masks" in sample:
+                mask = masks[i]
+                contours = find_contours(mask, 0.5)
+                for verts in contours:
+                    verts = np.fliplr(verts)
+                    p = patches.Polygon(
+                        verts, facecolor=color, alpha=mask_alpha, edgecolor="white"
+                    )
+                    axs[0].add_patch(p)
+
+            if show_titles:
+                axs[0].set_title("Ground Truth")
+
+        if show_predictions:
+            axs[1].imshow(image)
+            axs[1].axis("off")
+            for i in range(n_pred):
+                score = prediction_scores[i]
+                if score < 0.5:
+                    continue
+
+                class_num = prediction_labels[i]
+                color = cm(class_num / len(self.categories))
+
+                if show_pred_boxes:
+                    # Add bounding boxes
+                    x1, y1, x2, y2 = prediction_boxes[i]
+                    p = patches.Rectangle(
+                        (x1, y1),
+                        x2 - x1,
+                        y2 - y1,
+                        linewidth=2,
+                        alpha=box_alpha,
+                        linestyle="dashed",
+                        edgecolor=color,
+                        facecolor="none",
+                    )
+                    axs[1].add_patch(p)
+
+                    # Add labels
+                    label = self.categories[class_num]
+                    caption = f"{label} {score:.3f}"
+                    axs[1].text(
+                        x1,
+                        y1 - 8,
+                        caption,
+                        color="white",
+                        size=11,
+                        backgroundcolor="none",
+                    )
+
+                # Add masks
+                if show_pred_masks:
+                    mask = prediction_masks[i]
+                    contours = find_contours(mask, 0.5)
+                    for verts in contours:
+                        verts = np.fliplr(verts)
+                        p = patches.Polygon(
+                            verts, facecolor=color, alpha=mask_alpha, edgecolor="white"
+                        )
+                        axs[1].add_patch(p)
+
+            if show_titles:
+                axs[1].set_title("Prediction")
+
+        plt.tight_layout()
+
+        if suptitle is not None:
+            plt.suptitle(suptitle)
+
+        return fig


Test patch:
diff --git a/tests/data/nwpu/NWPU VHR-10 dataset.rar b/tests/data/nwpu/NWPU VHR-10 dataset.rar
new file mode 100644
Binary files /dev/null and b/tests/data/nwpu/NWPU VHR-10 dataset.rar differ
diff --git a/tests/data/nwpu/annotations.json b/tests/data/nwpu/annotations.json
new file mode 100644
--- /dev/null
+++ b/tests/data/nwpu/annotations.json
@@ -0,0 +1 @@
+{"images": [{"file_name": "001.jpg", "height": 8, "width": 8, "id": 0}, {"file_name": "002.jpg", "height": 8, "width": 8, "id": 1}, {"file_name": "003.jpg", "height": 8, "width": 8, "id": 2}, {"file_name": "004.jpg", "height": 8, "width": 8, "id": 3}, {"file_name": "005.jpg", "height": 8, "width": 8, "id": 4}], "annotations": [{"id": 0, "image_id": 0, "category_id": 1, "area": 4.0, "bbox": [4, 4, 2, 2], "iscrowd": 0}, {"id": 1, "image_id": 1, "category_id": 1, "area": 4.0, "bbox": [4, 4, 2, 2], "segmentation": [[1, 1, 2, 2, 3, 3, 4, 5, 5]], "iscrowd": 0}, {"id": 2, "image_id": 2, "category_id": 1, "area": 4.0, "bbox": [4, 4, 2, 2], "segmentation": [[1, 1, 2, 2, 3, 3, 4, 5, 5]], "iscrowd": 0}, {"id": 3, "image_id": 3, "category_id": 1, "area": 4.0, "bbox": [4, 4, 2, 2], "segmentation": [[1, 1, 2, 2, 3, 3, 4, 5, 5]], "iscrowd": 0}, {"id": 4, "image_id": 4, "category_id": 1, "area": 4.0, "bbox": [4, 4, 2, 2], "segmentation": [[1, 1, 2, 2, 3, 3, 4, 5, 5]], "iscrowd": 0}]}
\ No newline at end of file
diff --git a/tests/data/nwpu/data.py b/tests/data/nwpu/data.py
new file mode 100644
--- /dev/null
+++ b/tests/data/nwpu/data.py
@@ -0,0 +1,87 @@
+import json
+import os
+import shutil
+import subprocess
+from copy import deepcopy
+
+import numpy as np
+from PIL import Image
+from torchvision.datasets.utils import calculate_md5
+
+ANNOTATION_FILE = {"images": [], "annotations": []}
+
+
+def write_data(path: str, img: np.ndarray) -> None:
+    img = np.repeat(img[:, :, np.newaxis], 3, axis=2)
+    img = Image.fromarray(img)
+    img.save(path)
+
+
+def generate_test_data(root: str, n_imgs: int = 3) -> str:
+    folder_path = os.path.join(root, "NWPU VHR-10 dataset")
+    pos_img_dir = os.path.join(folder_path, "positive image set")
+    neg_img_dir = os.path.join(folder_path, "negative image set")
+    ann_file = os.path.join(folder_path, "annotations.json")
+    ann_file2 = os.path.join(root, "annotations.json")
+
+    if not os.path.exists(pos_img_dir):
+        os.makedirs(pos_img_dir)
+    if not os.path.exists(neg_img_dir):
+        os.makedirs(neg_img_dir)
+
+    for img_id in range(1, n_imgs + 1):
+        pos_img_name = os.path.join(pos_img_dir, f"00{img_id}.jpg")
+        neg_img_name = os.path.join(neg_img_dir, f"00{img_id}.jpg")
+
+        img = np.random.randint(255, size=(8, 8), dtype=np.dtype("uint8"))
+        write_data(pos_img_name, img)
+        write_data(neg_img_name, img)
+
+        img_name = os.path.basename(pos_img_name)
+
+        ANNOTATION_FILE["images"].append(
+            {"file_name": img_name, "height": 8, "width": 8, "id": img_id - 1}
+        )
+
+    ann = 0
+    for i, img in enumerate(ANNOTATION_FILE["images"]):
+        annot = {
+            "id": ann,
+            "image_id": img["id"],
+            "category_id": 1,
+            "area": 4.0,
+            "bbox": [4, 4, 2, 2],
+            "segmentation": [[1, 1, 2, 2, 3, 3, 4, 5, 5]],
+            "iscrowd": 0,
+        }
+        if i != 0:
+            ANNOTATION_FILE["annotations"].append(annot)
+        else:
+            noseg_annot = deepcopy(annot)
+            del noseg_annot["segmentation"]
+            ANNOTATION_FILE["annotations"].append(noseg_annot)
+        ann += 1
+
+    with open(ann_file, "w") as j:
+        json.dump(ANNOTATION_FILE, j)
+
+    with open(ann_file2, "w") as j:
+        json.dump(ANNOTATION_FILE, j)
+
+    # Create rar file
+    subprocess.run(
+        ["rar", "a", "NWPU VHR-10 dataset.rar", "-m5", "NWPU VHR-10 dataset"],
+        capture_output=True,
+        check=True,
+    )
+
+    annotations_md5 = calculate_md5(ann_file)
+    archive_md5 = calculate_md5("NWPU VHR-10 dataset.rar")
+    shutil.rmtree(folder_path)
+
+    return f"archive md5: {archive_md5}, annotation md5: {annotations_md5}"
+
+
+if __name__ == "__main__":
+    md5 = generate_test_data(os.getcwd(), 5)
+    print(md5)
diff --git a/tests/data/vhr10/NWPU VHR-10 dataset.rar b/tests/data/vhr10/NWPU VHR-10 dataset.rar
deleted file mode 100644
Binary files a/tests/data/vhr10/NWPU VHR-10 dataset.rar and /dev/null differ
diff --git a/tests/data/vhr10/annotations.json b/tests/data/vhr10/annotations.json
deleted file mode 100644
--- a/tests/data/vhr10/annotations.json
+++ /dev/null
@@ -1,134 +0,0 @@
-{
-    "info": {
-        "description": null,
-        "url": null,
-        "version": null,
-        "year": 2021,
-        "contributor": null,
-        "date_created": "2021-01-01 00:00:00"
-    },
-    "licenses": [
-        {
-            "url": null,
-            "id": 0,
-            "name": null
-        }
-    ],
-    "images": [
-        {
-            "license": 0,
-            "url": null,
-            "file_name": "001.jpg",
-            "height": 1,
-            "width": 1,
-            "date_captured": null,
-            "id": 0
-        },
-        {
-            "license": 0,
-            "url": null,
-            "file_name": "002.jpg",
-            "height": 1,
-            "width": 1,
-            "date_captured": null,
-            "id": 1
-        }
-    ],
-    "type": "instances",
-    "annotations": [
-        {
-            "id": 0,
-            "image_id": 0,
-            "category_id": 1,
-            "segmentation": [
-                [
-                    1,
-                    2,
-                    3,
-                    4
-                ]
-            ],
-            "area": 1.0,
-            "bbox": [
-                1,
-                2,
-                3,
-                4
-            ],
-            "iscrowd": 0
-        },
-        {
-            "id": 1,
-            "image_id": 1,
-            "category_id": 1,
-            "segmentation": [
-                [
-                    1,
-                    2,
-                    3,
-                    4
-                ]
-            ],
-            "area": 1.0,
-            "bbox": [
-                1,
-                2,
-                3,
-                4
-            ],
-            "iscrowd": 0
-        }
-    ],
-    "categories": [
-        {
-            "supercategory": null,
-            "id": 1,
-            "name": "airplane"
-        },
-        {
-            "supercategory": null,
-            "id": 2,
-            "name": "ship"
-        },
-        {
-            "supercategory": null,
-            "id": 3,
-            "name": "storage_tank"
-        },
-        {
-            "supercategory": null,
-            "id": 4,
-            "name": "baseball_diamond"
-        },
-        {
-            "supercategory": null,
-            "id": 5,
-            "name": "tennis_court"
-        },
-        {
-            "supercategory": null,
-            "id": 6,
-            "name": "basketball_court"
-        },
-        {
-            "supercategory": null,
-            "id": 7,
-            "name": "ground_track_field"
-        },
-        {
-            "supercategory": null,
-            "id": 8,
-            "name": "harbor"
-        },
-        {
-            "supercategory": null,
-            "id": 9,
-            "name": "bridge"
-        },
-        {
-            "supercategory": null,
-            "id": 10,
-            "name": "vehicle"
-        }
-    ]
-}
diff --git a/tests/datasets/test_advance.py b/tests/datasets/test_advance.py
--- a/tests/datasets/test_advance.py
+++ b/tests/datasets/test_advance.py
@@ -49,7 +49,7 @@ def mocked_import(name: str, *args: Any, **kwargs: Any) -> Any:
         monkeypatch.setattr(builtins, "__import__", mocked_import)
 
     def test_getitem(self, dataset: ADVANCE) -> None:
-        pytest.importorskip("scipy", minversion="1.2")
+        pytest.importorskip("scipy", minversion="1.6.2")
         x = dataset[0]
         assert isinstance(x, dict)
         assert isinstance(x["image"], torch.Tensor)
diff --git a/tests/datasets/test_nwpu.py b/tests/datasets/test_nwpu.py
--- a/tests/datasets/test_nwpu.py
+++ b/tests/datasets/test_nwpu.py
@@ -7,6 +7,7 @@
 from pathlib import Path
 from typing import Any
 
+import matplotlib.pyplot as plt
 import pytest
 import torch
 import torch.nn as nn
@@ -32,13 +33,13 @@ def dataset(
         pytest.importorskip("rarfile", minversion="3")
         monkeypatch.setattr(torchgeo.datasets.nwpu, "download_url", download_url)
         monkeypatch.setattr(torchgeo.datasets.utils, "download_url", download_url)
-        url = os.path.join("tests", "data", "vhr10", "NWPU VHR-10 dataset.rar")
+        url = os.path.join("tests", "data", "nwpu", "NWPU VHR-10 dataset.rar")
         monkeypatch.setitem(VHR10.image_meta, "url", url)
-        md5 = "e5c38351bd948479fe35a71136aedbc4"
+        md5 = "5fddb0dfd56a80638831df9f90cbf37a"
         monkeypatch.setitem(VHR10.image_meta, "md5", md5)
-        url = os.path.join("tests", "data", "vhr10", "annotations.json")
+        url = os.path.join("tests", "data", "nwpu", "annotations.json")
         monkeypatch.setitem(VHR10.target_meta, "url", url)
-        md5 = "16fc6aa597a19179dad84151cc221873"
+        md5 = "833899cce369168e0d4ee420dac326dc"
         monkeypatch.setitem(VHR10.target_meta, "md5", md5)
         root = str(tmp_path)
         split = request.param
@@ -46,25 +47,30 @@ def dataset(
         return VHR10(root, split, transforms, download=True, checksum=True)
 
     @pytest.fixture
-    def mock_missing_module(self, monkeypatch: MonkeyPatch) -> None:
+    def mock_missing_modules(self, monkeypatch: MonkeyPatch) -> None:
         import_orig = builtins.__import__
 
         def mocked_import(name: str, *args: Any, **kwargs: Any) -> Any:
-            if name == "pycocotools.coco":
+            if name in {"pycocotools.coco", "skimage.measure"}:
                 raise ImportError()
             return import_orig(name, *args, **kwargs)
 
         monkeypatch.setattr(builtins, "__import__", mocked_import)
 
     def test_getitem(self, dataset: VHR10) -> None:
-        x = dataset[0]
-        assert isinstance(x, dict)
-        assert isinstance(x["image"], torch.Tensor)
-        assert isinstance(x["label"], dict)
+        for i in range(2):
+            x = dataset[i]
+            assert isinstance(x, dict)
+            assert isinstance(x["image"], torch.Tensor)
+            if dataset.split == "positive":
+                assert isinstance(x["labels"], torch.Tensor)
+                assert isinstance(x["boxes"], torch.Tensor)
+                if "masks" in x:
+                    assert isinstance(x["masks"], torch.Tensor)
 
     def test_len(self, dataset: VHR10) -> None:
         if dataset.split == "positive":
-            assert len(dataset) == 650
+            assert len(dataset) == 5
         elif dataset.split == "negative":
             assert len(dataset) == 150
 
@@ -72,7 +78,7 @@ def test_add(self, dataset: VHR10) -> None:
         ds = dataset + dataset
         assert isinstance(ds, ConcatDataset)
         if dataset.split == "positive":
-            assert len(ds) == 1300
+            assert len(ds) == 10
         elif dataset.split == "negative":
             assert len(ds) == 300
 
@@ -88,7 +94,7 @@ def test_not_downloaded(self, tmp_path: Path) -> None:
             VHR10(str(tmp_path))
 
     def test_mock_missing_module(
-        self, dataset: VHR10, mock_missing_module: None
+        self, dataset: VHR10, mock_missing_modules: None
     ) -> None:
         if dataset.split == "positive":
             with pytest.raises(
@@ -96,3 +102,29 @@ def test_mock_missing_module(
                 match="pycocotools is not installed and is required to use this datase",
             ):
                 VHR10(dataset.root, dataset.split)
+
+            with pytest.raises(
+                ImportError,
+                match="scikit-image is not installed and is required to plot masks",
+            ):
+                x = dataset[0]
+                dataset.plot(x)
+
+    def test_plot(self, dataset: VHR10) -> None:
+        pytest.importorskip("skimage", minversion="0.18")
+        x = dataset[1].copy()
+        dataset.plot(x, suptitle="Test")
+        plt.close()
+        dataset.plot(x, show_titles=False)
+        plt.close()
+        if dataset.split == "positive":
+            scores = [0.7, 0.3, 0.7]
+            for i in range(3):
+                x = dataset[i]
+                x["prediction_labels"] = x["labels"]
+                x["prediction_boxes"] = x["boxes"]
+                x["prediction_scores"] = torch.Tensor([scores[i]])
+                if "masks" in x:
+                    x["prediction_masks"] = x["masks"]
+                    dataset.plot(x, show_feats="masks")
+                    plt.close()
diff --git a/tests/datasets/test_utils.py b/tests/datasets/test_utils.py
--- a/tests/datasets/test_utils.py
+++ b/tests/datasets/test_utils.py
@@ -87,7 +87,7 @@ def test_mock_missing_module(mock_missing_module: None) -> None:
     [
         os.path.join("cowc_detection", "COWC_Detection_Columbus_CSUAV_AFRL.tbz"),
         os.path.join("cowc_detection", "COWC_test_list_detection.txt.bz2"),
-        os.path.join("vhr10", "NWPU VHR-10 dataset.rar"),
+        os.path.join("nwpu", "NWPU VHR-10 dataset.rar"),
         os.path.join("landcoverai", "landcover.ai.v1.zip"),
         os.path.join("chesapeake", "BAYWIDE", "Baywide_13Class_20132014.zip"),
         os.path.join("sen12ms", "ROIs1158_spring_lc.tar.gz"),
@@ -107,7 +107,7 @@ def test_missing_rarfile(mock_missing_module: None) -> None:
         match="rarfile is not installed and is required to extract this dataset",
     ):
         extract_archive(
-            os.path.join("tests", "data", "vhr10", "NWPU VHR-10 dataset.rar")
+            os.path.join("tests", "data", "nwpu", "NWPU VHR-10 dataset.rar")
         )