Update structure to support subset

Author: sooahleex

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/training/calib/000001.txt renamed to tests/assets/kitti_dataset/kitti_3d/calib/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