diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index c6f98a3..523e990 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -1,10 +1,10 @@ -# Contributing to OpenSceneFlow +# Contributing to [OpenSceneFlow](https://github.com/KTH-RPL/OpenSceneFlow) We want to make contributing to this project as easy and transparent as possible. We welcome any contributions, from bug fixes to new features. If you're interested in adding your own scene flow method, this guide will walk you through the process. ## Adding a New Method -Here is a quick guide to integrating a new method into the OpenSceneFlow codebase. +Here is a quick guide to integrating a new method into the [OpenSceneFlow](https://github.com/KTH-RPL/OpenSceneFlow) codebase. ### 1. Data Preparation diff --git a/README.md b/README.md index 51d5307..41f9f6e 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,7 @@ 💞 If you find [*OpenSceneFlow*](https://github.com/KTH-RPL/OpenSceneFlow) useful to your research, please cite [**our works** 📖](#cite-us) and [give a star 🌟](https://github.com/KTH-RPL/OpenSceneFlow) as encouragement. (੭ˊ꒳​ˋ)੭✧ OpenSceneFlow is a codebase for point cloud scene flow estimation. -It is also an official implementation of the following papers (sored by the time of publication): +It is also an official implementation of the following papers (sorted by the time of publication): - **HiMo: High-Speed Objects Motion Compensation in Point Clouds** (SeFlow++) *Qingwen Zhang, Ajinkya Khoche, Yi Yang, Li Ling, Sina Sharif Mansouri, Olov Andersson, Patric Jensfelt* @@ -101,7 +101,7 @@ If you prefer to build the Docker image by yourself, Check [build-docker-image]( ## 1. Data Preparation -Refer to [dataprocess/README.md](dataprocess/README.md) for dataset download instructions. Currently, we support **Argoverse 2**, **Waymo**, and **custom datasets** (more datasets will be added in the future). +Refer to [dataprocess/README.md](dataprocess/README.md) for dataset download instructions. Currently, we support **Argoverse 2**, **Waymo**, **nuScenes** and **custom datasets** (more datasets will be added in the future). After downloading, convert the raw data to `.h5` format for easy training, evaluation, and visualization. Follow the steps in [dataprocess/README.md#process](dataprocess/README.md#process). diff --git a/assets/docs/nuscenes.png b/assets/docs/nuscenes.png new file mode 100644 index 0000000..3ed1ba0 Binary files /dev/null and b/assets/docs/nuscenes.png differ diff --git a/conf/others/nuscenes.toml b/conf/others/nuscenes.toml new file mode 100644 index 0000000..f77b06e --- /dev/null +++ b/conf/others/nuscenes.toml @@ -0,0 +1,31 @@ + +[important] +height = 1.85 # sensor height. 雷达传感器高度，主要是+offset + +# 整个雷点点云以自己为中心 分为多少个segment，每个segment又分成多少个bin +[segments] +r_min = 0.1 # minimum point distance. 感兴趣的区域 +r_max = 80 # maximum point distance. +n_segments = 360 # number of radial segments. +n_bins = 160 # number of radial bins. + + +[ground] +# 直线的 slope 斜率，一个seg里拟合的线斜率 +min_slope = 0.0 # minimum slope of a ground line. [T_m_small] +max_slope = 0.1 # maximum slope of a ground line. [T_m] +# 如果拟合点超过这个 long threshold，判断其 height 变化是否超过 max_long_height +long_threshold = 2.0 # Distance at which points are considered far from each other. [T_d_prev] +max_long_height = 0.3 # maximum height change to previous point in long line. [T_b] +# 如果没有超过 long threshold，判断现在拟合的点与fitline高度变化是否超过 max_start_height +max_start_height = 0.2 # Maximum heigh of starting line to be labelled ground. +# 最大的拟合误差，超过这个误差的点不会被加入 line 拟合 +max_fit_error = 0.1 # maximum error of a point during line fit. [T_RMSE: will sqaure inside code.] + +# 判断是否是地面点的时候，需要满足的条件： +max_dist_to_line = 0.2 # maximum vertical distance of point to line to be considered ground. [T_d_ground] +line_search_angle = 2.0 # How far to search for a line in angular direction [rad]. + +[general] +n_threads = 8 # number of threads for parallel processing. +verbose = false # if you don't want to see every output, set this to false. \ No newline at end of file diff --git a/dataprocess/README.md b/dataprocess/README.md index 7c57411..893af64 100644 --- a/dataprocess/README.md +++ b/dataprocess/README.md @@ -10,7 +10,7 @@ We've updated the process dataset for: - [x] Argoverse 2.0: check [here](#argoverse-20). The process script Involved from [DeFlow](https://github.com/KTH-RPL/DeFlow). - [x] Waymo: check [here](#waymo-dataset). The process script was involved from [SeFlow](https://github.com/KTH-RPL/SeFlow). -- [ ] nuScenes: done coding, public after review. Will be involved later by another paper. +- [x] nuScenes: check [here](#nuscenes), The process script was involved from [DeltaFlow](https://github.com/Kin-Zhang/DeltaFlow). - [ ] TruckScene: done coding, public after review. Will be involved later by another paper. - [ ] ZOD (w/o gt): done coding, public after review. Will be involved later by another paper. @@ -77,6 +77,20 @@ You need sign up an account at [nuScenes](https://www.nuscenes.org/) to download ![](../assets/docs/nuscenes.png) +Extracting & processing nuScenes require special handling: + +* Frame Rate: The raw LiDAR data is captured at 20Hz, while ground truth (GT) annotations are only available at 2Hz. +* Resampling: To standardize the data for consistent evaluation, we downsample the LiDAR point clouds to 10Hz. It is a GT-guided process that guarantees all annotated 2Hz frames are preserved within the final 10Hz sequence. +* The ground truth scene flow is generated using the official per-object velocity labels provided in the dataset, calculated between the resampled 10Hz frames. + + +#### Dataset frames + +| Dataset | # Total Scene | # Total Frames | +| ------- | ------------- | -------------- | +| train | 700 | 137575 / 27392 (w. gt) | +| val | 150 | 29126 / 5798 (w.gt) | + ### Waymo Dataset To download the Waymo dataset, you need to register an account at [Waymo Open Dataset](https://waymo.com/open/). You also need to install gcloud SDK and authenticate your account. Please refer to [this page](https://cloud.google.com/sdk/docs/install) for more details. @@ -123,6 +137,9 @@ python dataprocess/extract_av2.py --av2_type sensor --data_mode train --argo_dir # waymo: python dataprocess/extract_waymo.py --mode train --flow_data_dir /home/kin/data/waymo/flowlabel --map_dir /home/kin/data/waymo/flowlabel/map --output_dir /home/kin/data/waymo/preprocess --nproc 48 + +# nus: +python dataprocess/extract_nus.py --mode v1.0-trainval --output_dir /home/kin/data/nus/h5py/full --nproc 24 ``` diff --git a/dataprocess/extract_nus.py b/dataprocess/extract_nus.py index 6736e4c..4fcd2ad 100644 --- a/dataprocess/extract_nus.py +++ b/dataprocess/extract_nus.py @@ -6,19 +6,23 @@ # # Description: Preprocess Data, save as h5df format for faster loading # This one is for nuScenes dataset -# +# +# NOTE: nuscenes dataset need resample to 10Hz to align all other evaluations (e.g metric and frequency). +# and NOT ALL frames are annotated! So the SL training might be not that effective. +# """ import os -os.environ["OMP_NUM_THREADS"] = "1" +# os.environ["OMP_NUM_THREADS"] = "1" import multiprocessing from pathlib import Path from multiprocessing import Pool, current_process -from typing import Optional, Tuple, Dict, Union, Final +from typing import Optional, Final from tqdm import tqdm import numpy as np import fire, time, h5py +from collections import defaultdict from nuscenes.nuscenes import NuScenes from nuscenes.utils import splits @@ -30,9 +34,315 @@ import os, sys PARENT_DIR = os.path.abspath(os.path.join( os.path.dirname( __file__ ), '..')) sys.path.append(PARENT_DIR) -from src.utils.mics import create_reading_index -from src.utils.av2_eval import CATEGORY_TO_INDEX, NusNamMap +from dataprocess.misc_data import create_reading_index, check_h5py_file_exists, NusNamMap +from src.utils import npcal_pose0to1 +from src.utils.av2_eval import CATEGORY_TO_INDEX from linefit import ground_seg -BOUNDING_BOX_EXPANSION: Final = 0.2 -GROUNDSEG_config = f"{PARENT_DIR}/conf/others/nuscenes.toml" \ No newline at end of file +GROUNDSEG_config = f"{PARENT_DIR}/conf/others/nuscenes.toml" + +def remove_ego_points(pc: np.ndarray, + length_threshold: float = 4.084 / 2, + width_threshold: float = 1.730 / 2) -> np.ndarray: + """ + Remove ego points from a point cloud. + :param pc: point cloud to remove ego points from, shape: (N, 4) + :param length_threshold: Length threshold. + :param width_threshold: Width threshold. + :return: point cloud without ego points, shape: (N, 4). + """ + # NuScenes LiDAR position + # x: left --> width_threshold; y: front --> length_threshold; z: up --> height_threshold + x_filt = np.logical_and(pc[:, 0] > -width_threshold, pc[:, 0] < width_threshold) + y_filt = np.logical_and(pc[:, 1] > -length_threshold, pc[:, 1] < length_threshold) + not_close = np.logical_not(np.logical_and(x_filt, y_filt)) + not_close = not_close.astype(bool) + return pc[not_close] + +def _load_points_from_file(filename: str) -> np.ndarray: + """ + Private function to load point cloud from file. + :param filename: Path of the point cloud file. + :return: Point cloud as numpy array. + """ + pc = np.fromfile(filename, dtype=np.float32) + pc = pc.reshape((-1, 5))[:, :4] + return pc + +def get_pose(nusc, sweep_data): + # pose from lidar to world + ego2lidar = nusc.get('calibrated_sensor', sweep_data['calibrated_sensor_token']) + world2ego = nusc.get('ego_pose', sweep_data['ego_pose_token']) + ego2lidar_np = transform_matrix(ego2lidar['translation'], Quaternion(ego2lidar['rotation'])) + world2ego_np = transform_matrix(world2ego['translation'], Quaternion(world2ego['rotation'])) + return np.dot(world2ego_np, ego2lidar_np) + +def if_annotated_frame(sample_ann_dict, ts0): + gt_flow_flag = False + for key, anno_value in sample_ann_dict.items(): + if ts0 in anno_value.keys(): + gt_flow_flag = True # if there is one anno have gt vel then set to compute flow. + break + return gt_flow_flag + +def _resample_data(nusc, sample_data, sample_ann_dict, resample2frequency=10): + """ + NOTE(Qingwen) - 2025-05-18: + We always want to start from the first GT frame, and then resample the data! So we have as many GT frames as possible... + """ + sweep_data_lst, timestamps_lst = [], [] + skipFrame = int(20 / resample2frequency) # since nuscenes sweep at 20Hz, we want to resample to 10Hz + cnt = 0 + + # Find the first GT frame + while sample_data['next'] != '': + ts0 = sample_data['timestamp'] + gt_flow_flag = if_annotated_frame(sample_ann_dict, ts0) + + if gt_flow_flag: + # Found first GT frame, add it and break + sweep_data_lst.append(sample_data) + timestamps_lst.append(ts0) + sample_data = nusc.get('sample_data', sample_data['next']) + cnt = 1 # Reset counter after finding first GT + break + + sample_data = nusc.get('sample_data', sample_data['next']) + + # Continue processing from first GT frame + while sample_data['next'] != '': + ts0 = sample_data['timestamp'] + gt_flow_flag = if_annotated_frame(sample_ann_dict, ts0) + + # Always include GT frames + if gt_flow_flag: + sweep_data_lst.append(sample_data) + timestamps_lst.append(ts0) + sample_data = nusc.get('sample_data', sample_data['next']) + cnt = 1 # Reset counter after each GT frame + elif cnt % skipFrame == 0: + # Include non-GT frames according to sampling rate + sweep_data_lst.append(sample_data) + timestamps_lst.append(ts0) + sample_data = nusc.get('sample_data', sample_data['next']) + cnt += 1 + else: + # Skip this frame + sample_data = nusc.get('sample_data', sample_data['next']) + cnt += 1 + return sweep_data_lst, timestamps_lst + +def process_log(nusc_mode, data_dir: Path, scene_num_id: int, output_dir: Path, resample2frequency=10, n: Optional[int] = None) : + + def create_group_data(group, pc, pose, gm = None, flow_0to1=None, flow_valid=None, flow_category=None, flow_instance=None, ego_motion=None): + group.create_dataset('lidar', data=pc.astype(np.float32)) + group.create_dataset('pose', data=pose.astype(np.float64)) + group.create_dataset('ground_mask', data=gm.astype(bool)) + if ego_motion is not None: + group.create_dataset('ego_motion', data=ego_motion.astype(np.float32)) + if flow_0to1 is not None: + # ground truth flow information + group.create_dataset('flow', data=flow_0to1.astype(np.float32)) + group.create_dataset('flow_is_valid', data=flow_valid.astype(bool)) + group.create_dataset('flow_category_indices', data=flow_category.astype(np.uint8)) + group.create_dataset('flow_instance_id', data=flow_instance.astype(np.int16)) + + def compute_flow_simple(pc0, pose0, pose1, ts0, ts1, sample_ann_dict, dclass): + # compute delta transform between pose0 and pose1 + ego1_SE3_ego0 = npcal_pose0to1(pose0, pose1) + # flow due to ego motion + flow = np.zeros_like(pc0[:,:3]) + flow = pc0[:,:3] @ ego1_SE3_ego0[:3,:3].T + ego1_SE3_ego0[:3,3] - pc0[:,:3] # pose flow + + valid = np.ones(len(pc0), dtype=np.bool_) + classes = np.zeros(len(pc0), dtype=np.uint8) + instances = np.zeros(len(pc0), dtype=np.int16) + + delta_t = (ts1 - ts0) * 1e-6 + world_pc0 = pc0[:,:3] @ pose0[:3,:3].T + pose0[:3,3] + id_ = 0 + for key, anno_value in sample_ann_dict.items(): + # return: . Velocity in x/y/z direction in m/s. + if ts0 not in anno_value.keys(): + continue + ann_vel = anno_value[ts0]['vel'] + if np.isnan(ann_vel).any(): + continue + # previous x/y/z velocity is on world frame, need to transform to ego frame + ann_vel = ann_vel @ pose0[:3,:3] + box0 = anno_value[ts0]['bbx'] + cls = anno_value[ts0]['class'] + + # FIXME: compute points_in_box mask, expansion factor 1.1 here. + points_in_box_mask = points_in_box(box0, world_pc0[:,:3].T, wlh_factor=1.1) + classes[points_in_box_mask] = CATEGORY_TO_INDEX[NusNamMap[cls]] + if np.sum(points_in_box_mask) > 5: + obj_flow = np.ones_like(pc0[points_in_box_mask,:3]) * ann_vel * delta_t + flow[points_in_box_mask] += obj_flow + instances[points_in_box_mask] = (dclass[id_]+1) + id_ += 1 + else: + valid[points_in_box_mask] = False + + return {'flow_0_1': flow, + 'valid_0': valid, 'classes_0': classes, + 'ego_motion': ego1_SE3_ego0, + 'flow_instance_id': instances} + + nusc = NuScenes(dataroot=data_dir, version=nusc_mode, verbose=False) + scene = nusc.scene[scene_num_id] + log_id = scene['name'] + + # In nuscenes, samples are annotated at 2 Hz and sweeps at 20 Hz. + sample_data_lst = [] + for sample in nusc.sample: + if sample['scene_token'] != scene['token']: + continue + else: + sample_data_lst.append(sample) + sample_ann_dict = dict() + + # step: create a dictionary, with keys denoting the annotation token, + # and value denoting a list with annotation value and timestamp + for sample in sample_data_lst: + for ann_token in sample['anns']: + annotation = nusc.get('sample_annotation', ann_token) + if annotation['instance_token'] not in sample_ann_dict.keys(): + sample_ann_dict[annotation['instance_token']] = {} + # each annotation is list with attr + sample_ann_dict[annotation['instance_token']][sample['timestamp']]= { + 'bbx': Box(annotation['translation'], annotation['size'], Quaternion(annotation['rotation'])), \ + 'vel': nusc.box_velocity(ann_token).tolist(), \ + 'class': annotation['category_name']} + + # step note down timestamps for sweeps (interpolation points) and samples (at which data for interpolation exists) + now_sample_token_str = scene['first_sample_token'] + sample = nusc.get('sample', now_sample_token_str) + sample_data = nusc.get('sample_data', sample['data']['LIDAR_TOP']) + sweep_data_lst, timestamps = _resample_data(nusc, sample_data, sample_ann_dict, resample2frequency) + + if check_h5py_file_exists(output_dir/f'{log_id}.h5', timestamps): + print(f'{log_id} already exists and all timestamps are , skip...') + return + + dclass = defaultdict(lambda: len(dclass)) + mygroundseg = ground_seg(GROUNDSEG_config) + with h5py.File(output_dir/f'{log_id}.h5', 'a') as f: + for cnt, sweep_data in enumerate(sweep_data_lst): + ts0 = sweep_data['timestamp'] + # lidar point cloud + pc0 = _load_points_from_file(os.path.join(nusc.dataroot, sweep_data['filename'])) + pc0 = remove_ego_points(pc0) + if pc0.shape[0] < 10: + print(f'{log_id}/{ts0} has no points....') + continue + is_ground_0 = mygroundseg.run(pc0[:, :3]) + pose0 = get_pose(nusc, sweep_data) + + if cnt == len(sweep_data_lst) - 1: + group = f.create_group(str(ts0)) + create_group_data(group=group, pc=pc0, gm=is_ground_0.astype(np.bool_), pose=pose0.astype(np.float32)) + else: + sweep_data_next = sweep_data_lst[cnt+1] + ts1 = sweep_data_next['timestamp'] + pose1 = get_pose(nusc, sweep_data_next) + gt_flow_flag = if_annotated_frame(sample_ann_dict, ts0) + + group = f.create_group(str(ts0)) + if not gt_flow_flag: # no annotations, only save data + create_group_data(group=group, pc=pc0, gm=is_ground_0.astype(np.bool_), pose=pose0.astype(np.float32), \ + ego_motion=(np.linalg.inv(pose1) @ pose0).astype(np.float32)) + else: # compute sceneflow + scene_flow = compute_flow_simple(pc0, pose0, pose1, ts0, ts1, sample_ann_dict, dclass) + create_group_data(group=group, pc=pc0, gm=is_ground_0.astype(np.bool_), pose=pose0.astype(np.float32), + flow_0to1=scene_flow['flow_0_1'], flow_valid=scene_flow['valid_0'], flow_category=scene_flow['classes_0'], \ + flow_instance=scene_flow['flow_instance_id'], + ego_motion=scene_flow['ego_motion'].astype(np.float32)) + + +def proc(x, ignore_current_process=False): + if not ignore_current_process: + current=current_process() + pos = current._identity[0] + else: + pos = 1 + process_log(*x, n=pos) + +def process_logs(nusc_mode, data_dir: Path, scene_list: list, output_dir: Path, nproc: int): + """Compute sceneflow for all logs in the dataset. Logs are processed in parallel. + Args: + data_dir: Argoverse 2.0 directory + output_dir: Output directory. + """ + + if not (data_dir).exists(): + print(f'{data_dir} not found') + return + + args = sorted([(nusc_mode, data_dir, scene_num_id, output_dir) for scene_num_id in range(len(scene_list))]) + print(f'Using {nproc} processes') + + # # for debug + # for x in tqdm(args): + # proc(x, ignore_current_process=True) + # break + + if nproc <= 1: + for x in tqdm(args): + proc(x, ignore_current_process=True) + else: + with Pool(processes=nproc) as p: + res = list(tqdm(p.imap_unordered(proc, args), total=len(scene_list), ncols=100)) + +def main( + data_dir: str = "/home/kin/data/nus/v1.0", + mode: str = "v1.0-mini", + output_dir: str ="/home/kin/data/nus/h5py/demo", + nproc: int = (multiprocessing.cpu_count() - 1), + only_index: bool = False, + split_name = None +): + available_vers = ['v1.0-trainval', 'v1.0-test', 'v1.0-mini'] # defined by nus. + assert mode in available_vers + # nusc = NuScenes(dataroot=data_dir, version=mode, verbose=False) + # print(f"Processing {mode} dataset with {len(nusc.scene)} scenes...") + if mode == 'v1.0-trainval': + train_scenes = splits.train + val_scenes = splits.val + if split_name is not None and split_name == 'train': + input_dict = {'train': train_scenes} + elif split_name is not None and split_name == 'val': + input_dict = {'val': val_scenes} + else: + input_dict = {'train': train_scenes, 'val': val_scenes} + elif mode == 'v1.0-test': + test_scenes = splits.test + input_dict = {'test': test_scenes} + elif mode == 'v1.0-mini': + train_scenes = splits.mini_train + val_scenes = splits.mini_val + input_dict = { + 'train': train_scenes, + 'val': val_scenes + } + # NOTE(Qingwen): or if you don't want to split mini, use below + # input_dict = {'mini': train_scenes + val_scenes} + else: + raise ValueError('unknown') + + for input_key, input_val in input_dict.items(): + output_dir_ = Path(output_dir) / input_key + print("[INFO] We are processing data to ", output_dir_) + if only_index: + create_reading_index(Path(output_dir_)) + create_reading_index(Path(output_dir_), flow_inside_check=True) + return + output_dir_.mkdir(exist_ok=True, parents=True) + process_logs(mode, Path(data_dir), input_val, output_dir_, nproc) + create_reading_index(output_dir_) + create_reading_index(Path(output_dir_), flow_inside_check=True) + +if __name__ == '__main__': + start_time = time.time() + fire.Fire(main) + print(f"\nTime used: {(time.time() - start_time)/60:.2f} mins") \ No newline at end of file diff --git a/dataprocess/misc_data.py b/dataprocess/misc_data.py index 9698a77..b749a3c 100644 --- a/dataprocess/misc_data.py +++ b/dataprocess/misc_data.py @@ -3,7 +3,32 @@ from pathlib import Path from tqdm import tqdm -def create_reading_index(data_dir: Path): +def check_h5py_file_exists(h5py_file: Path, timestamps: list, verbose: bool = False) -> bool: + if not h5py_file.exists(): + return False + + log_id = h5py_file.name.split(".")[0] + try: + with h5py.File(h5py_file, 'r') as f: + for timestamp in timestamps: + if str(timestamp) not in f.keys(): + # delete the file if the timestamp is not in the file + # and it will reprocess this scene file + if verbose: + print(f"\n--- WARNING [data]: {log_id} has no {timestamp}, will be deleted the scene h5py.") + os.remove(h5py_file) + return False + except Exception as e: + if verbose: + print(f"\n--- WARNING [data]: {log_id} has error: {e}, will be deleted the scene h5py.") + os.remove(h5py_file) + return False + if verbose: + print(f'\n--- INFO [data]: {log_id} has been processed with total {len(timestamps)} timestamps.') + return True + +def create_reading_index(data_dir: Path, flow_inside_check=False): + pkl_file_name = 'index_total.pkl' if not flow_inside_check else 'index_flow.pkl' start_time = time.time() data_index = [] for file_name in tqdm(os.listdir(data_dir), ncols=100): @@ -12,14 +37,20 @@ def create_reading_index(data_dir: Path): scene_id = file_name.split(".")[0] timestamps = [] with h5py.File(data_dir/file_name, 'r') as f: - timestamps.extend(f.keys()) + if flow_inside_check: + for key in f.keys(): + if 'flow' in f[key]: + # print(f"Found flow in {scene_id} at {key}") + timestamps.append(key) + else: + timestamps.extend(f.keys()) timestamps.sort(key=lambda x: int(x)) # make sure the timestamps are in order for timestamp in timestamps: data_index.append([scene_id, timestamp]) - with open(data_dir/'index_total.pkl', 'wb') as f: + with open(data_dir/pkl_file_name, 'wb') as f: pickle.dump(data_index, f) - print(f"Create reading index Successfully, cost: {time.time() - start_time:.2f} s") + print(f"Create {pkl_file_name} index Successfully, cost: {time.time() - start_time:.2f} s") class SE2: @@ -88,4 +119,40 @@ def compose(self, right_se2: "SE2") -> "SE2": translation=chained_transform_matrix[:2, 2], ) return chained_se2 - \ No newline at end of file + + +## ====> nuScenes to Argoverse Mapping +NusNamMap = { + 'noise': 'NONE', + 'animal': 'ANIMAL', + 'human.pedestrian.adult': 'PEDESTRIAN', + 'human.pedestrian.child': 'PEDESTRIAN', + 'human.pedestrian.construction_worker': 'PEDESTRIAN', + 'human.pedestrian.personal_mobility': 'PEDESTRIAN', + 'human.pedestrian.police_officer': 'PEDESTRIAN', + 'human.pedestrian.stroller': 'STROLLER', + 'human.pedestrian.wheelchair': 'WHEELCHAIR', + 'movable_object.barrier': 'NONE', + 'movable_object.debris': 'NONE', + 'movable_object.pushable_pullable': 'NONE', + 'movable_object.trafficcone': 'CONSTRUCTION_CONE', + 'static_object.bicycle_rack': 'NONE', + 'vehicle.bicycle': 'BICYCLE', + 'vehicle.bus.bendy': 'ARTICULATED_BUS', + 'vehicle.bus.rigid': 'BUS', + 'vehicle.car': 'REGULAR_VEHICLE', + 'vehicle.construction': 'LARGE_VEHICLE', + 'vehicle.emergency.ambulance': 'LARGE_VEHICLE', + 'vehicle.emergency.police': 'REGULAR_VEHICLE', + 'vehicle.motorcycle': 'MOTORCYCLE', + 'vehicle.trailer': 'VEHICULAR_TRAILER', + 'vehicle.truck': 'TRUCK', + 'flat.driveable_surface': 'NONE', + 'flat.other': 'NONE', + 'flat.sidewalk': 'NONE', + 'flat.terrain': 'NONE', + 'static.manmade': 'NONE', + 'static.other': 'NONE', + 'static.vegetation': 'NONE', + 'vehicle.ego': 'NONE' +} \ No newline at end of file diff --git a/src/models/basic/__init__.py b/src/models/basic/__init__.py index 244c04c..6803e63 100644 --- a/src/models/basic/__init__.py +++ b/src/models/basic/__init__.py @@ -1,6 +1,21 @@ import torch import torch.nn as nn +import math +from torch.optim.lr_scheduler import _LRScheduler +class BaseModel(nn.Module): + def __init__(self): + super(BaseModel, self).__init__() + + def load_from_checkpoint(self, ckpt_path): + ckpt = torch.load(ckpt_path, map_location="cpu")["state_dict"] + state_dict = { + k[len("model.") :]: v for k, v in ckpt.items() if k.startswith("model.") + } + print("\nLoading... model weight from: ", ckpt_path, "\n") + return self.load_state_dict(state_dict=state_dict, strict=False) + + @torch.no_grad() def cal_pose0to1(pose0: torch.Tensor, pose1: torch.Tensor): """ @@ -76,4 +91,56 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: batchnorm_res = conv_res else: batchnorm_res = self.batchnorm(conv_res) - return self.nonlinearity(batchnorm_res) \ No newline at end of file + return self.nonlinearity(batchnorm_res) + +class WarmupCosLR(_LRScheduler): + def __init__( + self, optimizer, min_lr, lr, \ + warmup_epochs, epochs, \ + last_epoch=-1, verbose=False + ) -> None: + self.min_lr = min_lr + self.lr = lr + self.epochs = epochs + self.warmup_epochs = warmup_epochs + super(WarmupCosLR, self).__init__(optimizer, last_epoch, verbose) + + def state_dict(self): + """Returns the state of the scheduler as a :class:`dict`. + + It contains an entry for every variable in self.__dict__ which + is not the optimizer. + """ + return { + key: value for key, value in self.__dict__.items() if key != "optimizer" + } + + def load_state_dict(self, state_dict): + """Loads the schedulers state. + + Args: + state_dict (dict): scheduler state. Should be an object returned + from a call to :meth:`state_dict`. + """ + self.__dict__.update(state_dict) + + def get_init_lr(self): + lr = self.lr / self.warmup_epochs + return lr + + def get_lr(self): + if self.last_epoch < self.warmup_epochs: + lr = self.lr * (self.last_epoch + 1) / self.warmup_epochs + else: + lr = self.min_lr + 0.5 * (self.lr - self.min_lr) * ( + 1 + + math.cos( + math.pi + * (self.last_epoch - self.warmup_epochs) + / (self.epochs - self.warmup_epochs) + ) + ) + if "lr_scale" in self.optimizer.param_groups[0]: + return [lr * group["lr_scale"] for group in self.optimizer.param_groups] + + return [lr for _ in self.optimizer.param_groups] diff --git a/src/models/deflow.py b/src/models/deflow.py index ebb1629..d624690 100644 --- a/src/models/deflow.py +++ b/src/models/deflow.py @@ -15,9 +15,9 @@ from .basic.unet import FastFlow3DUNet from .basic.encoder import DynamicEmbedder from .basic.decoder import LinearDecoder, ConvGRUDecoder -from .basic import cal_pose0to1 +from .basic import cal_pose0to1, BaseModel -class DeFlow(nn.Module): +class DeFlow(BaseModel): def __init__(self, voxel_size = [0.2, 0.2, 6], point_cloud_range = [-51.2, -51.2, -3, 51.2, 51.2, 3], grid_feature_size = [512, 512], @@ -38,14 +38,6 @@ def __init__(self, voxel_size = [0.2, 0.2, 6], self.timer = dztimer.Timing() self.timer.start("Total") - def load_from_checkpoint(self, ckpt_path): - ckpt = torch.load(ckpt_path, map_location="cpu")["state_dict"] - state_dict = { - k[len("model.") :]: v for k, v in ckpt.items() if k.startswith("model.") - } - print("\nLoading... model weight from: ", ckpt_path, "\n") - return self.load_state_dict(state_dict=state_dict, strict=False) - def forward(self, batch): """ input: using the batch from dataloader, which is a dict diff --git a/src/models/fastflow3d.py b/src/models/fastflow3d.py index a08090f..e4d5774 100644 --- a/src/models/fastflow3d.py +++ b/src/models/fastflow3d.py @@ -12,10 +12,10 @@ from .basic.unet import FastFlow3DUNet from .basic.encoder import DynamicEmbedder from .basic.decoder import LinearDecoder -from .basic import cal_pose0to1 +from .basic import cal_pose0to1, BaseModel -class FastFlow3D(nn.Module): +class FastFlow3D(BaseModel): def __init__(self, voxel_size = [0.2, 0.2, 6], point_cloud_range = [-51.2, -51.2, -3, 51.2, 51.2, 3], grid_feature_size = [512, 512]): @@ -30,15 +30,7 @@ def __init__(self, voxel_size = [0.2, 0.2, 6], self.timer = dztimer.Timing() self.timer.start("Total") - - def load_from_checkpoint(self, ckpt_path): - ckpt = torch.load(ckpt_path, map_location="cpu")["state_dict"] - state_dict = { - k[len("model.") :]: v for k, v in ckpt.items() if k.startswith("model.") - } - print("\nLoading... model weight from: ", ckpt_path, "\n") - return self.load_state_dict(state_dict=state_dict, strict=False) - + def _model_forward(self, pc0s, pc1s): pc0_before_pseudoimages, pc0_voxel_infos_lst = self.embedder(pc0s) diff --git a/src/models/flow4d.py b/src/models/flow4d.py index b0f52b5..cbe5d33 100644 --- a/src/models/flow4d.py +++ b/src/models/flow4d.py @@ -9,11 +9,11 @@ import torch.nn as nn import dztimer, torch -from .basic import wrap_batch_pcs +from .basic import wrap_batch_pcs, BaseModel from .basic.flow4d_module import DynamicEmbedder_4D from .basic.flow4d_module import Network_4D, Seperate_to_3D, Point_head -class Flow4D(nn.Module): +class Flow4D(BaseModel): def __init__(self, voxel_size = [0.2, 0.2, 0.2], point_cloud_range = [-51.2, -51.2, -2.2, 51.2, 51.2, 4.2], grid_feature_size = [512, 512, 32], @@ -38,14 +38,6 @@ def __init__(self, voxel_size = [0.2, 0.2, 0.2], self.timer = dztimer.Timing() self.timer.start("Total") - def load_from_checkpoint(self, ckpt_path): - ckpt = torch.load(ckpt_path, map_location="cpu")["state_dict"] - state_dict = { - k[len("model.") :]: v for k, v in ckpt.items() if k.startswith("model.") - } - print("\nLoading... model weight from: ", ckpt_path, "\n") - return self.load_state_dict(state_dict=state_dict, strict=False) - def forward(self, batch): """ input: using the batch from dataloader, which is a dict diff --git a/src/models/ssf.py b/src/models/ssf.py index 8061180..8fb9799 100644 --- a/src/models/ssf.py +++ b/src/models/ssf.py @@ -16,9 +16,9 @@ from .basic.encoder import DynamicVoxelizer from .basic.ssf_module import DynamicScatterVFE from .basic.decoder import SimpleLinearDecoder -from .basic import cal_pose0to1 +from .basic import cal_pose0to1, BaseModel -class SSF(nn.Module): +class SSF(BaseModel): def __init__(self, voxel_size = [0.2, 0.2, 6], point_cloud_range = [-51.2, -51.2, -3, 51.2, 51.2, 3], grid_feature_size = [512, 512], @@ -57,14 +57,6 @@ def __init__(self, voxel_size = [0.2, 0.2, 6], self.timer = dztimer.Timing() self.timer.start("Total") - def load_from_checkpoint(self, ckpt_path): - ckpt = torch.load(ckpt_path, map_location="cpu")["state_dict"] - state_dict = { - k[len("model.") :]: v for k, v in ckpt.items() if k.startswith("model.") - } - print("\nLoading... model weight from: ", ckpt_path, "\n") - return self.load_state_dict(state_dict=state_dict, strict=False) - def forward(self, batch): """ input: using the batch from dataloader, which is a dict diff --git a/src/trainer.py b/src/trainer.py index d779388..fe2cb16 100644 --- a/src/trainer.py +++ b/src/trainer.py @@ -24,7 +24,8 @@ import os, sys, time, h5py, pickle BASE_DIR = os.path.abspath(os.path.join( os.path.dirname( __file__ ), '..' )) sys.path.append(BASE_DIR) -from src.utils.mics import import_func, weights_init, zip_res +from src.utils import import_func +from src.utils.mics import weights_init, zip_res from src.utils.av2_eval import write_output_file from src.models.basic import cal_pose0to1 from src.utils.eval_metric import OfficialMetrics, evaluate_leaderboard, evaluate_leaderboard_v2, evaluate_ssf diff --git a/src/utils/__init__.py b/src/utils/__init__.py index a9a803f..5111258 100644 --- a/src/utils/__init__.py +++ b/src/utils/__init__.py @@ -9,13 +9,24 @@ class bc: BOLD = '\033[1m' UNDERLINE = '\033[4m' +# ====> import func through string, ref: https://stackoverflow.com/a/19393328 +import importlib +def import_func(path: str): + function_string = path + mod_name, func_name = function_string.rsplit('.',1) + mod = importlib.import_module(mod_name) + func = getattr(mod, func_name) + return func -def hex_to_rgb(hex_color): - hex_color = hex_color.lstrip("#") - return tuple(int(hex_color[i:i + 2], 16) / 255.0 for i in (0, 2, 4)) - -color_map_hex = ['#a6cee3', '#de2d26', '#1f78b4','#b2df8a','#33a02c','#fb9a99','#e31a1c','#fdbf6f','#ff7f00',\ - '#cab2d6','#6a3d9a','#ffff99','#b15928', '#8dd3c7','#ffffb3','#bebada','#fb8072','#80b1d3',\ - '#fdb462','#b3de69','#fccde5','#d9d9d9','#bc80bd','#ccebc5','#ffed6f'] - -color_map = [hex_to_rgb(color) for color in color_map_hex] \ No newline at end of file +import numpy as np +def npcal_pose0to1(pose0, pose1): + """ + Note(Qingwen 2023-12-05 11:09): + Don't know why but it needed set the pose to float64 to calculate the inverse + otherwise it will be not expected result.... + """ + pose1_inv = np.eye(4, dtype=np.float64) + pose1_inv[:3,:3] = pose1[:3,:3].T + pose1_inv[:3,3] = (pose1[:3,:3].T * -pose1[:3,3]).sum(axis=1) + pose_0to1 = pose1_inv @ pose0.astype(np.float64) + return pose_0to1.astype(np.float32) \ No newline at end of file diff --git a/tools/visualization_rerun.py b/tools/visualization_rerun.py index 57331c5..6cc6ffc 100644 --- a/tools/visualization_rerun.py +++ b/tools/visualization_rerun.py @@ -19,7 +19,7 @@ BASE_DIR = os.path.abspath(os.path.join( os.path.dirname( __file__ ), '..' )) sys.path.append(BASE_DIR) from src.utils.mics import HDF5Data, flow_to_rgb -from src.utils import color_map +from src.utils.o3d_view import color_map import rerun as rr import rerun.blueprint as rrb import argparse