Updated code for cuda 11 and torch 1.10+

.gitignore

@@ -0,0 +1,8 @@
+.vscode/launch.json
+models/setup.log
+__pycache__/
+*.so
+models/correlation_package/build/
+*.egg-info/
+output/
+results_for_paper/

DatasetLidarCamera.py

@@ -23,8 +23,8 @@
 from PIL import Image
 from torch.utils.data import Dataset
 from torchvision import transforms
-
-from utils import invert_pose, rotate_forward, quaternion_from_matrix, read_calib_file
+from tqdm import tqdm
+from utils import invert_pose, rotate_forward, quaternion_from_matrix#, read_calib_file
 from pykitti import odometry
 import pykitti
 
@@ -50,8 +50,7 @@ def __init__(self, dataset_dir, transform=None, augmentation=False, use_reflecta
         self.suf = suf
 
         self.all_files = []
-        self.sequence_list = ['00', '01', '02', '03', '04', '05', '06', '07', '08', '09', '10',
-                              '11', '12', '13', '14', '15', '16', '17', '18', '19', '20', '21']
+        self.sequence_list = [f'{s:02}' for s in range(0, 21)]
         # self.model = CameraModel()
         # self.model.focal_length = [7.18856e+02, 7.18856e+02]
         # self.model.principal_point = [6.071928e+02, 1.852157e+02]
@@ -71,19 +70,17 @@ def __init__(self, dataset_dir, transform=None, augmentation=False, use_reflecta
             image_list = os.listdir(os.path.join(dataset_dir, 'sequences', seq, 'image_2'))
             image_list.sort()
 
-            for image_name in image_list:
-                if not os.path.exists(os.path.join(dataset_dir, 'sequences', seq, 'velodyne',
-                                                   str(image_name.split('.')[0])+'.bin')):
+            for image_name in tqdm(image_list, desc=f'Loading KITTI odometry dataset {seq}'):
+                if not os.path.exists(os.path.join(dataset_dir, 'sequences', seq, 'velodyne', str(image_name.split('.')[0])+'.bin')):
                     continue
-                if not os.path.exists(os.path.join(dataset_dir, 'sequences', seq, 'image_2',
-                                                   str(image_name.split('.')[0])+suf)):
+                if not os.path.exists(os.path.join(dataset_dir, 'sequences', seq, 'image_2', str(image_name.split('.')[0])+suf)):
                     continue
                 if seq == val_sequence:
                     if split.startswith('val') or split == 'test':
                         self.all_files.append(os.path.join(seq, image_name.split('.')[0]))
                 elif (not seq == val_sequence) and split == 'train':
                     self.all_files.append(os.path.join(seq, image_name.split('.')[0]))
-
+        assert self.all_files, 'No pcd files found'
         self.val_RT = []
         if split == 'val' or split == 'test':
             # val_RT_file = os.path.join(dataset_dir, 'sequences',

evaluate_calib.py

@@ -26,8 +26,10 @@
 import torch.utils.data
 from sacred import Experiment
 from sacred.utils import apply_backspaces_and_linefeeds
+from sacred import SETTINGS
+SETTINGS.CONFIG.READ_ONLY_CONFIG = False
 from skimage import io
-from tqdm import tqdm
+from tqdm import tqdm, trange
 import time
 
 from models.LCCNet import LCCNet
@@ -64,7 +66,7 @@
 @ex.config
 def config():
     dataset = 'kitti/odom'
-    data_folder = '/home/wangshuo/Datasets/KITTI/odometry_color/'
+    data_folder = './kitti_odometry_color/dataset/'
     test_sequence = 0
     use_prev_output = False
     max_t = 1.5
@@ -80,22 +82,22 @@ def config():
     # Set to True only if you use two network, the first for rotation and the second for translation
     rot_transl_separated = False
     random_initial_pose = False
-    save_log = False
+    save_log = True
     dropout = 0.0
     max_depth = 80.
     iterative_method = 'multi_range' # ['multi_range', 'single_range', 'single']
-    output = '../output'
+    output = './output'
     save_image = False
     outlier_filter = True
     outlier_filter_th = 10
     out_fig_lg = 'EN' # [EN, CN]
 
 weights = [
-   './pretrained/kitti_iter1.tar',
-   './pretrained/kitti_iter2.tar',
-   './pretrained/kitti_iter3.tar',
-   './pretrained/kitti_iter4.tar',
-   './pretrained/kitti_iter5.tar',
+   './LCCNet_pretrained/kitti_iter1.tar',
+   './LCCNet_pretrained/kitti_iter2.tar',
+   './LCCNet_pretrained/kitti_iter3.tar',
+   './LCCNet_pretrained/kitti_iter4.tar',
+   './LCCNet_pretrained/kitti_iter5.tar',
 ]
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -174,7 +176,7 @@ def main(_config, seed):
     def init_fn(x):
         return _init_fn(x, seed)
 
-    num_worker = 6
+    num_worker = 8
     batch_size = 1
 
     TestImgLoader = torch.utils.data.DataLoader(dataset=dataset_val,
@@ -189,7 +191,7 @@ def init_fn(x):
     print(len(TestImgLoader))
 
     models = [] # iterative model
-    for i in range(len(weights)):
+    for weight in tqdm(weights, desc='Loading weights'):
         # network choice and settings
         if _config['network'].startswith('Res'):
             feat = 1
@@ -207,7 +209,7 @@ def init_fn(x):
         else:
             raise TypeError("Network unknown")
 
-        checkpoint = torch.load(weights[i], map_location='cpu')
+        checkpoint = torch.load(weight, map_location='cpu')
         saved_state_dict = checkpoint['state_dict']
         model.load_state_dict(saved_state_dict)
         model = model.to(device)
@@ -216,8 +218,8 @@ def init_fn(x):
 
 
     if _config['save_log']:
+        os.makedirs('results_for_paper', exist_ok=True)
         log_file = f'./results_for_paper/log_seq{_config["test_sequence"]}.csv'
-        log_file = open(log_file, 'w')
         log_file = csv.writer(log_file)
         header = ['frame']
         for i in range(len(weights) + 1):
@@ -321,12 +323,12 @@ def init_fn(x):
 
             if _config['save_image']:
                 # save the Lidar pointcloud
-                pcl_lidar = o3.PointCloud()
+                pcl_lidar = o3.geometry.PointCloud()
                 pc_lidar = pc_lidar.detach().cpu().numpy()
-                pcl_lidar.points = o3.Vector3dVector(pc_lidar.T[:, :3])
+                pcl_lidar.points = o3.utility.Vector3dVector(pc_lidar.T[:, :3])
 
                 # o3.draw_geometries(downpcd)
-                o3.write_point_cloud(pc_lidar_path + '/{}.pcd'.format(batch_idx), pcl_lidar)
+                o3.io.write_point_cloud(pc_lidar_path + '/{}.pcd'.format(batch_idx), pcl_lidar)
 
 
             R = quat2mat(sample['rot_error'][idx])
@@ -353,12 +355,12 @@ def init_fn(x):
                 R = img[uv_input[:, 1], uv_input[:, 0], 0] / 255
                 G = img[uv_input[:, 1], uv_input[:, 0], 1] / 255
                 B = img[uv_input[:, 1], uv_input[:, 0], 2] / 255
-                pcl_input = o3.PointCloud()
-                pcl_input.points = o3.Vector3dVector(pc_input_valid[:, :3])
-                pcl_input.colors = o3.Vector3dVector(np.vstack((R, G, B)).T)
+                pcl_input = o3.geometry.PointCloud()
+                pcl_input.points = o3.utility.Vector3dVector(pc_input_valid[:, :3])
+                pcl_input.colors = o3.utility.Vector3dVector(np.vstack((R, G, B)).T)
 
                 # o3.draw_geometries(downpcd)
-                o3.write_point_cloud(pc_input_path + '/{}.pcd'.format(batch_idx), pcl_input)
+                o3.io.write_point_cloud(pc_input_path + '/{}.pcd'.format(batch_idx), pcl_input)
 
             # PAD ONLY ON RIGHT AND BOTTOM SIDE
             rgb = sample['rgb'][idx].cuda()
@@ -384,8 +386,8 @@ def init_fn(x):
 
         lidar_input = torch.stack(lidar_input)
         rgb_input = torch.stack(rgb_input)
-        rgb_resize = F.interpolate(rgb_input, size=[256, 512], mode="bilinear")
-        lidar_resize = F.interpolate(lidar_input, size=[256, 512], mode="bilinear")
+        rgb_resize = F.interpolate(rgb_input, size=[256, 512], mode="bilinear", align_corners=True)
+        lidar_resize = F.interpolate(lidar_input, size=[256, 512], mode="bilinear", align_corners=True)
 
 
         if _config['save_image']:
@@ -475,20 +477,20 @@ def init_fn(x):
                 depth_img_pred /= _config['max_depth']
                 depth_pred = F.pad(depth_img_pred, shape_pad_input[0])
                 lidar = depth_pred.unsqueeze(0)
-                lidar_resize = F.interpolate(lidar, size=[256, 512], mode="bilinear")
+                lidar_resize = F.interpolate(lidar, size=[256, 512], mode="bilinear", align_corners=True)
 
                 if iteration == len(weights)-1 and _config['save_image']:
                     # save the RGB pointcloud
                     img = cv2.imread(sample['img_path'][0])
                     R = img[uv_pred[:, 1], uv_pred[:, 0], 0] / 255
                     G = img[uv_pred[:, 1], uv_pred[:, 0], 1] / 255
                     B = img[uv_pred[:, 1], uv_pred[:, 0], 2] / 255
-                    pcl_pred = o3.PointCloud()
-                    pcl_pred.points = o3.Vector3dVector(pc_pred_valid[:, :3])
-                    pcl_pred.colors = o3.Vector3dVector(np.vstack((R, G, B)).T)
+                    pcl_pred = o3.geometry.PointCloud()
+                    pcl_pred.points = o3.utility.Vector3dVector(pc_pred_valid[:, :3])
+                    pcl_pred.colors = o3.utility.Vector3dVector(np.vstack((R, G, B)).T)
 
                     # o3.draw_geometries(downpcd)
-                    o3.write_point_cloud(pc_pred_path + '/{}.pcd'.format(batch_idx), pcl_pred)
+                    o3.io.write_point_cloud(pc_pred_path + '/{}.pcd'.format(batch_idx), pcl_pred)
 
 
                 if _config['save_image']:
@@ -537,10 +539,9 @@ def init_fn(x):
     # mis_calib_input[transl_x, transl_y, transl_z, rotx, roty, rotz] Nx6
     mis_calib_input = torch.stack(mis_calib_list)[:, :, 0]
 
-    if _config['save_log']:
-        log_file.close()
+
     print("Iterative refinement: ")
-    for i in range(len(weights) + 1):
+    for i in trange(len(weights) + 1, desc='Iterative refinement'):
         errors_r[i] = torch.tensor(errors_r[i]).abs() * (180.0 / 3.141592)
         errors_t[i] = torch.tensor(errors_t[i]).abs() * 100
 

models/correlation_package/correlation.py

@@ -16,8 +16,9 @@ def __init__(self, pad_size=3, kernel_size=3, max_displacement=20, stride1=1, st
         self.corr_multiply = corr_multiply
         # self.out_channel = ((max_displacement/stride2)*2 + 1) * ((max_displacement/stride2)*2 + 1)
 
-    def forward(self, input1, input2):
-        self.save_for_backward(input1, input2)
+    @staticmethod
+    def forward(ctx, self, input1, input2):
+        ctx.save_for_backward(self, input1, input2)
 
         with torch.cuda.device_of(input1):
             rbot1 = input1.new()
@@ -29,8 +30,9 @@ def forward(self, input1, input2):
 
         return output
 
-    def backward(self, grad_output):
-        input1, input2 = self.saved_tensors
+    @staticmethod
+    def backward(ctx, grad_output):
+        self, input1, input2 = ctx.saved_tensors
 
         with torch.cuda.device_of(input1):
             rbot1 = input1.new()
@@ -56,7 +58,7 @@ def __init__(self, pad_size=0, kernel_size=0, max_displacement=0, stride1=1, str
         self.corr_multiply = corr_multiply
 
     def forward(self, input1, input2):
-
-        result = CorrelationFunction(self.pad_size, self.kernel_size, self.max_displacement,self.stride1, self.stride2, self.corr_multiply)(input1, input2)
+        # result = CorrelationFunction(self.pad_size, self.kernel_size, self.max_displacement,self.stride1, self.stride2, self.corr_multiply)(input1, input2)
+        result = CorrelationFunction.apply(self, input1, input2)
 
         return result

models/correlation_package/pyproject.toml

@@ -1,3 +1,3 @@
 [build-system]
 # Minimum requirements for the build system to execute.
-requires = ["setuptools", "wheel", "numpy", "torch==1.0.1.post2"]  # PEP 508 specifications.
+requires = ["setuptools", "wheel", "numpy", "torch>=1.10.0"]  # PEP 508 specifications.

models/correlation_package/setup.py

@@ -5,15 +5,18 @@
 from setuptools import find_packages, setup
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension
 
-cxx_args = ['-std=c++11']
+# cxx_args = ['-std=c++11']
+cxx_args = ['-std=c++14']
 
 nvcc_args = [
-    '-gencode', 'arch=compute_50,code=sm_50',
-    '-gencode', 'arch=compute_52,code=sm_52',
-    '-gencode', 'arch=compute_60,code=sm_60',
-    '-gencode', 'arch=compute_61,code=sm_61',
-    '-gencode', 'arch=compute_70,code=sm_70',
-    '-gencode', 'arch=compute_70,code=compute_70'
+    # '-gencode', 'arch=compute_50,code=sm_50',
+    # '-gencode', 'arch=compute_52,code=sm_52',
+    # '-gencode', 'arch=compute_60,code=sm_60',
+    # '-gencode', 'arch=compute_61,code=sm_61',
+    # '-gencode', 'arch=compute_70,code=sm_70',
+    # '-gencode', 'arch=compute_70,code=compute_70',
+    '-gencode', 'arch=compute_75,code=compute_75',
+    '-gencode', 'arch=compute_86,code=compute_86'
 ]
 
 setup(

requirements.txt

@@ -1,20 +1,33 @@
 scikit_image
-git+https://gitlab.com/m1lhaus/blender-mathutils.git
+git+https://gitlab.com/ideasman42/blender-mathutils.git
+# 会在Python3.8遇到两个问题：
+# 1. PyModule_AddType()函数在Python3.8中不存在，建议方案为改成以下方式
+#   PyModule_AddObject(mod, vector_Type.tp_name,     (PyObject *)&vector_Type);
+#	PyModule_AddObject(mod, matrix_Type.tp_name,     (PyObject *)&matrix_Type);
+#	PyModule_AddObject(mod, euler_Type.tp_name,      (PyObject *)&euler_Type);
+#	PyModule_AddObject(mod, quaternion_Type.tp_name, (PyObject *)&quaternion_Type);
+#	PyModule_AddObject(mod, color_Type.tp_name,      (PyObject *)&color_Type);
+# 2. y = _Py_HashDouble(NULL, (double)(array[i++])); 这里的NULL会报错，建议方案为将NULL删除
+# mathutils==2.81.2 # 早期版本可以安装，但是会出问题：ValueError: mathutils.Euler(): invalid euler order 'XYZ'
 tqdm==4.19.9
 pandas
 h5py
 matplotlib
 scipy
 pyquaternion
 opencv-python
-cupy_cuda90==6.0.0
+# cupy_cuda90==6.0.0 # cuda11使用cupy-cuda11x
+cupy-cuda11x
 pykitti
 numpy
 tensorboardX
-open3d-python==0.7.0.0
+open3d-python
 Pillow
 scikit-image
-torch==1.0.1.post2
-torchvision==0.2.2
-sacred==0.7.4
--e ./models/correlation_package
+torch>=1.10.0
+torchvision>=0.11.0
+scared
+-e ./models/correlation_package --no-build-isolation
+# 注意，这里编译了一个本地库："correlation_package"(-e ./models/correlation_package)。
+# 因为pyproject.toml使用的torch文件比较古老，不能匹配本地cuda版本，同时会下载一个版本安装，会非常慢，
+# 因此我们这里使用pip install -e ./models/correlation_package --no-build-isolation 这个选项使得编译环境使用本地环境。