Use category name in label format; Update README

README.md

@@ -21,6 +21,8 @@ It is also the official code release of [`[PointRCNN]`](https://arxiv.org/abs/18
 
 
 ## Changelog
+[2022-07-05] Added support for [custom dataset tutorial and template](docs/CUSTOM_DATASET_TUTORIAL.md) 
+
 [2022-07-05] Added support for the 3D object detection backbone network [`Focals Conv`](https://openaccess.thecvf.com/content/CVPR2022/papers/Chen_Focal_Sparse_Convolutional_Networks_for_3D_Object_Detection_CVPR_2022_paper.pdf).
 
 [2022-02-12] Added support for using docker. Please refer to the guidance in [./docker](./docker).

docs/CUSTOM_DATASET_TUTORIAL.md

@@ -6,13 +6,11 @@ their corresponding annotations. Point clouds are supposed to be stored in `.npy
 We only consider the most basic information -- category and bounding box in the label template.
 Annotations are stored in the `.txt`. Each line represents a box in a given scene as below:
 ```
-[x y z dx dy dz heading_angle category_id]
-1.50 1.46 0.10 5.12 1.85 4.13 1.56 0
-5.54 0.57 0.41 1.08 0.74 1.95 1.57 1
+[x y z dx dy dz heading_angle category_name]
+1.50 1.46 0.10 5.12 1.85 4.13 1.56 Vehicle
+5.54 0.57 0.41 1.08 0.74 1.95 1.57 Pedestrian
 ```
 The box should in the unified 3D box definition (see [README](../README.md))
-The correspondence between `category_id` and `category_name` need to be pre-defined. 
-
 
 ## Files structure
 Files should be placed as the following folder structure:

pcdet/datasets/custom/custom_dataset.py

@@ -53,8 +53,14 @@ class CustomDataset(DatasetTemplate):
             lines = f.readlines()
 
         # [N, 8]: (x y z dx dy dz heading_angle category_id)
-        gt_boxes = [line.strip().split(' ') for line in lines]
-        return np.array(gt_boxes, dtype=np.float32)
+        gt_boxes = []
+        gt_names = []
+        for line in lines:
+            line_list = line.strip().split(' ')
+            gt_boxes.append(line_list[:-1])
+            gt_names.append(line_list[-1])
+
+        return np.array(gt_boxes, dtype=np.float32), np.array(gt_names)
 
     def get_lidar(self, idx):
         lidar_file = self.root_path / 'points' / ('%s.npy' % idx)
@@ -136,8 +142,6 @@ class CustomDataset(DatasetTemplate):
     def get_infos(self, class_names, num_workers=4, has_label=True, sample_id_list=None, num_features=4):
         import concurrent.futures as futures
 
-        class_names = np.array(class_names)
-
         def process_single_scene(sample_idx):
             print('%s sample_idx: %s' % (self.split, sample_idx))
             info = {}
@@ -146,8 +150,8 @@ class CustomDataset(DatasetTemplate):
 
             if has_label:
                 annotations = {}
-                gt_boxes_lidar = self.get_label(sample_idx)
-                annotations['name'] = class_names[gt_boxes_lidar[:, -1].astype(np.int64)]
+                gt_boxes_lidar, name = self.get_label(sample_idx)
+                annotations['name'] = name
                 annotations['gt_boxes_lidar'] = gt_boxes_lidar[:, :7]
                 info['annos'] = annotations
 
@@ -219,10 +223,9 @@ class CustomDataset(DatasetTemplate):
                 name = gt_names[idx]
                 if name not in class_names:
                     continue
-                category_id = class_names.index(name)
-                line = "{x} {y} {z} {l} {w} {h} {angle} {category_id}\n".format(
+                line = "{x} {y} {z} {l} {w} {h} {angle} {name}\n".format(
                     x=boxes[0], y=boxes[1], z=(boxes[2]), l=boxes[3],
-                    w=boxes[4], h=boxes[5], angle=boxes[6], category_id=category_id
+                    w=boxes[4], h=boxes[5], angle=boxes[6], name=name
                 )
                 f.write(line)