Support subset for KITTI 3D format

CHANGELOG.md

@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### New features
 - Support KITTI 3D format
   (<https://github.com/openvinotoolkit/datumaro/pull/1619>)
+  (<https://github.com/openvinotoolkit/datumaro/pull/1621>)
 - Add PseudoLabeling transform for unlabeled dataset
   (<https://github.com/openvinotoolkit/datumaro/pull/1594>)
 

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,26 +31,37 @@ 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)
         }
 
-        ann_dir = osp.join(self._path, Kitti3dPath.LABEL_DIR)
+        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]
 
-        for labels_path in sorted(glob.glob(osp.join(ann_dir, "*.txt"), recursive=True)):
+        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 = []
 
@@ -116,17 +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=self._subset,
-                    media=PointCloud.from_file(
-                        path=osp.join(self._path, Kitti3dPath.PCD_DIR, item_id + ".bin"),
-                        extra_images=[image],
-                    ),
-                    attributes={
-                        "calib_path": osp.join(self._path, Kitti3dPath.CALIB_DIR, item_id + ".txt")
-                    },
+                    media=image,
+                    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

@@ -2,6 +2,7 @@
 #
 # SPDX-License-Identifier: MIT
 
+import os.path as osp
 from typing import List
 
 from datumaro.components.errors import DatasetImportError
@@ -16,7 +17,7 @@ class Kitti3dImporter(Importer):
 
     @classmethod
     def detect(cls, context: FormatDetectionContext) -> FormatDetectionConfidence:
-        context.require_file(f"{Kitti3dPath.PCD_DIR}/*.bin")
+        context.require_file(f"{Kitti3dPath.CALIB_DIR}/*.txt")
         cls._check_ann_file(context.require_file(f"{Kitti3dPath.LABEL_DIR}/*.txt"), context)
         return FormatDetectionConfidence.MEDIUM
 
@@ -42,4 +43,11 @@ def get_file_extensions(cls) -> List[str]:
 
     @classmethod
     def find_sources(cls, path):
-        return [{"url": path, "format": "kitti3d"}]
+        # 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