diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..27d2120 --- /dev/null +++ b/.gitignore @@ -0,0 +1,8 @@ +.vscode/launch.json +models/setup.log +__pycache__/ +*.so +models/correlation_package/build/ +*.egg-info/ +output/ +results_for_paper/ \ No newline at end of file diff --git a/DatasetLidarCamera.py b/DatasetLidarCamera.py index 828ba4d..78917e6 100644 --- a/DatasetLidarCamera.py +++ b/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', diff --git a/evaluate_calib.py b/evaluate_calib.py index b4e477f..09cf6d3 100644 --- a/evaluate_calib.py +++ b/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,7 +477,7 @@ 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 @@ -483,12 +485,12 @@ def init_fn(x): 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 diff --git a/models/correlation_package/correlation.py b/models/correlation_package/correlation.py index f164a85..5a9023c 100644 --- a/models/correlation_package/correlation.py +++ b/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 diff --git a/models/correlation_package/pyproject.toml b/models/correlation_package/pyproject.toml index 3635a2d..9f6a83b 100644 --- a/models/correlation_package/pyproject.toml +++ b/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. diff --git a/models/correlation_package/setup.py b/models/correlation_package/setup.py index bfe8121..873d66d 100644 --- a/models/correlation_package/setup.py +++ b/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( diff --git a/requirements.txt b/requirements.txt index 7f7dd5b..f29dbce 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,14 @@ 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 @@ -7,14 +16,18 @@ 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 这个选项使得编译环境使用本地环境。 \ No newline at end of file