Update structure to support subset

src/datumaro/plugins/data_formats/kitti_3d/base.py

@@ -4,17 +4,18 @@
 
 import glob
 import logging
+import os
 import os.path as osp
 from typing import List, Optional, Type, TypeVar
 
-from datumaro.components.annotation import AnnotationType, Bbox, LabelCategories
+from datumaro.components.annotation import AnnotationType, Bbox
 from datumaro.components.dataset_base import DatasetItem, SubsetBase
 from datumaro.components.errors import InvalidAnnotationError
 from datumaro.components.importer import ImportContext
-from datumaro.components.media import Image, PointCloud
+from datumaro.components.media import Image
 from datumaro.util.image import find_images
 
-from .format import Kitti3dPath
+from .format import Kitti3DLabelMap, Kitti3dPath, make_kitti3d_categories
 
 T = TypeVar("T")
 
@@ -30,27 +31,38 @@ def __init__(
         ctx: Optional[ImportContext] = None,
     ):
         assert osp.isdir(path), path
-        super().__init__(subset=subset, media_type=PointCloud, ctx=ctx)
 
         self._path = path
 
-        common_attrs = {"truncated", "occluded", "alpha", "dimensions", "location", "rotation_y"}
-        self._categories = {AnnotationType.label: LabelCategories(attributes=common_attrs)}
+        if not subset:
+            folder_path = path.rsplit(Kitti3dPath.LABEL_DIR, 1)[0]
+            img_dir = osp.join(folder_path, Kitti3dPath.IMAGE_DIR)
+            if any(os.path.isdir(os.path.join(img_dir, item)) for item in os.listdir(img_dir)):
+                subset = osp.split(path)[-1]
+                self._path = folder_path
+        super().__init__(subset=subset, ctx=ctx)
+
+        self._categories = make_kitti3d_categories(Kitti3DLabelMap)
         self._items = self._load_items()
 
     def _load_items(self) -> List[DatasetItem]:
         items = []
+
         image_dir = osp.join(self._path, Kitti3dPath.IMAGE_DIR)
         image_path_by_id = {
-            osp.splitext(osp.relpath(p, image_dir))[0]: p
+            osp.split(osp.splitext(osp.relpath(p, image_dir))[0])[-1]: p
             for p in find_images(image_dir, recursive=True)
         }
 
+        if self._subset == "default":
+            ann_dir = osp.join(self._path, Kitti3dPath.LABEL_DIR)
+        else:
+            ann_dir = osp.join(self._path, Kitti3dPath.LABEL_DIR, self._subset)
+
         label_categories = self._categories[AnnotationType.label]
-        subset = osp.split(self._path)[-1]
 
-        for labels_path in sorted(glob.glob(osp.join(self._path, "*.txt"), recursive=True)):
-            item_id = osp.splitext(osp.relpath(labels_path, self._path))[0]
+        for labels_path in sorted(glob.glob(osp.join(ann_dir, "**", "*.txt"), recursive=True)):
+            item_id = osp.splitext(osp.relpath(labels_path, ann_dir))[0]
             anns = []
 
             try:
@@ -116,14 +128,18 @@ def _load_items(self) -> List[DatasetItem]:
             if image:
                 image = Image.from_file(path=image)
 
+            if self._subset == "default":
+                calib_path = osp.join(self._path, Kitti3dPath.CALIB_DIR, item_id + ".txt")
+            else:
+                calib_path = osp.join(
+                    self._path, Kitti3dPath.CALIB_DIR, self._subset, item_id + ".txt"
+                )
             items.append(
                 DatasetItem(
                     id=item_id,
-                    subset=subset,
+                    subset=self._subset,
                     media=image,
-                    attributes={
-                        "calib_path": osp.join(self._path, Kitti3dPath.CALIB_DIR, item_id + ".txt")
-                    },
+                    attributes={"calib_path": calib_path},
                     annotations=anns,
                 )
             )

src/datumaro/plugins/data_formats/kitti_3d/format.py

@@ -4,9 +4,40 @@
 
 import os.path as osp
 
+from datumaro.components.annotation import AnnotationType, LabelCategories
+
 
 class Kitti3dPath:
     PCD_DIR = osp.join("velodyne")
     IMAGE_DIR = "image_2"
     LABEL_DIR = "label_2"
     CALIB_DIR = "calib"
+
+
+Kitti3DLabelMap = [
+    "DontCare",
+    "Car",
+    "Pedestrian",
+    "Van",
+    "Truck",
+    "Cyclist",
+    "Sitter",
+    "Train",
+    "Motorcycle",
+    "Bus",
+    "Misc",
+]
+
+
+def make_kitti3d_categories(label_map=None):
+    if label_map is None:
+        label_map = Kitti3DLabelMap
+
+    categories = {}
+    common_attrs = {"truncated", "occluded", "alpha", "dimensions", "location", "rotation_y"}
+    label_categories = LabelCategories(attributes=common_attrs)
+    for label in label_map:
+        label_categories.add(label)
+    categories[AnnotationType.label] = label_categories
+
+    return categories

src/datumaro/plugins/data_formats/kitti_3d/importer.py

@@ -17,7 +17,6 @@ class Kitti3dImporter(Importer):
 
     @classmethod
     def detect(cls, context: FormatDetectionContext) -> FormatDetectionConfidence:
-        context.require_file(f"{Kitti3dPath.PCD_DIR}/*.bin")
         cls._check_ann_file(context.require_file(f"{Kitti3dPath.LABEL_DIR}/*.txt"), context)
         return FormatDetectionConfidence.MEDIUM
 
@@ -43,6 +42,11 @@ def get_file_extensions(cls) -> List[str]:
 
     @classmethod
     def find_sources(cls, path):
-        return cls._find_sources_recursive(
+        # return [{"url": path, "format": "kitti3d"}]
+        sources = cls._find_sources_recursive(
             path, "", "kitti3d", dirname=Kitti3dPath.LABEL_DIR, file_filter=lambda p: osp.isdir(p)
         )
+        if len(sources) == 0:
+            return [{"url": path, "format": "kitti3d"}]
+        else:
+            return sources

tests/assets/kitti_dataset/kitti_3d_with_subset/calib/test/000002.txt

@@ -0,0 +1,7 @@
+P0: 7.215377000000e+02 0.000000000000e+00 6.095593000000e+02 0.000000000000e+00 0.000000000000e+00 7.215377000000e+02 1.728540000000e+02 0.000000000000e+00 0.000000000000e+00 0.000000000000e+00 1.000000000000e+00 0.000000000000e+00
+P1: 7.215377000000e+02 0.000000000000e+00 6.095593000000e+02 -3.875744000000e+02 0.000000000000e+00 7.215377000000e+02 1.728540000000e+02 0.000000000000e+00 0.000000000000e+00 0.000000000000e+00 1.000000000000e+00 0.000000000000e+00
+P2: 7.215377000000e+02 0.000000000000e+00 6.095593000000e+02 4.485728000000e+01 0.000000000000e+00 7.215377000000e+02 1.728540000000e+02 2.163791000000e-01 0.000000000000e+00 0.000000000000e+00 1.000000000000e+00 2.745884000000e-03
+P3: 7.215377000000e+02 0.000000000000e+00 6.095593000000e+02 -3.395242000000e+02 0.000000000000e+00 7.215377000000e+02 1.728540000000e+02 2.199936000000e+00 0.000000000000e+00 0.000000000000e+00 1.000000000000e+00 2.729905000000e-03
+R0_rect: 9.999239000000e-01 9.837760000000e-03 -7.445048000000e-03 -9.869795000000e-03 9.999421000000e-01 -4.278459000000e-03 7.402527000000e-03 4.351614000000e-03 9.999631000000e-01
+Tr_velo_to_cam: 7.533745000000e-03 -9.999714000000e-01 -6.166020000000e-04 -4.069766000000e-03 1.480249000000e-02 7.280733000000e-04 -9.998902000000e-01 -7.631618000000e-02 9.998621000000e-01 7.523790000000e-03 1.480755000000e-02 -2.717806000000e-01
+Tr_imu_to_velo: 9.999976000000e-01 7.553071000000e-04 -2.035826000000e-03 -8.086759000000e-01 -7.854027000000e-04 9.998898000000e-01 -1.482298000000e-02 3.195559000000e-01 2.024406000000e-03 1.482454000000e-02 9.998881000000e-01 -7.997231000000e-01

tests/assets/kitti_dataset/kitti_3d_with_subset/calib/train/000000.txt

@@ -0,0 +1,7 @@
+P0: 7.070493000000e+02 0.000000000000e+00 6.040814000000e+02 0.000000000000e+00 0.000000000000e+00 7.070493000000e+02 1.805066000000e+02 0.000000000000e+00 0.000000000000e+00 0.000000000000e+00 1.000000000000e+00 0.000000000000e+00
+P1: 7.070493000000e+02 0.000000000000e+00 6.040814000000e+02 -3.797842000000e+02 0.000000000000e+00 7.070493000000e+02 1.805066000000e+02 0.000000000000e+00 0.000000000000e+00 0.000000000000e+00 1.000000000000e+00 0.000000000000e+00
+P2: 7.070493000000e+02 0.000000000000e+00 6.040814000000e+02 4.575831000000e+01 0.000000000000e+00 7.070493000000e+02 1.805066000000e+02 -3.454157000000e-01 0.000000000000e+00 0.000000000000e+00 1.000000000000e+00 4.981016000000e-03
+P3: 7.070493000000e+02 0.000000000000e+00 6.040814000000e+02 -3.341081000000e+02 0.000000000000e+00 7.070493000000e+02 1.805066000000e+02 2.330660000000e+00 0.000000000000e+00 0.000000000000e+00 1.000000000000e+00 3.201153000000e-03
+R0_rect: 9.999128000000e-01 1.009263000000e-02 -8.511932000000e-03 -1.012729000000e-02 9.999406000000e-01 -4.037671000000e-03 8.470675000000e-03 4.123522000000e-03 9.999556000000e-01
+Tr_velo_to_cam: 6.927964000000e-03 -9.999722000000e-01 -2.757829000000e-03 -2.457729000000e-02 -1.162982000000e-03 2.749836000000e-03 -9.999955000000e-01 -6.127237000000e-02 9.999753000000e-01 6.931141000000e-03 -1.143899000000e-03 -3.321029000000e-01
+Tr_imu_to_velo: 9.999976000000e-01 7.553071000000e-04 -2.035826000000e-03 -8.086759000000e-01 -7.854027000000e-04 9.998898000000e-01 -1.482298000000e-02 3.195559000000e-01 2.024406000000e-03 1.482454000000e-02 9.998881000000e-01 -7.997231000000e-01

tests/assets/kitti_dataset/kitti_3d_with_subset/calib/val/000001.txt

@@ -0,0 +1,7 @@
+P0: 7.215377000000e+02 0.000000000000e+00 6.095593000000e+02 0.000000000000e+00 0.000000000000e+00 7.215377000000e+02 1.728540000000e+02 0.000000000000e+00 0.000000000000e+00 0.000000000000e+00 1.000000000000e+00 0.000000000000e+00
+P1: 7.215377000000e+02 0.000000000000e+00 6.095593000000e+02 -3.875744000000e+02 0.000000000000e+00 7.215377000000e+02 1.728540000000e+02 0.000000000000e+00 0.000000000000e+00 0.000000000000e+00 1.000000000000e+00 0.000000000000e+00
+P2: 7.215377000000e+02 0.000000000000e+00 6.095593000000e+02 4.485728000000e+01 0.000000000000e+00 7.215377000000e+02 1.728540000000e+02 2.163791000000e-01 0.000000000000e+00 0.000000000000e+00 1.000000000000e+00 2.745884000000e-03
+P3: 7.215377000000e+02 0.000000000000e+00 6.095593000000e+02 -3.395242000000e+02 0.000000000000e+00 7.215377000000e+02 1.728540000000e+02 2.199936000000e+00 0.000000000000e+00 0.000000000000e+00 1.000000000000e+00 2.729905000000e-03
+R0_rect: 9.999239000000e-01 9.837760000000e-03 -7.445048000000e-03 -9.869795000000e-03 9.999421000000e-01 -4.278459000000e-03 7.402527000000e-03 4.351614000000e-03 9.999631000000e-01
+Tr_velo_to_cam: 7.533745000000e-03 -9.999714000000e-01 -6.166020000000e-04 -4.069766000000e-03 1.480249000000e-02 7.280733000000e-04 -9.998902000000e-01 -7.631618000000e-02 9.998621000000e-01 7.523790000000e-03 1.480755000000e-02 -2.717806000000e-01
+Tr_imu_to_velo: 9.999976000000e-01 7.553071000000e-04 -2.035826000000e-03 -8.086759000000e-01 -7.854027000000e-04 9.998898000000e-01 -1.482298000000e-02 3.195559000000e-01 2.024406000000e-03 1.482454000000e-02 9.998881000000e-01 -7.997231000000e-01

tests/assets/kitti_dataset/kitti_3d_with_subset/label_2/test/000002.txt

@@ -0,0 +1,2 @@
+Car 0.88 3 -0.69 0 190 400 380 1.60 1.57 3.23 -2.70 1.74 3.68 -1.29
+DontCare -1 -1 -10 800 160 825 185 -1 -1 -1 -1000 -1000 -1000 -10

tests/assets/kitti_dataset/kitti_3d_with_subset/label_2/train/000000.txt

@@ -0,0 +1 @@
+Pedestrian 0.00 0 -0.20 700 150 800 300 1.89 0.48 1.20 1.84 1.47 8.41 0.01

tests/assets/kitti_dataset/kitti_3d_with_subset/label_2/val/000001.txt

@@ -0,0 +1,2 @@
+Pedestrian 0.00 0 1.94 330 180 360 240 1.87 0.96 0.65 -8.50 2.07 23.02 1.59
+DontCare -1 -1 -10 600 170 620 185 -1 -1 -1 -1000 -1000 -1000 -10

tests/unit/test_kitti_3d_format.py

@@ -14,8 +14,8 @@
 from tests.utils.assets import get_test_asset_path
 from tests.utils.test_utils import compare_datasets_3d
 
-DUMMY_DATASET_DIR = get_test_asset_path("kitti_dataset", "kitti_3d", "training")
-DUMMY_SUBSET_DATASET_DIR = get_test_asset_path("kitti_dataset", "kitti_3d", "subset")
+DUMMY_DATASET_DIR = get_test_asset_path("kitti_dataset", "kitti_3d")
+DUMMY_SUBSET_DATASET_DIR = get_test_asset_path("kitti_dataset", "kitti_3d_with_subset")
 
 
 class Kitti3DImporterTest(TestCase):
@@ -40,16 +40,27 @@ def test_can_load(self):
         2. Load the dataset from the KITTI3D format.
         3. Compare the loaded dataset with the expected dataset.
         """
-        pcd1 = osp.join(DUMMY_DATASET_DIR, "velodyne", "000001.bin")
 
         image1 = Image.from_file(path=osp.join(DUMMY_DATASET_DIR, "image_2", "000001.png"))
 
         expected_label_cat = LabelCategories(
             attributes={"occluded", "truncated", "alpha", "dimensions", "location", "rotation_y"}
         )
-        expected_label_cat.add("Truck")
-        expected_label_cat.add("Car")
-        expected_label_cat.add("DontCare")
+        expected_label_list = [
+            "DontCare",
+            "Car",
+            "Pedestrian",
+            "Van",
+            "Truck",
+            "Cyclist",
+            "Sitter",
+            "Train",
+            "Motorcycle",
+            "Bus",
+            "Misc",
+        ]
+        for label in expected_label_list:
+            expected_label_cat.add(label)
         expected_dataset = Dataset.from_iterable(
             [
                 DatasetItem(
@@ -60,7 +71,7 @@ def test_can_load(self):
                             150,  # y1
                             30,  # x2-x1
                             40,  # y2-y1
-                            label=0,
+                            label=4,
                             id=0,
                             attributes={
                                 "truncated": 0.0,
@@ -92,7 +103,7 @@ def test_can_load(self):
                             170,  # y1
                             90,  # x2-x1
                             20,  # y2-y1
-                            label=2,
+                            label=0,
                             id=2,
                             attributes={
                                 "truncated": -1.0,
@@ -104,12 +115,11 @@ def test_can_load(self):
                             },
                         ),
                     ],
-                    media=PointCloud.from_file(path=pcd1, extra_images=[image1]),
+                    media=image1,
                     attributes={"calib_path": osp.join(DUMMY_DATASET_DIR, "calib", "000001.txt")},
                 ),
             ],
             categories={AnnotationType.label: expected_label_cat},
-            media_type=PointCloud,
         )
 
         parsed_dataset = Dataset.import_from(DUMMY_DATASET_DIR, "kitti3d")
@@ -134,20 +144,33 @@ def test_can_load_with_subset(self):
         expected_label_cat = LabelCategories(
             attributes={"occluded", "truncated", "alpha", "dimensions", "location", "rotation_y"}
         )
-        expected_label_cat.add("Pedestrian")
-        expected_label_cat.add("DontCare")
-        expected_label_cat.add("Car")
+        expected_label_list = [
+            "DontCare",
+            "Car",
+            "Pedestrian",
+            "Van",
+            "Truck",
+            "Cyclist",
+            "Sitter",
+            "Train",
+            "Motorcycle",
+            "Bus",
+            "Misc",
+        ]
+        for label in expected_label_list:
+            expected_label_cat.add(label)
         expected_dataset = Dataset.from_iterable(
             [
                 DatasetItem(
                     id="000000",
+                    subset="train",
                     annotations=[
                         Bbox(
-                            100,  # x1
+                            700,  # x1
                             150,  # y1
                             100,  # x2-x1
                             150,  # y2-y1
-                            label=0,
+                            label=2,
                             id=0,
                             attributes={
                                 "truncated": 0.0,
@@ -159,20 +182,25 @@ def test_can_load_with_subset(self):
                             },
                         ),
                     ],
-                    media=Image.from_numpy(data=np.ones((1, 5, 3))),
+                    media=Image.from_file(
+                        path=osp.join(DUMMY_SUBSET_DATASET_DIR, "image_2", "train", "000000.png")
+                    ),
                     attributes={
-                        "calib_path": osp.join(DUMMY_SUBSET_DATASET_DIR, "calib", "000000.txt")
+                        "calib_path": osp.join(
+                            DUMMY_SUBSET_DATASET_DIR, "calib", "train", "000000.txt"
+                        )
                     },
                 ),
                 DatasetItem(
                     id="000001",
+                    subset="val",
                     annotations=[
                         Bbox(
                             330,  # x1
                             180,  # y1
                             30,  # x2-x1
                             60,  # y2-y1
-                            label=0,
+                            label=2,
                             id=0,
                             attributes={
                                 "truncated": 0.0,
@@ -188,8 +216,8 @@ def test_can_load_with_subset(self):
                             170,  # y1
                             20,  # x2-x1
                             15,  # y2-y1
-                            label=1,
-                            id=0,
+                            label=0,
+                            id=1,
                             attributes={
                                 "truncated": -1,
                                 "occluded": -1,
@@ -200,20 +228,25 @@ def test_can_load_with_subset(self):
                             },
                         ),
                     ],
-                    media=Image.from_numpy(data=np.ones((1, 5, 3))),
+                    media=Image.from_file(
+                        path=osp.join(DUMMY_SUBSET_DATASET_DIR, "image_2", "val", "000001.png")
+                    ),
                     attributes={
-                        "calib_path": osp.join(DUMMY_SUBSET_DATASET_DIR, "calib", "000001.txt")
+                        "calib_path": osp.join(
+                            DUMMY_SUBSET_DATASET_DIR, "calib", "val", "000001.txt"
+                        )
                     },
                 ),
                 DatasetItem(
                     id="000002",
+                    subset="test",
                     annotations=[
                         Bbox(
                             0,  # x1
                             190,  # y1
                             400,  # x2-x1
                             190,  # y2-y1
-                            label=2,
+                            label=1,
                             id=0,
                             attributes={
                                 "truncated": 0.88,
@@ -228,9 +261,9 @@ def test_can_load_with_subset(self):
                             800,  # x1
                             160,  # y1
                             25,  # x2-x1
-                            15,  # y2-y1
-                            label=1,
-                            id=0,
+                            25,  # y2-y1
+                            label=0,
+                            id=1,
                             attributes={
                                 "truncated": -1,
                                 "occluded": -1,
@@ -241,9 +274,13 @@ def test_can_load_with_subset(self):
                             },
                         ),
                     ],
-                    media=Image.from_numpy(data=np.ones((1, 5, 3))),
+                    media=Image.from_file(
+                        path=osp.join(DUMMY_SUBSET_DATASET_DIR, "image_2", "test", "000002.png")
+                    ),
                     attributes={
-                        "calib_path": osp.join(DUMMY_SUBSET_DATASET_DIR, "calib", "000002.txt")
+                        "calib_path": osp.join(
+                            DUMMY_SUBSET_DATASET_DIR, "calib", "test", "000002.txt"
+                        )
                     },
                 ),
             ],