Adjust the hook function and nuscenes_dataset.py

Author: chenshi3

pcdet/datasets/augmentor/data_augmentor.py

@@ -24,7 +24,7 @@ def __init__(self, root_path, augmentor_configs, class_names, logger=None):
             cur_augmentor = getattr(self, cur_cfg.NAME)(config=cur_cfg)
             self.data_augmentor_queue.append(cur_augmentor)
 
-    def disableAugmentation(self, augmentor_configs):
+    def disable_augmentation(self, augmentor_configs):
         self.data_augmentor_queue = []
         aug_config_list = augmentor_configs if isinstance(augmentor_configs, list) \
             else augmentor_configs.AUG_CONFIG_LIST

pcdet/datasets/nuscenes/nuscenes_dataset.py

@@ -151,6 +151,63 @@ def crop_image(self, input_dict):
         input_dict['img_process_infos'] = img_process_infos
         input_dict['camera_imgs'] = crop_images
         return input_dict
+    
+    def load_camera_info(self, input_dict, info):
+        input_dict["image_paths"] = []
+        input_dict["lidar2camera"] = []
+        input_dict["lidar2image"] = []
+        input_dict["camera2ego"] = []
+        input_dict["camera_intrinsics"] = []
+        input_dict["camera2lidar"] = []
+
+        for _, camera_info in info["cams"].items():
+            input_dict["image_paths"].append(camera_info["data_path"])
+
+            # lidar to camera transform
+            lidar2camera_r = np.linalg.inv(camera_info["sensor2lidar_rotation"])
+            lidar2camera_t = (
+                camera_info["sensor2lidar_translation"] @ lidar2camera_r.T
+            )
+            lidar2camera_rt = np.eye(4).astype(np.float32)
+            lidar2camera_rt[:3, :3] = lidar2camera_r.T
+            lidar2camera_rt[3, :3] = -lidar2camera_t
+            input_dict["lidar2camera"].append(lidar2camera_rt.T)
+
+            # camera intrinsics
+            camera_intrinsics = np.eye(4).astype(np.float32)
+            camera_intrinsics[:3, :3] = camera_info["camera_intrinsics"]
+            input_dict["camera_intrinsics"].append(camera_intrinsics)
+
+            # lidar to image transform
+            lidar2image = camera_intrinsics @ lidar2camera_rt.T
+            input_dict["lidar2image"].append(lidar2image)
+
+            # camera to ego transform
+            camera2ego = np.eye(4).astype(np.float32)
+            camera2ego[:3, :3] = Quaternion(
+                camera_info["sensor2ego_rotation"]
+            ).rotation_matrix
+            camera2ego[:3, 3] = camera_info["sensor2ego_translation"]
+            input_dict["camera2ego"].append(camera2ego)
+
+            # camera to lidar transform
+            camera2lidar = np.eye(4).astype(np.float32)
+            camera2lidar[:3, :3] = camera_info["sensor2lidar_rotation"]
+            camera2lidar[:3, 3] = camera_info["sensor2lidar_translation"]
+            input_dict["camera2lidar"].append(camera2lidar)
+        # read image
+        filename = input_dict["image_paths"]
+        images = []
+        for name in filename:
+            images.append(Image.open(str(self.root_path / name)))
+        
+        input_dict["camera_imgs"] = images
+        input_dict["ori_shape"] = images[0].size
+        
+        # resize and crop image
+        input_dict = self.crop_image(input_dict)
+
+        return input_dict
 
     def __len__(self):
         if self._merge_all_iters_to_one_epoch:
@@ -182,59 +239,7 @@ def __getitem__(self, index):
                 'gt_boxes': info['gt_boxes'] if mask is None else info['gt_boxes'][mask]
             })
         if self.use_camera:
-            input_dict["image_paths"] = []
-            input_dict["lidar2camera"] = []
-            input_dict["lidar2image"] = []
-            input_dict["camera2ego"] = []
-            input_dict["camera_intrinsics"] = []
-            input_dict["camera2lidar"] = []
-
-            for _, camera_info in info["cams"].items():
-                input_dict["image_paths"].append(camera_info["data_path"])
-
-                # lidar to camera transform
-                lidar2camera_r = np.linalg.inv(camera_info["sensor2lidar_rotation"])
-                lidar2camera_t = (
-                    camera_info["sensor2lidar_translation"] @ lidar2camera_r.T
-                )
-                lidar2camera_rt = np.eye(4).astype(np.float32)
-                lidar2camera_rt[:3, :3] = lidar2camera_r.T
-                lidar2camera_rt[3, :3] = -lidar2camera_t
-                input_dict["lidar2camera"].append(lidar2camera_rt.T)
-
-                # camera intrinsics
-                camera_intrinsics = np.eye(4).astype(np.float32)
-                camera_intrinsics[:3, :3] = camera_info["camera_intrinsics"]
-                input_dict["camera_intrinsics"].append(camera_intrinsics)
-
-                # lidar to image transform
-                lidar2image = camera_intrinsics @ lidar2camera_rt.T
-                input_dict["lidar2image"].append(lidar2image)
-
-                # camera to ego transform
-                camera2ego = np.eye(4).astype(np.float32)
-                camera2ego[:3, :3] = Quaternion(
-                    camera_info["sensor2ego_rotation"]
-                ).rotation_matrix
-                camera2ego[:3, 3] = camera_info["sensor2ego_translation"]
-                input_dict["camera2ego"].append(camera2ego)
-
-                # camera to lidar transform
-                camera2lidar = np.eye(4).astype(np.float32)
-                camera2lidar[:3, :3] = camera_info["sensor2lidar_rotation"]
-                camera2lidar[:3, 3] = camera_info["sensor2lidar_translation"]
-                input_dict["camera2lidar"].append(camera2lidar)
-            # read image
-            filename = input_dict["image_paths"]
-            images = []
-            for name in filename:
-                images.append(Image.open(str(self.root_path / name)))
-            
-            input_dict["camera_imgs"] = images
-            input_dict["ori_shape"] = images[0].size
-            
-            # resize and crop image
-            input_dict = self.crop_image(input_dict)
+            input_dict = self.load_camera_info(input_dict, info)
 
         data_dict = self.prepare_data(data_dict=input_dict)
 

tools/train_utils/train_utils.py

@@ -267,6 +267,6 @@ def disable_augmentation_hook(hook_config, dataloader, total_epochs, cur_epoch,
                 dataset_cfg=cfg.DATA_CONFIG
                 logger.info(f'Disable augmentations: {DISABLE_AUG_LIST}')
                 dataset_cfg.DATA_AUGMENTOR.DISABLE_AUG_LIST = DISABLE_AUG_LIST
-                dataloader._dataset.data_augmentor.disableAugmentation(dataset_cfg.DATA_AUGMENTOR)
+                dataloader._dataset.data_augmentor.disable_augmentation(dataset_cfg.DATA_AUGMENTOR)
                 flag = True
     return flag