Fix some autolabel bugs

Author: TheCodez

autolabeling/autolabel.py

@@ -11,9 +11,8 @@ def get_points_in_fov_90(pc, clip_distance=0.0, max_distance=500):
     return pc[fov_inds, :]
 
 
-def pinhole_projection(points, img, T_cam0, K_cam0):
+def pinhole_projection(points, T_cam0, K_cam0):
     points_2d = np.zeros((points.shape[0], 2), dtype=np.int32)
-    height, width = img.shape
     points = np.dot(points, np.transpose(T_cam0))
 
     x = points[:, 0]
@@ -24,27 +23,26 @@ def pinhole_projection(points, img, T_cam0, K_cam0):
     x_img = (np.dot(K_cam0[0, 0], x / z) + K_cam0[0, 2]).astype(np.int32)
     y_img = (np.dot(K_cam0[1, 1], y / z) + K_cam0[1, 2]).astype(np.int32)
 
-    # clamp values
-    x_img = np.clip(x_img, 0, width - 1)
-    y_img = np.clip(y_img, 0, height - 1)
-
     points_2d[:, 0] = x_img
     points_2d[:, 1] = y_img
 
     return points_2d
 
 
 def transfer_labels(points, semantic, T_cam0, K_cam0):
-    labels = np.ones((points.shape[0], 1), dtype=points.dtype) * 255
+    labels = np.zeros((points.shape[0], 1), dtype=points.dtype)
+    height, width = semantic.shape
 
-    points_2d = pinhole_projection(points, semantic, T_cam0, K_cam0)
+    points_2d = pinhole_projection(points, T_cam0, K_cam0)
     semantic = convert_classes_to_lidar_classes(semantic)
 
     x_img = points_2d[:, 0]
     y_img = points_2d[:, 1]
 
-    label = semantic[y_img, x_img]
-    labels[:, 0] = label
+    for i in range(y_img.shape[0]):
+        if (0 <= x_img[i] < width) and (0 <= y_img[i] < height):
+            label = semantic[y_img[i], x_img[i]]
+            labels[i] = label
 
     return np.concatenate((points, labels), axis=1)
 
@@ -53,7 +51,7 @@ def spherical_projection(points, height=64, width=512):
     x = points[:, 0]
     y = points[:, 1]
     z = points[:, 2]
-    reflectivity = points[:, 3]
+    i = points[:, 3]
     label = points[:, 4]
     d = np.sqrt(x ** 2 + y ** 2 + z ** 2)
 
@@ -63,15 +61,15 @@ def spherical_projection(points, height=64, width=512):
     idx_h = height - 1 - ((height - 1) * (theta - theta.min()) / (theta.max() - theta.min()))
     idx_w = width - 1 - ((width - 1) * (phi - phi.min()) / (phi.max() - phi.min()))
 
-    idx_h = np.floor(idx_h).astype(np.int32)
-    idx_w = np.floor(idx_w).astype(np.int32)
+    idx_h = np.round(idx_h).astype(np.int32)
+    idx_w = np.round(idx_w).astype(np.int32)
 
     projected_img = np.zeros((height, width, 6), dtype=np.float32)
     projected_img[idx_h, idx_w, 0] = x
     projected_img[idx_h, idx_w, 1] = y
     projected_img[idx_h, idx_w, 2] = z
-    projected_img[idx_h, idx_w, 3] = d
-    projected_img[idx_h, idx_w, 4] = reflectivity
+    projected_img[idx_h, idx_w, 3] = i
+    projected_img[idx_h, idx_w, 4] = d
     projected_img[idx_h, idx_w, 5] = label
 
     return projected_img

autolabeling/classes.py

@@ -4,28 +4,56 @@
 
 LidarClass = namedtuple('LidarClass', ['name', 'train_id', 'lidar_name', 'lidar_id',
                                        'color'])
+
+Class = namedtuple('Class', ['lidar_name', 'lidar_id', 'color'])
+
+classes = [
+    Class('road', 19, (128, 64, 128)),
+    Class('sidewalk', 20, (244, 35, 232)),
+    Class('construction', 21, (70, 70, 70)),
+    Class('construction', 21, (70, 70, 70)),
+    Class('unlabeled', 0, (0, 0, 0)),
+    Class('pole', 22, (153, 153, 153)),
+    Class('construction', 21, (70, 70, 70)),
+    Class('traffic sign', 23, (220, 220, 0)),
+    Class('vegetation', 24, (104, 131, 15)),
+    Class('terrain', 25, (148, 255, 144)),
+    Class('sky', 26, (0, 0, 0)),
+    Class('person', 27, (220, 20, 60)),
+    Class('rider', 28, (255, 0, 0)),
+    Class('small vehicle', 29, (0, 0, 142)),
+    Class('large vehicle', 30, (0, 0, 70)),
+    Class('large vehicle', 30, (0, 0, 70)),
+    Class('large vehicle', 30, (0, 0, 70)),
+    Class('two wheeler', 31, (119, 11, 32)),
+    Class('two wheeler', 31, (119, 11, 32)),
+]
+
 lidar_classes = [
-    LidarClass('road', 0, 'road', 0, (128, 64, 128)),
-    LidarClass('sidewalk', 1, 'sidewalk', 1, (244, 35, 232)),
-    LidarClass('building', 2, 'construction', 2, (70, 70, 70)),
-    LidarClass('wall', 3, 'construction', 2, (70, 70, 70)),
-    LidarClass('fence', 4, 'unlabeled', 255, (70, 70, 70)),
-    LidarClass('pole', 5, 'pole', 3, (153, 153, 153)),
-    LidarClass('traffic light', 6, 'construction', 2, (70, 70, 70)),
-    LidarClass('traffic sign', 7, 'traffic sign', 4, (220, 220, 0)),
-    LidarClass('vegetation', 8, 'vegetation', 5, (104, 131, 15)),
-    LidarClass('terrain', 9, 'terrain', 6, (148, 255, 144)),
-    LidarClass('sky', 10, 'sky', 7, (0, 0, 0)),
-    LidarClass('person', 11, 'person', 8, (220, 20, 60)),
-    LidarClass('rider', 12, 'rider', 9, (255, 0, 0)),
-    LidarClass('car', 13, 'small vehicle', 10, (0, 0, 142)),
-    LidarClass('truck', 14, 'large vehicle', 11, (0, 0, 70)),
-    LidarClass('bus', 15, 'large vehicle', 11, (0, 0, 70)),
-    LidarClass('train', 16, 'large vehicle', 11, (0, 0, 70)),
-    LidarClass('motorcycle', 17, 'two wheeler', 12, (119, 11, 32)),
-    LidarClass('bicycle', 18, 'two wheeler', 12, (119, 11, 32)),
+    LidarClass('road', 0, 'road', 19, (128, 64, 128)),
+    LidarClass('sidewalk', 1, 'sidewalk', 20, (244, 35, 232)),
+    LidarClass('building', 2, 'construction', 21, (70, 70, 70)),
+    LidarClass('wall', 3, 'construction', 21, (70, 70, 70)),
+    LidarClass('fence', 4, 'unlabeled', 0, (0, 0, 0)),
+    LidarClass('pole', 5, 'pole', 22, (153, 153, 153)),
+    LidarClass('traffic light', 6, 'construction', 21, (70, 70, 70)),
+    LidarClass('traffic sign', 7, 'traffic sign', 23, (220, 220, 0)),
+    LidarClass('vegetation', 8, 'vegetation', 24, (104, 131, 15)),
+    LidarClass('terrain', 9, 'terrain', 25, (148, 255, 144)),
+    LidarClass('sky', 10, 'sky', 26, (0, 0, 0)),
+    LidarClass('person', 11, 'person', 27, (220, 20, 60)),
+    LidarClass('rider', 12, 'rider', 28, (255, 0, 0)),
+    LidarClass('car', 13, 'small vehicle', 29, (0, 0, 142)),
+    LidarClass('truck', 14, 'large vehicle', 30, (0, 0, 70)),
+    LidarClass('bus', 15, 'large vehicle', 30, (0, 0, 70)),
+    LidarClass('train', 16, 'large vehicle', 30, (0, 0, 70)),
+    LidarClass('motorcycle', 17, 'two wheeler', 31, (119, 11, 32)),
+    LidarClass('bicycle', 18, 'two wheeler', 31, (119, 11, 32)),
 ]
 
+train_id_to_lidar_id = {0: 19, 1: 20, 2: 21, 3: 21, 4: 0, 5: 22, 6: 21, 7: 23, 8: 24, 9: 25, 10: 26, 11: 27, 12: 28,
+                        13: 29, 14: 30, 15: 30, 16: 30, 17: 31, 18: 31, 255: 0}
+
 
 def convert_classes_to_lidar_classes(target):
     target_copy = target.clone()

autolabeling/kitti_demo.py

@@ -33,23 +33,27 @@ def convert_train_id_to_id(target):
 
 
 def show_lidar_on_image(points, image, segmentation, T_cam0, K_cam0):
-    points_2d = autolabel.pinhole_projection(points, segmentation, T_cam0, K_cam0)
+    points_2d = autolabel.pinhole_projection(points, T_cam0, K_cam0)
 
     cmap = get_cityscapes_colormap()
     segmentation = convert_train_id_to_id(segmentation)
     vis = colorize_seg(segmentation.cpu(), cmap)
+    height, width = segmentation.shape
 
     for i in range(points.shape[0]):
         img_x = points_2d[i, 0]
         img_y = points_2d[i, 1]
+        img_x = np.clip(img_x, 0, width - 1)
+        img_y = np.clip(img_y, 0, height - 1)
+
         color = vis[:, img_y, img_x].tolist()
         cv2.circle(image, (img_x, img_y), 2, color=tuple(color), thickness=-1)
 
     return image
 
 
-def show_lidar_depth_on_image(pc_velo, img, segmentation, T_cam0, K_cam0):
-    points_2d = autolabel.pinhole_projection(pc_velo, segmentation, T_cam0, K_cam0)
+def show_lidar_depth_on_image(pc_velo, img, T_cam0, K_cam0):
+    points_2d = autolabel.pinhole_projection(pc_velo, T_cam0, K_cam0)
 
     cmap = plt.cm.get_cmap('hsv', 256)
     cmap = np.array([cmap(i) for i in range(256)])[:, :3] * 255
@@ -143,8 +147,8 @@ def _normalize(x):
     proj_img = show_lidar_on_image(pc_velo, np.array(img), pred, dataset.calib.T_cam0_velo, dataset.calib.K_cam0)
 
     record = torch.as_tensor(lidar, dtype=torch.float32).permute(2, 0, 1).contiguous()
-    distance = record[3, :, :]
-    reflectivity = record[4, :, :]
+    reflectivity = record[3, :, :]
+    distance = record[4, :, :]
     label = record[5, :, :]
 
     plot_images(file_name, distance, reflectivity, label, pred, img, proj_img)