diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..06bb0e9 --- /dev/null +++ b/.gitignore @@ -0,0 +1,22 @@ +.mypy_cache/ +.vscode/ +**/__pycache__/ +docker_history.txt +work_dirs/ +model_checkpoints/ +model_checkpoints_old/ +tb_logs/ +tb_logs*/ +nsfp_split_configs/ +lightning_logs/ +kylevedder_offline_sceneflow_latest.sqsh +chamfer_investigation/merged_dicts.pkl +validation_results/ +validation_results*/ +**.pkl +**.png +job_dir/ +.pl_auto_save.ckpt +cluster_tb_logs/ +*_tb_logs/ +**.zip \ No newline at end of file diff --git a/GETTING_STARTED.md b/GETTING_STARTED.md new file mode 100644 index 0000000..7fdb9e1 --- /dev/null +++ b/GETTING_STARTED.md @@ -0,0 +1,75 @@ +# How to use this code + +## File system assumptions + +### Argoverse 2 + +Somewhere on disk, have an `argoverse2/` folder so that the downloaded files live inside + +``` +argoverse2/train +argoverse2/val +argoverse2/test +``` + +and generate the train and val supervision labels to + +``` +argoverse2/train_sceneflow +argoverse2/val_sceneflow +``` + + +The [Argoverse 2 Scene Flow generation script](https://github.com/nchodosh/argoverse2-sf) to compute ground truth flows for both `train/` and `val/`. + +### Waymo Open + +Download the Scene Flow labels contributed by _Scalable Scene Flow from Point Clouds in the Real World_. We preprocess these files, both to convert them from an annoying proto file format to a standard Python format and to remove the ground points. + +Do this using the `data_prep_scripts/waymo/extract_flow_and_remove_ground.py` file in the Waymo Open docker container. + +## Docker Images + +This project has three different docker images for different functions. + +Each image has an associated convinence script. You must edit this script to modify the mount commands to point to your Argoverse 2 / Waymo Open data locations. The `-v` commands are these mount commands. As an example, + +``` +-v `pwd`:/project +``` + +runs `pwd` inside the script, getting the current directory, and ensures that it's mounted as `/project` inside the container. You must edit the `/efs/` mount, i.e. + +``` +-v /efs:/efs +``` + +so that the source points to the containing folder of your `argoverse2/` and `waymo_open_processed_flow/` directories. As an example, on our cluster I have `~/datasets/argoverse2` and `~/datasets/waymo_open_processed_flow`, so if I were to run this on our cluster I would modify these mounts to be + +``` +-v $HOME/datasets:/efs +``` + +It's important that, once inside the docker container, the path to the Argoverse 2 dataset is `/efs/argoverse2/...` and the path to Waymo Open is `/efs/waymo_open_processed_flow/...` + +### Main image: + +Built with `docker/Dockerfile` [[dockerhub](https://hub.docker.com/repository/docker/kylevedder/zeroflow)] + +Convenience launch script is `./launch.sh`. Make sure that Argoverse 2 and Waymo Open preprocessed are mounted inside a folder in the container as `/efs`. + +### Waymo preprocessing image: + +Built with `docker/Dockerfilewaymo` [[dockerhub](https://hub.docker.com/repository/docker/kylevedder/zeroflow_waymo)]. Includes Tensorflow and other dependencies to preprocess the raw Waymo Open format and convert it to a standard format readable in the main image. + +Convenience launch script is `./launch_waymo.sh`. + +### AV2 challenge submission image: + +Built with `docker/Dockerav2` [[dockerhub](https://hub.docker.com/repository/docker/kylevedder/zeroflow_av2)]. Based on the main image, but includes the [AV2 API](https://github.com/argoverse/av2-api). + +Convenience launch script is `./launch_av2.sh`. + +## Setting up the base system + +The base system must have CUDA 11.3+ and [NVidia Docker](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html#docker) installed. diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..8ffb6f0 --- /dev/null +++ b/LICENSE @@ -0,0 +1,20 @@ +Copyright (c) 2023 Kyle Vedder + +Permission is hereby granted, free of charge, to any person obtaining +a copy of this software and associated documentation files (the +"Software"), to deal in the Software without restriction, including +without limitation the rights to use, copy, modify, merge, publish, +distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, subject to +the following conditions: + +The above copyright notice and this permission notice shall be +included in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE +LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION +WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..550dcb5 --- /dev/null +++ b/README.md @@ -0,0 +1,52 @@ +# ZeroFlow: Fast Zero Label Scene Flow via Distillation + +[Kyle Vedder](http://vedder.io), [Neehar Peri](http://www.neeharperi.com/), [Nathaniel Chodosh](https://scholar.google.com/citations?user=b4qKr7gAAAAJ&hl=en), [Ishan Khatri](https://ishan.khatri.io/), [Eric Eaton](https://www.seas.upenn.edu/~eeaton/), [Dinesh Jayaraman](https://www.seas.upenn.edu/~dineshj/), [Yang Liu](https://youngleox.github.io/), [Deva Ramanan](https://www.cs.cmu.edu/~deva/), and [James Hays](https://faculty.cc.gatech.edu/~hays/) + +Project webpage: [vedder.io/zeroflow](http://vedder.io/zeroflow) + +arXiv link: [arxiv.org/abs/2305.10424](http://arxiv.org/abs/2305.10424) + +**Citation:** + +``` +@article{Vedder2023zeroflow, + author = {Kyle Vedder and Neehar Peri and Nathaniel Chodosh and Ishan Khatri and Eric Eaton and Dinesh Jayaraman and Yang Liu Deva Ramanan and James Hays}, + title = {{ZeroFlow: Fast Zero Label Scene Flow via Distillation}}, + journal = {arXiv}, + year = {2023}, +} +``` + +## Pre-requisites / Getting Started + +Read the [Getting Started](./GETTING_STARTED.md) doc for detailed instructions to setup the AV2 and Waymo Open datasets and use the prepared docker environments. + +## Training a model + + Inside the main container (`./launch.sh`), run the `train_pl.py` with a path to a config (inside `configs/`) and optionally specify any number of GPUs (defaults to all GPUs on the system). + +``` +python train_pl.py --gpus +``` + +The script will start by verifying the val dataloader works, and then launch the train job. + +## Testing a model + +Inside the main (`./launch.sh`), run the `train_pl.py` with a path to a config (inside `configs/`), a path to a checkpoint, and the number of GPUs (defaults to a single GPU). + +``` +python test_pl.py --gpus +``` + +## Generating paper plots + +After all relevant checkpoints have been tested, thus generating result files in `validation_results/configs/...`, run `plot_performance.py` to generate the figures and tables used in the paper. + +## Submitting to the AV2 Scene Flow competition + +1. Dump the outputs of the model + a. `configs/fastflow3d/argo/nsfp_distilatation_dump_output.py` to dump the `val` set result + b. `configs/fastflow3d/argo/nsfp_distilatation_dump_output_test.py` to dump the `test` set result +2. Convert to the competition submission format (`av2_scene_flow_competition_submit.py`) +3. Use official zip `make_submission_archive.py` file (`python /av2-api/src/av2/evaluation/scene_flow/make_submission_archive.py /efs/argoverse2/test_official_masks.zip`) \ No newline at end of file diff --git a/av2_scene_flow_competition_submit.py b/av2_scene_flow_competition_submit.py new file mode 100644 index 0000000..027aaa6 --- /dev/null +++ b/av2_scene_flow_competition_submit.py @@ -0,0 +1,316 @@ +import os +# Set OMP env num threads to 1 to avoid deadlock in multithreading +os.environ["OMP_NUM_THREADS"] = "1" +import argparse +import pandas as pd +import numpy as np +from pathlib import Path +from loader_utils import load_npz +import multiprocessing as mp +import tqdm +from typing import Dict, Any, Tuple, List, Union +import open3d as o3d +from dataloaders import ArgoverseRawSequenceLoader, ArgoverseRawSequence +from pointclouds import PointCloud +import copy + +# Parse arguments from command line for scene flow masks and scene flow outputs +parser = argparse.ArgumentParser() +parser.add_argument("scene_flow_mask_dir", + type=Path, + help="Path to scene flow mask directory") +parser.add_argument("scene_flow_output_dir", + type=Path, + help="Path to scene flow output directory") +parser.add_argument("argoverse_dir", + type=Path, + help="Path to scene flow output directory") +parser.add_argument("output_dir", type=Path, help="Path to output directory") +# Num CPUs to use for multiprocessing +parser.add_argument("--num_cpus", + type=int, + default=mp.cpu_count(), + help="Number of cpus to use for multiprocessing") +# bool flag for subset of data to use +parser.add_argument("--subset", + action="store_true", + help="Whether to use a subset of the data") +parser.add_argument("--visualize", + action="store_true", + help="Whether to visualize the data") +args = parser.parse_args() + +assert args.scene_flow_mask_dir.is_dir( +), f"{args.scene_flow_mask_dir} is not a directory" +assert args.scene_flow_output_dir.is_dir( +), f"{args.scene_flow_output_dir} is not a directory" +assert args.argoverse_dir.is_dir(), f"{args.argoverse_dir} is not a directory" + +args.output_dir.mkdir(parents=True, exist_ok=True) + + +def load_feather(filepath: Path): + filepath = Path(filepath) + assert filepath.exists(), f'{filepath} does not exist' + return pd.read_feather(filepath) + + +def load_scene_flow_mask_from_folder(sequence_folder: Path) -> Dict[int, Any]: + files = sorted(sequence_folder.glob("*.feather")) + masks = [load_feather(file)['mask'].to_numpy() for file in files] + return sequence_folder.stem, { + int(file.stem): mask + for file, mask in zip(files, masks) + } + + +def load_scene_flow_predictions_from_folder( + sequence_folder: Path) -> List[Any]: + files = sorted(sequence_folder.glob("*.npz")) + return sequence_folder.stem, [ + dict(load_npz(file, verbose=False)) for file in files + ] + + +def multiprocess_load(folder: Path, worker_fn) -> Dict[str, Any]: + sequence_folders = sorted(e for e in folder.glob("*") if e.is_dir()) + if args.subset: + sequence_folders = sequence_folders[1:2] + sequence_lookup = {} + if args.num_cpus > 1: + with mp.Pool(processes=args.num_cpus) as pool: + for k, v in tqdm.tqdm(pool.imap_unordered(worker_fn, + sequence_folders), + total=len(sequence_folders)): + sequence_lookup[k] = v + else: + for sequence_folder in tqdm.tqdm(sequence_folders): + k, v = worker_fn(sequence_folder) + sequence_lookup[k] = v + return sequence_lookup + + +print("Loading scene flow masks...") +sequence_mask_lookup = multiprocess_load(args.scene_flow_mask_dir, + load_scene_flow_mask_from_folder) +print("Loading scene flow outputs...") +sequence_prediction_lookup = multiprocess_load( + args.scene_flow_output_dir, load_scene_flow_predictions_from_folder) +print("Done loading scene flow masks and outputs") + +sequence_loader = ArgoverseRawSequenceLoader(args.argoverse_dir) +mask_keys = set(sequence_mask_lookup.keys()) +output_keys = set(sequence_prediction_lookup.keys()) +assert mask_keys == output_keys, f"Mask keys {mask_keys} != output keys {output_keys}" + + +def make_lineset(pc, flowed_pc, draw_color): + line_set = o3d.geometry.LineSet() + assert len(pc) == len( + flowed_pc + ), f'pc and flowed_pc must have same length, but got {len(pc)} and {len(flowed_pc)}' + line_set_points = np.concatenate([pc, flowed_pc], axis=0) + + lines = np.array([[i, i + len(pc)] for i in range(len(pc))]) + line_set.points = o3d.utility.Vector3dVector(line_set_points) + line_set.lines = o3d.utility.Vector2iVector(lines) + draw_color = { + 'green': [0, 1, 0], + 'blue': [0, 0, 1], + 'red': [1, 0, 0], + 'yellow': [1, 1, 0], + 'orange': [1, 0.5, 0], + 'purple': [0.5, 0, 1], + 'pink': [1, 0, 1], + 'cyan': [0, 1, 1], + 'white': [1, 1, 1], + }[draw_color] + line_set.colors = o3d.utility.Vector3dVector( + [draw_color for _ in range(len(lines))]) + return line_set + + +def merge_predictions_and_mask_data(sequence: ArgoverseRawSequence, + timestamp_to_mask: Dict[int, Any], + predictions: List[Any]) -> Dict[int, Any]: + + # Predictions are over frame pairs, so there are N - 1 predictions for N frames in a sequence. + assert len(predictions) == len( + sequence + ) - 1, f"len(id_to_output) {len(predictions)} != len(sequence) - 1 {len(sequence) - 1}" + + timestamp_to_sequence_frame_idx = { + timestamp: idx + for idx, timestamp in enumerate(sequence.timestamp_list) + } + + timestamp_to_prediction = { + t: v + for t, v in zip(sequence.timestamp_list, predictions) + } + + timestamp_to_masked_output = {} + + # Check that the masks are corresponding to the lidar files. + for mask_timestamp in timestamp_to_mask: + frame_idx = timestamp_to_sequence_frame_idx[mask_timestamp] + frame_dict = sequence.load(frame_idx, frame_idx) + next_frame_ego_dict = sequence.load(frame_idx + 1, frame_idx + 1) + next_frame_curr_dict = sequence.load(frame_idx + 1, frame_idx) + submission_mask_data = timestamp_to_mask[mask_timestamp] + prediction_data = copy.deepcopy(timestamp_to_prediction[mask_timestamp]) + + assert len(frame_dict['relative_pc_with_ground']) == len( + submission_mask_data + ), f"len(frame_dict['relative_pc_with_ground']) {len(frame_dict['relative_pc_with_ground'])} != len(submission_mask_data) {len(submission_mask_data)}" + + prediction_data['pc'] = frame_dict['relative_pc'] + prediction_data['pc_with_ground'] = frame_dict[ + 'relative_pc_with_ground'] + prediction_data['pc_with_ground_is_ground_points'] = frame_dict[ + 'is_ground_points'] + prediction_data[ + 'pc_with_ground_is_submission_points'] = submission_mask_data + prediction_data['next_pc'] = next_frame_ego_dict['relative_pc'] + prediction_data['next_pc_with_ground'] = next_frame_ego_dict[ + 'relative_pc_with_ground'] + prediction_data[ + 'next_pc_with_ground_is_ground_points'] = next_frame_ego_dict[ + 'is_ground_points'] + prediction_data['relative_pose'] = next_frame_curr_dict[ + 'relative_pose'] + + timestamp_to_masked_output[mask_timestamp] = prediction_data + + return timestamp_to_masked_output + + +def save_sequence(sequence: ArgoverseRawSequence, + timestamp_to_mask: Dict[int, Any], predictions: List[Any]): + sequence_name = sequence.log_id + timestamp_to_masked_output = merge_predictions_and_mask_data( + sequence, timestamp_to_mask, predictions) + # print(f"Saving sequence {sequence_name}") + for timestamp, data in sorted(timestamp_to_masked_output.items()): + save_path = args.output_dir / sequence_name / f"{timestamp}.feather" + save_path.parent.mkdir(parents=True, exist_ok=True) + + # Sized to the full point cloud, including ground points. + full_size_flow_array = np.zeros((len(data['pc_with_ground']), 3), + dtype=data['flow'].dtype) + no_ground_flow_array = np.zeros((len(data['pc']), 3), + dtype=data['flow'].dtype) + + # Assign the flow to the non-ground and valid flow points. + no_ground_flow_array[data['valid_idxes']] = data['flow'] + + # Transform the flow from being in second frame to being in first frame. + # The relative pose describes the transform from first to second frame, + # so we need to rotate by the inverse of the relative pose. + no_ground_flow_array = (data['relative_pose'].inverse( + ).rotation_matrix @ no_ground_flow_array.T).T + + full_size_flow_array[ + ~data['pc_with_ground_is_ground_points']] = no_ground_flow_array + full_size_is_dynamic = (np.linalg.norm(full_size_flow_array, axis=1) > + 0.05) + + # Add ego motion to the flow -- evaluation is performed on uncompenstated flow. + # This is done by taking the point cloud and transforming it with the relative pose, and then computing the point differences. + + # Transform the point cloud with the relative pose from t to t+1. + pc_with_ground_transformed = data['pc_with_ground'].transform( + data['relative_pose'].inverse()) + point_diffs = pc_with_ground_transformed.matched_point_diffs( + data['pc_with_ground']) + # Add the point differences to the flow. + full_size_flow_array += point_diffs + + # Extract the submission mask points from scratch array. + output_flow_array = full_size_flow_array[ + data["pc_with_ground_is_submission_points"]] + output_is_dynamic = full_size_is_dynamic[ + data["pc_with_ground_is_submission_points"]] + + output_pc = data['pc_with_ground'][ + data["pc_with_ground_is_submission_points"]] + + def resize_output_pc_next(): + output_pc_next = data['next_pc_with_ground'].points + if len(output_pc_next) > len(output_pc): + return output_pc_next[:len(output_pc)] + else: + output_scratch = np.zeros_like(output_pc) + output_scratch[:len(output_pc_next)] = output_pc_next + return output_scratch + + output_pc_next = resize_output_pc_next() + + df = pd.DataFrame(output_flow_array.astype(np.float16), + columns=['flow_tx_m', 'flow_ty_m', 'flow_tz_m']) + df['is_dynamic'] = output_is_dynamic + # Debug info not needed for final submission + # df[['point_x', 'point_y', 'point_z']] = output_pc + # df[['point_x_next', 'point_y_next', 'point_z_next']] = output_pc_next + if save_path.exists(): + save_path.unlink() + df.to_feather(save_path) + + if args.visualize: + print("Sequence", sequence_name, "Timestamp", timestamp) + + def get_gt_flow(): + gt_path = Path("/efs/argoverse2/val_official_sceneflow/" + ) / sequence.log_id / f"{timestamp}.feather" + assert gt_path.exists(), f"gt_path {gt_path} does not exist" + gt_df = pd.read_feather(gt_path) + gt_flow = gt_df[['flow_tx_m', 'flow_ty_m', + 'flow_tz_m']].to_numpy() + return gt_flow + + # Visualize the lidar data and the flow data with Open3D + pc = o3d.geometry.PointCloud() + pc.points = o3d.utility.Vector3dVector(output_pc) + pc.paint_uniform_color([0.1, 0, 0]) + + next_pc = o3d.geometry.PointCloud() + next_pc.points = o3d.utility.Vector3dVector( + data['next_pc_with_ground']) + next_pc.paint_uniform_color([0, 0, 1]) + + pred_flow_lineset = make_lineset(output_pc, + output_pc + output_flow_array, + 'red') + gt_flow_lineset = make_lineset(output_pc, + output_pc + get_gt_flow(), 'green') + + print("Num dynamic points", output_is_dynamic.sum(), + "Percent dynamic", + output_is_dynamic.sum() / len(output_is_dynamic)) + + # visualize + o3d.visualization.draw_geometries( + [pc, next_pc, pred_flow_lineset, gt_flow_lineset]) + + +def save_wrapper(args): + save_sequence(*args) + + +print("Preparing args...") + +save_args = [(sequence_loader.load_sequence(sequence_name), + sequence_mask_lookup[sequence_name], + sequence_prediction_lookup[sequence_name]) + for sequence_name in tqdm.tqdm(sorted(sequence_mask_lookup))] + +print("Saving sequences...") + +if args.num_cpus > 1: + with mp.Pool(processes=args.num_cpus) as pool: + for _ in tqdm.tqdm(pool.imap_unordered(save_wrapper, save_args), + total=len(save_args)): + pass +else: + for arg in tqdm.tqdm(save_args): + save_wrapper(arg) diff --git a/configs/chodosh/argo/unsupervised_sensor_train_dumper.py b/configs/chodosh/argo/unsupervised_sensor_train_dumper.py new file mode 100644 index 0000000..068c3d4 --- /dev/null +++ b/configs/chodosh/argo/unsupervised_sensor_train_dumper.py @@ -0,0 +1,45 @@ +_base_ = "../../pseudoimage.py" + +is_trainable = False +has_labels = False + +test_sequence_dir = "/efs/argoverse2/train/" +nsfp_save_folder = "/efs/argoverse2/train_nsfp_flow/" +chodosh_save_folder = "/efs/argoverse2/train_chodosh_flow_problem_setup/" + +precision = 32 + + +def get_max_sequence_length(sequence_dir): + if "argoverse2" in sequence_dir: + return 146 + else: + return 296 + + +max_test_sequence_length = get_max_sequence_length(test_sequence_dir) + +epochs = 20 +learning_rate = 2e-6 +save_every = 500 +validate_every = 500 + +SEQUENCE_LENGTH = 2 + +model = dict(name="ChodoshProblemDumper", + args=dict(VOXEL_SIZE={{_base_.VOXEL_SIZE}}, + POINT_CLOUD_RANGE={{_base_.POINT_CLOUD_RANGE}}, + SEQUENCE_LENGTH=SEQUENCE_LENGTH, + nsfp_save_folder=nsfp_save_folder, + chodosh_save_folder=chodosh_save_folder)) + +test_loader = dict(name="ArgoverseRawSequenceLoader", + args=dict(sequence_dir=test_sequence_dir, verbose=True)) + +test_dataloader = dict( + args=dict(batch_size=1, num_workers=1, shuffle=False, pin_memory=True)) + +test_dataset = dict(name="SubsequenceRawDataset", + args=dict(subsequence_length=SEQUENCE_LENGTH, + max_sequence_length=max_test_sequence_length, + origin_mode="FIRST_ENTRY")) diff --git a/configs/chodosh/argo/unsupervised_sensor_val_cached.py b/configs/chodosh/argo/unsupervised_sensor_val_cached.py new file mode 100644 index 0000000..c37d588 --- /dev/null +++ b/configs/chodosh/argo/unsupervised_sensor_val_cached.py @@ -0,0 +1,17 @@ +_base_ = "./unsupervised_sensor_val_dumper.py" + +has_labels = True +test_sequence_dir = "/efs/argoverse2/val/" +test_flow_dir = "/efs/argoverse2/val_sceneflow/" + +flow_save_folder = "/efs/argoverse2/val_chodosh_flow/" + +test_loader = dict(name="ArgoverseSupervisedFlowSequenceLoader", + args=dict(_delete_=True, + raw_data_path=test_sequence_dir, + flow_data_path=test_flow_dir)) +test_dataset = dict(name="SubsequenceSupervisedFlowDataset") +model = dict(name="NSFPCached", args=dict(_delete_=True, VOXEL_SIZE={{_base_.VOXEL_SIZE}}, + POINT_CLOUD_RANGE={{_base_.POINT_CLOUD_RANGE}}, + SEQUENCE_LENGTH={{_base_.SEQUENCE_LENGTH}}, + flow_save_folder=flow_save_folder)) \ No newline at end of file diff --git a/configs/chodosh/argo/unsupervised_sensor_val_dumper.py b/configs/chodosh/argo/unsupervised_sensor_val_dumper.py new file mode 100644 index 0000000..775fe63 --- /dev/null +++ b/configs/chodosh/argo/unsupervised_sensor_val_dumper.py @@ -0,0 +1,45 @@ +_base_ = "../../pseudoimage.py" + +is_trainable = False +has_labels = False + +test_sequence_dir = "/efs/argoverse2/val/" +nsfp_save_folder = "/efs/argoverse2/val_nsfp_flow/" +chodosh_save_folder = "/efs/argoverse2/val_chodosh_flow_problem_setup/" + +precision = 32 + + +def get_max_sequence_length(sequence_dir): + if "argoverse2" in sequence_dir: + return 146 + else: + return 296 + + +max_test_sequence_length = get_max_sequence_length(test_sequence_dir) + +epochs = 20 +learning_rate = 2e-6 +save_every = 500 +validate_every = 500 + +SEQUENCE_LENGTH = 2 + +model = dict(name="ChodoshProblemDumper", + args=dict(VOXEL_SIZE={{_base_.VOXEL_SIZE}}, + POINT_CLOUD_RANGE={{_base_.POINT_CLOUD_RANGE}}, + SEQUENCE_LENGTH=SEQUENCE_LENGTH, + nsfp_save_folder=nsfp_save_folder, + chodosh_save_folder=chodosh_save_folder)) + +test_loader = dict(name="ArgoverseRawSequenceLoader", + args=dict(sequence_dir=test_sequence_dir, verbose=True)) + +test_dataloader = dict( + args=dict(batch_size=1, num_workers=1, shuffle=False, pin_memory=True)) + +test_dataset = dict(name="SubsequenceRawDataset", + args=dict(subsequence_length=SEQUENCE_LENGTH, + max_sequence_length=max_test_sequence_length, + origin_mode="FIRST_ENTRY")) diff --git a/configs/fastflow3d/argo/chodosh_distilatation.py b/configs/fastflow3d/argo/chodosh_distilatation.py new file mode 100644 index 0000000..2d32e88 --- /dev/null +++ b/configs/fastflow3d/argo/chodosh_distilatation.py @@ -0,0 +1,5 @@ +_base_ = "./nsfp_distilatation.py" + +train_flow_dir = "/efs/argoverse2/train_chodosh_flow/" + +loader = dict(args=dict(flow_data_path=train_flow_dir)) diff --git a/configs/fastflow3d/argo/chodosh_distilatation_dump_output.py b/configs/fastflow3d/argo/chodosh_distilatation_dump_output.py new file mode 100644 index 0000000..3bff175 --- /dev/null +++ b/configs/fastflow3d/argo/chodosh_distilatation_dump_output.py @@ -0,0 +1,6 @@ +_base_ = "./chodosh_distilatation.py" + +save_output_folder = "/efs/argoverse2/val_chodosh_distilation_out/" + +test_dataloader = dict( + args=dict(batch_size=16, num_workers=16, shuffle=False, pin_memory=True)) \ No newline at end of file diff --git a/configs/fastflow3d/argo/chodosh_distilatation_speed_scaled.py b/configs/fastflow3d/argo/chodosh_distilatation_speed_scaled.py new file mode 100644 index 0000000..4c0028e --- /dev/null +++ b/configs/fastflow3d/argo/chodosh_distilatation_speed_scaled.py @@ -0,0 +1,3 @@ +_base_ = "./chodosh_distilatation.py" + +loss_fn = dict(args=dict(fast_mover_scale=True)) \ No newline at end of file diff --git a/configs/fastflow3d/argo/nsfp_distilatation.py b/configs/fastflow3d/argo/nsfp_distilatation.py new file mode 100644 index 0000000..66898ea --- /dev/null +++ b/configs/fastflow3d/argo/nsfp_distilatation.py @@ -0,0 +1,58 @@ +_base_ = "../../pseudoimage.py" + +train_sequence_dir = "/efs/argoverse2/train/" +train_flow_dir = "/efs/argoverse2/train_nsfp_flow/" + +test_sequence_dir = "/efs/argoverse2/val/" +test_flow_dir = "/efs/argoverse2/val_sceneflow/" + + +def get_max_sequence_length(sequence_dir): + if "argoverse2" in sequence_dir: + return 145 + else: + return 296 + + +max_train_sequence_length = get_max_sequence_length(train_sequence_dir) +max_test_sequence_length = get_max_sequence_length(test_sequence_dir) + +epochs = 50 +learning_rate = 2e-6 +save_every = 500 +validate_every = 500 + +SEQUENCE_LENGTH = 2 + +model = dict(name="FastFlow3D", + args=dict(VOXEL_SIZE={{_base_.VOXEL_SIZE}}, + PSEUDO_IMAGE_DIMS={{_base_.PSEUDO_IMAGE_DIMS}}, + POINT_CLOUD_RANGE={{_base_.POINT_CLOUD_RANGE}}, + FEATURE_CHANNELS=32, + SEQUENCE_LENGTH=SEQUENCE_LENGTH)) + +loader = dict(name="ArgoverseUnsupervisedFlowSequenceLoader", + args=dict(raw_data_path=train_sequence_dir, + flow_data_path=train_flow_dir)) + +dataloader = dict( + args=dict(batch_size=16, num_workers=16, shuffle=True, pin_memory=False)) + +dataset = dict(name="SubsequenceUnsupervisedFlowDataset", + args=dict(subsequence_length=SEQUENCE_LENGTH, + max_sequence_length=max_train_sequence_length, + origin_mode="FIRST_ENTRY")) + +loss_fn = dict(name="FastFlow3DDistillationLoss", args=dict()) + +test_loader = dict(name="ArgoverseSupervisedFlowSequenceLoader", + args=dict(raw_data_path=test_sequence_dir, + flow_data_path=test_flow_dir)) + +test_dataloader = dict( + args=dict(batch_size=8, num_workers=8, shuffle=False, pin_memory=True)) + +test_dataset = dict(name="SubsequenceSupervisedFlowDataset", + args=dict(subsequence_length=SEQUENCE_LENGTH, + max_sequence_length=max_test_sequence_length, + origin_mode="FIRST_ENTRY")) diff --git a/configs/fastflow3d/argo/nsfp_distilatation_0.01x_sequential.py b/configs/fastflow3d/argo/nsfp_distilatation_0.01x_sequential.py new file mode 100644 index 0000000..8caeef4 --- /dev/null +++ b/configs/fastflow3d/argo/nsfp_distilatation_0.01x_sequential.py @@ -0,0 +1,6 @@ +_base_ = "./nsfp_distilatation.py" + +epochs = 500 +check_val_every_n_epoch = 5 +validate_every = None +dataset = dict(args=dict(subset_fraction=0.01, subset_mode='sequential')) \ No newline at end of file diff --git a/configs/fastflow3d/argo/nsfp_distilatation_0.1x_random.py b/configs/fastflow3d/argo/nsfp_distilatation_0.1x_random.py new file mode 100644 index 0000000..320953b --- /dev/null +++ b/configs/fastflow3d/argo/nsfp_distilatation_0.1x_random.py @@ -0,0 +1,5 @@ +_base_ = "./nsfp_distilatation.py" + +epochs = 500 +validate_every = 158 # Number of training batches +dataset = dict(args=dict(subset_fraction=0.1)) \ No newline at end of file diff --git a/configs/fastflow3d/argo/nsfp_distilatation_0.1x_sequential.py b/configs/fastflow3d/argo/nsfp_distilatation_0.1x_sequential.py new file mode 100644 index 0000000..e613c39 --- /dev/null +++ b/configs/fastflow3d/argo/nsfp_distilatation_0.1x_sequential.py @@ -0,0 +1,6 @@ +_base_ = "./nsfp_distilatation.py" + +epochs = 200 +check_val_every_n_epoch = 5 +validate_every = None +dataset = dict(args=dict(subset_fraction=0.1, subset_mode='sequential')) diff --git a/configs/fastflow3d/argo/nsfp_distilatation_0.2x_sequential.py b/configs/fastflow3d/argo/nsfp_distilatation_0.2x_sequential.py new file mode 100644 index 0000000..7bce333 --- /dev/null +++ b/configs/fastflow3d/argo/nsfp_distilatation_0.2x_sequential.py @@ -0,0 +1,6 @@ +_base_ = "./nsfp_distilatation.py" + +epochs = 100 +check_val_every_n_epoch = 1 +validate_every = None +dataset = dict(args=dict(subset_fraction=0.2, subset_mode='sequential')) \ No newline at end of file diff --git a/configs/fastflow3d/argo/nsfp_distilatation_0.5x_sequential.py b/configs/fastflow3d/argo/nsfp_distilatation_0.5x_sequential.py new file mode 100644 index 0000000..a81399f --- /dev/null +++ b/configs/fastflow3d/argo/nsfp_distilatation_0.5x_sequential.py @@ -0,0 +1,4 @@ +_base_ = "./nsfp_distilatation.py" + +epochs = 100 +dataset = dict(args=dict(subset_fraction=0.5, subset_mode='sequential')) \ No newline at end of file diff --git a/configs/fastflow3d/argo/nsfp_distilatation_dump_output.py b/configs/fastflow3d/argo/nsfp_distilatation_dump_output.py new file mode 100644 index 0000000..e324ec6 --- /dev/null +++ b/configs/fastflow3d/argo/nsfp_distilatation_dump_output.py @@ -0,0 +1,22 @@ +_base_ = "./nsfp_distilatation.py" + +has_labels = False + +test_sequence_dir = "/efs/argoverse2/val/" +test_flow_dir = "/efs/argoverse2/val_sceneflow/" + +save_output_folder = "/efs/argoverse2/val_distilation_out/" +SEQUENCE_LENGTH = 2 + +test_loader = dict(_delete_=True, + name="ArgoverseRawSequenceLoader", + args=dict(sequence_dir=test_sequence_dir, verbose=True)) + +test_dataset = dict(name="VarLenSubsequenceRawDataset", + args=dict(_delete_=True, + subsequence_length=SEQUENCE_LENGTH, + origin_mode="FIRST_ENTRY", + max_pc_points=120000)) + +test_dataloader = dict( + args=dict(batch_size=16, num_workers=16, shuffle=False, pin_memory=True)) \ No newline at end of file diff --git a/configs/fastflow3d/argo/nsfp_distilatation_dump_output_test.py b/configs/fastflow3d/argo/nsfp_distilatation_dump_output_test.py new file mode 100644 index 0000000..00eeeb8 --- /dev/null +++ b/configs/fastflow3d/argo/nsfp_distilatation_dump_output_test.py @@ -0,0 +1,22 @@ +_base_ = "./nsfp_distilatation.py" + +has_labels = False + +test_sequence_dir = "/efs/argoverse2/test/" +test_flow_dir = "/efs/argoverse2/test_sceneflow/" + +save_output_folder = "/efs/argoverse2/test_distilation_out/" +SEQUENCE_LENGTH = 2 + +test_loader = dict(_delete_=True, + name="ArgoverseRawSequenceLoader", + args=dict(sequence_dir=test_sequence_dir, verbose=True)) + +test_dataset = dict(name="VarLenSubsequenceRawDataset", + args=dict(_delete_=True, + subsequence_length=SEQUENCE_LENGTH, + origin_mode="FIRST_ENTRY", + max_pc_points=120000)) + +test_dataloader = dict( + args=dict(batch_size=16, num_workers=16, shuffle=False, pin_memory=True)) \ No newline at end of file diff --git a/configs/fastflow3d/argo/nsfp_distilatation_run3.py b/configs/fastflow3d/argo/nsfp_distilatation_run3.py new file mode 100644 index 0000000..4e88841 --- /dev/null +++ b/configs/fastflow3d/argo/nsfp_distilatation_run3.py @@ -0,0 +1,4 @@ +_base_ = "./nsfp_distilatation.py" + +test_dataloader = dict( + args=dict(batch_size=16, num_workers=16, shuffle=False, pin_memory=True)) \ No newline at end of file diff --git a/configs/fastflow3d/argo/nsfp_distilatation_speed_scaled_updated.py b/configs/fastflow3d/argo/nsfp_distilatation_speed_scaled_updated.py new file mode 100644 index 0000000..05f3956 --- /dev/null +++ b/configs/fastflow3d/argo/nsfp_distilatation_speed_scaled_updated.py @@ -0,0 +1,4 @@ +_base_ = "./nsfp_distilatation.py" + +loss_fn = dict(args=dict(fast_mover_scale=True)) + diff --git a/configs/fastflow3d/argo/nsfp_distilatation_speed_scaled_updated_batch8_train.py b/configs/fastflow3d/argo/nsfp_distilatation_speed_scaled_updated_batch8_train.py new file mode 100644 index 0000000..acee72a --- /dev/null +++ b/configs/fastflow3d/argo/nsfp_distilatation_speed_scaled_updated_batch8_train.py @@ -0,0 +1,5 @@ +_base_ = "./nsfp_distilatation_speed_scaled_updated.py" + +dataloader = dict(args=dict(batch_size=8, num_workers=8)) + +test_dataloader = dict(args=dict(batch_size=4, num_workers=4)) diff --git a/configs/fastflow3d/argo/nsfp_distilatation_speed_scaled_updated_run2.py b/configs/fastflow3d/argo/nsfp_distilatation_speed_scaled_updated_run2.py new file mode 100644 index 0000000..ec3ef03 --- /dev/null +++ b/configs/fastflow3d/argo/nsfp_distilatation_speed_scaled_updated_run2.py @@ -0,0 +1,2 @@ +_base_ = "./nsfp_distilatation_speed_scaled_updated.py" +test_dataloader = dict(args=dict(batch_size=16, num_workers=16)) diff --git a/configs/fastflow3d/argo/nsfp_distilatation_speed_scaled_updated_run3.py b/configs/fastflow3d/argo/nsfp_distilatation_speed_scaled_updated_run3.py new file mode 100644 index 0000000..ec3ef03 --- /dev/null +++ b/configs/fastflow3d/argo/nsfp_distilatation_speed_scaled_updated_run3.py @@ -0,0 +1,2 @@ +_base_ = "./nsfp_distilatation_speed_scaled_updated.py" +test_dataloader = dict(args=dict(batch_size=16, num_workers=16)) diff --git a/configs/fastflow3d/argo/old/nsfp_distilatation_0.1x.py b/configs/fastflow3d/argo/old/nsfp_distilatation_0.1x.py new file mode 100644 index 0000000..6b7121b --- /dev/null +++ b/configs/fastflow3d/argo/old/nsfp_distilatation_0.1x.py @@ -0,0 +1,6 @@ +_base_ = "./nsfp_distilatation.py" + +epochs = 500 +validate_every = 158 # Number of training batches + +dataset = dict(args=dict(subset_fraction=0.1)) diff --git a/configs/fastflow3d/argo/old/nsfp_distilatation_2x.py b/configs/fastflow3d/argo/old/nsfp_distilatation_2x.py new file mode 100644 index 0000000..aaa2623 --- /dev/null +++ b/configs/fastflow3d/argo/old/nsfp_distilatation_2x.py @@ -0,0 +1,53 @@ +_base_ = "./nsfp_distilatation.py" + +sensor_train_sequence_dir = "/efs/argoverse2/train/" +sensor_train_flow_dir = "/efs/argoverse2/train_nsfp_flow/" + +lidar_train_sequence_dir = "/efs/argoverse_lidar/train/" +lidar_train_flow_dir = "/efs/argoverse_lidar_subset/train_nsfp_flow/" + + +def get_max_sequence_length(sequence_dir): + if "argoverse2" in sequence_dir: + return 145 + else: + return 294 + + +max_sensor_train_sequence_length = get_max_sequence_length( + sensor_train_sequence_dir) +max_lidar_train_sequence_length = get_max_sequence_length( + lidar_train_sequence_dir) + +epochs = 50 +SUBSEQUENCE_LENGTH = 2 + +loader = [ + dict(name="ArgoverseUnsupervisedFlowSequenceLoader", + args=dict(raw_data_path=sensor_train_sequence_dir, + flow_data_path=sensor_train_flow_dir)), + dict(name="ArgoverseUnsupervisedFlowSequenceLoader", + args=dict(raw_data_path=lidar_train_sequence_dir, + flow_data_path=lidar_train_flow_dir)) +] + +# The sensor dataset is 3600 batches from 750 * 145 sequences +# 3600 / (750 * 145) = 0.0331034483 +# 3600 = 0.0331034483 * (750 * 145) + +# We want to match this with the lidar dataset. +# 3600 = X * (2500 * 294) * 0.0331034483 +# 3600 / (2500 * 294 * 0.0331034483) = 0.1479 + + +dataset = [ + dict(name="SubsequenceUnsupervisedFlowDataset", + args=dict(subsequence_length=SUBSEQUENCE_LENGTH, + max_sequence_length=max_sensor_train_sequence_length, + origin_mode="FIRST_ENTRY")), + dict(name="SubsequenceUnsupervisedFlowDataset", + args=dict(subsequence_length=SUBSEQUENCE_LENGTH, + max_sequence_length=max_lidar_train_sequence_length, + origin_mode="FIRST_ENTRY", + subset_fraction=0.1479)) +] diff --git a/configs/fastflow3d/argo/old/nsfp_distilatation_half.py b/configs/fastflow3d/argo/old/nsfp_distilatation_half.py new file mode 100644 index 0000000..6ffd6f0 --- /dev/null +++ b/configs/fastflow3d/argo/old/nsfp_distilatation_half.py @@ -0,0 +1,4 @@ +_base_ = "./nsfp_distilatation.py" + +epochs = 100 +dataset = dict(args=dict(subset_fraction=0.5)) diff --git a/configs/fastflow3d/argo/old/nsfp_distilatation_tiny.py b/configs/fastflow3d/argo/old/nsfp_distilatation_tiny.py new file mode 100644 index 0000000..9eb8157 --- /dev/null +++ b/configs/fastflow3d/argo/old/nsfp_distilatation_tiny.py @@ -0,0 +1,7 @@ +_base_ = "./nsfp_distilatation.py" + +epochs = 100 +validate_every = 2 +save_every = 2 +dataset = dict(args=dict(subset_fraction=0.003)) +test_dataset = dict(args=dict(subset_fraction=0.003)) diff --git a/configs/fastflow3d/argo/old/unsupervised.py b/configs/fastflow3d/argo/old/unsupervised.py new file mode 100644 index 0000000..6e8eea4 --- /dev/null +++ b/configs/fastflow3d/argo/old/unsupervised.py @@ -0,0 +1,58 @@ +train_sequence_dir = "/efs/argoverse_lidar/train/" + +test_sequence_dir = "/efs/argoverse2/val/" +test_flow_dir = "/efs/argoverse2/val_sceneflow/" + + +def get_max_sequence_length(sequence_dir): + if "argoverse2" in sequence_dir: + return 146 + else: + return 296 + + +max_train_sequence_length = get_max_sequence_length(train_sequence_dir) +max_test_sequence_length = get_max_sequence_length(test_sequence_dir) + +gradient_clip_val = 0 +epochs = 10 +accumulate_grad_batches = 1 +learning_rate = 2e-6 +save_every = 500 +validate_every = 500 + +SEQUENCE_LENGTH = 2 + +model = dict(name="FastFlow3D", + args=dict(VOXEL_SIZE=(0.2, 0.2, 4), + PSEUDO_IMAGE_DIMS=(512, 512), + POINT_CLOUD_RANGE=(-51.2, -51.2, -3, 51.2, 51.2, 1), + FEATURE_CHANNELS=32, + SEQUENCE_LENGTH=SEQUENCE_LENGTH)) +train_forward_args = dict(compute_cycle=False, compute_symmetry_x=False, compute_symmetry_y=False) + + +loader = dict(name="ArgoverseRawSequenceLoader", + args=dict(sequence_dir=train_sequence_dir)) + +dataloader = dict(args=dict(batch_size=16, num_workers=16, shuffle=True, pin_memory=False)) + + + +dataset = dict(name="SubsequenceRawDataset", + args=dict(subsequence_length=SEQUENCE_LENGTH, + max_sequence_length=max_train_sequence_length, + origin_mode="FIRST_ENTRY")) + +loss_fn = dict(name="FastFlow3DSelfSupervisedLoss", args=dict(warp_upscale=1)) + +test_loader = dict(name="ArgoverseSupervisedFlowSequenceLoader", + args=dict(raw_data_path=test_sequence_dir, + flow_data_path=test_flow_dir)) + +test_dataloader = dict(args=dict(batch_size=8, num_workers=8, shuffle=False, pin_memory=True)) + +test_dataset = dict(name="SubsequenceSupervisedFlowDataset", + args=dict(subsequence_length=SEQUENCE_LENGTH, + max_sequence_length=max_test_sequence_length, + origin_mode="FIRST_ENTRY")) diff --git a/configs/fastflow3d/argo/old/unsupervised_bottleneck.py b/configs/fastflow3d/argo/old/unsupervised_bottleneck.py new file mode 100644 index 0000000..f1c206e --- /dev/null +++ b/configs/fastflow3d/argo/old/unsupervised_bottleneck.py @@ -0,0 +1,3 @@ +_base_ = "./unsupervised.py" + +model = dict(name="FastFlow3D", args=dict(bottleneck_head=True)) diff --git a/configs/fastflow3d/argo/old/unsupervised_symmetry.py b/configs/fastflow3d/argo/old/unsupervised_symmetry.py new file mode 100644 index 0000000..a88488f --- /dev/null +++ b/configs/fastflow3d/argo/old/unsupervised_symmetry.py @@ -0,0 +1,4 @@ +_base_ = './unsupervised.py' + +dataloader = dict(args=dict(batch_size=8, num_workers=8)) +train_forward_args = dict(compute_symmetry_x=True, compute_symmetry_y=True) diff --git a/configs/fastflow3d/argo/supervised.py b/configs/fastflow3d/argo/supervised.py new file mode 100644 index 0000000..bee5f2d --- /dev/null +++ b/configs/fastflow3d/argo/supervised.py @@ -0,0 +1,58 @@ +_base_ = "../../pseudoimage.py" + +train_sequence_dir = "/efs/argoverse2/train/" +train_flow_dir = "/efs/argoverse2/train_sceneflow/" + +test_sequence_dir = "/efs/argoverse2/val/" +test_flow_dir = "/efs/argoverse2/val_sceneflow/" + + +def get_max_sequence_length(sequence_dir): + if "argoverse2" in sequence_dir: + return 145 + else: + return 296 + + +max_train_sequence_length = get_max_sequence_length(train_sequence_dir) +max_test_sequence_length = get_max_sequence_length(test_sequence_dir) + +epochs = 50 +learning_rate = 2e-6 +save_every = 500 +validate_every = 500 + +SEQUENCE_LENGTH = 2 + +model = dict(name="FastFlow3D", + args=dict(VOXEL_SIZE={{_base_.VOXEL_SIZE}}, + PSEUDO_IMAGE_DIMS={{_base_.PSEUDO_IMAGE_DIMS}}, + POINT_CLOUD_RANGE={{_base_.POINT_CLOUD_RANGE}}, + FEATURE_CHANNELS=32, + SEQUENCE_LENGTH=SEQUENCE_LENGTH)) + +loader = dict(name="ArgoverseSupervisedFlowSequenceLoader", + args=dict(raw_data_path=train_sequence_dir, + flow_data_path=train_flow_dir)) + +dataloader = dict( + args=dict(batch_size=16, num_workers=16, shuffle=True, pin_memory=False)) + +dataset = dict(name="SubsequenceSupervisedFlowDataset", + args=dict(subsequence_length=SEQUENCE_LENGTH, + max_sequence_length=max_train_sequence_length, + origin_mode="FIRST_ENTRY")) + +loss_fn = dict(name="FastFlow3DSupervisedLoss", args=dict()) + +test_loader = dict(name="ArgoverseSupervisedFlowSequenceLoader", + args=dict(raw_data_path=test_sequence_dir, + flow_data_path=test_flow_dir)) + +test_dataloader = dict( + args=dict(batch_size=8, num_workers=8, shuffle=False, pin_memory=True)) + +test_dataset = dict(name="SubsequenceSupervisedFlowDataset", + args=dict(subsequence_length=SEQUENCE_LENGTH, + max_sequence_length=max_test_sequence_length, + origin_mode="FIRST_ENTRY")) diff --git a/configs/fastflow3d/argo/supervised_0.01x_finetune_from_base.py b/configs/fastflow3d/argo/supervised_0.01x_finetune_from_base.py new file mode 100644 index 0000000..54d3934 --- /dev/null +++ b/configs/fastflow3d/argo/supervised_0.01x_finetune_from_base.py @@ -0,0 +1,7 @@ +_base_ = "./supervised.py" + +epochs = 200 +check_val_every_n_epoch = 1 +learning_rate = 2e-8 +validate_every = None +dataset = dict(args=dict(subset_fraction=0.01, subset_mode='sequential')) diff --git a/configs/fastflow3d/argo/supervised_0.01x_sequential.py b/configs/fastflow3d/argo/supervised_0.01x_sequential.py new file mode 100644 index 0000000..6d97d0d --- /dev/null +++ b/configs/fastflow3d/argo/supervised_0.01x_sequential.py @@ -0,0 +1,6 @@ +_base_ = "./supervised.py" + +epochs = 500 +check_val_every_n_epoch = 500 +validate_every = None +dataset = dict(args=dict(subset_fraction=0.01, subset_mode='sequential')) diff --git a/configs/fastflow3d/argo/supervised_0.1x_finetune_from_base.py b/configs/fastflow3d/argo/supervised_0.1x_finetune_from_base.py new file mode 100644 index 0000000..92cf6d5 --- /dev/null +++ b/configs/fastflow3d/argo/supervised_0.1x_finetune_from_base.py @@ -0,0 +1,7 @@ +_base_ = "./supervised.py" + +epochs = 100 +check_val_every_n_epoch = 1 +learning_rate = 2e-8 +validate_every = None +dataset = dict(args=dict(subset_fraction=0.1, subset_mode='sequential')) \ No newline at end of file diff --git a/configs/fastflow3d/argo/supervised_0.1x_random.py b/configs/fastflow3d/argo/supervised_0.1x_random.py new file mode 100644 index 0000000..041890e --- /dev/null +++ b/configs/fastflow3d/argo/supervised_0.1x_random.py @@ -0,0 +1,5 @@ +_base_ = "./supervised.py" + +epochs = 500 +validate_every = 158 # Number of training batches +dataset = dict(args=dict(subset_fraction=0.1, subset_mode='sequential')) \ No newline at end of file diff --git a/configs/fastflow3d/argo/supervised_0.1x_sequential.py b/configs/fastflow3d/argo/supervised_0.1x_sequential.py new file mode 100644 index 0000000..9010202 --- /dev/null +++ b/configs/fastflow3d/argo/supervised_0.1x_sequential.py @@ -0,0 +1,5 @@ +_base_ = "./supervised.py" + +epochs = 500 +validate_every = 158 # Number of training batches +dataset = dict(args=dict(subset_fraction=0.1, subset_mode='sequential')) \ No newline at end of file diff --git a/configs/fastflow3d/argo/supervised_0.2x_finetune_from_base.py b/configs/fastflow3d/argo/supervised_0.2x_finetune_from_base.py new file mode 100644 index 0000000..221af05 --- /dev/null +++ b/configs/fastflow3d/argo/supervised_0.2x_finetune_from_base.py @@ -0,0 +1,7 @@ +_base_ = "./supervised.py" + +epochs = 100 +check_val_every_n_epoch = 1 +learning_rate = 2e-8 +validate_every = None +dataset = dict(args=dict(subset_fraction=0.2, subset_mode='sequential')) \ No newline at end of file diff --git a/configs/fastflow3d/argo/supervised_0.2x_finetune_from_base_random.py b/configs/fastflow3d/argo/supervised_0.2x_finetune_from_base_random.py new file mode 100644 index 0000000..bf5df9b --- /dev/null +++ b/configs/fastflow3d/argo/supervised_0.2x_finetune_from_base_random.py @@ -0,0 +1,2 @@ +_base_ = "./supervised_0.2x_finetune_from_base.py" +dataset = dict(args=dict(subset_mode='random')) \ No newline at end of file diff --git a/configs/fastflow3d/argo/supervised_0.2x_sequential.py b/configs/fastflow3d/argo/supervised_0.2x_sequential.py new file mode 100644 index 0000000..e9692c3 --- /dev/null +++ b/configs/fastflow3d/argo/supervised_0.2x_sequential.py @@ -0,0 +1,6 @@ +_base_ = "./supervised.py" + +epochs = 100 +check_val_every_n_epoch = 1 +validate_every = None +dataset = dict(args=dict(subset_fraction=0.2, subset_mode='sequential')) \ No newline at end of file diff --git a/configs/fastflow3d/argo/supervised_0.5x_finetune_from_base.py b/configs/fastflow3d/argo/supervised_0.5x_finetune_from_base.py new file mode 100644 index 0000000..8c5dc02 --- /dev/null +++ b/configs/fastflow3d/argo/supervised_0.5x_finetune_from_base.py @@ -0,0 +1,7 @@ +_base_ = "./supervised.py" + +epochs = 300 +check_val_every_n_epoch = 1 +learning_rate = 2e-8 +validate_every = None +dataset = dict(args=dict(subset_fraction=0.5, subset_mode='sequential')) \ No newline at end of file diff --git a/configs/fastflow3d/argo/supervised_0.5x_finetune_from_base_random.py b/configs/fastflow3d/argo/supervised_0.5x_finetune_from_base_random.py new file mode 100644 index 0000000..b3de88d --- /dev/null +++ b/configs/fastflow3d/argo/supervised_0.5x_finetune_from_base_random.py @@ -0,0 +1,2 @@ +_base_ = "./supervised_0.5x_finetune_from_base.py" +dataset = dict(args=dict(subset_mode='random')) \ No newline at end of file diff --git a/configs/fastflow3d/argo/supervised_0.5x_random.py b/configs/fastflow3d/argo/supervised_0.5x_random.py new file mode 100644 index 0000000..d0edb81 --- /dev/null +++ b/configs/fastflow3d/argo/supervised_0.5x_random.py @@ -0,0 +1,6 @@ +_base_ = "./supervised.py" + +epochs = 100 +check_val_every_n_epoch = 1 +validate_every = None +dataset = dict(args=dict(subset_fraction=0.5, subset_mode='random')) \ No newline at end of file diff --git a/configs/fastflow3d/argo/supervised_0.5x_sequential.py b/configs/fastflow3d/argo/supervised_0.5x_sequential.py new file mode 100644 index 0000000..823722d --- /dev/null +++ b/configs/fastflow3d/argo/supervised_0.5x_sequential.py @@ -0,0 +1,6 @@ +_base_ = "./supervised.py" + +epochs = 100 +check_val_every_n_epoch = 1 +validate_every = None +dataset = dict(args=dict(subset_fraction=0.5, subset_mode='sequential')) \ No newline at end of file diff --git a/configs/fastflow3d/argo/supervised_batch8_train.py b/configs/fastflow3d/argo/supervised_batch8_train.py new file mode 100644 index 0000000..0d2eb6a --- /dev/null +++ b/configs/fastflow3d/argo/supervised_batch8_train.py @@ -0,0 +1,5 @@ +_base_ = "./supervised.py" + +dataloader = dict(args=dict(batch_size=8, num_workers=8)) + +test_dataloader = dict(args=dict(batch_size=4, num_workers=4)) diff --git a/configs/fastflow3d/argo/supervised_dump_output.py b/configs/fastflow3d/argo/supervised_dump_output.py new file mode 100644 index 0000000..de2b0b1 --- /dev/null +++ b/configs/fastflow3d/argo/supervised_dump_output.py @@ -0,0 +1,6 @@ +_base_ = "./supervised.py" + +save_output_folder = "/efs/argoverse2/val_supervised_out/" + +test_dataloader = dict( + args=dict(batch_size=16, num_workers=16, shuffle=False, pin_memory=True)) \ No newline at end of file diff --git a/configs/fastflow3d/argo/supervised_no_scale.py b/configs/fastflow3d/argo/supervised_no_scale.py new file mode 100644 index 0000000..55d9207 --- /dev/null +++ b/configs/fastflow3d/argo/supervised_no_scale.py @@ -0,0 +1,3 @@ +_base_ = "./supervised.py" +epochs = 100 +loss_fn = dict(name="FastFlow3DSupervisedLoss", args=dict(scale_background=False)) diff --git a/configs/fastflow3d/argo/supervised_perf_test.py b/configs/fastflow3d/argo/supervised_perf_test.py new file mode 100644 index 0000000..62ca1d0 --- /dev/null +++ b/configs/fastflow3d/argo/supervised_perf_test.py @@ -0,0 +1,10 @@ +_base_ = "./supervised.py" + +test_dataset = dict( + name="VarLenSubsequenceSupervisedFlowSpecificSubsetDataset", + args=dict(_delete_=True, + subsequence_length={{_base_.SEQUENCE_LENGTH}}, + origin_mode="FIRST_ENTRY", + subset_file="data_prep_scripts/argo/eval_sequence_id_index.txt")) + +test_dataloader = dict(args=dict(batch_size=1)) diff --git a/configs/fastflow3d/argo/supervised_run2.py b/configs/fastflow3d/argo/supervised_run2.py new file mode 100644 index 0000000..22fb8ed --- /dev/null +++ b/configs/fastflow3d/argo/supervised_run2.py @@ -0,0 +1,4 @@ +_base_ = "./supervised.py" + +test_dataloader = dict( + args=dict(batch_size=16, num_workers=16, shuffle=False, pin_memory=True)) \ No newline at end of file diff --git a/configs/fastflow3d/argo/supervised_run3.py b/configs/fastflow3d/argo/supervised_run3.py new file mode 100644 index 0000000..22fb8ed --- /dev/null +++ b/configs/fastflow3d/argo/supervised_run3.py @@ -0,0 +1,4 @@ +_base_ = "./supervised.py" + +test_dataloader = dict( + args=dict(batch_size=16, num_workers=16, shuffle=False, pin_memory=True)) \ No newline at end of file diff --git a/configs/fastflow3d/argo/supervised_val_debug.py b/configs/fastflow3d/argo/supervised_val_debug.py new file mode 100644 index 0000000..3bbb5eb --- /dev/null +++ b/configs/fastflow3d/argo/supervised_val_debug.py @@ -0,0 +1,4 @@ +_base_ = "./supervised.py" + +validate_every = 1 +gradient_clip_val = 0.000001 diff --git a/configs/fastflow3d/waymo/nsfp_distilatation.py b/configs/fastflow3d/waymo/nsfp_distilatation.py new file mode 100644 index 0000000..8bb813b --- /dev/null +++ b/configs/fastflow3d/waymo/nsfp_distilatation.py @@ -0,0 +1,45 @@ +_base_ = "../../pseudoimage.py" + +train_sequence_dir = "/efs/waymo_open_processed_flow/training/" +train_flow_dir = "/efs/waymo_open_processed_flow/train_nsfp_flow/" + +test_sequence_dir = "/efs/waymo_open_processed_flow/validation/" + +epochs = 50 +learning_rate = 2e-6 +save_every = 500 +validate_every = 500 + +SEQUENCE_LENGTH = 2 + +model = dict(name="FastFlow3D", + args=dict(VOXEL_SIZE={{_base_.VOXEL_SIZE}}, + PSEUDO_IMAGE_DIMS={{_base_.PSEUDO_IMAGE_DIMS}}, + POINT_CLOUD_RANGE={{_base_.POINT_CLOUD_RANGE}}, + FEATURE_CHANNELS=32, + SEQUENCE_LENGTH=SEQUENCE_LENGTH)) + +loader = dict(name="WaymoUnsupervisedFlowSequenceLoader", + args=dict(raw_data_path=train_sequence_dir, + flow_data_path=train_flow_dir)) + +dataloader = dict( + args=dict(batch_size=16, num_workers=16, shuffle=True, pin_memory=False)) + +dataset = dict(name="VarLenSubsequenceUnsupervisedFlowDataset", + args=dict(subsequence_length=SEQUENCE_LENGTH, + max_pc_points=150000, + origin_mode="FIRST_ENTRY")) + +loss_fn = dict(name="FastFlow3DDistillationLoss", args=dict()) + +test_loader = dict(name="WaymoSupervisedFlowSequenceLoader", + args=dict(sequence_dir=test_sequence_dir)) + +test_dataloader = dict( + args=dict(batch_size=8, num_workers=8, shuffle=False, pin_memory=True)) + +test_dataset = dict(name="VarLenSubsequenceSupervisedFlowDataset", + args=dict(subsequence_length=SEQUENCE_LENGTH, + max_pc_points=150000, + origin_mode="FIRST_ENTRY")) diff --git a/configs/fastflow3d/waymo/nsfp_distilatation_dump_output.py b/configs/fastflow3d/waymo/nsfp_distilatation_dump_output.py new file mode 100644 index 0000000..9cdcc06 --- /dev/null +++ b/configs/fastflow3d/waymo/nsfp_distilatation_dump_output.py @@ -0,0 +1,6 @@ +_base_ = "./nsfp_distilatation.py" + +save_output_folder = "/bigdata/waymo_open_processed_flow/val_distillation_out/" + +test_dataloader = dict( + args=dict(batch_size=16, num_workers=16, shuffle=False, pin_memory=True)) \ No newline at end of file diff --git a/configs/fastflow3d/waymo/nsfp_distilatation_scaled.py b/configs/fastflow3d/waymo/nsfp_distilatation_scaled.py new file mode 100644 index 0000000..2534769 --- /dev/null +++ b/configs/fastflow3d/waymo/nsfp_distilatation_scaled.py @@ -0,0 +1,3 @@ +_base_ = "./nsfp_distilatation.py" + +loss_fn = dict(args=dict(fast_mover_scale=True)) diff --git a/configs/fastflow3d/waymo/supervised.py b/configs/fastflow3d/waymo/supervised.py new file mode 100644 index 0000000..c64ba49 --- /dev/null +++ b/configs/fastflow3d/waymo/supervised.py @@ -0,0 +1,43 @@ +_base_ = "../../pseudoimage.py" + +train_sequence_dir = "/efs/waymo_open_processed_flow/training/" + +test_sequence_dir = "/efs/waymo_open_processed_flow/validation/" + +max_train_sequence_length = 192 +max_test_sequence_length = 192 + +epochs = 50 +learning_rate = 2e-6 +save_every = 500 +validate_every = 500 + +SEQUENCE_LENGTH = 2 + +model = dict(name="FastFlow3D", + args=dict(VOXEL_SIZE={{_base_.VOXEL_SIZE}}, + PSEUDO_IMAGE_DIMS={{_base_.PSEUDO_IMAGE_DIMS}}, + POINT_CLOUD_RANGE={{_base_.POINT_CLOUD_RANGE}}, + FEATURE_CHANNELS=32, + SEQUENCE_LENGTH=SEQUENCE_LENGTH)) + +loader = dict(name="WaymoSupervisedFlowSequenceLoader", + args=dict(sequence_dir=train_sequence_dir)) + +dataloader = dict( + args=dict(batch_size=16, num_workers=16, shuffle=True, pin_memory=False)) + +dataset = dict(name="VarLenSubsequenceSupervisedFlowDataset", + args=dict(subsequence_length=SEQUENCE_LENGTH, + max_pc_points=150000, + origin_mode="FIRST_ENTRY")) + +loss_fn = dict(name="FastFlow3DSupervisedLoss", args=dict()) + +test_loader = dict(name="WaymoSupervisedFlowSequenceLoader", + args=dict(sequence_dir=test_sequence_dir)) + +test_dataset = dict(name="VarLenSubsequenceSupervisedFlowDataset", + args=dict(subsequence_length=SEQUENCE_LENGTH, + max_pc_points=150000, + origin_mode="FIRST_ENTRY")) diff --git a/configs/fastflow3d/waymo/supervised_dump_output.py b/configs/fastflow3d/waymo/supervised_dump_output.py new file mode 100644 index 0000000..4d4f1c5 --- /dev/null +++ b/configs/fastflow3d/waymo/supervised_dump_output.py @@ -0,0 +1,6 @@ +_base_ = "./supervised.py" + +save_output_folder = "/bigdata/waymo_open_processed_flow/val_supervised_out/" + +test_dataloader = dict( + args=dict(batch_size=16, num_workers=16, shuffle=False, pin_memory=True)) \ No newline at end of file diff --git a/configs/fastflow3d/waymo/supervised_perf_test.py b/configs/fastflow3d/waymo/supervised_perf_test.py new file mode 100644 index 0000000..9f82300 --- /dev/null +++ b/configs/fastflow3d/waymo/supervised_perf_test.py @@ -0,0 +1,3 @@ +_base_ = "./supervised.py" + +test_dataloader = dict(args=dict(batch_size=1)) diff --git a/configs/nearest_neighbor/argo/sensor_val.py b/configs/nearest_neighbor/argo/sensor_val.py new file mode 100644 index 0000000..52f5072 --- /dev/null +++ b/configs/nearest_neighbor/argo/sensor_val.py @@ -0,0 +1,43 @@ +_base_ = "../../pseudoimage.py" + +is_trainable = False +has_labels = False + +test_sequence_dir = "/efs/argoverse2/val/" +flow_save_folder = "/efs/argoverse2/val_nearest_neighbor_flow/" + +precision = 32 + + +def get_max_sequence_length(sequence_dir): + if "argoverse2" in sequence_dir: + return 146 + else: + return 296 + + +max_test_sequence_length = get_max_sequence_length(test_sequence_dir) + +epochs = 20 +learning_rate = 2e-6 +save_every = 500 +validate_every = 500 + +SEQUENCE_LENGTH = 2 + +model = dict(name="NearestNeighborFlow", + args=dict(VOXEL_SIZE={{_base_.VOXEL_SIZE}}, + POINT_CLOUD_RANGE={{_base_.POINT_CLOUD_RANGE}}, + SEQUENCE_LENGTH=SEQUENCE_LENGTH, + flow_save_folder=flow_save_folder, skip_existing=True)) + +test_loader = dict(name="ArgoverseRawSequenceLoader", + args=dict(sequence_dir=test_sequence_dir, verbose=True)) + +test_dataloader = dict( + args=dict(batch_size=1, num_workers=1, shuffle=False, pin_memory=True)) + +test_dataset = dict(name="SubsequenceRawDataset", + args=dict(subsequence_length=SEQUENCE_LENGTH, + max_sequence_length=max_test_sequence_length, + origin_mode="FIRST_ENTRY")) diff --git a/configs/nsfp/argo/old/unsupervised_flow_data.py b/configs/nsfp/argo/old/unsupervised_flow_data.py new file mode 100644 index 0000000..b4e436a --- /dev/null +++ b/configs/nsfp/argo/old/unsupervised_flow_data.py @@ -0,0 +1,42 @@ + +test_sequence_dir = "/efs/argoverse2/val/" +test_flow_dir = "/efs/argoverse2/val_sceneflow/" + +flow_save_folder = "/efs/argoverse_lidar/train_flow/" + + +def get_max_sequence_length(sequence_dir): + if "argoverse2" in sequence_dir: + return 146 + else: + return 296 + + +max_test_sequence_length = get_max_sequence_length(test_sequence_dir) + +epochs = 20 +learning_rate = 2e-6 +save_every = 500 +validate_every = 500 + +SEQUENCE_LENGTH = 2 + +is_trainable = False + +model = dict(name="NSFP", + args=dict(VOXEL_SIZE=(0.2, 0.2, 4), + POINT_CLOUD_RANGE=(-51.2, -51.2, -3, 51.2, 51.2, 1), + SEQUENCE_LENGTH=SEQUENCE_LENGTH, + flow_save_folder=flow_save_folder)) + +test_loader = dict(name="ArgoverseSupervisedFlowSequenceLoader", + args=dict(raw_data_path=test_sequence_dir, + flow_data_path=test_flow_dir)) + +test_dataloader = dict( + args=dict(batch_size=1, num_workers=1, shuffle=False, pin_memory=True)) + +test_dataset = dict(name="SubsequenceSupervisedFlowDataset", + args=dict(subsequence_length=SEQUENCE_LENGTH, + max_sequence_length=max_test_sequence_length, + origin_mode="FIRST_ENTRY")) diff --git a/configs/nsfp/argo/old/unsupervised_flow_data_cached.py b/configs/nsfp/argo/old/unsupervised_flow_data_cached.py new file mode 100644 index 0000000..c23af95 --- /dev/null +++ b/configs/nsfp/argo/old/unsupervised_flow_data_cached.py @@ -0,0 +1,4 @@ +_base_ = "./unsupervised_flow_data.py" + +flow_save_folder = "/efs/argoverse2/val_nsfp_flow/" +model = dict(name="NSFPCached", args=dict(flow_save_folder=flow_save_folder)) \ No newline at end of file diff --git a/configs/nsfp/argo/old/unsupervised_lidar_data.py b/configs/nsfp/argo/old/unsupervised_lidar_data.py new file mode 100644 index 0000000..17808c9 --- /dev/null +++ b/configs/nsfp/argo/old/unsupervised_lidar_data.py @@ -0,0 +1,42 @@ +is_trainable = False +has_labels = False + + +test_sequence_dir = "/efs/argoverse_lidar/train/" +flow_save_folder = "/efs/argoverse_lidar/train_nsfp_flow/" + +precision = 32 + + +def get_max_sequence_length(sequence_dir): + if "argoverse2" in sequence_dir: + return 146 + else: + return 296 + + +max_test_sequence_length = get_max_sequence_length(test_sequence_dir) + +epochs = 20 +learning_rate = 2e-6 +save_every = 500 +validate_every = 500 + +SEQUENCE_LENGTH = 2 + +model = dict(name="NSFP", + args=dict(VOXEL_SIZE=(0.2, 0.2, 4), + POINT_CLOUD_RANGE=(-51.2, -51.2, -3, 51.2, 51.2, 1), + SEQUENCE_LENGTH=SEQUENCE_LENGTH, + flow_save_folder=flow_save_folder)) + +test_loader = dict(name="ArgoverseRawSequenceLoader", + args=dict(sequence_dir=test_sequence_dir)) + +test_dataloader = dict( + args=dict(batch_size=1, num_workers=1, shuffle=False, pin_memory=True)) + +test_dataset = dict(name="SubsequenceRawDataset", + args=dict(subsequence_length=SEQUENCE_LENGTH, + max_sequence_length=max_test_sequence_length, + origin_mode="FIRST_ENTRY")) diff --git a/configs/nsfp/argo/old/unsupervised_lidar_data_subset.py b/configs/nsfp/argo/old/unsupervised_lidar_data_subset.py new file mode 100644 index 0000000..d03a54c --- /dev/null +++ b/configs/nsfp/argo/old/unsupervised_lidar_data_subset.py @@ -0,0 +1,44 @@ +_base_ = './unsupervised_lidar_data.py' + +is_trainable = False +has_labels = False + + +test_sequence_dir = "/efs/argoverse_lidar_subset/train/" +flow_save_folder = "/efs/argoverse_lidar_subset/train_nsfp_flow/" + +precision = 32 + + +def get_max_sequence_length(sequence_dir): + if "argoverse2" in sequence_dir: + return 146 + else: + return 296 + + +max_test_sequence_length = get_max_sequence_length(test_sequence_dir) + +epochs = 20 +learning_rate = 2e-6 +save_every = 500 +validate_every = 500 + +SEQUENCE_LENGTH = 2 + +model = dict(name="NSFP", + args=dict(VOXEL_SIZE=(0.2, 0.2, 4), + POINT_CLOUD_RANGE=(-51.2, -51.2, -3, 51.2, 51.2, 1), + SEQUENCE_LENGTH=SEQUENCE_LENGTH, + flow_save_folder=flow_save_folder)) + +test_loader = dict(name="ArgoverseRawSequenceLoader", + args=dict(sequence_dir=test_sequence_dir)) + +test_dataloader = dict( + args=dict(batch_size=1, num_workers=1, shuffle=False, pin_memory=True)) + +test_dataset = dict(name="SubsequenceRawDataset", + args=dict(subsequence_length=SEQUENCE_LENGTH, + max_sequence_length=max_test_sequence_length, + origin_mode="FIRST_ENTRY")) diff --git a/configs/nsfp/argo/old/unsupervised_lidar_data_subset_compliment.py b/configs/nsfp/argo/old/unsupervised_lidar_data_subset_compliment.py new file mode 100644 index 0000000..bdb7c1e --- /dev/null +++ b/configs/nsfp/argo/old/unsupervised_lidar_data_subset_compliment.py @@ -0,0 +1,8 @@ +_base_ = './unsupervised_lidar_data_subset.py' + +test_sequence_dir = "/efs/argoverse_lidar_subset_compliment/train/" +flow_save_folder = "/efs/argoverse_lidar_subset_compliment/train_nsfp_flow/" + +model = dict(args=dict(flow_save_folder=flow_save_folder)) + +test_loader = dict(args=dict(sequence_dir=test_sequence_dir)) diff --git a/configs/nsfp/argo/unsupervised_sensor_test.py b/configs/nsfp/argo/unsupervised_sensor_test.py new file mode 100644 index 0000000..37d0b60 --- /dev/null +++ b/configs/nsfp/argo/unsupervised_sensor_test.py @@ -0,0 +1,9 @@ +_base_ = "./unsupervised_sensor_train.py" + +test_sequence_dir = "/efs/argoverse2/test/" +flow_save_folder = "/efs/argoverse2/test_nsfp_flow/" +save_output_folder = "/efs/argoverse2/test_nsfp_out/" + +model = dict(args=dict(flow_save_folder=flow_save_folder)) + +test_loader = dict(args=dict(sequence_dir=test_sequence_dir)) diff --git a/configs/nsfp/argo/unsupervised_sensor_train.py b/configs/nsfp/argo/unsupervised_sensor_train.py new file mode 100644 index 0000000..87729bb --- /dev/null +++ b/configs/nsfp/argo/unsupervised_sensor_train.py @@ -0,0 +1,33 @@ +_base_ = "../../pseudoimage.py" + +is_trainable = False +has_labels = False + +test_sequence_dir = "/efs/argoverse2/train/" +flow_save_folder = "/efs/argoverse2/train_nsfp_flow/" + +precision = 32 + + +epochs = 20 +learning_rate = 2e-6 +save_every = 500 +validate_every = 500 + +SEQUENCE_LENGTH = 2 + +model = dict(name="NSFP", + args=dict(VOXEL_SIZE={{_base_.VOXEL_SIZE}}, + POINT_CLOUD_RANGE={{_base_.POINT_CLOUD_RANGE}}, + SEQUENCE_LENGTH=SEQUENCE_LENGTH, + flow_save_folder=flow_save_folder, skip_existing=True)) + +test_loader = dict(name="ArgoverseRawSequenceLoader", + args=dict(sequence_dir=test_sequence_dir, verbose=True)) + +test_dataloader = dict( + args=dict(batch_size=1, num_workers=1, shuffle=False, pin_memory=True)) + +test_dataset = dict(name="VarLenSubsequenceRawDataset", + args=dict(subsequence_length=SEQUENCE_LENGTH, + origin_mode="FIRST_ENTRY")) diff --git a/configs/nsfp/argo/unsupervised_sensor_val.py b/configs/nsfp/argo/unsupervised_sensor_val.py new file mode 100644 index 0000000..218cc39 --- /dev/null +++ b/configs/nsfp/argo/unsupervised_sensor_val.py @@ -0,0 +1,8 @@ +_base_ = "./unsupervised_sensor_train.py" + +test_sequence_dir = "/efs/argoverse2/val/" +flow_save_folder = "/efs/argoverse2/val_nsfp_flow/" + +model = dict(args=dict(flow_save_folder=flow_save_folder)) + +test_loader = dict(args=dict(sequence_dir=test_sequence_dir)) diff --git a/configs/nsfp/argo/unsupervised_sensor_val_cached.py b/configs/nsfp/argo/unsupervised_sensor_val_cached.py new file mode 100644 index 0000000..6a9558d --- /dev/null +++ b/configs/nsfp/argo/unsupervised_sensor_val_cached.py @@ -0,0 +1,14 @@ +_base_ = "./unsupervised_sensor_val.py" + +has_labels = True +test_sequence_dir = "/efs/argoverse2/val/" +test_flow_dir = "/efs/argoverse2/val_sceneflow/" + +flow_save_folder = "/efs/argoverse2/val_nsfp_flow/" + +test_loader = dict(name="ArgoverseSupervisedFlowSequenceLoader", + args=dict(_delete_=True, + raw_data_path=test_sequence_dir, + flow_data_path=test_flow_dir)) +test_dataset = dict(name="SubsequenceSupervisedFlowDataset") +model = dict(name="NSFPCached", args=dict(flow_save_folder=flow_save_folder)) \ No newline at end of file diff --git a/configs/nsfp/argo/unsupervised_sensor_val_fix_broken.py b/configs/nsfp/argo/unsupervised_sensor_val_fix_broken.py new file mode 100644 index 0000000..2a4190f --- /dev/null +++ b/configs/nsfp/argo/unsupervised_sensor_val_fix_broken.py @@ -0,0 +1,3 @@ +_base_ = "./unsupervised_sensor_val.py" + +test_loader = dict(args=dict(log_subset=['02a00399-3857-444e-8db3-a8f58489c394'])) diff --git a/configs/nsfp/argo/unsupervised_sensor_val_perf_test.py b/configs/nsfp/argo/unsupervised_sensor_val_perf_test.py new file mode 100644 index 0000000..83f7b66 --- /dev/null +++ b/configs/nsfp/argo/unsupervised_sensor_val_perf_test.py @@ -0,0 +1,21 @@ +_base_ = "./unsupervised_sensor_val.py" + +flow_save_folder = "/tmp/nsfp_perf_test_results/" + +has_labels = True +test_sequence_dir = "/efs/argoverse2/val/" +test_flow_dir = "/efs/argoverse2/val_sceneflow/" + +test_loader = dict(name="ArgoverseSupervisedFlowSequenceLoader", + args=dict(_delete_=True, + raw_data_path=test_sequence_dir, + flow_data_path=test_flow_dir)) + +test_dataset = dict( + name="VarLenSubsequenceSupervisedFlowSpecificSubsetDataset", + args=dict(_delete_=True, + subsequence_length={{_base_.SEQUENCE_LENGTH}}, + origin_mode="FIRST_ENTRY", + subset_file="data_prep_scripts/argo/eval_sequence_id_index.txt")) + +model = dict(args=dict(flow_save_folder=flow_save_folder)) diff --git a/configs/nsfp/waymo/unsupervised_sensor_train.py b/configs/nsfp/waymo/unsupervised_sensor_train.py new file mode 100644 index 0000000..f44869f --- /dev/null +++ b/configs/nsfp/waymo/unsupervised_sensor_train.py @@ -0,0 +1,34 @@ +_base_ = "../../pseudoimage.py" + +is_trainable = False +has_labels = False + +test_sequence_dir = "/efs/waymo_open_processed_flow/training/" +flow_save_folder = "/efs/waymo_open_processed_flow/train_nsfp_flow/" + +precision = 32 + +epochs = 20 +learning_rate = 2e-6 +save_every = 500 +validate_every = 500 + +SEQUENCE_LENGTH = 2 + +model = dict(name="NSFP", + args=dict(VOXEL_SIZE={{_base_.VOXEL_SIZE}}, + POINT_CLOUD_RANGE={{_base_.POINT_CLOUD_RANGE}}, + SEQUENCE_LENGTH=SEQUENCE_LENGTH, + flow_save_folder=flow_save_folder, + skip_existing=True)) + +test_loader = dict(name="WaymoSupervisedFlowSequenceLoader", + args=dict(sequence_dir=test_sequence_dir, verbose=True)) + +test_dataloader = dict( + args=dict(batch_size=1, num_workers=1, shuffle=False, pin_memory=True)) + +test_dataset = dict(name="VarLenSubsequenceRawDataset", + args=dict(subsequence_length=SEQUENCE_LENGTH, + max_pc_points=150000, + origin_mode="FIRST_ENTRY")) diff --git a/configs/nsfp/waymo/unsupervised_sensor_val.py b/configs/nsfp/waymo/unsupervised_sensor_val.py new file mode 100644 index 0000000..33ba06f --- /dev/null +++ b/configs/nsfp/waymo/unsupervised_sensor_val.py @@ -0,0 +1,8 @@ +_base_ = "./unsupervised_sensor_train.py" + +test_sequence_dir = "/efs/waymo_open_processed_flow/validation/" +flow_save_folder = "/efs/waymo_open_preprocessed/validation_nsfp_flow/" + +model = dict(args=dict(flow_save_folder=flow_save_folder)) + +test_loader = dict(args=dict(sequence_dir=test_sequence_dir)) diff --git a/configs/pseudoimage.py b/configs/pseudoimage.py new file mode 100644 index 0000000..9313cdd --- /dev/null +++ b/configs/pseudoimage.py @@ -0,0 +1,3 @@ +POINT_CLOUD_RANGE = (-51.2, -51.2, -3, 51.2, 51.2, 3) +VOXEL_SIZE = (0.2, 0.2, POINT_CLOUD_RANGE[5] - POINT_CLOUD_RANGE[2]) +PSEUDO_IMAGE_DIMS = (512, 512) diff --git a/data_prep_scripts/argo/convert_test_outputs_to_flow_format.py b/data_prep_scripts/argo/convert_test_outputs_to_flow_format.py new file mode 100644 index 0000000..acf2041 --- /dev/null +++ b/data_prep_scripts/argo/convert_test_outputs_to_flow_format.py @@ -0,0 +1,75 @@ +import os + +# Set Open MP's num threads to 1 +os.environ["OMP_NUM_THREADS"] = "1" + +from loader_utils import load_npy, save_npz +from pathlib import Path +import multiprocessing +import tqdm + +import argparse + +parser = argparse.ArgumentParser() +# Path to the input folder +parser.add_argument("input_folder", type=Path, help="Path to the input folder") +# Path to the output folder +parser.add_argument("output_folder", + type=Path, + help="Path to the output folder") +# Number of CPUs to use for parallel processing +parser.add_argument("--num_cpus", + type=int, + default=multiprocessing.cpu_count(), + help="Number of CPUs to use for parallel processing") +args = parser.parse_args() + +assert args.input_folder.exists(), f"Error: {args.input_folder} does not exist" +assert args.input_folder.is_dir( +), f"Error: {args.input_folder} is not a directory" + +args.output_folder.mkdir(parents=True, exist_ok=True) + + +def prepare_jobs(): + sequence_folders = sorted(args.input_folder.glob("*/")) + return sequence_folders + + +jobs = prepare_jobs() + + +def sequence_to_npz_format(data: dict): + flow = data['est_flow'] + valid_idxes = data['pc1_flows_valid_idx'] + delta_time = -1.0 + return {'flow': flow, 'valid_idxes': valid_idxes, 'delta_time': delta_time} + + +def process_sequence_folder(folder: Path): + # Load sequence files + sequence_files = sorted(folder.glob("*.npy")) + assert len(sequence_files) > 0, f"Error: {folder} has no .npy files" + + for sequence_file in sequence_files: + # Load the sequence file + sequence = load_npy(sequence_file, verbose=False).item() + # Convert the sequence to the npz format + sequence = sequence_to_npz_format(sequence) + # Save the sequence to the output folder + save_npz( + args.output_folder / folder.name / (sequence_file.stem + ".npz"), + sequence, verbose=False) + + +# If CPUs <= 1, don't use multiprocessing +if args.num_cpus <= 1: + for sequence_folder in tqdm.tqdm(jobs): + process_sequence_folder(sequence_folder) +else: + # Otherwise, use multiprocessing + with multiprocessing.Pool(args.num_cpus) as p: + # Use tqdm to show progress + for _ in tqdm.tqdm(p.imap_unordered(process_sequence_folder, jobs), + total=len(jobs)): + pass diff --git a/data_prep_scripts/argo/create_chodosh.py b/data_prep_scripts/argo/create_chodosh.py new file mode 100644 index 0000000..79db7c8 --- /dev/null +++ b/data_prep_scripts/argo/create_chodosh.py @@ -0,0 +1,161 @@ +import os + +# Set Open MP's num threads to 1 +os.environ["OMP_NUM_THREADS"] = "1" + +import argparse +from pathlib import Path +import numpy as np +import multiprocessing +from multiprocessing import Pool, cpu_count +from sklearn.cluster import DBSCAN +import time +from tqdm import tqdm + +# Take in input and output folder with no default arguments +parser = argparse.ArgumentParser() +parser.add_argument("input_folder", help="Path to the input folder") +parser.add_argument("output_folder", help="Path to the output folder") +parser.add_argument("--num_cpus", + type=int, + default=multiprocessing.cpu_count(), + help="Number of CPUs to use for parallel processing") + +# Parse the arguments +args = parser.parse_args() + +# Verify that the input folder is a folder and exists +input_folder_path = Path(args.input_folder) +assert input_folder_path.exists(), f"Error: {input_folder_path} does not exist" +assert input_folder_path.is_dir( +), f"Error: {input_folder_path} is not a directory" + +# Create the output folder if it doesn't exist +output_folder_path = Path(args.output_folder) +output_folder_path.mkdir(parents=True, exist_ok=True) + +# Load all the .npz file paths in the input folder into a list +npz_file_list = list(input_folder_path.glob("*.npz")) + +######################################################### + + +def fit_rigid(A, B): + """ + Fit Rigid transformation A @ R.T + t = B + """ + assert A.shape == B.shape + num_rows, num_cols = A.shape + + # find mean column wise + centroid_A = np.mean(A, axis=0) + centroid_B = np.mean(B, axis=0) + + # ensure centroids are 1x3 + centroid_A = centroid_A.reshape(1, -1) + centroid_B = centroid_B.reshape(1, -1) + + # subtract mean + Am = A - centroid_A + Bm = B - centroid_B + + H = Am.T @ Bm + # find rotation + U, S, Vt = np.linalg.svd(H) + R = Vt.T @ U.T + + # special reflection case + if np.linalg.det(R) < 0: + Vt[2, :] *= -1 + R = Vt.T @ U.T + + t = -centroid_A @ R.T + centroid_B + + return R, t + + +def ransac_fit_rigid(pts1, pts2, inlier_thresh, ntrials): + best_R = np.eye(3) + best_t = np.zeros((3, )) + best_inliers = (np.linalg.norm(pts1 - pts2, axis=-1) < inlier_thresh) + best_inliers_sum = best_inliers.sum() + for i in range(ntrials): + choice = np.random.choice(len(pts1), 3) + R, t = fit_rigid(pts1[choice], pts2[choice]) + inliers = (np.linalg.norm(pts1 @ R.T + t - pts2, axis=-1) < + inlier_thresh) + if inliers.sum() > best_inliers_sum: + best_R = R + best_t = t + best_inliers = inliers + best_inliers_sum = best_inliers.sum() + if best_inliers_sum / len(pts1) > 0.5: + break + best_R, best_t = fit_rigid(pts1[best_inliers], pts2[best_inliers]) + return best_R, best_t + + +def refine_flow(pc0, + flow_pred, + eps=0.4, + min_points=10, + motion_threshold=0.05, + inlier_thresh=0.2, + ntrials=250): + labels = DBSCAN(eps=eps, min_samples=min_points).fit_predict(pc0) + max_label = labels.max() + refined_flow = np.zeros_like(flow_pred) + for l in range(max_label + 1): + label_mask = labels == l + cluster_pts = pc0[label_mask] + cluster_flows = flow_pred[label_mask] + R, t = ransac_fit_rigid(cluster_pts, cluster_pts + cluster_flows, + inlier_thresh, ntrials) + if np.linalg.norm(t) < motion_threshold: + R = np.eye(3) + t = np.zeros((3, )) + refined_flow[label_mask] = (cluster_pts @ R.T + t) - cluster_pts + + refined_flow[labels == -1] = flow_pred[labels == -1] + return refined_flow + + +######################################################### + + +# Define the worker function for parallel processing +def process_npz_file(npz_path): + # Load the data from the .npz file + data = dict(np.load(npz_path)) + before_time = time.time() + refined_flow = refine_flow(data["pc"], data["flow"]) + after_time = time.time() + assert refined_flow.shape == data["flow"].shape + data["flow"] = refined_flow + delta_time = after_time - before_time + data["refine_delta_time"] = delta_time + data["delta_time"] = delta_time + data["nsfp_delta_time"] + npz_name = npz_path.stem + folder_name, index_name, minibatch_name = npz_name.split("_") + # Convert index_name to int + index_name = int(index_name) + # Convert minibatch_name to int + minibatch_name = int(minibatch_name) + # Convert + output_file_path = output_folder_path / folder_name / f"{index_name:010d}_{minibatch_name:03d}.npz" + output_file_path.parent.mkdir(parents=True, exist_ok=True) + if output_file_path.exists(): + output_file_path.unlink() + np.savez(output_file_path, **data) + + +# Process the .npz files either in parallel or with a standard for loop +if args.num_cpus <= 1: + for npz_path in tqdm(npz_file_list): + process_npz_file(npz_path) +else: + print(f"Using {args.num_cpus} CPUs for parallel processing...") + with Pool(processes=args.num_cpus) as pool: + for _ in tqdm(pool.imap_unordered(process_npz_file, npz_file_list), + total=len(npz_file_list)): + pass diff --git a/data_prep_scripts/argo/create_odom_only_flow.py b/data_prep_scripts/argo/create_odom_only_flow.py new file mode 100644 index 0000000..d3c96c1 --- /dev/null +++ b/data_prep_scripts/argo/create_odom_only_flow.py @@ -0,0 +1,37 @@ +import argparse +from pathlib import Path +import numpy as np +from loader_utils import load_npz, save_npz +import joblib +import multiprocessing + +# Get path to existing nsfp flow folder +# Get path to new odom flow folder + +parser = argparse.ArgumentParser() +parser.add_argument('nsfp_dir', type=Path, help='path to the nsfp flow folder') +parser.add_argument('odom_dir', type=Path, help='path to the odom flow folder') +args = parser.parse_args() + +assert args.nsfp_dir.exists(), 'nsfp directory does not exist' +args.odom_dir.mkdir(exist_ok=True) + +sequence_dirs = sorted(args.nsfp_dir.glob('*/')) + + +def process_sequence(sequence_dir: Path): + # Collect the npz files + npz_files = sorted(sequence_dir.glob('*.npz')) + # Create a new folder for the odom flows + odom_dir = args.odom_dir / sequence_dir.name + odom_dir.mkdir(exist_ok=True) + # For each npz file, load the data and save the odom flow + for npz_file in npz_files: + data = dict(load_npz(npz_file, verbose=False)) + data['flow'] = np.zeros_like(data['flow']) + save_npz(odom_dir / npz_file.name, data, verbose=False) + + +# Parallel process the sequences +joblib.Parallel(n_jobs=multiprocessing.cpu_count() - 1)(joblib.delayed(process_sequence)(sequence_dir) + for sequence_dir in sequence_dirs) diff --git a/data_prep_scripts/argo/eval_frames.txt b/data_prep_scripts/argo/eval_frames.txt new file mode 100644 index 0000000..9e02e42 --- /dev/null +++ b/data_prep_scripts/argo/eval_frames.txt @@ -0,0 +1,100 @@ +0b5142c1-420b-3fea-9e98-b87327ae22c6,315968122260198000,315968122359731000 +0b5142c1-420b-3fea-9e98-b87327ae22c6,315968127559929000,315968127660126000 +11ba4e81-c26f-3cd1-827d-b6913bcef64e,315966066760060000,315966066860256000 +11ba4e81-c26f-3cd1-827d-b6913bcef64e,315966068560241000,315966068659774000 +15ec0778-826e-3ed7-9775-54fbf66997f4,315970268859920000,315970268960117000 +19350c96-623d-4d77-af96-f8c23f00c358,315972066960005000,315972067060185000 +1f434d15-8745-3fba-9c3e-ccb026688397,315972878259718000,315972878359914000 +20d47f81-46e8-3adf-a0ca-564fbb5c599d,315968561859653000,315968561959849000 +20dd185d-b4eb-3024-a17a-b4e5d8b15b65,315969734159653000,315969734259849000 +22052525-4f85-3fe8-9d7d-000a9fffce36,315970673059898000,315970673160095000 +24642607-2a51-384a-90a7-228067956d05,315970925960088000,315970926060273000 +24642607-2a51-384a-90a7-228067956d05,315970927460357000,315970927559889000 +25e5c600-36fe-3245-9cc0-40ef91620c22,315966117659687000,315966117759883000 +27c03d98-6ac3-38a3-ba5e-102b184d01ef,315975015659864000,315975015759397000 +27c03d98-6ac3-38a3-ba5e-102b184d01ef,315975023059688000,315975023159884000 +27c03d98-6ac3-38a3-ba5e-102b184d01ef,315975027360115000,315975027459648000 +2ff4f798-78d9-3384-87e9-61928aa4cb6d,315974474560006000,315974474660202000 +335aabef-269e-3211-a99d-2c3a3a8f8475,315975285359460000,315975285458993000 +335aabef-269e-3211-a99d-2c3a3a8f8475,315975286360750000,315975286460283000 +335aabef-269e-3211-a99d-2c3a3a8f8475,315975286959938000,315975287060123000 +3bffdcff-c3a7-38b6-a0f2-64196d130958,315975593060303000,315975593159836000 +44adf4c4-6064-362f-94d3-323ed42cfda9,315966152059704000,315966152159900000 +44adf4c4-6064-362f-94d3-323ed42cfda9,315966157660052000,315966157760248000 +47286726-5dd4-4e26-bd2d-5324f429e445,315972129559829000,315972129660026000 +52971a8a-ed62-3bfd-bcd4-ca3308b594e0,315968616859955000,315968616960151000 +544a8102-0ef5-3044-921e-dc0544370376,315968326059706000,315968326159903000 +5589de60-1727-3e3f-9423-33437fc5da4b,315967919259399000,315967919359595000 +5589de60-1727-3e3f-9423-33437fc5da4b,315967921659456000,315967921759653000 +5589de60-1727-3e3f-9423-33437fc5da4b,315967932259644000,315967932359840000 +58e82365-03bc-3b2f-b55a-a4ad0e3e792d,315968793660007000,315968793760204000 +58e82365-03bc-3b2f-b55a-a4ad0e3e792d,315968797059999000,315968797160195000 +58fed0d4-97d5-469b-89a4-4394838e10c7,315974258359783000,315974258459979000 +5c0584a3-52a6-3029-b6ff-ca45a19d8aa6,315967118960063000,315967119059582000 +5f278cdd-ca28-3c53-8f5c-04e62308811d,315969638659746000,315969638759942000 +5f278cdd-ca28-3c53-8f5c-04e62308811d,315969642360336000,315969642459869000 +5f8f4a26-59b1-3f70-bcab-b5e3e615d3bc,315966660659605000,315966660759801000 +65387aee-4490-38b9-8f4f-1fc43bd4ac06,315969317559880000,315969317660076000 +65387aee-4490-38b9-8f4f-1fc43bd4ac06,315969324260352000,315969324359885000 +6803104a-bb06-402e-8471-e5af492db0a8,315977902159939000,315977902260135000 +6803104a-bb06-402e-8471-e5af492db0a8,315977908859734000,315977908959931000 +6aaf5b08-9f84-3a2e-8a32-2e50e5e11a3c,315970475560177000,315970475659710000 +6f128f23-ee40-3ea9-8c50-c9cdb9d3e8b6,315967388960111000,315967389059636000 +78683234-e6f1-3e4e-af52-6f839254e4c0,315966766260097000,315966766360293000 +78da7b7e-8ddf-3c7d-8716-eaa890106dd3,315974834959662000,315974835059850000 +7a2c222d-addc-30b2-aac6-596cb65a22e3,315971812359328000,315971812459524000 +7dbc2eac-5871-3480-b322-246e03d954d2,315968146260158000,315968146360354000 +7e4d67b3-c3cc-3288-afe5-043602ea3c70,315969976560031000,315969976659564000 +7e4d67b3-c3cc-3288-afe5-043602ea3c70,315969985460128000,315969985560324000 +7fab2350-7eaf-3b7e-a39d-6937a4c1bede,315966253760553000,315966253860086000 +7fab2350-7eaf-3b7e-a39d-6937a4c1bede,315966256459865000,315966256560061000 +8749f79f-a30b-3c3f-8a44-dbfa682bbef1,315971055560254000,315971055659787000 +88f47a10-87b4-3ea8-a0c7-a07d825b647d,315967650060104000,315967650159636000 +91aab547-1912-3b8e-8e7f-df3b202147bf,315968307860021000,315968307960217000 +9e9bcfb7-601d-3d80-bc12-ef7025174beb,315966293260056000,315966293359589000 +9e9bcfb7-601d-3d80-bc12-ef7025174beb,315966297060217000,315966297159750000 +a33a44fb-6008-3dc2-b7c5-2d27b70741e8,315967270060068000,315967270160264000 +adf9a841-e0db-30ab-b5b3-bf0b61658e1e,315972240960065000,315972241060252000 +adf9a841-e0db-30ab-b5b3-bf0b61658e1e,315972246760122000,315972246860319000 +adf9a841-e0db-30ab-b5b3-bf0b61658e1e,315972246959851000,315972247060039000 +b50c4763-5d1e-37f4-a009-2244aeebabcd,315974150359697000,315974150459893000 +b5a7ff7e-d74a-3be6-b95d-3fc0042215f6,315977333460219000,315977333559751000 +b5a7ff7e-d74a-3be6-b95d-3fc0042215f6,315977334760109000,315977334859642000 +b6e967f6-92bc-3bf5-99c9-1b0c4649fd67,315966561759650000,315966561859847000 +b87683ae-14c5-321f-8af3-623e7bafc3a7,315972338260322000,315972338360518000 +b87683ae-14c5-321f-8af3-623e7bafc3a7,315972343260124000,315972343359657000 +b87683ae-14c5-321f-8af3-623e7bafc3a7,315972345760365000,315972345859898000 +bf360aeb-1bbd-3c1e-b143-09cf83e4f2e4,315969622960010000,315969623060197000 +bf382949-3515-3c16-b505-319442937a43,315967943859754000,315967943959950000 +c2d44a70-9fd4-3298-ad0a-c4c9712e6f1e,315976595360581000,315976595460114000 +c85a88a8-c916-30a7-923c-0c66bd3ebbd3,315967360160262000,315967360259794000 +c865c156-0f26-411c-a16c-be985333f675,315979577860048000,315979577959581000 +cae56e40-8470-3c9c-af75-6e444189488f,315967283060282000,315967283159815000 +d1395998-7e8a-417d-91e9-5ca6ec045ee1,315968264560197000,315968264660393000 +d1395998-7e8a-417d-91e9-5ca6ec045ee1,315968268059722000,315968268159919000 +d3ca0450-2167-38fb-b34b-449741cb38f3,315968249559651000,315968249659847000 +d89f80be-76d0-3853-8daa-76605cf4ce5e,315970802759878000,315970802860074000 +d89f80be-76d0-3853-8daa-76605cf4ce5e,315970805060382000,315970805159914000 +da036982-92bf-36a8-b880-4ccf4e20b74e,315976671259685000,315976671359882000 +dafe14f5-825c-4e7a-9009-6dfdfdd5b030,315977628359965000,315977628459498000 +dc9077b9-2fe0-3d18-9b97-8067ff090874,315971086159609000,315971086259805000 +de56b100-508b-3479-81fe-735349f8e8de,315966307259783000,315966307359979000 +ded5ef6e-46ea-3a66-9180-18a6fa0a2db4,315976259559848000,315976259660045000 +dfc32963-1524-34f4-9f5e-0292f0f223ae,315970541559826000,315970541660023000 +dfc6d65f-20f5-389d-a5cd-81c1c7ecb11f,315969778159965000,315969778260162000 +e13c06cb-cd01-380e-946f-6d92ac1af49d,315966598859962000,315966598960158000 +e13c06cb-cd01-380e-946f-6d92ac1af49d,315966602759658000,315966602859854000 +e1d68dde-22a9-3918-a526-0850b21ff2eb,315969776959918000,315969777060105000 +e2e921fe-e489-3656-a0a2-5e17bd399ddf,315966472060002000,315966472160199000 +eec284b2-840a-3c75-aa42-04d2e309bbe1,315972311159687000,315972311259883000 +f1275002-842e-3571-8f7d-05816bc7cf56,315977585560157000,315977585659690000 +f2576c8a-da9b-450e-88cf-a70af1b0eadf,315980512759794000,315980512859990000 +f2576c8a-da9b-450e-88cf-a70af1b0eadf,315980513259434000,315980513359630000 +f2576c8a-da9b-450e-88cf-a70af1b0eadf,315980514459785000,315980514559981000 +f3cd0d0d-8b71-3266-9732-d9f0d5778eb6,315968818559990000,315968818660187000 +f3cd0d0d-8b71-3266-9732-d9f0d5778eb6,315968823060155000,315968823159688000 +f4c94798-4d77-36ab-bdc5-c1194e5e7aff,315968015459869000,315968015560066000 +f4c94798-4d77-36ab-bdc5-c1194e5e7aff,315968016060376000,315968016159909000 +f4c94798-4d77-36ab-bdc5-c1194e5e7aff,315968025859646000,315968025959843000 +fbee355f-8878-31fa-8ac8-b9a45a3f130a,315967341660050000,315967341759583000 +ff0dbfc5-8a7b-3a6e-8936-e5e812e45408,315972914759754000,315972914859951000 \ No newline at end of file diff --git a/data_prep_scripts/argo/eval_sequence_id_index.txt b/data_prep_scripts/argo/eval_sequence_id_index.txt new file mode 100644 index 0000000..712dba7 --- /dev/null +++ b/data_prep_scripts/argo/eval_sequence_id_index.txt @@ -0,0 +1,100 @@ +0b5142c1-420b-3fea-9e98-b87327ae22c6,10 +0b5142c1-420b-3fea-9e98-b87327ae22c6,63 +11ba4e81-c26f-3cd1-827d-b6913bcef64e,62 +11ba4e81-c26f-3cd1-827d-b6913bcef64e,80 +15ec0778-826e-3ed7-9775-54fbf66997f4,89 +19350c96-623d-4d77-af96-f8c23f00c358,5 +1f434d15-8745-3fba-9c3e-ccb026688397,53 +20d47f81-46e8-3adf-a0ca-564fbb5c599d,80 +20dd185d-b4eb-3024-a17a-b4e5d8b15b65,124 +22052525-4f85-3fe8-9d7d-000a9fffce36,43 +24642607-2a51-384a-90a7-228067956d05,101 +24642607-2a51-384a-90a7-228067956d05,116 +25e5c600-36fe-3245-9cc0-40ef91620c22,134 +27c03d98-6ac3-38a3-ba5e-102b184d01ef,13 +27c03d98-6ac3-38a3-ba5e-102b184d01ef,87 +27c03d98-6ac3-38a3-ba5e-102b184d01ef,130 +2ff4f798-78d9-3384-87e9-61928aa4cb6d,0 +335aabef-269e-3211-a99d-2c3a3a8f8475,86 +335aabef-269e-3211-a99d-2c3a3a8f8475,96 +335aabef-269e-3211-a99d-2c3a3a8f8475,102 +3bffdcff-c3a7-38b6-a0f2-64196d130958,120 +44adf4c4-6064-362f-94d3-323ed42cfda9,2 +44adf4c4-6064-362f-94d3-323ed42cfda9,58 +47286726-5dd4-4e26-bd2d-5324f429e445,153 +52971a8a-ed62-3bfd-bcd4-ca3308b594e0,130 +544a8102-0ef5-3044-921e-dc0544370376,94 +5589de60-1727-3e3f-9423-33437fc5da4b,0 +5589de60-1727-3e3f-9423-33437fc5da4b,24 +5589de60-1727-3e3f-9423-33437fc5da4b,130 +58e82365-03bc-3b2f-b55a-a4ad0e3e792d,29 +58e82365-03bc-3b2f-b55a-a4ad0e3e792d,63 +58fed0d4-97d5-469b-89a4-4394838e10c7,50 +5c0584a3-52a6-3029-b6ff-ca45a19d8aa6,89 +5f278cdd-ca28-3c53-8f5c-04e62308811d,18 +5f278cdd-ca28-3c53-8f5c-04e62308811d,55 +5f8f4a26-59b1-3f70-bcab-b5e3e615d3bc,45 +65387aee-4490-38b9-8f4f-1fc43bd4ac06,42 +65387aee-4490-38b9-8f4f-1fc43bd4ac06,109 +6803104a-bb06-402e-8471-e5af492db0a8,29 +6803104a-bb06-402e-8471-e5af492db0a8,96 +6aaf5b08-9f84-3a2e-8a32-2e50e5e11a3c,5 +6f128f23-ee40-3ea9-8c50-c9cdb9d3e8b6,90 +78683234-e6f1-3e4e-af52-6f839254e4c0,128 +78da7b7e-8ddf-3c7d-8716-eaa890106dd3,77 +7a2c222d-addc-30b2-aac6-596cb65a22e3,2 +7dbc2eac-5871-3480-b322-246e03d954d2,32 +7e4d67b3-c3cc-3288-afe5-043602ea3c70,22 +7e4d67b3-c3cc-3288-afe5-043602ea3c70,111 +7fab2350-7eaf-3b7e-a39d-6937a4c1bede,1 +7fab2350-7eaf-3b7e-a39d-6937a4c1bede,28 +8749f79f-a30b-3c3f-8a44-dbfa682bbef1,114 +88f47a10-87b4-3ea8-a0c7-a07d825b647d,10 +91aab547-1912-3b8e-8e7f-df3b202147bf,72 +9e9bcfb7-601d-3d80-bc12-ef7025174beb,8 +9e9bcfb7-601d-3d80-bc12-ef7025174beb,46 +a33a44fb-6008-3dc2-b7c5-2d27b70741e8,27 +adf9a841-e0db-30ab-b5b3-bf0b61658e1e,37 +adf9a841-e0db-30ab-b5b3-bf0b61658e1e,95 +adf9a841-e0db-30ab-b5b3-bf0b61658e1e,97 +b50c4763-5d1e-37f4-a009-2244aeebabcd,126 +b5a7ff7e-d74a-3be6-b95d-3fc0042215f6,123 +b5a7ff7e-d74a-3be6-b95d-3fc0042215f6,136 +b6e967f6-92bc-3bf5-99c9-1b0c4649fd67,127 +b87683ae-14c5-321f-8af3-623e7bafc3a7,39 +b87683ae-14c5-321f-8af3-623e7bafc3a7,89 +b87683ae-14c5-321f-8af3-623e7bafc3a7,114 +bf360aeb-1bbd-3c1e-b143-09cf83e4f2e4,49 +bf382949-3515-3c16-b505-319442937a43,137 +c2d44a70-9fd4-3298-ad0a-c4c9712e6f1e,84 +c85a88a8-c916-30a7-923c-0c66bd3ebbd3,88 +c865c156-0f26-411c-a16c-be985333f675,147 +cae56e40-8470-3c9c-af75-6e444189488f,154 +d1395998-7e8a-417d-91e9-5ca6ec045ee1,72 +d1395998-7e8a-417d-91e9-5ca6ec045ee1,107 +d3ca0450-2167-38fb-b34b-449741cb38f3,73 +d89f80be-76d0-3853-8daa-76605cf4ce5e,122 +d89f80be-76d0-3853-8daa-76605cf4ce5e,145 +da036982-92bf-36a8-b880-4ccf4e20b74e,109 +dafe14f5-825c-4e7a-9009-6dfdfdd5b030,25 +dc9077b9-2fe0-3d18-9b97-8067ff090874,11 +de56b100-508b-3479-81fe-735349f8e8de,141 +ded5ef6e-46ea-3a66-9180-18a6fa0a2db4,4 +dfc32963-1524-34f4-9f5e-0292f0f223ae,37 +dfc6d65f-20f5-389d-a5cd-81c1c7ecb11f,8 +e13c06cb-cd01-380e-946f-6d92ac1af49d,20 +e13c06cb-cd01-380e-946f-6d92ac1af49d,59 +e1d68dde-22a9-3918-a526-0850b21ff2eb,113 +e2e921fe-e489-3656-a0a2-5e17bd399ddf,35 +eec284b2-840a-3c75-aa42-04d2e309bbe1,110 +f1275002-842e-3571-8f7d-05816bc7cf56,115 +f2576c8a-da9b-450e-88cf-a70af1b0eadf,132 +f2576c8a-da9b-450e-88cf-a70af1b0eadf,137 +f2576c8a-da9b-450e-88cf-a70af1b0eadf,149 +f3cd0d0d-8b71-3266-9732-d9f0d5778eb6,107 +f3cd0d0d-8b71-3266-9732-d9f0d5778eb6,152 +f4c94798-4d77-36ab-bdc5-c1194e5e7aff,38 +f4c94798-4d77-36ab-bdc5-c1194e5e7aff,44 +f4c94798-4d77-36ab-bdc5-c1194e5e7aff,142 +fbee355f-8878-31fa-8ac8-b9a45a3f130a,4 +ff0dbfc5-8a7b-3a6e-8936-e5e812e45408,115 diff --git a/data_prep_scripts/argo/load_sequence_index.py b/data_prep_scripts/argo/load_sequence_index.py new file mode 100644 index 0000000..03e4e93 --- /dev/null +++ b/data_prep_scripts/argo/load_sequence_index.py @@ -0,0 +1,55 @@ +import argparse +from pathlib import Path +import numpy as np +from typing import List, Tuple, Dict, Optional + +# Get path to the eval frames file + +parser = argparse.ArgumentParser() +parser.add_argument('argoverse_dir', + type=Path, + help='path to the Argoverse dataset') +parser.add_argument('eval_frames_file', + type=Path, + help='path to the eval frames file') +parser.add_argument('output_file', type=Path, help='path to the out file') +args = parser.parse_args() + +assert args.argoverse_dir.exists(), 'Argoverse directory does not exist' +assert args.eval_frames_file.exists(), 'Eval frames file does not exist' + + +def process_lines(lines: List[str]) -> List[Tuple[str, int, int]]: + + def process_line(line: str) -> Tuple[str, int, int]: + line = line.strip() + line = line.split(',') + return line[0], int(line[1]), int(line[2]) + + return [process_line(line) for line in lines] + + +# Load txt file +with open(args.eval_frames_file, 'r') as f: + lines = f.readlines() +configs = process_lines(lines) + + +def config_to_id_index(config: Tuple[str, int, int]) -> Tuple[str, int]: + id, start, end = config + feather_folder = args.argoverse_dir / id / 'sensors' / 'lidar' + feather_files = sorted(feather_folder.glob('*.feather')) + feather_names = [int(f.stem) for f in feather_files] + index = feather_names.index(start) + assert feather_names[index] == start, 'Start frame not found' + assert feather_names[index + 1] == end, 'End frame not found' + return id, index + + +# Convert configs to id, index +configs = [config_to_id_index(config) for config in configs] + +# Write to file +with open(args.output_file, 'w') as f: + for config in configs: + f.write(f'{config[0]},{config[1]}\n') diff --git a/data_prep_scripts/argo/rasterize_heightmap.py b/data_prep_scripts/argo/rasterize_heightmap.py new file mode 100644 index 0000000..c7cd8f6 --- /dev/null +++ b/data_prep_scripts/argo/rasterize_heightmap.py @@ -0,0 +1,235 @@ +import numpy as np +from PIL import Image, ImageDraw, ImageFilter + +from sklearn.neighbors import NearestNeighbors +import numpy as np +import json +from pathlib import Path +import multiprocessing +import matplotlib.pyplot as plt +from functools import partial +from joblib import Parallel, delayed +from tqdm import tqdm + +import argparse + + +def build_knn(points, num_neighbors): + return NearestNeighbors(n_neighbors=num_neighbors, + radius=20, + leaf_size=num_neighbors).fit(points) + + +def build_global_grid(points, cells_per_meter): + xs, ys, _ = zip(*points) + min_x, max_x = min(xs), max(xs) + min_y, max_y = min(ys), max(ys) + + grid_max_global_frame = np.array([max_x, max_y]) + grid_min_global_frame = np.array([min_x, min_y]) + + area_grid_global_frame = grid_max_global_frame - grid_min_global_frame + grid_shape = np.ceil( + area_grid_global_frame * cells_per_meter).astype(int) + 1 + + def global_to_grid_float(pts): + assert (pts <= grid_max_global_frame + ).all(), f"({pts} <= {grid_max_global_frame})" + assert (pts >= grid_min_global_frame + ).all(), f"({pts} >= {grid_min_global_frame})" + relative_to_grid_origin = pts - grid_min_global_frame + floating_point_grid_coordinate = relative_to_grid_origin * cells_per_meter + return floating_point_grid_coordinate + + def global_to_grid_index(pts): + coords = global_to_grid_float(pts) + return np.round(coords).astype(int) + + return grid_shape, global_to_grid_index, global_to_grid_float, grid_min_global_frame + + +def collect_points(drivable_areas): + all_points = [] + + for key in drivable_areas.keys(): + points = [(e['x'], e['y'], e['z']) + for e in drivable_areas[key]['area_boundary']] + all_points.append(points) + + return all_points + + +def get_road_polygon_mask_arr(polygon_list, mask_shape, global_to_grid_float, + cells_per_meter, meters_beyond_poly_edge): + mask_img = Image.new('1', (mask_shape[1], mask_shape[0]), 0) + mask_draw = ImageDraw.Draw(mask_img) + for polygon in polygon_list: + grid_pts = [tuple(global_to_grid_float(p[:2])) for p in polygon] + mask_draw.polygon(grid_pts, fill=1, outline=1) + + blur_radius = int(np.ceil(meters_beyond_poly_edge * cells_per_meter)) + mask_img = mask_img.convert('L') + mask_img = mask_img.filter(ImageFilter.GaussianBlur(radius=blur_radius)) + return (np.array(mask_img) > 0).astype(int) + + +def render_heightmap(drivable_areas, cells_per_meter, meters_beyond_poly_edge, + num_neighbors): + + polygon_list = collect_points(drivable_areas) + + flattened_poly_points = [e for e_list in polygon_list for e in e_list] + + # We construct this full heightmap with 0, 0 at xy_min_offset in the global coordinate frame. + grid_shape, global_to_grid_index, global_to_grid_float, grid_min_global_frame = build_global_grid( + flattened_poly_points, cells_per_meter) + + polygon_grid_points = np.array([(*global_to_grid_index(p[:2]), p[2]) + for p in flattened_poly_points]) + mean_z = np.mean([e[2] for e in polygon_grid_points]) + knn = build_knn(polygon_grid_points, num_neighbors) + + # Construct a grid shaped array whose last axis holds the X, Y index value for that square, + # with the average Z value for the purposes of querying. + xs_lin = np.arange(0, grid_shape[0], 1) + ys_lin = np.arange(0, grid_shape[1], 1) + xs_square = np.expand_dims(np.tile(xs_lin, (grid_shape[1], 1)), 2) + ys_square = np.expand_dims(np.tile(ys_lin, (grid_shape[0], 1)).T, 2) + zs_square = np.ones_like(xs_square) * mean_z + pts_square = np.concatenate((xs_square, ys_square, zs_square), 2) + + # Flatten the pts square into an N x 3 array for querying KNN + pts_square_shape = pts_square.shape + pts_line = pts_square.reshape(pts_square_shape[0] * pts_square_shape[1], 3) + + _, indices = knn.kneighbors(pts_line) + neighbor_values = polygon_grid_points[indices] + avg_neighbor_z_values = np.mean(neighbor_values[:, :, 2], axis=1) + + # Reshape flattened average Z values back into grid shape. + grid = avg_neighbor_z_values.reshape(pts_square_shape[0], + pts_square_shape[1]) + + if meters_beyond_poly_edge is not None: + road_mask_array = get_road_polygon_mask_arr(polygon_list, grid.shape, + global_to_grid_float, + cells_per_meter, + meters_beyond_poly_edge) + grid[~road_mask_array] = np.NaN + return grid, grid_min_global_frame + + +def save_grid_global_offset(sequence_folder: Path, grid, grid_min_global_frame, + cells_per_meter): + se2 = { + "R": [1.0, 0.0, 0.0, 1.0], # Identity rotation matrix flattened + "t": [-grid_min_global_frame[0], -grid_min_global_frame[1]], + "s": cells_per_meter + } + + se2_name = f"{sequence_folder.name}___img_Sim2_city.json" + height_map_name = f"{sequence_folder.name}_ground_height_surface____.npy" + + height_map_file = sequence_folder / "map" / height_map_name + if height_map_file.exists(): + height_map_file.unlink() + + se2_file = sequence_folder / "map" / se2_name + if se2_file.exists(): + se2_file.unlink() + + np.save(height_map_file, grid.astype(np.float16)) + with open(se2_file, 'w') as fp: + json.dump(se2, fp) + + +def process_sequence_folder(sequence_folder: Path, + cells_per_meter=10.0 / 3.0, + meters_beyond_poly_edge=None, + num_neighbors=20): + sequence_folder = Path(sequence_folder) + map_files = list((sequence_folder / "map").glob("log_map_archive_*.json")) + assert len( + map_files + ) == 1, f"Expected 1 map file, got {len(map_files)} in {sequence_folder / 'map'}" + map_file = map_files[0] + with open(map_file) as f: + map_json = json.load(f) + + grid, grid_min_global_frame = render_heightmap(map_json['drivable_areas'], + cells_per_meter, + meters_beyond_poly_edge, + num_neighbors) + save_grid_global_offset(sequence_folder, grid, grid_min_global_frame, + cells_per_meter) + + +def build_work_queue(root_path: Path): + root_path = Path(root_path) + root_folders = [ + e for e in root_path.glob("*/") + if e.is_dir() and e.name in ["train", "val"] + ] + work_queue = [] + for rf in root_folders: + work_queue.extend([e for e in rf.glob("*/") if e.is_dir()]) + return work_queue + + +parser = argparse.ArgumentParser() +parser.add_argument("argoverse_directory", + help="Path to Argoverse 2 directory.") +parser.add_argument( + '--meters_beyond_poly_edge', + type=float, + default=5.0, + help="Meters beyond the poly road edge to project out the height map." + "For the Argoverse 2 sensor dataset, the heightmap extends 5 meters outside the road geometry polygon. " + "For values less than 0, this will generate the height map for the entire scene." +) +parser.add_argument( + '--cells_per_meter', + type=int, + default=10.0 / 3.0, + help= + "Cells per meter for discritization. Default is Argoverse 2 Sensor dataset default is 30cm (10/3 cells per meter)" +) +parser.add_argument('--num_neighbors', + type=int, + default=20, + help="Number of neighbors to use to compute height") +parser.add_argument('--cpus', + type=int, + default=multiprocessing.cpu_count(), + help="Number of cpus to use for parallel processing") + +args = parser.parse_args() + +argoverse_directory = Path(args.argoverse_directory) +assert argoverse_directory.is_dir( +), f"{argoverse_directory} is not a directory" + +cells_per_meter = args.cells_per_meter +assert cells_per_meter > 0, f"Cells per meter must be positive (got {cells_per_meter})" + +meters_beyond_poly_edge = args.meters_beyond_poly_edge if args.meters_beyond_poly_edge >= 0 else None + +num_neighbors = args.num_neighbors +assert num_neighbors > 0, f"Number of neighbors must be greater than zero (got {num_neighbors})" + +work_queue = build_work_queue(argoverse_directory) + +print("Work queue size:", len(work_queue)) + +Parallel(args.cpus)(delayed(process_sequence_folder)( + sequence_folder=e, + cells_per_meter=cells_per_meter, + meters_beyond_poly_edge=meters_beyond_poly_edge, + num_neighbors=num_neighbors) for e in tqdm(work_queue)) + +# with multiprocessing.Pool(multiprocessing.cpu_count()) as pool: +# pool.map( +# partial(process_sequence_folder, +# cells_per_meter=cells_per_meter, +# meters_beyond_poly_edge=meters_beyond_poly_edge, +# num_neighbors=num_neighbors), work_queue) diff --git a/data_prep_scripts/check_nsfp_status.py b/data_prep_scripts/check_nsfp_status.py new file mode 100644 index 0000000..1d7d397 --- /dev/null +++ b/data_prep_scripts/check_nsfp_status.py @@ -0,0 +1,22 @@ +from pathlib import Path +import argparse +import numpy as np + +# Path to the NSFP flows +parser = argparse.ArgumentParser() +parser.add_argument('nsfp_flows_dir', type=Path, help='path to the NSFP flows') +parser.add_argument('expected_num_flows', + type=int, + help='expected number of flows') +args = parser.parse_args() + +assert args.nsfp_flows_dir.exists(), 'NSFP flows dir does not exist' +assert args.expected_num_flows > 0, 'Expected number of flows must be positive' + +subfolders = sorted([f for f in args.nsfp_flows_dir.glob('*') if f.is_dir()]) + +# Check that all subfolders have the expected number of flows +for subfolder in subfolders: + num_flows = len([f for f in subfolder.glob('*.npz') if f.is_file()]) + if num_flows != args.expected_num_flows: + print(subfolder.name, num_flows) diff --git a/data_prep_scripts/compute_total_runtime.py b/data_prep_scripts/compute_total_runtime.py new file mode 100644 index 0000000..0bba999 --- /dev/null +++ b/data_prep_scripts/compute_total_runtime.py @@ -0,0 +1,58 @@ +import argparse +from loader_utils import load_npz +from pathlib import Path +import tqdm +from multiprocessing import Pool, cpu_count +from functools import partial +import numpy as np + +parser = argparse.ArgumentParser(description='Process some folders.') +parser.add_argument('--cpus', type=int, default=cpu_count(), + help='Number of CPUs to use for multiprocessing') +parser.add_argument('--key', type=str, default="delta_time", + help='Key to use for measuring delta time') +args = parser.parse_args() + +cpus = args.cpus + + +argoverse_train_dir = Path('/efs/argoverse2/train_nsfp_flow') +argoverse_val_dir = Path('/efs/argoverse2/val_nsfp_flow') +argoverse_chodosh_train_dir = Path('/efs/argoverse2/train_chodosh_flow/') +waymo_open_train_dir = Path('/efs/waymo_open_processed_flow/train_nsfp_flow') + + +def process_file(file, key="delta_time"): + npz_data = dict(load_npz(file, verbose=False)) + return npz_data[key] + +def process_folder(path : Path): + time_list = [] + + process_function = partial(process_file, key=args.key) + if cpus <= 1: + # Use a for loop instead of multiprocessing + subfolders = path.iterdir() + file_list = [subfolder.glob("*.npz") for subfolder in subfolders] + file_list = [file for sublist in file_list for file in sublist] + for file in tqdm.tqdm(file_list): + delta_time = process_function(file) + time_list.append(delta_time) + + else: + with Pool(processes=cpus) as pool: + subfolders = path.iterdir() + file_list = [subfolder.glob("*.npz") for subfolder in subfolders] + delta_time_list = pool.map(process_function, [file for sublist in file_list for file in sublist]) + time_list.extend(delta_time_list) + return time_list + + +# waymo_train = process_folder(waymo_open_train_dir) +# print("Waymo train: ", waymo_train) +# argo_train = process_folder(argoverse_train_dir) +# print("Argo train: ", argo_train) +# argo_val = process_folder(argoverse_val_dir) +# print("Argo val: ", argo_val) +argo_chodosh_train = process_folder(argoverse_chodosh_train_dir) +print("Argo chodosh train: ", np.sum(argo_chodosh_train), np.mean(argo_chodosh_train), np.std(argo_chodosh_train)) \ No newline at end of file diff --git a/data_prep_scripts/investigate_pointclouds.py b/data_prep_scripts/investigate_pointclouds.py new file mode 100644 index 0000000..8d527bc --- /dev/null +++ b/data_prep_scripts/investigate_pointclouds.py @@ -0,0 +1,153 @@ +import os +# set OMP_NUM_THREADS before importing torch or numpy +os.environ["OMP_NUM_THREADS"] = "1" +import argparse +from pathlib import Path +import numpy as np +import tqdm +import joblib +import multiprocessing +from loader_utils import save_pickle +import open3d as o3d +from pointclouds import PointCloud +import torch + +from dataloaders import ArgoverseSupervisedFlowSequenceLoader, ArgoverseUnsupervisedFlowSequenceLoader, WaymoSupervisedFlowSequenceLoader, WaymoUnsupervisedFlowSequenceLoader + +# cli argument to pick argo or waymo +parser = argparse.ArgumentParser() +parser.add_argument('dataset', + type=str, + help='argo or waymo', + choices=['argo', 'waymo']) +parser.add_argument('--cpus', + type=int, + help='number of cpus to use', + default=multiprocessing.cpu_count() - 1) +args = parser.parse_args() + +argoverse_data = Path("/efs/argoverse2") +waymo_data = Path("/efs/waymo_open_processed_flow") + +SLOW_THRESHOLD = 0.4 +FAST_THRESHOLD = 1.0 + + +def get_pc_infos(seq_idx, supervised_seq, unsupervised_seq): + infos_lst = [] + max_range = min(len(supervised_seq), len(unsupervised_seq)) + for idx in range(max_range): + if idx == max_range - 1: + continue + supervised_frame = supervised_seq.load(idx, idx) + unsupervised_frame = unsupervised_seq.load(idx, idx) + + assert ( + supervised_frame['relative_pc'] == + unsupervised_frame['relative_pc'] + ), f"supervised_frame['relative_pc'] {supervised_frame['relative_pc']} != unsupervised_frame['relative_pc'] {unsupervised_frame['relative_pc']}" + + valid_indices = unsupervised_frame['valid_idxes'] + classes = supervised_frame['pc_classes'][valid_indices] + pc = supervised_frame['relative_pc'][valid_indices] + pseudoflow = unsupervised_frame['flow'] + if pseudoflow.ndim == 3 and pseudoflow.shape[0] == 1: + pseudoflow = pseudoflow.squeeze(0) + + assert len(valid_indices) == len( + pseudoflow + ), f"len(valid_indices) {len(valid_indices)} != len(pseudoflow) {len(pseudoflow)}" + + pseudoflow_speed_meters_per_second = np.linalg.norm(pseudoflow, + axis=1) * 10.0 + + assert len(pc) == len( + pseudoflow_speed_meters_per_second + ), f"len(pc) {len(pc)} != len(pseudoflow_speed) {len(pseudoflow_speed_meters_per_second)}" + + slow_point_count = np.sum( + pseudoflow_speed_meters_per_second <= SLOW_THRESHOLD) + medium_point_count = np.sum( + (pseudoflow_speed_meters_per_second > SLOW_THRESHOLD) + & (pseudoflow_speed_meters_per_second < FAST_THRESHOLD)) + fast_point_count = np.sum( + pseudoflow_speed_meters_per_second >= FAST_THRESHOLD) + + foreground_count = np.sum(classes >= 0) + background_count = np.sum(classes == -1) + assert foreground_count + background_count == len( + pc + ), f"foreground_count {foreground_count} + background_count {background_count} != len(pc) {len(pc)}" + + num_points = len(pc) + if num_points < 100: + print("Seq idx", seq_idx, "idx", idx, "num_points", num_points) + infos_lst.append({ + "num_points": num_points, + "foreground_count": foreground_count, + "background_count": background_count, + "slow_point_count": slow_point_count, + "medium_point_count": medium_point_count, + "fast_point_count": fast_point_count, + }) + print("Finished sequence", + seq_idx, + "with", + len(supervised_seq), + "frames", + flush=True) + return infos_lst + + +if args.dataset == 'waymo': + supervised_seq_loader = WaymoSupervisedFlowSequenceLoader(waymo_data / + "training") + unsupervised_seq_loader = WaymoUnsupervisedFlowSequenceLoader( + waymo_data / "training", waymo_data / "train_nsfp_flow") +else: + supervised_seq_loader = ArgoverseSupervisedFlowSequenceLoader( + argoverse_data / "val", argoverse_data / "val_sceneflow") + unsupervised_seq_loader = ArgoverseUnsupervisedFlowSequenceLoader( + argoverse_data / "val", argoverse_data / "val_nsfp_flow") +seq_ids = supervised_seq_loader.get_sequence_ids() + +if args.cpus > 1: + # Use joblib to parallelize the loading of the sequences + pc_infos_lst = joblib.Parallel(n_jobs=args.cpus)( + joblib.delayed(get_pc_infos)( + idx, supervised_seq_loader.load_sequence(seq_id), + unsupervised_seq_loader.load_sequence(seq_id)) + for idx, seq_id in enumerate(seq_ids)) + pc_infos_lst = [item for sublist in pc_infos_lst for item in sublist] +else: + pc_infos_lst = [] + for idx, seq_id in enumerate(tqdm.tqdm(seq_ids)): + supervised_seq = supervised_seq_loader.load_sequence(seq_id) + unsupervised_seq = unsupervised_seq_loader.load_sequence(seq_id) + pc_infos = get_pc_infos(idx, supervised_seq, unsupervised_seq) + pc_infos_lst.extend(pc_infos) + +save_pickle( + f"validation_results/{args.dataset}_validation_pointcloud_pc_infos.pkl", + pc_infos_lst) + +pc_sizes_lst = [e["num_points"] for e in pc_infos_lst] +foreground_counts = [e["foreground_count"] for e in pc_infos_lst] +background_counts = [e["background_count"] for e in pc_infos_lst] + +slow_counts = [e["slow_point_count"] for e in pc_infos_lst] +medium_counts = [e["medium_point_count"] for e in pc_infos_lst] +fast_counts = [e["fast_point_count"] for e in pc_infos_lst] + +print("PC Point Count Mean:", np.mean(pc_sizes_lst), "Std:", + np.std(pc_sizes_lst), "Median:", np.median(pc_sizes_lst)) +print("Foreground Mean:", np.mean(foreground_counts), "Std:", + np.std(foreground_counts), "Median:", np.median(foreground_counts)) +print("Background Mean:", np.mean(background_counts), "Std:", + np.std(background_counts), "Median:", np.median(background_counts)) +print("Slow Mean:", np.mean(slow_counts), "Std:", np.std(slow_counts), + "Median:", np.median(slow_counts)) +print("Medium Mean:", np.mean(medium_counts), "Std:", np.std(medium_counts), + "Median:", np.median(medium_counts)) +print("Fast Mean:", np.mean(fast_counts), "Std:", np.std(fast_counts), + "Median:", np.median(fast_counts)) diff --git a/data_prep_scripts/split_nsfp_jobs.py b/data_prep_scripts/split_nsfp_jobs.py new file mode 100644 index 0000000..8311528 --- /dev/null +++ b/data_prep_scripts/split_nsfp_jobs.py @@ -0,0 +1,118 @@ +import argparse +from pathlib import Path +import math +import shutil + +# Get path to argoverse lidar dataset and number of sequences per job +parser = argparse.ArgumentParser() +parser.add_argument('lidar_path', type=Path) +parser.add_argument('sequences_per_job', type=int) +parser.add_argument('base_config', type=Path) +parser.add_argument('--reset_config_dir', action='store_true') +parser.add_argument('--configs_path', + type=Path, + default=Path("./nsfp_split_configs")) +parser.add_argument('--runtime_mins', type=int, default=180) +parser.add_argument('--job_prefix', type=str, default='np') +args = parser.parse_args() + +assert args.lidar_path.is_dir(), f"Path {args.lidar_path} is not a directory" +assert args.sequences_per_job > 0, f"Number of sequences per job must be positive" +assert args.base_config.is_file( +), f"Config file {args.base_config} does not exist" + +assert args.runtime_mins > 0, f"Runtime must be positive" + +sequence_folders = sorted([c for c in args.lidar_path.glob("*") if c.is_dir()], + key=lambda x: x.name.lower()) +num_jobs = math.ceil(len(sequence_folders) / args.sequences_per_job) + +print(f"Splitting {len(sequence_folders)} sequences into {num_jobs} jobs") + +job_sequence_names_lst = [] +for i in range(num_jobs): + start = i * args.sequences_per_job + end = min(start + args.sequences_per_job, len(sequence_folders)) + job_sequence_folders = sequence_folders[start:end] + job_sequence_names = [f.name for f in job_sequence_folders] + job_sequence_names_lst.append(job_sequence_names) + +sequence_names_set = set(f for seqs in job_sequence_names_lst for f in seqs) +assert len(sequence_names_set) == len( + sequence_folders), "Some sequences are missing from jobs" + +configs_path = args.configs_path +if args.reset_config_dir: + if configs_path.exists(): + shutil.rmtree(configs_path) + configs_path.mkdir(exist_ok=False) +else: + configs_path.mkdir(exist_ok=True, parents=True) + + +def get_runtime_format(runtime_mins): + hours = runtime_mins // 60 + minutes = runtime_mins % 60 + return f"{hours:02d}:{minutes:02d}:00" + + +def make_config(i, job_sequence_names): + config_path = configs_path / f"nsfp_split_{i:06d}.py" + config_file_content = f""" +_base_ = '../{args.base_config}' +test_loader = dict(args=dict(log_subset={job_sequence_names})) +""" + with open(config_path, "w") as f: + f.write(config_file_content) + + +def make_srun(i): + srun_path = configs_path / f"srun_{i:06d}.sh" + docker_image_path = Path( + "kylevedder_offline_sceneflow_latest.sqsh").absolute() + assert docker_image_path.is_file( + ), f"Docker image {docker_image_path} squash file does not exist" + srun_file_content = f"""#!/bin/bash +srun --gpus=1 --mem-per-gpu=12G --cpus-per-gpu=2 --time={get_runtime_format(args.runtime_mins)} --exclude=kd-2080ti-2.grasp.maas --job-name={args.job_prefix}{i:06d} --container-mounts=../../datasets/:/efs/,`pwd`:/project --container-image={docker_image_path} bash -c "python test_pl.py {configs_path}/nsfp_split_{i:06d}.py; echo 'done' > {configs_path}/nsfp_{i:06d}.done" +""" + with open(srun_path, "w") as f: + f.write(srun_file_content) + + +def make_screen(i): + screen_path = configs_path / f"screen_{i:06d}.sh" + screen_file_content = f"""#!/bin/bash +rm -f {configs_path}/nsfp_{i:06d}.out; + screen -L -Logfile {configs_path}/nsfp_{i:06d}.out -dmS nsfp_{i:06d} bash {configs_path}/srun_{i:06d}.sh +""" + with open(screen_path, "w") as f: + f.write(screen_file_content) + + +def make_runall(): + runall_path = configs_path / f"runall.sh" + runall_file_content = f"""#!/bin/bash +for i in {configs_path / "screen_*.sh"}; do + DONEFILE=nsfp_split_configs/nsfp_${{i:26:6}}.done + if [ -f $DONEFILE ]; then + echo "Skipping already completed0: $i" + continue + fi + echo "Launching: $i" + bash $i +done +""" + with open(runall_path, "w") as f: + f.write(runall_file_content) + + +for i, job_sequence_names in enumerate(job_sequence_names_lst): + make_config(i, job_sequence_names) + make_srun(i) + make_screen(i) + if i % 100 == 0: + print(f"Made configs for {i} jobs") + +make_runall() + +print(f"Config files written to {configs_path.absolute()}") diff --git a/data_prep_scripts/split_nsfp_jobs_sbatch.py b/data_prep_scripts/split_nsfp_jobs_sbatch.py new file mode 100644 index 0000000..9b019a2 --- /dev/null +++ b/data_prep_scripts/split_nsfp_jobs_sbatch.py @@ -0,0 +1,140 @@ +import argparse +from pathlib import Path +import math +import shutil + +# Get path to argoverse lidar dataset and number of sequences per job +parser = argparse.ArgumentParser() +parser.add_argument('lidar_path', type=Path) +parser.add_argument('sequences_per_job', type=int) +parser.add_argument('base_config', type=Path) +parser.add_argument('--reset_config_dir', action='store_true') +parser.add_argument('--configs_path', + type=Path, + default=Path("./nsfp_split_configs")) +parser.add_argument('--runtime_mins', type=int, default=180) +parser.add_argument('--job_prefix', type=str, default='nsfp') +parser.add_argument('--elevated', action='store_true') +parser.add_argument('--job_subset_file', type=Path, default=None) +args = parser.parse_args() + +assert args.lidar_path.is_dir(), f"Path {args.lidar_path} is not a directory" +assert args.sequences_per_job > 0, f"Number of sequences per job must be positive" +assert args.base_config.is_file( +), f"Config file {args.base_config} does not exist" + +assert args.runtime_mins > 0, f"Runtime must be positive" + +sequence_folders = sorted([c for c in args.lidar_path.glob("*") if c.is_dir()], + key=lambda x: x.name.lower()) +num_jobs = math.ceil(len(sequence_folders) / args.sequences_per_job) + +print(f"Splitting {len(sequence_folders)} sequences into {num_jobs} jobs") + +job_sequence_names_lst = [] +for i in range(num_jobs): + start = i * args.sequences_per_job + end = min(start + args.sequences_per_job, len(sequence_folders)) + job_sequence_folders = sequence_folders[start:end] + job_sequence_names = [f.name for f in job_sequence_folders] + job_sequence_names_lst.append(job_sequence_names) + +if args.job_subset_file is not None: + assert args.job_subset_file.is_file( + ), f"Job subset file {args.job_subset_file} does not exist" + print(f"Using job subset file {args.job_subset_file} to filter jobs") + + def parse_line(line: str) -> str: + if ' ' not in line: + return line.strip() + # Get the first column + return line.split(' ')[0].strip() + + with open(args.job_subset_file, "r") as f: + valid_jobs = [parse_line(l) for l in f.readlines()] + valid_jobs = set(valid_jobs) + + # Filter out jobs that are not in the subset + job_sequence_names_lst = [[job for job in sequence if job in valid_jobs] + for sequence in job_sequence_names_lst] + + # Filter out empty sequences + job_sequence_names_lst = [ + sequence for sequence in job_sequence_names_lst if len(sequence) > 0 + ] + + print(f"Filtered to {len(job_sequence_names_lst)} jobs") + +# sequence_names_set = set(f for seqs in job_sequence_names_lst for f in seqs) +# assert len(sequence_names_set) == len( +# sequence_folders), "Some sequences are missing from jobs" + +configs_path = args.configs_path +if args.reset_config_dir: + if configs_path.exists(): + shutil.rmtree(configs_path) + configs_path.mkdir(exist_ok=False) +else: + configs_path.mkdir(exist_ok=True, parents=True) + + +def get_runtime_format(runtime_mins): + hours = runtime_mins // 60 + minutes = runtime_mins % 60 + return f"{hours:02d}:{minutes:02d}:00" + + +def make_config(i, job_sequence_names): + config_path = configs_path / f"nsfp_split_{i:06d}.py" + config_file_content = f""" +_base_ = '../{args.base_config}' +test_loader = dict(args=dict(log_subset={job_sequence_names})) +""" + with open(config_path, "w") as f: + f.write(config_file_content) + + +def make_sbatch(): + current_working_dir = Path.cwd().absolute() + sbatch_path = configs_path / f"sbatch.bash" + if args.elevated: + qos = "ee-med" + partition = "eaton-compute" + else: + qos = "batch" + partition = "batch" + sbatch_file_content = f"""#!/bin/bash +#SBATCH --job-name={args.job_prefix} +#SBATCH --output={configs_path}/nsfp_%a.out +#SBATCH --error={configs_path}/nsfp_%a.err +#SBATCH --time={get_runtime_format(args.runtime_mins)}""" + if args.elevated: + sbatch_file_content += f""" +#SBATCH --qos={qos} +#SBATCH --partition={partition}""" + sbatch_file_content += f""" +#SBATCH --gpus=1 +#SBATCH --mem-per-gpu=12G +#SBATCH --cpus-per-gpu=2 +#SBATCH --exclude=kd-2080ti-2.grasp.maas +#SBATCH --array=0-{len(job_sequence_names_lst) - 1} +#SBATCH --container-mounts=../../datasets/:/efs/,{current_working_dir}:/project +#SBATCH --container-image={current_working_dir / "kylevedder_offline_sceneflow_latest.sqsh"} + +echo "Running job $SLURM_ARRAY_TASK_ID on $HOSTNAME" +export MY_RUN_ID=$(printf "%06d" $SLURM_ARRAY_TASK_ID) +python test_pl.py {configs_path}/nsfp_split_$MY_RUN_ID.py; echo 'done' > {configs_path}/nsfp_$MY_RUN_ID.done +""" + with open(sbatch_path, "w") as f: + f.write(sbatch_file_content) + + +for i, job_sequence_names in enumerate(job_sequence_names_lst): + make_config(i, job_sequence_names) + if i % 100 == 0: + print(f"Made configs for {i} jobs") +print(f"Made total of {len(job_sequence_names_lst)} jobs") + +make_sbatch() + +print(f"Config files written to {configs_path.absolute()}") diff --git a/data_prep_scripts/waymo/check_waymo_data.py b/data_prep_scripts/waymo/check_waymo_data.py new file mode 100644 index 0000000..4679cfd --- /dev/null +++ b/data_prep_scripts/waymo/check_waymo_data.py @@ -0,0 +1,54 @@ +import argparse +from pathlib import Path + +# Get path to the Waymo dataset + +parser = argparse.ArgumentParser() +parser.add_argument('waymo_dir', type=Path, help='path to the Waymo dataset') +args = parser.parse_args() + +training_dir = args.waymo_dir / 'training' +validation_dir = args.waymo_dir / 'validation' + +flow_label_dir = args.waymo_dir / 'train_nsfp_flow' + +assert training_dir.exists(), 'Training directory does not exist' +assert validation_dir.exists(), 'Validation directory does not exist' + +training_subfolders = sorted(training_dir.glob('segment*')) +validation_subfolders = sorted(validation_dir.glob('segment*')) + +flow_subfolders = sorted(flow_label_dir.glob('segment*')) + +assert len( + training_subfolders +) == 798, f'Expected 798 training subfolders, got {len(training_subfolders)}' +assert len( + validation_subfolders +) == 202, f'Expected 202 validation subfolders, got {len(validation_subfolders)}' + +# Check that all subfolders have the expected number of flows +for subfolder in training_subfolders: + num_flows = sorted(subfolder.glob('*.pkl')) + if len(num_flows) < 196: + print(str(subfolder), len(num_flows)) + +for subfolder in validation_subfolders: + num_flows = sorted(subfolder.glob('*.pkl')) + if len(num_flows) < 196: + print(str(subfolder), len(num_flows)) + +train_subfolder_names = [subfolder.name for subfolder in training_subfolders] +for subfolder in flow_subfolders: + num_flows = sorted(subfolder.glob('*.npz')) + if len(num_flows) < 196: + idx = train_subfolder_names.index(subfolder.name) + print(subfolder.name, len(num_flows), idx) + +# Get all the folders in the training subset that aren't in flow subfolders +missing_training_subfolders = sorted( + set(training_subfolders) - set(flow_subfolders)) +print("+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++") +print("Missing training subfolders:") +for subfolder in missing_training_subfolders: + print(subfolder.name) diff --git a/data_prep_scripts/waymo/check_waymo_nsfp_flows.py b/data_prep_scripts/waymo/check_waymo_nsfp_flows.py new file mode 100644 index 0000000..cb1150a --- /dev/null +++ b/data_prep_scripts/waymo/check_waymo_nsfp_flows.py @@ -0,0 +1,49 @@ +import argparse +from pathlib import Path +from collections import defaultdict + +# Get path to the Waymo dataset + +parser = argparse.ArgumentParser() +parser.add_argument('waymo_dir', type=Path, help='path to the Waymo dataset') +args = parser.parse_args() + +training_dir = args.waymo_dir / 'training' + +flow_label_dir = args.waymo_dir / 'train_nsfp_flow' + +assert training_dir.exists(), 'Training directory does not exist' + +training_subfolders = sorted(training_dir.glob('segment*')) + +flow_subfolders = sorted(flow_label_dir.glob('segment*')) + +assert len( + training_subfolders +) == 798, f'Expected 798 training subfolders, got {len(training_subfolders)}' + +train_subfolder_names = [e.name for e in training_subfolders] + +flow_folder_lookup_dict = defaultdict(int) +lidar_folder_lookup_dict = defaultdict(int) +valid_flow_folders = [] +for flow_subfolder in flow_subfolders: + num_flows = len(sorted(flow_subfolder.glob('*.npz'))) + train_subfolder_idx = train_subfolder_names.index(flow_subfolder.name) + train_subfolder = training_subfolders[train_subfolder_idx] + num_lidar = len(sorted(train_subfolder.glob('*.pkl'))) + + flow_folder_lookup_dict[flow_subfolder.name] = num_flows + lidar_folder_lookup_dict[flow_subfolder.name] = num_lidar + if num_flows == (num_lidar - 1): + valid_flow_folders.append(flow_subfolder) + +valid_flow_folder_names = sorted([e.name for e in valid_flow_folders]) + +# Get all the folders in the training subset that aren't in flow subfolders +missing_folder_names = sorted( + set(train_subfolder_names) - set(valid_flow_folder_names)) + +for folder_name in missing_folder_names: + print(folder_name, flow_folder_lookup_dict[folder_name], + lidar_folder_lookup_dict[folder_name]) diff --git a/data_prep_scripts/waymo/extract_flow_and_remove_ground.py b/data_prep_scripts/waymo/extract_flow_and_remove_ground.py new file mode 100644 index 0000000..8540a4a --- /dev/null +++ b/data_prep_scripts/waymo/extract_flow_and_remove_ground.py @@ -0,0 +1,435 @@ +from sklearn.neighbors import NearestNeighbors +import os + +os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' +import tensorflow as tf +import argparse +import multiprocessing +from pathlib import Path +from joblib import Parallel, delayed +import numpy as np +from typing import Tuple, List, Dict +from pointclouds import PointCloud, SE3, SE2 +from loader_utils import load_json, save_pickle + +from waymo_open_dataset import dataset_pb2 +from waymo_open_dataset.utils import frame_utils +from waymo_open_dataset.utils import plot_maps + +import open3d as o3d + +GROUND_HEIGHT_THRESHOLD = 0.4 # 40 centimeters + +parser = argparse.ArgumentParser() +parser.add_argument("flow_directory", + type=Path, + help="Path to Waymo flow data directory.") +parser.add_argument("heightmap_directory", + type=Path, + help="Path to rasterized heightmap.") +parser.add_argument("save_directory", + type=Path, + help="Path to save directory.") +parser.add_argument('--cpus', + type=int, + default=multiprocessing.cpu_count(), + help="Number of cpus to use for parallel processing") + +args = parser.parse_args() + + +def parse_range_image_and_camera_projection(frame): + """ + Parse range images and camera projections given a frame. + + Args: + frame: open dataset frame proto + + Returns: + range_images: A dict of {laser_name, + [range_image_first_return, range_image_second_return]}. + camera_projections: A dict of {laser_name, + [camera_projection_from_first_return, + camera_projection_from_second_return]}. + range_image_top_pose: range image pixel pose for top lidar. + """ + range_images = {} + camera_projections = {} + point_flows = {} + range_image_top_pose = None + for laser in frame.lasers: + if len(laser.ri_return1.range_image_compressed) > 0: # pylint: disable=g-explicit-length-test + range_image_str_tensor = tf.io.decode_compressed( + laser.ri_return1.range_image_compressed, 'ZLIB') + ri = dataset_pb2.MatrixFloat() + ri.ParseFromString(bytearray(range_image_str_tensor.numpy())) + range_images[laser.name] = [ri] + + if len(laser.ri_return1.range_image_flow_compressed) > 0: + range_image_flow_str_tensor = tf.io.decode_compressed( + laser.ri_return1.range_image_flow_compressed, 'ZLIB') + ri = dataset_pb2.MatrixFloat() + ri.ParseFromString( + bytearray(range_image_flow_str_tensor.numpy())) + point_flows[laser.name] = [ri] + + if laser.name == dataset_pb2.LaserName.TOP: + range_image_top_pose_str_tensor = tf.io.decode_compressed( + laser.ri_return1.range_image_pose_compressed, 'ZLIB') + range_image_top_pose = dataset_pb2.MatrixFloat() + range_image_top_pose.ParseFromString( + bytearray(range_image_top_pose_str_tensor.numpy())) + + camera_projection_str_tensor = tf.io.decode_compressed( + laser.ri_return1.camera_projection_compressed, 'ZLIB') + cp = dataset_pb2.MatrixInt32() + cp.ParseFromString(bytearray(camera_projection_str_tensor.numpy())) + camera_projections[laser.name] = [cp] + if len(laser.ri_return2.range_image_compressed) > 0: # pylint: disable=g-explicit-length-test + range_image_str_tensor = tf.io.decode_compressed( + laser.ri_return2.range_image_compressed, 'ZLIB') + ri = dataset_pb2.MatrixFloat() + ri.ParseFromString(bytearray(range_image_str_tensor.numpy())) + range_images[laser.name].append(ri) + + if len(laser.ri_return2.range_image_flow_compressed) > 0: + range_image_flow_str_tensor = tf.io.decode_compressed( + laser.ri_return2.range_image_flow_compressed, 'ZLIB') + ri = dataset_pb2.MatrixFloat() + ri.ParseFromString( + bytearray(range_image_flow_str_tensor.numpy())) + point_flows[laser.name].append(ri) + + camera_projection_str_tensor = tf.io.decode_compressed( + laser.ri_return2.camera_projection_compressed, 'ZLIB') + cp = dataset_pb2.MatrixInt32() + cp.ParseFromString(bytearray(camera_projection_str_tensor.numpy())) + camera_projections[laser.name].append(cp) + return range_images, camera_projections, point_flows, range_image_top_pose + + +def convert_range_image_to_point_cloud(frame, + range_images, + camera_projections, + point_flows, + range_image_top_pose, + ri_index=0, + keep_polar_features=False): + """Convert range images to point cloud. + + Args: + frame: open dataset frame + range_images: A dict of {laser_name, [range_image_first_return, + range_image_second_return]}. + camera_projections: A dict of {laser_name, + [camera_projection_from_first_return, + camera_projection_from_second_return]}. + range_image_top_pose: range image pixel pose for top lidar. + ri_index: 0 for the first return, 1 for the second return. + keep_polar_features: If true, keep the features from the polar range image + (i.e. range, intensity, and elongation) as the first features in the + output range image. + + Returns: + points: {[N, 3]} list of 3d lidar points of length 5 (number of lidars). + (NOTE: Will be {[N, 6]} if keep_polar_features is true. + cp_points: {[N, 6]} list of camera projections of length 5 + (number of lidars). + """ + calibrations = sorted(frame.context.laser_calibrations, + key=lambda c: c.name) + points = [] + cp_points = [] + flows = [] + + cartesian_range_images = frame_utils.convert_range_image_to_cartesian( + frame, range_images, range_image_top_pose, ri_index, + keep_polar_features) + + for c in calibrations: + range_image = range_images[c.name][ri_index] + range_image_tensor = tf.reshape( + tf.convert_to_tensor(value=range_image.data), + range_image.shape.dims) + range_image_mask = range_image_tensor[..., 0] > 0 + + range_image_cartesian = cartesian_range_images[c.name] + points_tensor = tf.gather_nd(range_image_cartesian, + tf.compat.v1.where(range_image_mask)) + + flow = point_flows[c.name][ri_index] + flow_tensor = tf.reshape(tf.convert_to_tensor(value=flow.data), + flow.shape.dims) + flow_points_tensor = tf.gather_nd(flow_tensor, + tf.compat.v1.where(range_image_mask)) + + cp = camera_projections[c.name][ri_index] + cp_tensor = tf.reshape(tf.convert_to_tensor(value=cp.data), + cp.shape.dims) + cp_points_tensor = tf.gather_nd(cp_tensor, + tf.compat.v1.where(range_image_mask)) + + points.append(points_tensor.numpy()) + cp_points.append(cp_points_tensor.numpy()) + flows.append(flow_points_tensor.numpy()) + + return points, cp_points, flows + + +def get_pc(frame: dataset_pb2.Frame) -> np.ndarray: + # Parse the frame lidar data into range images. + range_images, camera_projections, seg_labels, range_image_top_poses = ( + frame_utils.parse_range_image_and_camera_projection(frame)) + # Project the range images into points. + points, cp_points = frame_utils.convert_range_image_to_point_cloud( + frame, + range_images, + camera_projections, + range_image_top_poses, + keep_polar_features=True, + ) + car_frame_pc = points[0][:, 3:] + num_points = car_frame_pc.shape[0] + print(f'num_points: {num_points}') + + # Transform the points from the vehicle frame to the world frame. + world_frame_pc = np.concatenate( + [car_frame_pc, np.ones([num_points, 1])], axis=-1) + car_to_global_transform = np.reshape(np.array(frame.pose.transform), + [4, 4]) + world_frame_pc = np.transpose( + np.matmul(car_to_global_transform, np.transpose(world_frame_pc)))[:, + 0:3] + + # Transform the points from the world frame to the map frame. + offset = frame.map_pose_offset + points_offset = np.array([offset.x, offset.y, offset.z]) + world_frame_pc += points_offset + + return car_frame_pc, world_frame_pc + + +def get_car_pc_global_pc_flow_transform( + frame: dataset_pb2.Frame +) -> Tuple[PointCloud, PointCloud, np.ndarray, SE3]: + + # Parse the frame lidar data into range images. + range_images, camera_projections, point_flows, range_image_top_poses = parse_range_image_and_camera_projection( + frame) + + # Project the range images into points. + points_lst, cp_points, flows_lst = convert_range_image_to_point_cloud( + frame, + range_images, + camera_projections, + point_flows, + range_image_top_poses, + keep_polar_features=True) + + car_frame_pc = points_lst[0][:, 3:] + car_frame_flows = flows_lst[0][:, :3] + car_frame_labels = flows_lst[0][:, 3] + num_points = car_frame_pc.shape[0] + + # Transform the points from the vehicle frame to the world frame. + world_frame_pc = np.concatenate( + [car_frame_pc, np.ones([num_points, 1])], axis=-1) + car_to_global_transform = np.reshape(np.array(frame.pose.transform), + [4, 4]) + world_frame_pc = np.transpose( + np.matmul(car_to_global_transform, np.transpose(world_frame_pc)))[:, + 0:3] + + # Transform the points from the world frame to the map frame. + offset = frame.map_pose_offset + points_offset = np.array([offset.x, offset.y, offset.z]) + world_frame_pc += points_offset + return PointCloud(car_frame_pc), PointCloud( + world_frame_pc), car_frame_flows, car_frame_labels, SE3.from_array( + car_to_global_transform) + + +def flow_path_to_height_map_path(flow_path: Path): + heightmap_dir = args.heightmap_directory / flow_path.parent.name / ( + flow_path.stem + "_map") + assert heightmap_dir.is_dir(), f"{heightmap_dir} is not a directory" + return heightmap_dir + + +def flow_path_to_save_folder(flow_path: Path): + save_folder = args.save_directory / flow_path.parent.name / flow_path.stem + save_folder.mkdir(parents=True, exist_ok=True) + return save_folder + + +def load_ground_height_raster(map_path: Path): + raster_height_path = map_path / "ground_height.npy" + transform_path = map_path / "se2.json" + + raster_heightmap = np.load(raster_height_path) + transform = load_json(transform_path) + + transform_rotation = np.array(transform['R']).reshape(2, 2) + transform_translation = np.array(transform['t']) + transform_scale = np.array(transform['s']) + + transform_se2 = SE2(rotation=transform_rotation, + translation=transform_translation) + + return raster_heightmap, transform_se2, transform_scale + + +def get_ground_heights(raster_heightmap, global_to_raster_se2, + global_to_raster_scale, + global_point_cloud: PointCloud) -> np.ndarray: + """Get ground height for each of the xy locations in a point cloud. + Args: + point_cloud: Numpy array of shape (k,2) or (k,3) in global coordinates. + Returns: + ground_height_values: Numpy array of shape (k,) + """ + + global_points_xy = global_point_cloud.points[:, :2] + + raster_points_xy = global_to_raster_se2.transform_point_cloud( + global_points_xy) * global_to_raster_scale + + raster_points_xy = np.round(raster_points_xy).astype(np.int64) + + ground_height_values = np.full((raster_points_xy.shape[0]), np.nan) + # outside max X + outside_max_x = (raster_points_xy[:, 0] >= + raster_heightmap.shape[1]).astype(bool) + # outside max Y + outside_max_y = (raster_points_xy[:, 1] >= + raster_heightmap.shape[0]).astype(bool) + # outside min X + outside_min_x = (raster_points_xy[:, 0] < 0).astype(bool) + # outside min Y + outside_min_y = (raster_points_xy[:, 1] < 0).astype(bool) + ind_valid_pts = ~np.logical_or(np.logical_or(outside_max_x, outside_max_y), + np.logical_or(outside_min_x, outside_min_y)) + + ground_height_values[ind_valid_pts] = raster_heightmap[raster_points_xy[ + ind_valid_pts, 1], raster_points_xy[ind_valid_pts, 0]] + + return ground_height_values + + +def is_ground_points(raster_heightmap, global_to_raster_se2, + global_to_raster_scale, + global_point_cloud: PointCloud) -> np.ndarray: + """Remove ground points from a point cloud. + Args: + point_cloud: Numpy array of shape (k,3) in global coordinates. + Returns: + ground_removed_point_cloud: Numpy array of shape (k,3) in global coordinates. + """ + ground_height_values = get_ground_heights(raster_heightmap, + global_to_raster_se2, + global_to_raster_scale, + global_point_cloud) + is_ground_boolean_arr = ( + np.absolute(global_point_cloud[:, 2] - ground_height_values) <= + GROUND_HEIGHT_THRESHOLD) | ( + np.array(global_point_cloud[:, 2] - ground_height_values) < 0) + return is_ground_boolean_arr + + +def visualize_point_cloud_flow(point_cloud: PointCloud, flow: np.ndarray): + + print("Visualizing point cloud and flow") + # Use open3d to visualize the point cloud and flow. + + flowed_point_cloud = point_cloud.flow(flow[:, :3]) + + geometries = [] + + # Make base pointcloud + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(point_cloud) + geometries.append(pcd) + + # Add line set + line_set = o3d.geometry.LineSet() + line_set_points = np.concatenate( + [point_cloud.points, flowed_point_cloud.points], axis=0) + + lines = np.array([[i, i + len(point_cloud)] + for i in range(len(point_cloud))]) + line_set.points = o3d.utility.Vector3dVector(line_set_points) + line_set.lines = o3d.utility.Vector2iVector(lines) + # Line set is blue + line_set.colors = o3d.utility.Vector3dVector([[1, 0, 0] + for _ in range(len(lines))]) + geometries.append(line_set) + + # Visualize the pointcloud + o3d.visualization.draw_geometries(geometries) + + +def build_work_queue(waymo_directory): + waymo_directory = Path(waymo_directory) + assert waymo_directory.is_dir(), f"{waymo_directory} is not a directory" + + train_records = sorted((waymo_directory / 'training').glob('*.tfrecord')) + val_records = sorted((waymo_directory / 'validation').glob('*.tfrecord')) + + queue = train_records + val_records + for record in queue: + assert record.is_file(), f"{record} is not a file" + + return queue + + +def process_record(file_path: Path): + file_path = Path(file_path) + print("Processing", file_path) + height_map_path = flow_path_to_height_map_path(file_path) + save_folder = flow_path_to_save_folder(file_path) + raster_heightmap, transform_se2, transform_scale = load_ground_height_raster( + height_map_path) + + print("LOADING FILEPATH", file_path) + + dataset = tf.data.TFRecordDataset(file_path, compression_type='') + for idx, data in enumerate(dataset): + frame = dataset_pb2.Frame.FromString(bytearray(data.numpy())) + + car_frame_pc, global_frame_pc, flow, label, pose = get_car_pc_global_pc_flow_transform( + frame) + + if (flow == -1).all(): + continue + + keep_points_mask = ~is_ground_points(raster_heightmap, transform_se2, + transform_scale, global_frame_pc) + + masked_car_frame_pc = car_frame_pc.mask_points(keep_points_mask) + masked_flow = flow[keep_points_mask] / 10.0 + masked_label = label[keep_points_mask] + + # visualize_point_cloud_flow(masked_car_frame_pc, masked_flow) + + save_pickle( + save_folder / f"{idx:06d}.pkl", { + "car_frame_pc": masked_car_frame_pc.points, + "flow": masked_flow, + "label": masked_label, + "pose": pose.to_array(), + "fraction_kept": + np.sum(keep_points_mask) / len(keep_points_mask) + }) + + +work_queue = build_work_queue(args.flow_directory) +print("Work queue size:", len(work_queue)) + +assert args.cpus > 0, "Must have at least 1 CPU" +assert len(work_queue) > 0, "Work queue must have at least 1 element" + +# for idx, record in enumerate(work_queue): +# process_record(record) + +num_cores = min(args.cpus, len(work_queue)) +Parallel(num_cores)(delayed(process_record)(record) for record in work_queue) diff --git a/data_prep_scripts/waymo/extract_flow_and_remove_ground.sbatch b/data_prep_scripts/waymo/extract_flow_and_remove_ground.sbatch new file mode 100755 index 0000000..6373d92 --- /dev/null +++ b/data_prep_scripts/waymo/extract_flow_and_remove_ground.sbatch @@ -0,0 +1,15 @@ +#!/bin/bash + +## Resource Request +#SBATCH --job-name=waymo_remove_ground +#SBATCH --output=waymo_remove_ground.out +#SBATCH --qos=ee-med +#SBATCH --partition=eaton-compute +#SBATCH --cpus-per-task=120 +#SBATCH --mem-per-cpu=2GB +#SBATCH --time=04:00:00 +#SBATCH --container-mounts=../../datasets/:/efs/,./:/project +#SBATCH --container-image=/home/kvedder/kylevedder+offline_sceneflow_waymo+latest.sqsh + +## Job Steps +python extract_flow_and_remove_ground.py /efs/waymo_open_flow/ /efs/waymo_open_raster_heights/ /efs/waymo_open_processed_flow --cpus 120 \ No newline at end of file diff --git a/data_prep_scripts/waymo/investigate_pointclouds_waymo.sbatch b/data_prep_scripts/waymo/investigate_pointclouds_waymo.sbatch new file mode 100644 index 0000000..7a3c1f9 --- /dev/null +++ b/data_prep_scripts/waymo/investigate_pointclouds_waymo.sbatch @@ -0,0 +1,14 @@ +#!/bin/bash + +## Resource Request +#SBATCH --job-name=investigate_pointclouds_waymo +#SBATCH --output=investigate_pointclouds_waymo.out +#SBATCH --qos=ee-med +#SBATCH --partition=eaton-compute +#SBATCH --cpus-per-task=120 +#SBATCH --mem-per-cpu=2GB +#SBATCH --time=00:10:00 +#SBATCH --container-mounts=../../datasets/:/efs/,./:/project +#SBATCH --container-image=/home/kvedder/kylevedder+offline_sceneflow+latest.sqsh +## Job Steps +python data_prep_scripts/investigate_pointclouds.py waymo --cpus 120 \ No newline at end of file diff --git a/data_prep_scripts/waymo/old/data/preprocess.py b/data_prep_scripts/waymo/old/data/preprocess.py new file mode 100644 index 0000000..9f286f5 --- /dev/null +++ b/data_prep_scripts/waymo/old/data/preprocess.py @@ -0,0 +1,493 @@ +# Source code adapted from: +# https://github.com/waymo-research/waymo-open-dataset/blob/bbcd77fc503622a292f0928bfa455f190ca5946e/waymo_open_dataset/utils/frame_utils.py + +import glob +import os +import pickle +import re + +import cv2 +import numpy as np +import tensorflow as tf +from waymo_open_dataset import dataset_pb2, dataset_pb2 as open_dataset +from waymo_open_dataset.utils import frame_utils +import tqdm + + +def convert_range_image_to_point_cloud(frame, + range_images, + camera_projections, + point_flows, + range_image_top_pose, + ri_index=0, + keep_polar_features=False): + """Convert range images to point cloud. + + Args: + frame: open dataset frame + range_images: A dict of {laser_name, [range_image_first_return, + range_image_second_return]}. + camera_projections: A dict of {laser_name, + [camera_projection_from_first_return, + camera_projection_from_second_return]}. + range_image_top_pose: range image pixel pose for top lidar. + ri_index: 0 for the first return, 1 for the second return. + keep_polar_features: If true, keep the features from the polar range image + (i.e. range, intensity, and elongation) as the first features in the + output range image. + + Returns: + points: {[N, 3]} list of 3d lidar points of length 5 (number of lidars). + (NOTE: Will be {[N, 6]} if keep_polar_features is true. + cp_points: {[N, 6]} list of camera projections of length 5 + (number of lidars). + """ + calibrations = sorted(frame.context.laser_calibrations, key=lambda c: c.name) + points = [] + cp_points = [] + flows = [] + + cartesian_range_images = frame_utils.convert_range_image_to_cartesian( + frame, range_images, range_image_top_pose, ri_index, keep_polar_features) + + for c in calibrations: + range_image = range_images[c.name][ri_index] + range_image_tensor = tf.reshape( + tf.convert_to_tensor(value=range_image.data), range_image.shape.dims) + range_image_mask = range_image_tensor[..., 0] > 0 + + range_image_cartesian = cartesian_range_images[c.name] + points_tensor = tf.gather_nd(range_image_cartesian, + tf.compat.v1.where(range_image_mask)) + + flow = point_flows[c.name][ri_index] + flow_tensor = tf.reshape(tf.convert_to_tensor(value=flow.data), flow.shape.dims) + flow_points_tensor = tf.gather_nd(flow_tensor, + tf.compat.v1.where(range_image_mask)) + + cp = camera_projections[c.name][ri_index] + cp_tensor = tf.reshape(tf.convert_to_tensor(value=cp.data), cp.shape.dims) + cp_points_tensor = tf.gather_nd(cp_tensor, + tf.compat.v1.where(range_image_mask)) + + points.append(points_tensor.numpy()) + cp_points.append(cp_points_tensor.numpy()) + flows.append(flow_points_tensor.numpy()) + + return points, cp_points, flows + + +def parse_range_image_and_camera_projection(frame): + """ + Parse range images and camera projections given a frame. + + Args: + frame: open dataset frame proto + + Returns: + range_images: A dict of {laser_name, + [range_image_first_return, range_image_second_return]}. + camera_projections: A dict of {laser_name, + [camera_projection_from_first_return, + camera_projection_from_second_return]}. + range_image_top_pose: range image pixel pose for top lidar. + """ + range_images = {} + camera_projections = {} + point_flows = {} + range_image_top_pose = None + for laser in frame.lasers: + if len(laser.ri_return1.range_image_compressed) > 0: # pylint: disable=g-explicit-length-test + range_image_str_tensor = tf.io.decode_compressed( + laser.ri_return1.range_image_compressed, 'ZLIB') + ri = dataset_pb2.MatrixFloat() + ri.ParseFromString(bytearray(range_image_str_tensor.numpy())) + range_images[laser.name] = [ri] + + if len(laser.ri_return1.range_image_flow_compressed) > 0: + range_image_flow_str_tensor = tf.io.decode_compressed( + laser.ri_return1.range_image_flow_compressed, 'ZLIB') + ri = dataset_pb2.MatrixFloat() + ri.ParseFromString(bytearray(range_image_flow_str_tensor.numpy())) + point_flows[laser.name] = [ri] + + if laser.name == dataset_pb2.LaserName.TOP: + range_image_top_pose_str_tensor = tf.io.decode_compressed( + laser.ri_return1.range_image_pose_compressed, 'ZLIB') + range_image_top_pose = dataset_pb2.MatrixFloat() + range_image_top_pose.ParseFromString( + bytearray(range_image_top_pose_str_tensor.numpy())) + + camera_projection_str_tensor = tf.io.decode_compressed( + laser.ri_return1.camera_projection_compressed, 'ZLIB') + cp = dataset_pb2.MatrixInt32() + cp.ParseFromString(bytearray(camera_projection_str_tensor.numpy())) + camera_projections[laser.name] = [cp] + if len(laser.ri_return2.range_image_compressed) > 0: # pylint: disable=g-explicit-length-test + range_image_str_tensor = tf.io.decode_compressed( + laser.ri_return2.range_image_compressed, 'ZLIB') + ri = dataset_pb2.MatrixFloat() + ri.ParseFromString(bytearray(range_image_str_tensor.numpy())) + range_images[laser.name].append(ri) + + if len(laser.ri_return2.range_image_flow_compressed) > 0: + range_image_flow_str_tensor = tf.io.decode_compressed( + laser.ri_return2.range_image_flow_compressed, 'ZLIB') + ri = dataset_pb2.MatrixFloat() + ri.ParseFromString(bytearray(range_image_flow_str_tensor.numpy())) + point_flows[laser.name].append(ri) + + camera_projection_str_tensor = tf.io.decode_compressed( + laser.ri_return2.camera_projection_compressed, 'ZLIB') + cp = dataset_pb2.MatrixInt32() + cp.ParseFromString(bytearray(camera_projection_str_tensor.numpy())) + camera_projections[laser.name].append(cp) + return range_images, camera_projections, point_flows, range_image_top_pose + + +def save_point_cloud(compressed_frame, file_path): + """ + Compute the point cloud from a frame and stores it into disk. + :param compressed_frame: compressed frame from a TFRecord + :param file_path: name path that will have the stored point cloud + :returns: + - points - [N, 5] matrix which stores the [x, y, z, intensity, elongation] in the frame reference + - flows - [N, 4] matrix where each row is the flow for each point in the form [vx, vy, vz, label] + in the reference frame + - transform - [,16] flattened transformation matrix + """ + frame = get_uncompressed_frame(compressed_frame) + points, flows = compute_features(frame) + point_cloud = np.hstack((points, flows)) + np.savez_compressed(file_path, frame=point_cloud) + transform = list(frame.pose.transform) + return points, flows, transform + + +def preprocess(tfrecord_file, output_path, frames_per_segment = None): + """ + TFRecord file to store in a suitable form for training + in disk. A point cloud in disk has dimensions [N, 9] where N is the number of points + and per each point it stores [x, y, z, intensity, elongation, vx, vy, vz, label]. + It stores in a dictionary relevant metadata: + - look-up table: It has the form [[t_1, t_0], [t_2, t_1], ... , [t_n, t_(n-1)]], where t_i is + (file_path, transform), where file_path is the file where the point cloud is stored and transform the transformation + to apply to a point to change it reference frame from global to the car frame in that moment. + - flows information: min and max flows encountered for then visualize pointclouds properly + + :param tfrecord_file: TFRecord file. It should have the flow extension. + They can be downloaded from https://console.cloud.google.com/storage/browser/waymo_open_dataset_scene_flow + :param output_path: path where the processed point clouds will be saved. + """ + tfrecord_filename = os.path.basename(tfrecord_file) + tfrecord_filename = os.path.splitext(tfrecord_filename)[0] + + look_up_table = [] + metadata_path = os.path.join(output_path, f"metadata_{tfrecord_filename}") # look-up table and flows mins and maxs + loaded_file = tf.data.TFRecordDataset(tfrecord_file, compression_type='') + previous_frame = None + + # Needed of max and min flows for normalizing in visualization + min_vx_global, max_vx_global = np.inf, -np.inf + min_vy_global, max_vy_global = np.inf, -np.inf + min_vz_global, max_vz_global = np.inf, -np.inf + + for j, frame in enumerate(tqdm.tqdm(loaded_file)): + output_file_name = f"pointCloud_file_{tfrecord_filename}_frame_{j:06d}.npz" + point_cloud_path = os.path.join(output_path, output_file_name) + # Process frame and store point clouds into disk + _, flows, pose_transform = save_point_cloud(frame, point_cloud_path) + flows = flows[flows[:, -1] != -1] + if flows.size != 0: # May all the flows are invalid + min_vx, min_vy, min_vz = flows[:, :-1].min(axis=0) + max_vx, max_vy, max_vz = flows[:, :-1].max(axis=0) + min_vx_global = min(min_vx_global, min_vx) + min_vy_global = min(min_vy_global, min_vy) + min_vz_global = min(min_vz_global, min_vz) + max_vx_global = max(max_vx_global, max_vx) + max_vy_global = max(max_vy_global, max_vy) + max_vz_global = max(max_vz_global, max_vz) + + if j == 0: + previous_frame = (output_file_name, pose_transform) + else: + current_frame = (output_file_name, pose_transform) + look_up_table.append([current_frame, previous_frame]) + previous_frame = current_frame + if frames_per_segment is not None and j == frames_per_segment: + break + + # Save metadata into disk + flows_info = {'min_vx': min_vx_global, + 'max_vx': max_vx_global, + 'min_vy': min_vy_global, + 'max_vy': max_vy_global, + 'min_vz': min_vz_global, + 'max_vz': max_vz_global} + + metadata = {'look_up_table': look_up_table, + 'flows_information': flows_info} + + with open(metadata_path, 'wb') as metadata_file: + pickle.dump(metadata, metadata_file) + + +def get_uncompressed_frame(compressed_frame): + """ + :param compressed_frame: Compressed frame + :return: Uncompressed frame + """ + frame = open_dataset.Frame() + frame.ParseFromString(bytearray(compressed_frame.numpy())) + return frame + + +def compute_features(frame): + """ + :param frame: Uncompressed frame + :return: [N, F], [N, 4], where N is the number of points, F the number of features, + which is [x, y, z, intensity, elongation] and 4 in the second results stands for [vx, vy, vz, label], + which corresponds to the flow information + """ + range_images, camera_projections, point_flows, range_image_top_pose = parse_range_image_and_camera_projection( + frame) + + points, cp_points, flows = convert_range_image_to_point_cloud( + frame, + range_images, + camera_projections, + point_flows, + range_image_top_pose, + keep_polar_features=True) + + # 3D points in the vehicle reference frame + points_all = np.concatenate(points, axis=0) + flows_all = np.concatenate(flows, axis=0) + # We skip the range feature since pillars will account for it + # Note that first are features and then point coordinates + points_features, points_coord = points_all[:, 1:3], points_all[:, 3:points_all.shape[1]] + points_all = np.hstack((points_coord, points_features)) + return points_all, flows_all + + +def merge_metadata(input_path): + """ + Merge individual look-up table and flows mins and maxs and store it in the input_path with the name metadata + :param input_path: Path with the local metadata in the form metadata_[tfRecordName] + """ + look_up_table = [] + flows_info = None + + os.chdir(input_path) + for file in glob.glob("metadata_*"): + file_name = os.path.abspath(file) + try: + with open(file_name, 'rb') as metadata_file: + metadata_local = pickle.load(metadata_file) + look_up_table.extend(metadata_local['look_up_table']) + flows_information = metadata_local['flows_information'] + if flows_info is None: + flows_info = flows_information + else: + flows_info['min_vx'] = min(flows_info['min_vx'], flows_information['min_vx']) + flows_info['min_vx'] = min(flows_info['min_vx'], flows_information['min_vy']) + flows_info['min_vz'] = min(flows_info['min_vz'], flows_information['min_vz']) + flows_info['max_vx'] = max(flows_info['max_vx'], flows_information['max_vx']) + flows_info['max_vy'] = max(flows_info['max_vy'], flows_information['max_vy']) + flows_info['max_vz'] = max(flows_info['max_vz'], flows_information['max_vz']) + + except FileNotFoundError: + raise FileNotFoundError( + "Metadata not found when merging individual metadata") + + # Save metadata into disk + metadata = {'look_up_table': look_up_table, + 'flows_information': flows_info} + with open(os.path.join(input_path, "metadata"), 'wb') as metadata_file: + pickle.dump(metadata, metadata_file) + + +def load_pfm(filename): + file = open(filename, 'r', newline='', encoding='latin-1') + + header = file.readline().rstrip() + if header == 'PF': + color = True + elif header == 'Pf': + color = False + else: + raise Exception('Not a PFM file.') + + dim_match = re.match(r'^(\d+)\s(\d+)\s$', file.readline()) + if dim_match: + width, height = map(int, dim_match.groups()) + else: + raise Exception('Malformed PFM header.') + + scale = float(file.readline().rstrip()) + if scale < 0: # little-endian + endian = '<' + scale = -scale + else: + endian = '>' # big-endian + + data = np.fromfile(file, endian + 'f') + shape = (height, width, 3) if color else (height, width) + + file.close() + + return np.reshape(data, shape), scale + + +def bilinear_interp_val(vmap, y, x): + """ + bilinear interpolation on a 2D map + """ + h, w = vmap.shape + x1 = int(x) + x2 = x1 + 1 + x2 = w - 1 if x2 > (w - 1) else x2 + y1 = int(y) + y2 = y1 + 1 + y2 = h - 1 if y2 > (h - 1) else y2 + Q11 = vmap[y1, x1] + Q21 = vmap[y1, x2] + Q12 = vmap[y2, x1] + Q22 = vmap[y2, x2] + return Q11 * (x2 - x) * (y2 - y) + Q21 * (x - x1) * (y2 - y) + Q12 * (x2 - x) * (y - y1) + Q22 * (x - x1) * (y - y1) + + +def get_3d_pos_xy(y_prime, x_prime, depth, focal_length=1050., w=960, h=540): + """ depth pop up """ + y = (y_prime - h / 2.) * depth / focal_length + x = (x_prime - w / 2.) * depth / focal_length + return [x, y, depth] + + +def readFlow(name): + f = open(name, 'rb') + + header = f.read(4) + if header.decode("utf-8") != 'PIEH': + raise Exception('Flow file header does not contain PIEH') + + width = np.fromfile(f, np.int32, 1).squeeze() + height = np.fromfile(f, np.int32, 1).squeeze() + + flow = np.fromfile(f, np.float32, width * height * 2).reshape((height, width, 2)) + + return flow.astype(np.float32) + + +def generate_flying_things_point_cloud(fname_disparity, fname_disparity_next_frame, fname_disparity_change, + fname_optical_flow, image, image_next_frame, max_cut=35, focal_length=1050., + n = 2048, add_label=True): + # generate needed data + disparity_np, _ = load_pfm(fname_disparity) + disparity_next_frame_np, _ = load_pfm(fname_disparity_next_frame) + disparity_change_np, _ = load_pfm(fname_disparity_change) + optical_flow_np = readFlow(fname_optical_flow) + rgb_np = cv2.imread(image)[:, :, ::-1] / 255. + rgb_next_frame_np = cv2.imread(image_next_frame)[:, :, ::-1] / 255. + + depth_np = focal_length / disparity_np + depth_next_frame_np = focal_length / disparity_next_frame_np + future_depth_np = focal_length / (disparity_np + disparity_change_np) + h, w = disparity_np.shape + + try: + depth_requirement = depth_np < max_cut + except: + return None + + satisfy_pix1 = np.column_stack(np.where(depth_requirement)) + if satisfy_pix1.shape[0] < n: + return None + sample_choice1 = np.random.choice(satisfy_pix1.shape[0], size=n, replace=False) + sampled_pix1_x = satisfy_pix1[sample_choice1, 1] + sampled_pix1_y = satisfy_pix1[sample_choice1, 0] + n_1 = sampled_pix1_x.shape[0] + current_pos1 = np.array([get_3d_pos_xy(sampled_pix1_y[i], sampled_pix1_x[i], + depth_np[int(sampled_pix1_y[i]), + int(sampled_pix1_x[i])]) for i in range(n_1)]) + + current_rgb1 = np.array([[rgb_np[h - 1 - int(sampled_pix1_y[i]), int(sampled_pix1_x[i]), 0], + rgb_np[h - 1 - int(sampled_pix1_y[i]), int(sampled_pix1_x[i]), 1], + rgb_np[h - 1 - int(sampled_pix1_y[i]), int(sampled_pix1_x[i]), 2]] for i in range(n)]) + + # sampled_optical_flow_x = np.array( + # [optical_flow_np[int(sampled_pix1_y[i]), int(sampled_pix1_x[i])][0] for i in range(n_1)]) + # sampled_optical_flow_y = np.array( + # [optical_flow_np[int(sampled_pix1_y[i]), int(sampled_pix1_x[i])][1] for i in range(n_1)]) + # future_pix1_x = sampled_pix1_x + sampled_optical_flow_x + # future_pix1_y = sampled_pix1_y - sampled_optical_flow_y + # future_pos1 = np.array([get_3d_pos_xy(future_pix1_y[i], future_pix1_x[i], + # future_depth_np[int(sampled_pix1_y[i]), int(sampled_pix1_x[i])]) for i in + # range(n_1)]) + # flow = future_pos1 - current_pos1 + + try: + depth_requirement = depth_next_frame_np < max_cut + except: + return None + + satisfy_pix2 = np.column_stack(np.where(depth_next_frame_np < max_cut)) + if satisfy_pix2.shape[0] < n: + return None + sample_choice2 = np.random.choice(satisfy_pix2.shape[0], size=n, replace=False) + sampled_pix2_x = satisfy_pix2[sample_choice2, 1] + sampled_pix2_y = satisfy_pix2[sample_choice2, 0] + n_2 = sampled_pix2_x.shape[0] + + current_pos2 = np.array([get_3d_pos_xy(sampled_pix2_y[i], sampled_pix2_x[i], + depth_next_frame_np[int(sampled_pix2_y[i]), int(sampled_pix2_x[i])]) for i in + range(n_2)]) + + current_rgb2 = np.array([[rgb_next_frame_np[h-1-int(sampled_pix2_y[i]), int(sampled_pix2_x[i]), 0], + rgb_next_frame_np[h-1-int(sampled_pix2_y[i]), int(sampled_pix2_x[i]), 1], + rgb_next_frame_np[h-1-int(sampled_pix2_y[i]), int(sampled_pix2_x[i]), 2]] + for i in range(n)]) + + # Compute backward flow + sampled_optical_flow_x = np.array( + [optical_flow_np[int(sampled_pix2_y[i]), int(sampled_pix2_x[i])][0] for i in range(n_2)]) + sampled_optical_flow_y = np.array( + [optical_flow_np[int(sampled_pix2_y[i]), int(sampled_pix2_x[i])][1] for i in range(n_2)]) + future_pix2_x = sampled_pix2_x + sampled_optical_flow_x + future_pix2_y = sampled_pix2_y - sampled_optical_flow_y + future_pos2 = np.array([get_3d_pos_xy(future_pix2_y[i], future_pix2_x[i], + future_depth_np[int(sampled_pix2_y[i]), int(sampled_pix2_x[i])]) for i in + range(n_2)]) + flow = future_pos2 - current_pos2 + + # mask, judge whether point move out of fov or occluded by other object after motion + future_pos2_depth = future_depth_np[sampled_pix2_y, sampled_pix2_x] + future_pos2_foreground_depth = np.zeros_like(future_pos2_depth) + valid_mask_fov2 = np.ones_like(future_pos2_depth, dtype=bool) + for i in range(future_pos2_depth.shape[0]): + if 0 < future_pix2_y[i] < h and 0 < future_pix2_x[i] < w: + future_pos2_foreground_depth[i] = bilinear_interp_val(depth_next_frame_np, future_pix2_y[i], future_pix2_x[i]) + else: + valid_mask_fov2[i] = False + valid_mask_occ2 = (future_pos2_foreground_depth - future_pos2_depth) > -5e-1 + + mask2 = valid_mask_occ2 & valid_mask_fov2 + + mask2 = np.ones(shape=mask2.shape) + + # Add redundant zero labels for each flow in order to fit the shape + if add_label: + labelled_flow = np.ones(shape=(flow.shape[0], flow.shape[1] + 1)) + labelled_flow[:, :-1] = flow + flow = labelled_flow + + return current_pos1, current_pos2, current_rgb1, current_rgb2, flow, mask2 + + +def get_all_flying_things_frames(input_dir, disp_dir, opt_dir, disp_change_dir, img_dir): + all_files_disparity = glob.glob(os.path.join(input_dir, '{0}/*.pfm'.format(disp_dir))) + all_files_disparity_change = glob.glob(os.path.join(input_dir, '{0}/*.pfm'.format(disp_change_dir))) + all_files_opt_flow = glob.glob(os.path.join(input_dir, '{0}/*.flo'.format(opt_dir))) + all_files_img = glob.glob(os.path.join(input_dir, '{0}/*.png'.format(img_dir))) + + return all_files_disparity, all_files_disparity_change, all_files_opt_flow, all_files_img \ No newline at end of file diff --git a/data_prep_scripts/waymo/old/patchworkpp_read.py b/data_prep_scripts/waymo/old/patchworkpp_read.py new file mode 100644 index 0000000..c42f612 --- /dev/null +++ b/data_prep_scripts/waymo/old/patchworkpp_read.py @@ -0,0 +1,166 @@ +import sys + +sys.path.append('.') + +import sys +import open3d as o3d +import numpy as np +import pypatchworkpp + +from dataloaders import WaymoRawSequence, WaymoRawSequenceLoader +from pointclouds import PointCloud, SE3 +import numpy as np +import tqdm +import argparse +from pathlib import Path +import collections +import open3d as o3d +import sklearn + +from loader_utils import save_npy + +# Get path to waymo data from command line +parser = argparse.ArgumentParser() +parser.add_argument('path', type=Path) +args = parser.parse_args() + +params = pypatchworkpp.Parameters() +params.RNR_intensity_thr = 0.2 +params.RNR_ver_angle_thr = -15.0 +params.adaptive_seed_selection_margin = -1.2 +params.elevation_thr = [0.0, 0.0, 0.0, 0.0] +params.enable_RNR = True +params.enable_RVPF = True +params.enable_TGR = True +params.flatness_thr = [0.0, 0.0, 0.0, 0.0] +params.intensity_thr = -8.728169361654678e-12 +params.max_elevation_storage = 1000 +params.max_flatness_storage = 1000 +params.max_range = 80.0 +params.min_range = 1 +params.num_iter = 5 +params.num_lpr = 20 +params.num_min_pts = 10 +params.num_rings_each_zone = [2, 4, 4, 4] +params.num_rings_of_interest = 4 +params.num_sectors_each_zone = [16, 32, 54, 32] +params.num_zones = 4 +params.sensor_height = 1.9304 +params.th_dist = 0.125 +params.th_dist_v = 0.1 +params.th_seeds = 0.125 +params.th_seeds_v = 0.25 +params.uprightness_thr = 0.707 +params.verbose = True + +PatchworkPLUSPLUS = pypatchworkpp.patchworkpp(params) + + +def process_sequence(sequence: WaymoRawSequence): + for idx, entry_dict in enumerate(tqdm.tqdm(sequence.load_frame_list(None))): + pc = entry_dict['relative_pc'] + + pc_points = pc.points + + PatchworkPLUSPLUS.estimateGround(pc_points) + # Does not provide ground mask, only ground points. + ground_points = PatchworkPLUSPLUS.getGround() + # Requires a second step to get the mask by doing NN matching. + ground_point_idxes = sklearn.neighbors.NearestNeighbors( + n_neighbors=1).fit(pc_points).kneighbors( + ground_points, return_distance=False).squeeze(1) + ground_mask = np.zeros(pc_points.shape[0], dtype=np.bool) + ground_mask[ground_point_idxes] = True + + # Use open3d to visualize the pointcloud + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(pc) + # Use the ground mask to color the points + pcd.colors = o3d.utility.Vector3dVector( + np.array([[0, 0, 1], [1, 0, 0]])[ground_mask.astype(np.int)]) + + o3d.visualization.draw_geometries([pcd]) + + save_path = args.path / 'ground_removed' / sequence.sequence_name / f'{idx:06d}.npy' + save_npy(save_path, ground_point_idxes) + + +sequence_loader = WaymoRawSequenceLoader(args.path) +for sequence_id in sequence_loader.get_sequence_ids(): + process_sequence(sequence_loader.load_sequence(sequence_id)) + + +def read_bin(bin_path): + scan = np.fromfile(bin_path, dtype=np.float32) + scan = scan.reshape((-1, 4)) + + return scan + + +# if __name__ == "__main__": + +# # Patchwork++ initialization +# params = pypatchworkpp.Parameters() +# params.verbose = True + +# PatchworkPLUSPLUS = pypatchworkpp.patchworkpp(params) + +# # Load point cloud +# pointcloud = read_bin(input_cloud_filepath) + +# # Estimate Ground +# PatchworkPLUSPLUS.estimateGround(pointcloud) + +# # Get Ground and Nonground +# ground = PatchworkPLUSPLUS.getGround() +# nonground = PatchworkPLUSPLUS.getNonground() +# time_taken = PatchworkPLUSPLUS.getTimeTaken() + +# # Get centers and normals for patches +# centers = PatchworkPLUSPLUS.getCenters() +# normals = PatchworkPLUSPLUS.getNormals() + +# print("Origianl Points #: ", pointcloud.shape[0]) +# print("Ground Points #: ", ground.shape[0]) +# print("Nonground Points #: ", nonground.shape[0]) +# print("Time Taken : ", time_taken / 1000000, "(sec)") +# print("Press ... \n") +# print("\t H : help") +# print("\t N : visualize the surface normals") +# print("\tESC : close the Open3D window") + +# # Visualize +# vis = o3d.visualization.VisualizerWithKeyCallback() +# vis.create_window(width=600, height=400) + +# mesh = o3d.geometry.TriangleMesh.create_coordinate_frame() + +# ground_o3d = o3d.geometry.PointCloud() +# ground_o3d.points = o3d.utility.Vector3dVector(ground) +# ground_o3d.colors = o3d.utility.Vector3dVector( +# np.array([[0.0, 1.0, 0.0] for _ in range(ground.shape[0])], +# dtype=float) # RGB +# ) + +# nonground_o3d = o3d.geometry.PointCloud() +# nonground_o3d.points = o3d.utility.Vector3dVector(nonground) +# nonground_o3d.colors = o3d.utility.Vector3dVector( +# np.array([[1.0, 0.0, 0.0] for _ in range(nonground.shape[0])], +# dtype=float) #RGB +# ) + +# centers_o3d = o3d.geometry.PointCloud() +# centers_o3d.points = o3d.utility.Vector3dVector(centers) +# centers_o3d.normals = o3d.utility.Vector3dVector(normals) +# centers_o3d.colors = o3d.utility.Vector3dVector( +# np.array([[1.0, 1.0, 0.0] for _ in range(centers.shape[0])], +# dtype=float) #RGB +# ) + +# vis.add_geometry(mesh) +# vis.add_geometry(ground_o3d) +# vis.add_geometry(nonground_o3d) +# vis.add_geometry(centers_o3d) + +# vis.run() +# vis.destroy_window() \ No newline at end of file diff --git a/data_prep_scripts/waymo/old/read_raw_data.py b/data_prep_scripts/waymo/old/read_raw_data.py new file mode 100644 index 0000000..95e4336 --- /dev/null +++ b/data_prep_scripts/waymo/old/read_raw_data.py @@ -0,0 +1,23 @@ +import numpy as np +from loader_utils import load_npz +from pathlib import Path +from typing import List, Tuple, Dict, Any + +from dataloaders import WaymoSupervisedFlowSequenceLoader + + +root_path = Path('/efs/waymo_open_preprocessed/train') +sequence_loader = WaymoSupervisedFlowSequenceLoader(root_path) + +for sequence_id in sequence_loader.get_sequence_ids(): + sequence = sequence_loader.load_sequence(sequence_id) + for seq_idx in range(len(sequence)): + frame = sequence.load(seq_idx, seq_idx) + print(frame) + + + +# for (before_frame, after_frame) in look_up_table: +# before_frame = Frame(*before_frame) +# after_frame = Frame(*after_frame) +# breakpoint() \ No newline at end of file diff --git a/data_prep_scripts/waymo/old/remove_ground.py b/data_prep_scripts/waymo/old/remove_ground.py new file mode 100644 index 0000000..af6dc3c --- /dev/null +++ b/data_prep_scripts/waymo/old/remove_ground.py @@ -0,0 +1,133 @@ +import sys + +sys.path.append('.') + +import torch + +from loader_utils import save_npy + +from dataloaders import WaymoRawSequence, WaymoRawSequenceLoader +from pointclouds import PointCloud, SE3 +import numpy as np +import tqdm +import argparse +from pathlib import Path +import collections +import open3d as o3d +import copy + +# Get path to waymo data from command line +parser = argparse.ArgumentParser() +parser.add_argument('path', type=Path) +args = parser.parse_args() + +assert args.path.exists(), f"Path {args.path} does not exist" + + +class GroundRemoverNet(torch.nn.Module): + """ + 3 layer coordinate network of 16 hidden units each + """ + + def __init__(self): + super().__init__() + self.net = torch.nn.Sequential( + torch.nn.Linear(2, 16), + torch.nn.ReLU(), + torch.nn.Linear(16, 16), + torch.nn.ReLU(), + torch.nn.Linear(16, 16), + torch.nn.ReLU(), + torch.nn.Linear(16, 16), + torch.nn.ReLU(), + torch.nn.Linear(16, 1), + ) + + self.huber_loss = torch.nn.HuberLoss(delta=0.1) + + def forward(self, pc_points: torch.Tensor): + """ + Args: + pc: PointCloud of shape (N, 3) + Returns: + logits: Tensor of shape (N, 1) + """ + xy_points = pc_points[:, :2] + res = self.net(xy_points) + return res.squeeze(1) + + def loss(self, h, z, c=0.3): + case_1 = (h - z)**2 + case_2 = self.huber_loss(h, z) + case_3 = torch.ones_like(h) * c + + case_1_mask = (z < h) + case_2_mask = ((h <= z) & (z < h + c)) + + return torch.mean( + torch.where(case_1_mask, case_1, + torch.where(case_2_mask, case_2, case_3))) + + +def process_sequence(sequence: WaymoRawSequence): + + for idx, entry_dict in enumerate(tqdm.tqdm( + sequence.load_frame_list(None))): + pc = entry_dict['relative_pc'] + pose = entry_dict['relative_pose'] + + net = GroundRemoverNet().cuda() + # setup gradient descent optimizer + optimizer = torch.optim.Adam(net.parameters(), + lr=0.01, + weight_decay=0.0) + schedule = torch.optim.lr_scheduler.StepLR(optimizer, + step_size=750, + gamma=0.5) + max_iterations = 5000 + best_loss = None + best_params = None + best_idx = None + + c = 0.3 + + pc = torch.from_numpy(pc.points).float().cuda() + for i in tqdm.tqdm(range(max_iterations)): + optimizer.zero_grad() + logits = net(pc) + loss = net.loss(logits, pc[:, 2], c=c) + loss.backward() + optimizer.step() + schedule.step() + + loss_item = loss.item() + + if best_loss is None or loss_item < best_loss: + best_loss = loss_item + best_params = copy.deepcopy(net.state_dict()) + best_idx = i + + net.load_state_dict(best_params) + print("Best idx:", best_idx) + + # Points where z <= net(x, y) + c are ground + z_points = pc[:, 2] + ground_mask = ((net(pc) + c) >= z_points) + + # Use open3d to visualize the pointcloud + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(pc.cpu().numpy()) + # Use the ground mask to color the points + pcd.colors = o3d.utility.Vector3dVector( + np.array([[1, 0, 0], + [0, 1, 0]])[ground_mask.cpu().numpy().astype(np.int)]) + + o3d.visualization.draw_geometries([pcd]) + + save_path = args.path / 'ground_removed' / sequence.sequence_name / f'{idx:06d}.npy' + save_npy(save_path, ground_mask.cpu().numpy()) + + +sequence_loader = WaymoRawSequenceLoader(args.path) +for sequence_id in sequence_loader.get_sequence_ids(): + process_sequence(sequence_loader.load_sequence(sequence_id)) diff --git a/data_prep_scripts/waymo/old/tf_record_to_npz.py b/data_prep_scripts/waymo/old/tf_record_to_npz.py new file mode 100644 index 0000000..419481e --- /dev/null +++ b/data_prep_scripts/waymo/old/tf_record_to_npz.py @@ -0,0 +1,89 @@ +import glob +import multiprocessing as mp +import os +import time +from argparse import ArgumentParser +from pathlib import Path +import tensorflow as tf +import numpy as np +from tqdm import tqdm + +from data.preprocess import generate_flying_things_point_cloud, get_all_flying_things_frames +from data.preprocess import preprocess, merge_metadata + +global output_directory + + +def preprocess_wrap(tfrecord_file): + preprocess(tfrecord_file, output_directory, frames_per_segment=None) + + +# https://github.com/tqdm/tqdm/issues/484 +if __name__ == '__main__': + tf.config.set_visible_devices([], 'GPU') + visible_devices = tf.config.get_visible_devices() + for device in visible_devices: + assert device.device_type != 'GPU' + parser = ArgumentParser() + parser.add_argument('input_directory', type=str) + parser.add_argument('output_directory', type=str) + parser.add_argument('--n_cores', default=None, type=int) + args = parser.parse_args() + + print(f"Extracting frames from {args.input_directory} to {args.output_directory}") + + input_directory = Path(args.input_directory) + if not input_directory.exists() or not input_directory.is_dir(): + print("Input directory does not exist") + exit(1) + + output_directory = Path(args.output_directory) + if not output_directory.exists(): + output_directory.mkdir(parents=True) + #if list(output_directory.iterdir()): + # print("Output directory not empty! Please remove existing files as there is no merge.") + # exit(1) + output_directory = os.path.abspath(output_directory) + + n_cores = mp.cpu_count() + if args.n_cores is not None: + if args.n_cores <= 0: + print("Number of cores cannot be negative") + exit(1) + if args.n_cores > n_cores: + print(f"Number of cores cannot be more than{n_cores}") + exit(1) + else: + n_cores = args.n_cores + + print(f"{n_cores} number of cores available") + + tfrecord_filenames = [] + os.chdir(input_directory) + for file in glob.glob("*.tfrecord"): + file_name = os.path.abspath(file) + tfrecord_filenames.append(file_name) + + t = time.time() + + + if n_cores > 1: + pool = mp.Pool(min(n_cores, len(tfrecord_filenames))) + + for _ in tqdm(pool.imap_unordered(preprocess_wrap, tfrecord_filenames), total=len(tfrecord_filenames)): + pass + + # Close Pool and let all the processes complete + pool.close() + pool.join() # postpones the execution of next line of code until all processes in the queue are done. + else: + for tfrecord_file in tqdm(tfrecord_filenames): + preprocess_wrap(tfrecord_file) + + # Merge look up tables + print("Merging individual metadata...") + merge_metadata(os.path.abspath(output_directory)) + + print(f"Preprocessing duration: {(time.time() - t):.2f} s") + + diff --git a/data_prep_scripts/waymo/old/tutorial_maps.ipynb b/data_prep_scripts/waymo/old/tutorial_maps.ipynb new file mode 100644 index 0000000..4db2882 --- /dev/null +++ b/data_prep_scripts/waymo/old/tutorial_maps.ipynb @@ -0,0 +1,240 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "zy79JF35iDaE" + }, + "source": [ + "## Plot point clouds with maps" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": { + "id": "2MnJVU1jVe7a" + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "2023-03-23 18:41:26.609244: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations: AVX2 FMA\n", + "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n", + "2023-03-23 18:41:27.087125: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: :/home/k/.mujoco/mujoco200/bin\n", + "2023-03-23 18:41:27.087188: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: :/home/k/.mujoco/mujoco200/bin\n", + "2023-03-23 18:41:27.087194: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.\n" + ] + } + ], + "source": [ + "from typing import List\n", + "\n", + "import numpy as np\n", + "import plotly.graph_objs as go\n", + "from waymo_open_dataset import dataset_pb2\n", + "from waymo_open_dataset.utils import frame_utils\n", + "from waymo_open_dataset.utils import plot_maps\n", + "\n", + "def plot_point_clouds_with_maps(frames: List[dataset_pb2.Frame])->None:\n", + " \"\"\"Plot the point clouds within the given frames with map data.\n", + "\n", + " Map data must be populated in the first frame in the list.\n", + "\n", + " Args:\n", + " frames: A list of frames to be plotted, frames[0] must contain map data.\n", + " \"\"\"\n", + "\n", + " # Plot the map features.\n", + " if len(frames) == 0:\n", + " return\n", + " figure = plot_maps.plot_map_features(frames[0].map_features)\n", + "\n", + " for frame in frames:\n", + " # Parse the frame lidar data into range images.\n", + " range_images, camera_projections, seg_labels, range_image_top_poses = (\n", + " frame_utils.parse_range_image_and_camera_projection(frame)\n", + " )\n", + "\n", + " # Project the range images into points.\n", + " points, cp_points = frame_utils.convert_range_image_to_point_cloud(\n", + " frame,\n", + " range_images,\n", + " camera_projections,\n", + " range_image_top_poses,\n", + " keep_polar_features=True,\n", + " )\n", + " xyz = points[0][:, 3:]\n", + " num_points = xyz.shape[0]\n", + " \n", + " # Transform the points from the vehicle frame to the world frame.\n", + " xyz = np.concatenate([xyz, np.ones([num_points, 1])], axis=-1)\n", + " transform = np.reshape(np.array(frame.pose.transform), [4, 4])\n", + " xyz = np.transpose(np.matmul(transform, np.transpose(xyz)))[:, 0:3]\n", + "\n", + " # Correct the pose of the points into the coordinate system of the first\n", + " # frame to align with the map data.\n", + " offset = frame.map_pose_offset\n", + " points_offset = np.array([offset.x, offset.y, offset.z])\n", + " xyz += points_offset\n", + "\n", + " # Plot the point cloud for this frame aligned with the map data.\n", + " intensity = points[0][:, 0]\n", + " figure.add_trace(\n", + " go.Scatter3d(\n", + " x=xyz[:, 0],\n", + " y=xyz[:, 1],\n", + " z=xyz[:, 2],\n", + " mode='markers',\n", + " marker=dict(\n", + " size=1,\n", + " color=intensity, # set color to an array/list of desired values\n", + " colorscale='Pinkyl', # choose a colorscale\n", + " opacity=0.8,\n", + " ),\n", + " )\n", + " )\n", + "\n", + " figure.show()\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "5mhOt4geq12V" + }, + "source": [ + "## Load frames data" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": { + "id": "WfRqjpyEt_Px" + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "2023-03-23 18:41:28.778289: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n", + "2023-03-23 18:41:28.778587: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n", + "2023-03-23 18:41:28.778855: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n", + "2023-03-23 18:41:28.779123: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:981] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n", + "2023-03-23 18:41:28.784104: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudnn.so.8'; dlerror: libcudnn.so.8: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: :/home/k/.mujoco/mujoco200/bin\n", + "2023-03-23 18:41:28.784123: W tensorflow/core/common_runtime/gpu/gpu_device.cc:1934] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.\n", + "Skipping registering GPU devices...\n", + "2023-03-23 18:41:28.784941: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations: AVX2 FMA\n", + "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n" + ] + }, + { + "ename": "NotFoundError", + "evalue": "{{function_node __wrapped__IteratorGetNext_output_types_1_device_/job:localhost/replica:0/task:0/device:CPU:0}} /content/frames_with_maps.tfrecord; No such file or directory [Op:IteratorGetNext]", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mNotFoundError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[3], line 10\u001b[0m\n\u001b[1;32m 8\u001b[0m frames \u001b[39m=\u001b[39m []\n\u001b[1;32m 9\u001b[0m count \u001b[39m=\u001b[39m \u001b[39m0\u001b[39m\n\u001b[0;32m---> 10\u001b[0m \u001b[39mfor\u001b[39;00m data \u001b[39min\u001b[39;00m dataset:\n\u001b[1;32m 11\u001b[0m frame \u001b[39m=\u001b[39m dataset_pb2\u001b[39m.\u001b[39mFrame\u001b[39m.\u001b[39mFromString(\u001b[39mbytearray\u001b[39m(data\u001b[39m.\u001b[39mnumpy()))\n\u001b[1;32m 12\u001b[0m frames\u001b[39m.\u001b[39mappend(frame)\n", + "File \u001b[0;32m~/miniconda3/envs/waymo_open/lib/python3.9/site-packages/tensorflow/python/data/ops/iterator_ops.py:787\u001b[0m, in \u001b[0;36mOwnedIterator.__next__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 785\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39m__next__\u001b[39m(\u001b[39mself\u001b[39m):\n\u001b[1;32m 786\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[0;32m--> 787\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_next_internal()\n\u001b[1;32m 788\u001b[0m \u001b[39mexcept\u001b[39;00m errors\u001b[39m.\u001b[39mOutOfRangeError:\n\u001b[1;32m 789\u001b[0m \u001b[39mraise\u001b[39;00m \u001b[39mStopIteration\u001b[39;00m\n", + "File \u001b[0;32m~/miniconda3/envs/waymo_open/lib/python3.9/site-packages/tensorflow/python/data/ops/iterator_ops.py:770\u001b[0m, in \u001b[0;36mOwnedIterator._next_internal\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 767\u001b[0m \u001b[39m# TODO(b/77291417): This runs in sync mode as iterators use an error status\u001b[39;00m\n\u001b[1;32m 768\u001b[0m \u001b[39m# to communicate that there is no more data to iterate over.\u001b[39;00m\n\u001b[1;32m 769\u001b[0m \u001b[39mwith\u001b[39;00m context\u001b[39m.\u001b[39mexecution_mode(context\u001b[39m.\u001b[39mSYNC):\n\u001b[0;32m--> 770\u001b[0m ret \u001b[39m=\u001b[39m gen_dataset_ops\u001b[39m.\u001b[39;49miterator_get_next(\n\u001b[1;32m 771\u001b[0m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_iterator_resource,\n\u001b[1;32m 772\u001b[0m output_types\u001b[39m=\u001b[39;49m\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_flat_output_types,\n\u001b[1;32m 773\u001b[0m output_shapes\u001b[39m=\u001b[39;49m\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_flat_output_shapes)\n\u001b[1;32m 775\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[1;32m 776\u001b[0m \u001b[39m# Fast path for the case `self._structure` is not a nested structure.\u001b[39;00m\n\u001b[1;32m 777\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_element_spec\u001b[39m.\u001b[39m_from_compatible_tensor_list(ret) \u001b[39m# pylint: disable=protected-access\u001b[39;00m\n", + "File \u001b[0;32m~/miniconda3/envs/waymo_open/lib/python3.9/site-packages/tensorflow/python/ops/gen_dataset_ops.py:3017\u001b[0m, in \u001b[0;36miterator_get_next\u001b[0;34m(iterator, output_types, output_shapes, name)\u001b[0m\n\u001b[1;32m 3015\u001b[0m \u001b[39mreturn\u001b[39;00m _result\n\u001b[1;32m 3016\u001b[0m \u001b[39mexcept\u001b[39;00m _core\u001b[39m.\u001b[39m_NotOkStatusException \u001b[39mas\u001b[39;00m e:\n\u001b[0;32m-> 3017\u001b[0m _ops\u001b[39m.\u001b[39;49mraise_from_not_ok_status(e, name)\n\u001b[1;32m 3018\u001b[0m \u001b[39mexcept\u001b[39;00m _core\u001b[39m.\u001b[39m_FallbackException:\n\u001b[1;32m 3019\u001b[0m \u001b[39mpass\u001b[39;00m\n", + "File \u001b[0;32m~/miniconda3/envs/waymo_open/lib/python3.9/site-packages/tensorflow/python/framework/ops.py:7215\u001b[0m, in \u001b[0;36mraise_from_not_ok_status\u001b[0;34m(e, name)\u001b[0m\n\u001b[1;32m 7213\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mraise_from_not_ok_status\u001b[39m(e, name):\n\u001b[1;32m 7214\u001b[0m e\u001b[39m.\u001b[39mmessage \u001b[39m+\u001b[39m\u001b[39m=\u001b[39m (\u001b[39m\"\u001b[39m\u001b[39m name: \u001b[39m\u001b[39m\"\u001b[39m \u001b[39m+\u001b[39m name \u001b[39mif\u001b[39;00m name \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m \u001b[39melse\u001b[39;00m \u001b[39m\"\u001b[39m\u001b[39m\"\u001b[39m)\n\u001b[0;32m-> 7215\u001b[0m \u001b[39mraise\u001b[39;00m core\u001b[39m.\u001b[39m_status_to_exception(e) \u001b[39mfrom\u001b[39;00m \u001b[39mNone\u001b[39m\n", + "\u001b[0;31mNotFoundError\u001b[0m: {{function_node __wrapped__IteratorGetNext_output_types_1_device_/job:localhost/replica:0/task:0/device:CPU:0}} /content/frames_with_maps.tfrecord; No such file or directory [Op:IteratorGetNext]" + ] + } + ], + "source": [ + "import tensorflow as tf\n", + "\n", + "FILENAME = '/efs/waymo_open_with_maps/training/training/segment-10017090168044687777_6380_000_6400_000_with_camera_labels.tfrecord'\n", + "\n", + "dataset = tf.data.TFRecordDataset(FILENAME, compression_type='')\n", + "\n", + "# Load only 2 frames. Note that using too many frames may be slow to display.\n", + "frames = []\n", + "count = 0\n", + "for data in dataset:\n", + " frame = dataset_pb2.Frame.FromString(bytearray(data.numpy()))\n", + " frames.append(frame)\n", + " count += 1\n", + " if count == 2: \n", + " break\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "RzTEio68tE2B" + }, + "source": [ + "## Plot frames data" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "-Y59EuqTs8c1" + }, + "outputs": [], + "source": [ + "# Interactive plot of multiple point clouds aligned to the maps frame.\n", + "\n", + "# For most systems:\n", + "# left mouse button: rotate\n", + "# right mouse button: pan\n", + "# scroll wheel: zoom\n", + "\n", + "plot_point_clouds_with_maps(frames)" + ] + } + ], + "metadata": { + "colab": { + "last_runtime": { + "build_target": "", + "kind": "local" + }, + "private_outputs": true, + "provenance": [ + { + "file_id": "1ZOQUmLALJyfE-1zH4tynEQmSoYXUWSXF", + "timestamp": 1675474768372 + }, + { + "file_id": "1YdikcSqhy9NvprGaQzeSYfTe_UtPsMnv", + "timestamp": 1675452906901 + }, + { + "file_id": "1qK84E7Wug2MmhaqXTtkBzm9id9U1r1NC", + "timestamp": 1663010675571 + }, + { + "file_id": "13RtPddJ5Fzgi27_vv2v3Xm7-m1EV5dAL", + "timestamp": 1661916016522 + } + ] + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.9.16" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/data_prep_scripts/waymo/old/waymo_map.py b/data_prep_scripts/waymo/old/waymo_map.py new file mode 100644 index 0000000..79824cc --- /dev/null +++ b/data_prep_scripts/waymo/old/waymo_map.py @@ -0,0 +1,110 @@ +from typing import List + +import numpy as np +import plotly.graph_objs as go +from waymo_open_dataset import dataset_pb2 +from waymo_open_dataset.utils import frame_utils +from waymo_open_dataset.utils import plot_maps + +import open3d as o3d + + +def polygon_to_points(polygon) -> np.ndarray: + return [np.array([e.x, e.y, e.z]) for e in polygon] + + +def get_map_points(frame: dataset_pb2.Frame) -> np.ndarray: + map_features = frame.map_features + points = [] + for feature in map_features: + if feature.HasField('road_edge'): + points.extend(polygon_to_points(feature.road_edge.polyline)) + elif feature.HasField('crosswalk'): + points.extend(polygon_to_points(feature.crosswalk.polygon)) + elif feature.HasField('road_line'): + points.extend(polygon_to_points(feature.road_line.polyline)) + return np.array(points) + + +def get_pc(frame: dataset_pb2.Frame) -> np.ndarray: + # Parse the frame lidar data into range images. + range_images, camera_projections, seg_labels, range_image_top_poses = ( + frame_utils.parse_range_image_and_camera_projection(frame)) + # Project the range images into points. + points, cp_points = frame_utils.convert_range_image_to_point_cloud( + frame, + range_images, + camera_projections, + range_image_top_poses, + keep_polar_features=True, + ) + car_frame_pc = points[0][:, 3:] + num_points = car_frame_pc.shape[0] + print(f'num_points: {num_points}') + + # Transform the points from the vehicle frame to the world frame. + world_frame_pc = np.concatenate( + [car_frame_pc, np.ones([num_points, 1])], axis=-1) + car_to_global_transform = np.reshape(np.array(frame.pose.transform), + [4, 4]) + world_frame_pc = np.transpose( + np.matmul(car_to_global_transform, np.transpose(world_frame_pc)))[:, + 0:3] + + # Transform the points from the world frame to the map frame. + offset = frame.map_pose_offset + points_offset = np.array([offset.x, offset.y, offset.z]) + world_frame_pc += points_offset + + return car_frame_pc, world_frame_pc + + +def remove_ground(frame: dataset_pb2.Frame) -> None: + car_frame_pc, global_frame_pc = get_pc(frame) + # map_points = get_map_points(frame) + + # Use open3d to visualize the point cloud in xyz + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(car_frame_pc) + + # Use open3d to draw a 1 meter sphere at the origin + sphere = o3d.geometry.TriangleMesh.create_sphere(radius=1.0) + + geometry_list = [pcd, sphere] + + # for point in map_points: + # sphere = o3d.geometry.TriangleMesh.create_sphere(radius=0.1) + # sphere.translate(point) + # # Make the sphere red + # sphere.paint_uniform_color([1, 0, 0]) + # geometry_list.append(sphere) + + o3d.visualization.draw_geometries(geometry_list) + + + +import tensorflow as tf + +FILENAME = '/efs/waymo_open_with_maps/training/segment-10017090168044687777_6380_000_6400_000_with_camera_labels.tfrecord' + +dataset = tf.data.TFRecordDataset(FILENAME, compression_type='') + +# Load only 3 frames. Note that using too many frames may be slow to display. +frames = [] +count = 0 +for data in dataset: + frame = dataset_pb2.Frame.FromString(bytearray(data.numpy())) + remove_ground(frame) + # frames.append(frame) + # count += 1 + # if count == 3: + # break + +# Interactive plot of multiple point clouds aligned to the maps frame. + +# For most systems: +# left mouse button: rotate +# right mouse button: pan +# scroll wheel: zoom + +# plot_point_clouds_with_maps(frames) \ No newline at end of file diff --git a/data_prep_scripts/waymo/old/waymo_map_min.py b/data_prep_scripts/waymo/old/waymo_map_min.py new file mode 100644 index 0000000..51cacab --- /dev/null +++ b/data_prep_scripts/waymo/old/waymo_map_min.py @@ -0,0 +1,68 @@ +from typing import List + +import numpy as np +import plotly.graph_objs as go +from waymo_open_dataset import dataset_pb2 +from waymo_open_dataset.utils import frame_utils +from waymo_open_dataset.utils import plot_maps + +import tensorflow as tf + +import open3d as o3d + + +def get_pc(frame: dataset_pb2.Frame) -> np.ndarray: + # Parse the frame lidar data into range images. + range_images, camera_projections, seg_labels, range_image_top_poses = ( + frame_utils.parse_range_image_and_camera_projection(frame)) + # Project the range images into points. + points, cp_points = frame_utils.convert_range_image_to_point_cloud( + frame, + range_images, + camera_projections, + range_image_top_poses, + keep_polar_features=True, + ) + car_frame_pc = points[0][:, 3:] + num_points = car_frame_pc.shape[0] + print(f'num_points: {num_points}') + + # Transform the points from the vehicle frame to the world frame. + world_frame_pc = np.concatenate( + [car_frame_pc, np.ones([num_points, 1])], axis=-1) + car_to_global_transform = np.reshape(np.array(frame.pose.transform), + [4, 4]) + world_frame_pc = np.transpose( + np.matmul(car_to_global_transform, np.transpose(world_frame_pc)))[:, + 0:3] + + # Transform the points from the world frame to the map frame. + offset = frame.map_pose_offset + points_offset = np.array([offset.x, offset.y, offset.z]) + world_frame_pc += points_offset + + return car_frame_pc, world_frame_pc + + +def remove_ground(frame: dataset_pb2.Frame) -> None: + car_frame_pc, global_frame_pc = get_pc(frame) + + # Use open3d to visualize the point cloud in xyz + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(car_frame_pc) + + # Use open3d to draw a 1 meter sphere at the origin + sphere = o3d.geometry.TriangleMesh.create_sphere(radius=1.0) + + geometry_list = [pcd, sphere] + + o3d.visualization.draw_geometries(geometry_list) + + +FILENAME = '/efs/waymo_open_with_maps/training/segment-10017090168044687777_6380_000_6400_000_with_camera_labels.tfrecord' + +dataset = tf.data.TFRecordDataset(FILENAME, compression_type='') + +for data in dataset: + frame = dataset_pb2.Frame.FromString(bytearray(data.numpy())) + remove_ground(frame) diff --git a/data_prep_scripts/waymo/rasterize_heightmap.py b/data_prep_scripts/waymo/rasterize_heightmap.py new file mode 100644 index 0000000..2b0f6f1 --- /dev/null +++ b/data_prep_scripts/waymo/rasterize_heightmap.py @@ -0,0 +1,207 @@ +from sklearn.neighbors import NearestNeighbors +import os + +os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' +import tensorflow as tf +import argparse +import multiprocessing +from pathlib import Path +from joblib import Parallel, delayed +import numpy as np + +from waymo_open_dataset import dataset_pb2 +import json + +parser = argparse.ArgumentParser() +parser.add_argument("waymo_directory", + type=Path, + help="Path to Waymo Open directory.") +parser.add_argument("output_directory", + type=Path, + help="Path to output directory.") +parser.add_argument( + '--cells_per_meter', + type=int, + default=10.0 / 3.0, + help= + "Cells per meter for discritization. Default is Argoverse 2 Sensor dataset default is 30cm (10/3 cells per meter)" +) +parser.add_argument('--num_neighbors', + type=int, + default=20, + help="Number of neighbors to use to compute height") +parser.add_argument('--cpus', + type=int, + default=multiprocessing.cpu_count(), + help="Number of cpus to use for parallel processing") + +args = parser.parse_args() + +assert args.cells_per_meter > 0, "Cells per meter must be positive" +assert args.num_neighbors > 0, "Number of neighbors must be greater than zero" +assert args.waymo_directory.is_dir( +), f"{args.waymo_directory} is not a directory" + +print("Waymo directory:", args.waymo_directory) +print("Output directory:", args.output_directory) + + +def build_knn(points, num_neighbors): + return NearestNeighbors(n_neighbors=num_neighbors, + radius=20, + leaf_size=num_neighbors).fit(points) + + +def build_global_grid(points, cells_per_meter): + xs, ys, _ = zip(*points) + min_x, max_x = min(xs), max(xs) + min_y, max_y = min(ys), max(ys) + + grid_max_global_frame = np.array([max_x, max_y]) + grid_min_global_frame = np.array([min_x, min_y]) + + area_grid_global_frame = grid_max_global_frame - grid_min_global_frame + grid_shape = np.ceil( + area_grid_global_frame * cells_per_meter).astype(int) + 1 + + def global_to_grid_float(pts): + assert (pts <= grid_max_global_frame + ).all(), f"({pts} <= {grid_max_global_frame})" + assert (pts >= grid_min_global_frame + ).all(), f"({pts} >= {grid_min_global_frame})" + relative_to_grid_origin = pts - grid_min_global_frame + floating_point_grid_coordinate = relative_to_grid_origin * cells_per_meter + return floating_point_grid_coordinate + + def global_to_grid_index(pts): + coords = global_to_grid_float(pts) + return np.round(coords).astype(int) + + return grid_shape, global_to_grid_index, global_to_grid_float, grid_min_global_frame + + +def render_heightmap(points, cells_per_meter, num_neighbors): + + # We construct this full heightmap with 0, 0 at xy_min_offset in the global coordinate frame. + grid_shape, global_to_grid_index, global_to_grid_float, grid_min_global_frame = build_global_grid( + points, cells_per_meter) + + polygon_grid_points = np.array([(*global_to_grid_index(p[:2]), p[2]) + for p in points]) + mean_z = np.mean([e[2] for e in polygon_grid_points]) + knn = build_knn(polygon_grid_points, num_neighbors) + + # Construct a grid shaped array whose last axis holds the X, Y index value for that square, + # with the average Z value for the purposes of querying. + xs_lin = np.arange(0, grid_shape[0], 1) + ys_lin = np.arange(0, grid_shape[1], 1) + xs_square = np.expand_dims(np.tile(xs_lin, (grid_shape[1], 1)), 2) + ys_square = np.expand_dims(np.tile(ys_lin, (grid_shape[0], 1)).T, 2) + zs_square = np.ones_like(xs_square) * mean_z + pts_square = np.concatenate((xs_square, ys_square, zs_square), 2) + + # Flatten the pts square into an N x 3 array for querying KNN + pts_square_shape = pts_square.shape + pts_line = pts_square.reshape(pts_square_shape[0] * pts_square_shape[1], 3) + + _, indices = knn.kneighbors(pts_line) + neighbor_values = polygon_grid_points[indices] + avg_neighbor_z_values = np.mean(neighbor_values[:, :, 2], axis=1) + + # Reshape flattened average Z values back into grid shape. + grid = avg_neighbor_z_values.reshape(pts_square_shape[0], + pts_square_shape[1]) + + return grid, grid_min_global_frame + + +def save_grid_global_offset(file_path: Path, + grid, + grid_min_global_frame, + cells_per_meter, + verbose=False): + se2 = { + "R": [1.0, 0.0, 0.0, 1.0], # Identity rotation matrix flattened + "t": [-grid_min_global_frame[0], -grid_min_global_frame[1]], + "s": cells_per_meter + } + + save_folder = args.output_directory / file_path.parent.name / ( + file_path.stem + "_map") + save_folder.mkdir(parents=True, exist_ok=True) + + se2_name = "se2.json" + height_map_name = "ground_height.npy" + + height_map_file = save_folder / height_map_name + if height_map_file.exists(): + height_map_file.unlink() + + se2_file = save_folder / se2_name + if se2_file.exists(): + se2_file.unlink() + + np.save(height_map_file, grid.astype(np.float16)) + if verbose: + print(f"Saving heightmap to {height_map_file}") + with open(se2_file, 'w') as fp: + json.dump(se2, fp) + + +def polygon_to_points(polygon) -> np.ndarray: + return [np.array([e.x, e.y, e.z]) for e in polygon] + + +def collect_points(frame: dataset_pb2.Frame) -> np.ndarray: + map_features = frame.map_features + points = [] + for feature in map_features: + if feature.HasField('road_edge'): + points.extend(polygon_to_points(feature.road_edge.polyline)) + elif feature.HasField('crosswalk'): + points.extend(polygon_to_points(feature.crosswalk.polygon)) + elif feature.HasField('road_line'): + points.extend(polygon_to_points(feature.road_line.polyline)) + return np.array(points) + + +def build_work_queue(waymo_directory): + waymo_directory = Path(waymo_directory) + assert waymo_directory.is_dir(), f"{waymo_directory} is not a directory" + + train_records = sorted((waymo_directory / 'training').glob('*.tfrecord')) + val_records = sorted((waymo_directory / 'validation').glob('*.tfrecord')) + + queue = train_records + val_records + for record in queue: + assert record.is_file(), f"{record} is not a file" + + return queue + + +def process_record(file_path: Path): + print("Processing", file_path) + dataset = tf.data.TFRecordDataset(file_path, compression_type='') + + # Hack because I can't figure out how to extract the first frame from the dataset. + for data in dataset: + frame = dataset_pb2.Frame.FromString(bytearray(data.numpy())) + break + + points = collect_points(frame) + grid, grid_min_global_frame = render_heightmap(points, + args.cells_per_meter, + args.num_neighbors) + save_grid_global_offset(file_path, + grid, + grid_min_global_frame, + args.cells_per_meter, + verbose=True) + + +work_queue = build_work_queue(args.waymo_directory) +print("Work queue size:", len(work_queue)) + +num_processes = min(args.cpus, len(work_queue)) +Parallel(num_processes)(delayed(process_record)(record) + for record in work_queue) diff --git a/data_prep_scripts/waymo/rasterize_heightmap.sbatch b/data_prep_scripts/waymo/rasterize_heightmap.sbatch new file mode 100755 index 0000000..f27b36a --- /dev/null +++ b/data_prep_scripts/waymo/rasterize_heightmap.sbatch @@ -0,0 +1,14 @@ +#!/bin/bash + +## Resource Request +#SBATCH --job-name=waymo_rasterize +#SBATCH --output=waymo_rasterize.out +#SBATCH --qos=ee-med +#SBATCH --partition=eaton-compute +#SBATCH --cpus-per-task=120 +#SBATCH --mem-per-cpu=2GB +#SBATCH --container-mounts=../../datasets/:/efs/,./:/project +#SBATCH --container-image=/home/kvedder/kylevedder+offline_sceneflow_waymo+latest.sqsh + +## Job Steps +python rasterize_heightmap.py /efs/waymo_open_with_maps/ /efs/waymo_open_raster_heights/ --cpus 120 \ No newline at end of file diff --git a/dataloaders/__init__.py b/dataloaders/__init__.py new file mode 100644 index 0000000..bfc04cf --- /dev/null +++ b/dataloaders/__init__.py @@ -0,0 +1,29 @@ +from .argoverse_raw_lidar import ArgoverseRawSequenceLoader, ArgoverseRawSequence +from .argoverse_supervised_flow import ArgoverseSupervisedFlowSequenceLoader, ArgoverseSupervisedFlowSequence +from .argoverse_unsupervised_flow import ArgoverseUnsupervisedFlowSequenceLoader, ArgoverseUnsupervisedFlowSequence + +from .waymo_supervised_flow import WaymoSupervisedFlowSequenceLoader, WaymoSupervisedFlowSequence +from .waymo_unsupervised_flow import WaymoUnsupervisedFlowSequenceLoader, WaymoUnsupervisedFlowSequence + +from .sequence_dataset import OriginMode, SubsequenceRawDataset, SubsequenceSupervisedFlowDataset, SubsequenceSupervisedFlowSpecificSubsetDataset, SubsequenceUnsupervisedFlowDataset, ConcatDataset +from .var_len_sequence_dataset import VarLenSubsequenceRawDataset, VarLenSubsequenceSupervisedFlowDataset, VarLenSubsequenceSupervisedFlowSpecificSubsetDataset, VarLenSubsequenceUnsupervisedFlowDataset +from .pointcloud_dataset import PointCloudDataset + +__all__ = [ + 'ArgoverseRawSequenceLoader', + 'ArgoverseRawSequence', + 'SubsequenceRawDataset', + 'PointCloudDataset', + 'OriginMode', + 'ArgoverseSupervisedFlowSequenceLoader', + 'ArgoverseSupervisedFlowSequence', + 'SubsequenceSupervisedFlowDataset', + 'ArgoverseUnsupervisedFlowSequenceLoader', + 'ArgoverseUnsupervisedFlowSequence', + 'SubsequenceUnsupervisedFlowDataset', + 'ConcatDataset', + 'VarLenSubsequenceRawDataset', + 'VarLenSubsequenceSupervisedFlowDataset', + 'VarLenSubsequenceUnsupervisedFlowDataset', + 'WaymoSupervisedFlowSequenceLoader', +] diff --git a/dataloaders/argoverse_raw_lidar.py b/dataloaders/argoverse_raw_lidar.py new file mode 100644 index 0000000..49dbc8b --- /dev/null +++ b/dataloaders/argoverse_raw_lidar.py @@ -0,0 +1,255 @@ +import numpy as np +import pandas as pd +from pathlib import Path +from pointclouds import PointCloud, SE3, SE2 +from loader_utils import load_json +from typing import List, Tuple, Dict, Optional, Any +import time + +GROUND_HEIGHT_THRESHOLD = 0.4 # 40 centimeters + + +class ArgoverseRawSequence(): + + def __init__(self, + log_id: str, + dataset_dir: Path, + verbose: bool = False, + sample_every: Optional[int] = None): + self.log_id = log_id + + self.dataset_dir = Path(dataset_dir) + assert self.dataset_dir.is_dir( + ), f'dataset_dir {dataset_dir} does not exist' + self.frame_paths = sorted( + (self.dataset_dir / 'sensors' / 'lidar').glob('*.feather')) + assert len( + self.frame_paths) > 0, f'no frames found in {self.dataset_dir}' + self.frame_infos = pd.read_feather(self.dataset_dir / + 'city_SE3_egovehicle.feather') + + self.timestamp_to_file_map = {int(e.stem): e for e in self.frame_paths} + file_timestamps = set(self.timestamp_to_file_map.keys()) + + self.timestamp_to_info_idx_map = { + int(timestamp): idx + for idx, timestamp in enumerate( + self.frame_infos['timestamp_ns'].values) + } + info_timestamps = set(self.timestamp_to_info_idx_map.keys()) + + self.timestamp_list = sorted( + file_timestamps.intersection(info_timestamps)) + assert len(self.timestamp_list + ) > 0, f'no timestamps found in {self.dataset_dir}' + + if sample_every is not None: + self.timestamp_list = self.timestamp_list[::sample_every] + + self.raster_heightmap, self.global_to_raster_se2, self.global_to_raster_scale = self._load_ground_height_raster( + ) + + if verbose: + print( + f'Loaded {len(self.timestamp_list)} frames from {self.dataset_dir} at timestamp {time.time():.3f}' + ) + + def _load_ground_height_raster(self): + raster_height_paths = list( + (self.dataset_dir / + 'map').glob("*_ground_height_surface____*.npy")) + assert len( + raster_height_paths + ) == 1, f'Expected 1 raster, got {len(raster_height_paths)} in path {self.dataset_dir / "map"}' + raster_height_path = raster_height_paths[0] + + transform_paths = list( + (self.dataset_dir / 'map').glob("*img_Sim2_city.json")) + assert len(transform_paths + ) == 1, f'Expected 1 transform, got {len(transform_paths)}' + transform_path = transform_paths[0] + + raster_heightmap = np.load(raster_height_path) + transform = load_json(transform_path, verbose=False) + + transform_rotation = np.array(transform['R']).reshape(2, 2) + transform_translation = np.array(transform['t']) + transform_scale = np.array(transform['s']) + + transform_se2 = SE2(rotation=transform_rotation, + translation=transform_translation) + + return raster_heightmap, transform_se2, transform_scale + + def get_ground_heights(self, global_point_cloud: PointCloud) -> np.ndarray: + """Get ground height for each of the xy locations in a point cloud. + Args: + point_cloud: Numpy array of shape (k,2) or (k,3) in global coordinates. + Returns: + ground_height_values: Numpy array of shape (k,) + """ + + global_points_xy = global_point_cloud.points[:, :2] + + raster_points_xy = self.global_to_raster_se2.transform_point_cloud( + global_points_xy) * self.global_to_raster_scale + + raster_points_xy = np.round(raster_points_xy).astype(np.int64) + + ground_height_values = np.full((raster_points_xy.shape[0]), np.nan) + # outside max X + outside_max_x = (raster_points_xy[:, 0] >= + self.raster_heightmap.shape[1]).astype(bool) + # outside max Y + outside_max_y = (raster_points_xy[:, 1] >= + self.raster_heightmap.shape[0]).astype(bool) + # outside min X + outside_min_x = (raster_points_xy[:, 0] < 0).astype(bool) + # outside min Y + outside_min_y = (raster_points_xy[:, 1] < 0).astype(bool) + ind_valid_pts = ~np.logical_or( + np.logical_or(outside_max_x, outside_max_y), + np.logical_or(outside_min_x, outside_min_y)) + + ground_height_values[ind_valid_pts] = self.raster_heightmap[ + raster_points_xy[ind_valid_pts, 1], raster_points_xy[ind_valid_pts, + 0]] + + return ground_height_values + + def is_ground_points(self, global_point_cloud: PointCloud) -> np.ndarray: + """Remove ground points from a point cloud. + Args: + point_cloud: Numpy array of shape (k,3) in global coordinates. + Returns: + ground_removed_point_cloud: Numpy array of shape (k,3) in global coordinates. + """ + ground_height_values = self.get_ground_heights(global_point_cloud) + is_ground_boolean_arr = ( + np.absolute(global_point_cloud[:, 2] - ground_height_values) <= + GROUND_HEIGHT_THRESHOLD) | ( + np.array(global_point_cloud[:, 2] - ground_height_values) < 0) + return is_ground_boolean_arr + + def __repr__(self) -> str: + return f'ArgoverseSequence with {len(self)} frames' + + def __len__(self): + return len(self.timestamp_list) + + def _load_pc(self, idx) -> PointCloud: + assert idx < len( + self + ), f'idx {idx} out of range, len {len(self)} for {self.dataset_dir}' + timestamp = self.timestamp_list[idx] + frame_path = self.timestamp_to_file_map[timestamp] + frame_content = pd.read_feather(frame_path) + xs = frame_content['x'].values + ys = frame_content['y'].values + zs = frame_content['z'].values + points = np.stack([xs, ys, zs], axis=1) + return PointCloud(points) + + def _load_pose(self, idx) -> SE3: + assert idx < len( + self + ), f'idx {idx} out of range, len {len(self)} for {self.dataset_dir}' + timestamp = self.timestamp_list[idx] + infos_idx = self.timestamp_to_info_idx_map[timestamp] + frame_info = self.frame_infos.iloc[infos_idx] + se3 = SE3.from_rot_w_x_y_z_translation_x_y_z( + frame_info['qw'], frame_info['qx'], frame_info['qy'], + frame_info['qz'], frame_info['tx_m'], frame_info['ty_m'], + frame_info['tz_m']) + return se3 + + def load(self, idx: int, relative_to_idx: int) -> Dict[str, Any]: + assert idx < len( + self + ), f'idx {idx} out of range, len {len(self)} for {self.dataset_dir}' + pc = self._load_pc(idx) + start_pose = self._load_pose(relative_to_idx) + idx_pose = self._load_pose(idx) + relative_pose = start_pose.inverse().compose(idx_pose) + absolute_global_frame_pc = pc.transform(idx_pose) + is_ground_points = self.is_ground_points(absolute_global_frame_pc) + relative_global_frame_pc_with_ground = pc.transform(relative_pose) + relative_global_frame_pc_no_ground = relative_global_frame_pc_with_ground.mask_points( + ~is_ground_points) + return { + "relative_pc": relative_global_frame_pc_no_ground, + "relative_pc_with_ground": relative_global_frame_pc_with_ground, + "relative_pose": relative_pose, + "is_ground_points": is_ground_points, + "log_id": self.log_id, + "log_idx": idx, + } + + def load_frame_list( + self, relative_to_idx: Optional[int]) -> List[Dict[str, Any]]: + + return [ + self.load(idx, + relative_to_idx if relative_to_idx is not None else idx) + for idx in range(len(self)) + ] + + +class ArgoverseRawSequenceLoader(): + + def __init__(self, + sequence_dir: Path, + log_subset: Optional[List[str]] = None, + verbose: bool = False, + num_sequences: Optional[int] = None, + per_sequence_sample_every: Optional[int] = None): + self.dataset_dir = Path(sequence_dir) + self.verbose = verbose + self.per_sequence_sample_every = per_sequence_sample_every + assert self.dataset_dir.is_dir( + ), f'dataset_dir {sequence_dir} does not exist' + self.log_lookup = {e.name: e for e in self.dataset_dir.glob('*/')} + if log_subset is not None: + log_subset = set(log_subset) + log_keys = set(self.log_lookup.keys()) + assert log_subset.issubset( + log_keys + ), f'log_subset {log_subset} is not a subset of {log_keys}' + self.log_lookup = { + k: v + for k, v in self.log_lookup.items() if k in log_subset + } + + if num_sequences is not None: + self.log_lookup = { + k: v + for idx, (k, v) in enumerate(self.log_lookup.items()) + if idx < num_sequences + } + + if self.verbose: + print(f'Loaded {len(self.log_lookup)} logs') + + self.last_loaded_sequence = None + self.last_loaded_sequence_id = None + + def get_sequence_ids(self): + return sorted(self.log_lookup.keys()) + + def _raw_load_sequence(self, log_id: str) -> ArgoverseRawSequence: + assert log_id in self.log_lookup, f'log_id {log_id} is not in the {len(self.log_lookup)}' + log_dir = self.log_lookup[log_id] + assert log_dir.is_dir(), f'log_id {log_id} does not exist' + return ArgoverseRawSequence( + log_id, + log_dir, + verbose=self.verbose, + sample_every=self.per_sequence_sample_every) + + def load_sequence(self, log_id: str) -> ArgoverseRawSequence: + # Basic caching mechanism for repeated loads of the same sequence + if self.last_loaded_sequence_id != log_id: + self.last_loaded_sequence = self._raw_load_sequence(log_id) + self.last_loaded_sequence_id = log_id + + return self.last_loaded_sequence diff --git a/dataloaders/argoverse_supervised_flow.py b/dataloaders/argoverse_supervised_flow.py new file mode 100644 index 0000000..9872882 --- /dev/null +++ b/dataloaders/argoverse_supervised_flow.py @@ -0,0 +1,181 @@ +from pathlib import Path +from collections import defaultdict +from typing import List, Tuple, Dict, Optional, Any +import pandas as pd +from pointclouds import PointCloud, SE3, SE2 +import numpy as np + +from . import ArgoverseRawSequence + +CATEGORY_MAP = { + 0: 'ANIMAL', + 1: 'ARTICULATED_BUS', + 2: 'BICYCLE', + 3: 'BICYCLIST', + 4: 'BOLLARD', + 5: 'BOX_TRUCK', + 6: 'BUS', + 7: 'CONSTRUCTION_BARREL', + 8: 'CONSTRUCTION_CONE', + 9: 'DOG', + 10: 'LARGE_VEHICLE', + 11: 'MESSAGE_BOARD_TRAILER', + 12: 'MOBILE_PEDESTRIAN_CROSSING_SIGN', + 13: 'MOTORCYCLE', + 14: 'MOTORCYCLIST', + 15: 'OFFICIAL_SIGNALER', + 16: 'PEDESTRIAN', + 17: 'RAILED_VEHICLE', + 18: 'REGULAR_VEHICLE', + 19: 'SCHOOL_BUS', + 20: 'SIGN', + 21: 'STOP_SIGN', + 22: 'STROLLER', + 23: 'TRAFFIC_LIGHT_TRAILER', + 24: 'TRUCK', + 25: 'TRUCK_CAB', + 26: 'VEHICULAR_TRAILER', + 27: 'WHEELCHAIR', + 28: 'WHEELED_DEVICE', + 29: 'WHEELED_RIDER' +} + + +class ArgoverseSupervisedFlowSequence(ArgoverseRawSequence): + + def __init__(self, log_id: str, dataset_dir: Path, + flow_data_lst: List[Tuple[int, Path]]): + super().__init__(log_id, dataset_dir) + + # Each flow contains info for the t and t+1 timestamps. This means the last pointcloud in the sequence + # will not have a corresponding flow. + self.timestamp_to_flow_map = { + int(timestamp): flow_data_file + for timestamp, flow_data_file in flow_data_lst + } + # Make sure all of the timestamps in self.timestamp_list *other than the last timestamp* have a corresponding flow. + assert set(self.timestamp_list[:-1]) == set( + self.timestamp_to_flow_map.keys( + )), f'Flow data missing for some timestamps in {self.dataset_dir}' + + def _load_flow(self, idx): + assert idx < len( + self + ), f'idx {idx} out of range, len {len(self)} for {self.dataset_dir}' + # There is no flow information for the last pointcloud in the sequence. + if idx == len(self) - 1 or idx == -1: + return None, None, None, None, None, None + timestamp = self.timestamp_list[idx] + flow_data_file = self.timestamp_to_flow_map[timestamp] + flow_info = dict(np.load(flow_data_file)) + flow_0_1 = flow_info['flow_0_1'] + classes_0 = flow_info['classes_0'] + classes_1 = flow_info['classes_1'] + is_ground0 = flow_info['is_ground_0'] + is_ground1 = flow_info['is_ground_1'] + ego_motion = flow_info['ego_motion'] + return flow_0_1, classes_0, classes_1, is_ground0, is_ground1, ego_motion + + def load(self, idx: int, relative_to_idx: int) -> Dict[str, Any]: + assert idx < len( + self + ), f'idx {idx} out of range, len {len(self)} for {self.dataset_dir}' + raw_pc = self._load_pc(idx) + flow_0_1, classes_0, _, is_ground0, _, _ = self._load_flow(idx) + start_pose = self._load_pose(relative_to_idx) + idx_pose = self._load_pose(idx) + + relative_pose = start_pose.inverse().compose(idx_pose) + + # Global frame PC is needed to compute hte ground point mask. + absolute_global_frame_pc = raw_pc.transform(idx_pose) + is_ground_points = self.is_ground_points(absolute_global_frame_pc) + + relative_global_frame_pc_with_ground = raw_pc.transform(relative_pose) + relative_global_frame_pc = relative_global_frame_pc_with_ground.mask_points( + ~is_ground_points) + if flow_0_1 is not None: + ix_plus_one_pose = self._load_pose(idx + 1) + relative_pose_plus_one = start_pose.inverse().compose( + ix_plus_one_pose) + flowed_pc = raw_pc.flow(flow_0_1) + flowed_pc = flowed_pc.mask_points(~is_ground_points) + relative_global_frame_flowed_pc = flowed_pc.transform( + relative_pose_plus_one) + classes_0 = classes_0[~is_ground_points] + is_ground0 = is_ground0[~is_ground_points] + else: + relative_global_frame_flowed_pc = None + return { + "relative_pc": relative_global_frame_pc, + "relative_pc_with_ground": relative_global_frame_pc_with_ground, + "is_ground_points": is_ground_points, + "relative_pose": relative_pose, + "relative_flowed_pc": relative_global_frame_flowed_pc, + "pc_classes": classes_0, + "pc_is_ground": is_ground0, + "log_id": self.log_id, + "log_idx": idx, + } + + +class ArgoverseSupervisedFlowSequenceLoader(): + + def __init__(self, + raw_data_path: Path, + flow_data_path: Path, + log_subset: Optional[List[str]] = None, + num_sequences: Optional[int] = None): + self.raw_data_path = Path(raw_data_path) + self.flow_data_path = Path(flow_data_path) + assert self.raw_data_path.is_dir( + ), f'raw_data_path {raw_data_path} does not exist' + assert self.flow_data_path.is_dir( + ), f'flow_data_path {flow_data_path} does not exist' + + # Convert folder of flow NPZ files to a lookup table for different sequences + flow_data_files = sorted(self.flow_data_path.glob('*.npz')) + self.sequence_id_to_flow_lst = defaultdict(list) + for flow_data_file in flow_data_files: + sequence_id, timestamp = flow_data_file.stem.split('_') + timestamp = int(timestamp) + self.sequence_id_to_flow_lst[sequence_id].append( + (timestamp, flow_data_file)) + + self.sequence_id_to_raw_data = {} + for sequence_id in self.sequence_id_to_flow_lst.keys(): + sequence_folder = self.raw_data_path / sequence_id + assert sequence_folder.is_dir( + ), f'raw_data_path {raw_data_path} does not exist' + self.sequence_id_to_raw_data[sequence_id] = sequence_folder + + self.sequence_id_lst = sorted(self.sequence_id_to_flow_lst.keys()) + if log_subset is not None: + self.sequence_id_lst = [ + sequence_id for sequence_id in self.sequence_id_lst + if sequence_id in log_subset + ] + + if num_sequences is not None: + self.sequence_id_lst = self.sequence_id_lst[:num_sequences] + + self.last_loaded_sequence = None + self.last_loaded_sequence_id = None + + def get_sequence_ids(self): + return self.sequence_id_lst + + def _load_sequence_raw( + self, sequence_id: str) -> ArgoverseSupervisedFlowSequence: + assert sequence_id in self.sequence_id_to_flow_lst, f'sequence_id {sequence_id} does not exist' + return ArgoverseSupervisedFlowSequence( + sequence_id, self.sequence_id_to_raw_data[sequence_id], + self.sequence_id_to_flow_lst[sequence_id]) + + def load_sequence(self, + sequence_id: str) -> ArgoverseSupervisedFlowSequence: + # Basic caching mechanism for repeated loads of the same sequence + if self.last_loaded_sequence_id != sequence_id: + self.last_loaded_sequence = self._load_sequence_raw(sequence_id) + self.last_loaded_sequence_id = sequence_id + return self.last_loaded_sequence diff --git a/dataloaders/argoverse_unsupervised_flow.py b/dataloaders/argoverse_unsupervised_flow.py new file mode 100644 index 0000000..7f1a89d --- /dev/null +++ b/dataloaders/argoverse_unsupervised_flow.py @@ -0,0 +1,124 @@ +from pathlib import Path +from collections import defaultdict +from typing import List, Tuple, Dict, Optional, Any +import pandas as pd +from pointclouds import PointCloud, SE3, SE2 +import numpy as np + +from . import ArgoverseRawSequence + + +class ArgoverseUnsupervisedFlowSequence(ArgoverseRawSequence): + + def __init__(self, log_id: str, dataset_dir: Path, flow_data_lst: Path): + super().__init__(log_id, dataset_dir) + + # The flow data does not have a timestamp, so we need to just rely on the order of the files. + self.flow_data_files = sorted(flow_data_lst.glob('*.npz')) + + assert len(self.timestamp_list) > len( + self.flow_data_files + ), f"More flow data files in {flow_data_lst} than pointclouds in {self.dataset_dir}; {len(self.timestamp_list)} vs {len(self.flow_data_files)}" + + # The first len(self.flow_data_files) timestamps have flow data. + # We keep those timestamps, plus the final timestamp. + self.timestamp_list = self.timestamp_list[:len(self.flow_data_files) + + 1] + + def _load_flow(self, idx): + assert idx < len( + self + ), f'idx {idx} out of range, len {len(self)} for {self.dataset_dir}' + # There is no flow information for the last pointcloud in the sequence. + if idx == len(self) - 1 or idx == -1: + return None, None + + flow_data_file = self.flow_data_files[idx] + flow_info = dict(np.load(flow_data_file)) + flow_0_1, valid_idxes = flow_info['flow'], flow_info['valid_idxes'] + return flow_0_1, valid_idxes + + def load(self, idx, relative_to_idx) -> Dict[str, Any]: + assert idx < len( + self + ), f'idx {idx} out of range, len {len(self)} for {self.dataset_dir}' + raw_pc = self._load_pc(idx) + flow_0_1, valid_idxes = self._load_flow(idx) + start_pose = self._load_pose(relative_to_idx) + idx_pose = self._load_pose(idx) + + relative_pose = start_pose.inverse().compose(idx_pose) + + absolute_global_frame_pc = raw_pc.transform(idx_pose) + is_ground_points = self.is_ground_points(absolute_global_frame_pc) + relative_global_frame_pc_with_ground = raw_pc.transform(relative_pose) + relative_global_frame_pc = relative_global_frame_pc_with_ground.mask_points( + ~is_ground_points) + return { + "relative_pc": relative_global_frame_pc, + "relative_pc_with_ground": relative_global_frame_pc_with_ground, + "is_ground_points": is_ground_points, + "relative_pose": relative_pose, + "flow": flow_0_1, + "valid_idxes": valid_idxes, + "log_id": self.log_id, + "log_idx": idx, + } + + +class ArgoverseUnsupervisedFlowSequenceLoader(): + + def __init__(self, + raw_data_path: Path, + flow_data_path: Path, + log_subset: Optional[List[str]] = None): + self.raw_data_path = Path(raw_data_path) + self.flow_data_path = Path(flow_data_path) + assert self.raw_data_path.is_dir( + ), f'raw_data_path {raw_data_path} does not exist' + assert self.flow_data_path.is_dir( + ), f'flow_data_path {flow_data_path} does not exist' + + # Raw data folders + raw_data_folders = sorted(self.raw_data_path.glob('*/')) + self.sequence_id_to_raw_data = { + folder.stem: folder + for folder in raw_data_folders + } + # Flow data folders + flow_data_folders = sorted(self.flow_data_path.glob('*/')) + self.sequence_id_to_flow_data = { + folder.stem: folder + for folder in flow_data_folders + } + + self.sequence_id_list = sorted( + set(self.sequence_id_to_raw_data.keys()).intersection( + set(self.sequence_id_to_flow_data.keys()))) + + if log_subset is not None: + self.sequence_id_list = [ + sequence_id for sequence_id in self.sequence_id_list + if sequence_id in log_subset + ] + + self.last_loaded_sequence = None + self.last_loaded_sequence_id = None + + def get_sequence_ids(self): + return self.sequence_id_list + + def _load_sequence_raw( + self, sequence_id: str) -> ArgoverseUnsupervisedFlowSequence: + assert sequence_id in self.sequence_id_to_flow_data, f'sequence_id {sequence_id} does not exist' + return ArgoverseUnsupervisedFlowSequence( + sequence_id, self.sequence_id_to_raw_data[sequence_id], + self.sequence_id_to_flow_data[sequence_id]) + + def load_sequence(self, + sequence_id: str) -> ArgoverseUnsupervisedFlowSequence: + # Basic caching mechanism for repeated loads of the same sequence + if self.last_loaded_sequence_id != sequence_id: + self.last_loaded_sequence = self._load_sequence_raw(sequence_id) + self.last_loaded_sequence_id = sequence_id + return self.last_loaded_sequence diff --git a/dataloaders/pointcloud_dataset.py b/dataloaders/pointcloud_dataset.py new file mode 100644 index 0000000..7cbce03 --- /dev/null +++ b/dataloaders/pointcloud_dataset.py @@ -0,0 +1,69 @@ +import torch +import torch.nn as nn +import numpy as np +# Import rotation as R +from scipy.spatial.transform import Rotation as R + + +class PointCloudDataset(torch.utils.data.Dataset): + + def __init__(self, sequence_loader, max_pc_points: int = 90000): + self.sequence_loader = sequence_loader + self.max_pc_points = max_pc_points + + self.pcs_per_sequence = 295 + self.sequence_ids = sequence_loader.get_sequence_ids() + self.random_state = np.random.RandomState(len(self.sequence_ids)) + + def __len__(self): + return len(self.sequence_ids) * self.pcs_per_sequence + + def _random_rot(self): + # Uniform random between -pi and pi + random_angle = self.random_state.uniform(-np.pi, np.pi) + # Create random rotation matrix for PC + return R.from_euler('z', random_angle, + degrees=False).as_matrix()[:2, :2] + + def _random_sheer(self): + + # Uniform random scale between 0.95 and 1.05 + random_scale = self.random_state.uniform(0.95, 1.05) + # Binary random to select upper or lower off diagonal for sheer + is_upper = self.random_state.randint(2, dtype=bool) + + # Create shear amount between -0.6 and 0.6 + shear_amount = self.random_state.uniform(-0.6, 0.6) + + # Create random sheer matrix for PC + shear_matrix = np.eye(2) * random_scale + if is_upper: + shear_matrix[0, 1] = shear_amount + else: + shear_matrix[1, 0] = shear_amount + return shear_matrix, random_scale, shear_amount, is_upper + + def _random_translation(self): + # Uniform random scale between -100 and 100 for x, y, -20 and 20 for z + return self.random_state.uniform([-100, -100, -20], [100, 100, 20]) + + def __getitem__(self, index): + sequence_id = self.sequence_ids[index // self.pcs_per_sequence] + sequence = self.sequence_loader.load_sequence(sequence_id) + sequence_idx = index % self.pcs_per_sequence + pc, _ = sequence.load(sequence_idx, sequence_idx) + + # Create random rotation matrix for PC + rotation_matrix = self._random_rot() + pc.points[:2] = rotation_matrix @ pc.points[:2] + shear_matrix, random_scale, shear_amount, is_upper = self._random_sheer( + ) + translation = self._random_translation() + pc.points[:2] = shear_matrix @ pc.points[:2] + return { + "pc": pc.to_fixed_array(self.max_pc_points), + "translation": translation, + "random_scale": random_scale, + "shear_amount": shear_amount, + "is_upper": np.int64(is_upper), + } diff --git a/dataloaders/sequence_dataset.py b/dataloaders/sequence_dataset.py new file mode 100644 index 0000000..8864419 --- /dev/null +++ b/dataloaders/sequence_dataset.py @@ -0,0 +1,295 @@ +import torch +import numpy as np +from tqdm import tqdm +import enum +from typing import Union, List, Tuple, Dict, Optional, Any +from pointclouds import to_fixed_array +from pathlib import Path + + +class OriginMode(enum.Enum): + FIRST_ENTRY = 0 + LAST_ENTRY = 1 + + +class SubsequenceRawDataset(torch.utils.data.Dataset): + + def __init__(self, + sequence_loader, + subsequence_length: int, + max_sequence_length: int, + origin_mode: Union[OriginMode, str], + max_pc_points: int = 90000, + subset_fraction: float = 1.0, + subset_mode='random', + shuffle=False): + assert subset_mode in [ + 'random', 'sequential' + ], f"subset_mode must be 'random' or 'sequential', got {subset_mode}" + assert 1.0 >= subset_fraction > 0.0, f"subset_fraction must be in (0.0, 1.0], got {subset_fraction}" + self.sequence_loader = sequence_loader + assert subsequence_length > 0, f"subsequence_length must be > 0, got {subsequence_length}" + assert max_sequence_length > 0, f"max_sequence_length must be > 0, got {max_sequence_length}" + + self.subsequence_length = subsequence_length + self.max_sequence_length = max_sequence_length + + if isinstance(origin_mode, str): + origin_mode = OriginMode[origin_mode] + self.origin_mode = origin_mode + self.max_pc_points = max_pc_points + + self.subsequence_id_begin_index = [] + for id in sequence_loader.get_sequence_ids(): + num_subsequences = max_sequence_length - subsequence_length + 1 + assert num_subsequences > 0, f"num_subsequences must be > 0, got {num_subsequences}" + self.subsequence_id_begin_index.extend([ + (id, i) for i in range(num_subsequences) + ]) + + self.subsequence_id_shuffled_index = list( + range(len(self.subsequence_id_begin_index))) + + subsequence_max_index = int( + len(self.subsequence_id_shuffled_index) * subset_fraction) + + if subset_fraction < 1.0 and subset_mode == 'sequential': + print( + f"Using only {subsequence_max_index} of {len(self.subsequence_id_shuffled_index)} sequences." + ) + self.subsequence_id_shuffled_index = self.subsequence_id_shuffled_index[: + subsequence_max_index] + + if shuffle: + random_state = np.random.RandomState( + len(self.subsequence_id_begin_index)) + random_state.shuffle(self.subsequence_id_shuffled_index) + + if subset_fraction < 1.0 and subset_mode == 'random': + print( + f"Using only {subsequence_max_index} of {len(self.subsequence_id_shuffled_index)} sequences." + ) + self.subsequence_id_shuffled_index = self.subsequence_id_shuffled_index[: + subsequence_max_index] + + def __len__(self): + return len(self.subsequence_id_shuffled_index) + + def _get_subsequence(self, index): + assert index >= 0 and index < len( + self + ), f"index must be >= 0 and < len(self), got {index} and {len(self)}" + assert index < len( + self.subsequence_id_shuffled_index + ), f"index must be < len(self.subsequence_id_shuffled_index), got {index} and {len(self.subsequence_id_shuffled_index)}" + shuffled_index = self.subsequence_id_shuffled_index[index] + id, subsequence_begin_index = self.subsequence_id_begin_index[ + shuffled_index] + sequence = self.sequence_loader.load_sequence(id) + assert len( + sequence + ) >= self.max_sequence_length, f"actual len(sequence) must be >= promised self.max_sequence_length, got {len(sequence)} vs {self.max_sequence_length}" + + if self.origin_mode == OriginMode.FIRST_ENTRY: + origin_idx = subsequence_begin_index + elif self.origin_mode == OriginMode.LAST_ENTRY: + origin_idx = subsequence_begin_index + self.subsequence_length - 1 + else: + raise ValueError(f"Unknown origin mode {self.origin_mode}") + + assert origin_idx >= 0 and origin_idx < len( + sequence + ), f"origin_idx must be >= 0 and < len(sequence), got {origin_idx} and {len(sequence)}" + assert subsequence_begin_index >= 0 and subsequence_begin_index + self.subsequence_length <= len( + sequence + ), f"offset must be >= 0 and offset + self.subsequence_length <= len(sequence), got subsequence_begin_index {subsequence_begin_index} and len(sequence) {len(sequence)} for max sequence len {self.max_sequence_length} and a subsequence length {self.subsequence_length}" + subsequence_lst = [ + sequence.load(subsequence_begin_index + i, origin_idx) + for i in range(self.subsequence_length) + ] + return subsequence_lst + + def _subsequence_lst_to_output_dict(self, + subsequence_lst) -> Dict[str, Any]: + pc_arrays = [ + e['relative_pc'].to_fixed_array(self.max_pc_points) + for e in subsequence_lst + ] + pose_arrays = [e['relative_pose'].to_array() for e in subsequence_lst] + pc_with_ground_arrays = [ + e['relative_pc_with_ground'].to_fixed_array(self.max_pc_points) + for e in subsequence_lst + ] + is_ground_arrays = [ + to_fixed_array(e['is_ground_points'].astype(np.float32), + self.max_pc_points) for e in subsequence_lst + ] + log_ids = [e['log_id'] for e in subsequence_lst] + log_idxes = [e['log_idx'] for e in subsequence_lst] + pc_array_stack = np.stack(pc_arrays, axis=0).astype(np.float32) + pose_array_stack = np.stack(pose_arrays, axis=0).astype(np.float32) + pc_with_ground_array_stack = np.stack(pc_with_ground_arrays, + axis=0).astype(np.float32) + is_ground_array_stack = np.stack(is_ground_arrays, + axis=0).astype(np.float32) + + return { + "pc_array_stack": pc_array_stack, + "pose_array_stack": pose_array_stack, + "pc_with_ground_array_stack": pc_with_ground_array_stack, + "is_ground_array_stack": is_ground_array_stack, + "log_ids": log_ids, + "log_idxes": log_idxes, + } + + def __getitem__(self, index): + assert index >= 0 and index < len( + self + ), f"index must be >= 0 and < len(self), got {index} and {len(self)}" + + subsequence_lst = self._get_subsequence(index) + + ret_dict = self._subsequence_lst_to_output_dict(subsequence_lst) + ret_dict['data_index'] = index + return ret_dict + + +class SubsequenceSupervisedFlowDataset(SubsequenceRawDataset): + + def _subsequence_lst_to_output_dict(self, + subsequence_lst) -> Dict[str, Any]: + ret_dict = super()._subsequence_lst_to_output_dict(subsequence_lst) + flowed_pc_arrays = [ + e['relative_flowed_pc'].to_fixed_array(self.max_pc_points) + for e in subsequence_lst + ] + pc_class_masks = [ + to_fixed_array(e['pc_classes'].astype(np.float32), + self.max_pc_points) for e in subsequence_lst + ] + flowed_pc_array_stack = np.stack(flowed_pc_arrays, + axis=0).astype(np.float32) + pc_class_mask_stack = np.stack(pc_class_masks, + axis=0).astype(np.float32) + ret_dict['flowed_pc_array_stack'] = flowed_pc_array_stack + ret_dict['pc_class_mask_stack'] = pc_class_mask_stack + return ret_dict + + +class SubsequenceSupervisedFlowSpecificSubsetDataset( + SubsequenceSupervisedFlowDataset): + + def __init__(self, + sequence_loader, + subsequence_length: int, + max_sequence_length: int, + origin_mode: Union[OriginMode, str], + subset_file: Path, + max_pc_points: int = 90000): + super().__init__(sequence_loader, subsequence_length, + max_sequence_length, origin_mode, max_pc_points) + subset_file = Path(subset_file) + assert subset_file.exists( + ), f"subset file {self.subset_file} does not exist" + self.subset_list = self._parse_subset_file(subset_file) + + def _parse_subset_file(self, subset_file) -> List[Tuple[str, int]]: + # Load each file line by line and extract tuple of (log_id, log_idx) + with open(subset_file, 'r') as f: + lines = f.readlines() + res_list = [] + for line in lines: + log_id, log_idx = line.split(",") + res_list.append((log_id, int(log_idx))) + return res_list + + def __len__(self): + return len(self.subset_list) + + def _get_subsequence(self, index): + assert index >= 0 and index < len( + self + ), f"index must be >= 0 and < len(self), got {index} and {len(self)}" + log_id, log_idx = self.subset_list[index] + sequence = self.sequence_loader.load_sequence(log_id) + + if self.origin_mode == OriginMode.FIRST_ENTRY: + origin_idx = log_idx + elif self.origin_mode == OriginMode.LAST_ENTRY: + origin_idx = log_idx + self.subsequence_length - 1 + else: + raise ValueError(f"Unknown origin mode {self.origin_mode}") + + subsequence_lst = [ + sequence.load(log_idx + i, origin_idx) + for i in range(self.subsequence_length) + ] + + # Special process the last entry in the subsequence because it does not have a flow but we still + # want to use it for eval, so we need to shim in a flow of zeros and a pc_classes of -1 + + e = subsequence_lst[-1] + if e['relative_flowed_pc'] is None: + e['relative_flowed_pc'] = e['relative_pc'] + if e['pc_classes'] is None: + e['pc_classes'] = np.zeros(e['relative_pc'].points.shape[0]) * -1 + + return subsequence_lst + + +class SubsequenceUnsupervisedFlowDataset(SubsequenceRawDataset): + + def _subsequence_lst_to_output_dict(self, + subsequence_lst) -> Dict[str, Any]: + ret_dict = super()._subsequence_lst_to_output_dict(subsequence_lst) + + def _squeeze_flow(flow: np.ndarray) -> np.ndarray: + if flow.ndim == 3: + assert flow.shape[ + 0] == 1, f"Flow must have 1 channel, got {flow.shape[0]}" + return flow.squeeze(0) + elif flow.ndim == 2: + return flow + else: + raise ValueError( + f"Flow must have 2 or 3 dimensions, got {flow.ndim}") + + flow_arrays = [ + to_fixed_array(_squeeze_flow(e['flow']), self.max_pc_points) + for e in subsequence_lst + ] + flow_array_stack = np.stack(flow_arrays, axis=0).astype(np.float32) + ret_dict['flow_array_stack'] = flow_array_stack + return ret_dict + + +class ConcatDataset(torch.utils.data.Dataset): + r"""Dataset to concatenate multiple datasets. + + Args: + datasets (sequence): List of datasets to be concatenated + """ + + def __init__(self, datasets: List[torch.utils.data.Dataset]): + self.datasets = datasets + for d in self.datasets: + assert isinstance( + d, torch.utils.data.Dataset + ), f"ConcatDataset only supports datasets, got {type(d)}" + self._length = sum(len(d) for d in self.datasets) + print( + f"Concatenated {len(self.datasets)} datasets with total length {self._length})" + ) + + def __len__(self): + return self._length + + def __getitem__(self, idx): + assert idx >= 0, f"Index must be >= 0, got {idx}" + assert idx < len( + self), f"Index must be < len(self), got {idx} and {len(self)}" + for d in self.datasets: + if idx < len(d): + return d[idx] + idx -= len(d) + raise IndexError('Index out of range') \ No newline at end of file diff --git a/dataloaders/var_len_sequence_dataset.py b/dataloaders/var_len_sequence_dataset.py new file mode 100644 index 0000000..44ee59d --- /dev/null +++ b/dataloaders/var_len_sequence_dataset.py @@ -0,0 +1,191 @@ +import torch +import numpy as np +from tqdm import tqdm +import enum +from typing import Union, List, Tuple, Dict, Optional, Any +from dataloaders.sequence_dataset import OriginMode +from pathlib import Path +from .sequence_dataset import SubsequenceRawDataset, SubsequenceSupervisedFlowDataset, SubsequenceUnsupervisedFlowDataset +from functools import partial + + +class VarLenSubsequenceRawDataset(SubsequenceRawDataset): + + def __init__(self, + sequence_loader, + subsequence_length: int, + origin_mode: Union[OriginMode, str], + max_pc_points: int = 90000, + subset_fraction: float = 1.0, + shuffle=False): + self.sequence_loader = sequence_loader + + # Subsequence length is the number of pointclouds projected into the given frame. + assert subsequence_length > 0, f"subsequence_length must be > 0, got {subsequence_length}" + self.subsequence_length = subsequence_length + + if isinstance(origin_mode, str): + origin_mode = OriginMode[origin_mode] + self.origin_mode = origin_mode + self.max_pc_points = max_pc_points + + self.index_array_range, self.sequence_list = self._build_sequence_lookup( + ) + self.shuffled_idx_lookup = self._build_shuffle_lookup( + shuffle, subset_fraction) + + def _build_sequence_lookup(self): + ids = sorted(self.sequence_loader.get_sequence_ids()) + index_range_list = [0] + sequence_list = [] + for id in ids: + sequence = self.sequence_loader.load_sequence(id) + + # Get number of unique subsequences in this sequence. + sequence_length = len(sequence) + num_subsequences = sequence_length - self.subsequence_length + 1 + + index_range_list.append(index_range_list[-1] + num_subsequences) + sequence_list.append(sequence) + + index_range_array = np.array(index_range_list) + return index_range_array, sequence_list + + def _build_shuffle_lookup(self, shuffle, subset_fraction): + shuffled_idx_lookup = np.arange(self.index_array_range[-1]) + if shuffle: + np.random.shuffle(shuffled_idx_lookup) + + assert 1.0 >= subset_fraction > 0.0, f"subset_fraction must be in (0.0, 1.0], got {subset_fraction}" + if subset_fraction == 1.0: + return shuffled_idx_lookup + max_index = int(len(shuffled_idx_lookup) * subset_fraction) + print( + f"Using only {max_index} of {len(shuffled_idx_lookup)} sequences.") + return shuffled_idx_lookup[:max_index] + + def __len__(self): + return len(self.shuffled_idx_lookup) + + def _global_idx_to_seq_and_seq_idx(self, input_global_idx): + assert input_global_idx >= 0 and input_global_idx < len( + self + ), f"global_idx must be >= 0 and < len(self), got {input_global_idx} and {len(self)}" + + global_idx = self.shuffled_idx_lookup[input_global_idx] + + # Find the sequence that contains this index. self.index_array_range provides a + # sorted global index range table, whose index can extract the relevant sequence + # from self.sequence_list. + seq_idx = np.searchsorted( + self.index_array_range, global_idx, side='right') - 1 + assert seq_idx >= 0 and seq_idx < len( + self.sequence_list + ), f"seq_idx must be >= 0 and < len(self.sequence_list), got {seq_idx} and {len(self.sequence_list)}" + + sequence = self.sequence_list[seq_idx] + sequence_idx = global_idx - self.index_array_range[seq_idx] + + assert sequence_idx >= 0 and sequence_idx < len( + sequence + ), f"sequence_idx must be >= 0 and < len(sequence), got {sequence_idx} and {len(sequence)}" + return sequence, sequence_idx + + def _get_subsequence(self, global_idx): + + sequence, subsequence_begin_index = self._global_idx_to_seq_and_seq_idx( + global_idx) + + if self.origin_mode == OriginMode.FIRST_ENTRY: + origin_idx = subsequence_begin_index + elif self.origin_mode == OriginMode.LAST_ENTRY: + origin_idx = subsequence_begin_index + self.subsequence_length - 1 + else: + raise ValueError(f"Unknown origin mode {self.origin_mode}") + + assert origin_idx >= 0 and origin_idx < len( + sequence + ), f"origin_idx must be >= 0 and < len(sequence), got {origin_idx} and {len(sequence)}" + assert subsequence_begin_index >= 0 and subsequence_begin_index + self.subsequence_length <= len( + sequence + ), f"offset must be >= 0 and offset + self.subsequence_length <= len(sequence), got subsequence_begin_index {subsequence_begin_index} and len(sequence) {len(sequence)} for max sequence len {self.max_sequence_length} and a subsequence length {self.subsequence_length}" + subsequence_lst = [ + sequence.load(subsequence_begin_index + i, origin_idx) + for i in range(self.subsequence_length) + ] + return subsequence_lst + + +class VarLenSubsequenceSupervisedFlowDataset(VarLenSubsequenceRawDataset): + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.__getitem__ = partial( + SubsequenceSupervisedFlowDataset.__getitem__, self=self) + + +class VarLenSubsequenceSupervisedFlowSpecificSubsetDataset( + VarLenSubsequenceSupervisedFlowDataset): + + def __init__(self, + sequence_loader, + subsequence_length: int, + origin_mode: Union[OriginMode, str], + subset_file: Path, + max_pc_points: int = 90000): + super().__init__(sequence_loader, subsequence_length, origin_mode, + max_pc_points) + subset_file = Path(subset_file) + assert subset_file.exists( + ), f"subset file {self.subset_file} does not exist" + self.subset_list = self._parse_subset_file(subset_file) + + def _parse_subset_file(self, subset_file) -> List[Tuple[str, int]]: + # Load each file line by line and extract tuple of (log_id, log_idx) + with open(subset_file, 'r') as f: + lines = f.readlines() + res_list = [] + for line in lines: + log_id, log_idx = line.split(",") + res_list.append((log_id, int(log_idx))) + return res_list + + def __len__(self): + return len(self.subset_list) + + def _get_subsequence(self, index): + assert index >= 0 and index < len( + self + ), f"index must be >= 0 and < len(self), got {index} and {len(self)}" + log_id, log_idx = self.subset_list[index] + sequence = self.sequence_loader.load_sequence(log_id) + + if self.origin_mode == OriginMode.FIRST_ENTRY: + origin_idx = log_idx + elif self.origin_mode == OriginMode.LAST_ENTRY: + origin_idx = log_idx + self.subsequence_length - 1 + else: + raise ValueError(f"Unknown origin mode {self.origin_mode}") + + subsequence_lst = [ + sequence.load(log_idx + i, origin_idx) + for i in range(self.subsequence_length) + ] + + # Special process the last entry in the subsequence because it does not have a flow but we still + # want to use it for eval, so we need to shim in a flow of zeros and a pc_classes of -1 + e = subsequence_lst[-1] + if e['relative_flowed_pc'] is None: + e['relative_flowed_pc'] = e['relative_pc'] + if e['pc_classes'] is None: + e['pc_classes'] = np.zeros(e['relative_pc'].points.shape[0]) * -1 + + return subsequence_lst + + +class VarLenSubsequenceUnsupervisedFlowDataset(VarLenSubsequenceRawDataset): + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.__getitem__ = partial( + SubsequenceUnsupervisedFlowDataset.__getitem__, self=self) diff --git a/dataloaders/waymo_supervised_flow.py b/dataloaders/waymo_supervised_flow.py new file mode 100644 index 0000000..968db6d --- /dev/null +++ b/dataloaders/waymo_supervised_flow.py @@ -0,0 +1,112 @@ +from pathlib import Path +from collections import defaultdict +from typing import List, Tuple, Dict, Optional, Any +import pandas as pd +from pointclouds import PointCloud, SE3, SE2 +import numpy as np +from loader_utils import load_pickle + + +class WaymoSupervisedFlowSequence(): + + def __init__(self, sequence_folder: Path, verbose: bool = False): + self.sequence_folder = Path(sequence_folder) + self.sequence_files = sorted(self.sequence_folder.glob('*.pkl')) + assert len(self.sequence_files + ) > 0, f'no frames found in {self.sequence_folder}' + + def __repr__(self) -> str: + return f'WaymoRawSequence with {len(self)} frames' + + def __len__(self): + return len(self.sequence_files) + + def _load_idx(self, idx: int): + assert idx < len( + self + ), f'idx {idx} out of range, len {len(self)} for {self.dataset_dir}' + pickle_path = self.sequence_files[idx] + pkl = load_pickle(pickle_path, verbose=False) + pc = PointCloud(pkl['car_frame_pc']) + flow = pkl['flow'] + labels = pkl['label'] + pose = SE3.from_array(pkl['pose']) + return pc, flow, labels, pose + + def cleanup_flow(self, flow): + flow[np.isnan(flow)] = 0 + flow[np.isinf(flow)] = 0 + flow_speed = np.linalg.norm(flow, axis=1) + flow[flow_speed > 30] = 0 + return flow + + def load(self, idx: int, start_idx: int) -> Dict[str, Any]: + assert idx < len( + self + ), f'idx {idx} out of range, len {len(self)} for {self.dataset_dir}' + + idx_pc, idx_flow, idx_labels, idx_pose = self._load_idx(idx) + # Unfortunatly, the flow has some artifacts that we need to clean up. These are very adhoc and will + # need to be updated if the flow is updated. + idx_flow = self.cleanup_flow(idx_flow) + _, _, _, start_pose = self._load_idx(start_idx) + + relative_pose = start_pose.inverse().compose(idx_pose) + relative_global_frame_pc = idx_pc.transform(relative_pose) + car_frame_flowed_pc = idx_pc.flow(idx_flow) + relative_global_frame_flowed_pc = car_frame_flowed_pc.transform( + relative_pose) + + return { + "relative_pc": relative_global_frame_pc, + "relative_pc_with_ground": relative_global_frame_pc, + "is_ground_points": np.zeros(len(relative_global_frame_pc), dtype=bool), + "relative_pose": relative_pose, + "relative_flowed_pc": relative_global_frame_flowed_pc, + "pc_classes": idx_labels, + "pc_is_ground": (idx_labels == -1), + "log_id": self.sequence_folder.name, + "log_idx": idx, + } + + def load_frame_list( + self, relative_to_idx: Optional[int]) -> List[Dict[str, Any]]: + + return [ + self.load(idx, + relative_to_idx if relative_to_idx is not None else idx) + for idx in range(len(self)) + ] + + +class WaymoSupervisedFlowSequenceLoader(): + + def __init__(self, + sequence_dir: Path, + log_subset: Optional[List[str]] = None, + verbose: bool = False): + self.dataset_dir = Path(sequence_dir) + self.verbose = verbose + assert self.dataset_dir.is_dir( + ), f'dataset_dir {sequence_dir} does not exist' + + # Load list of sequences from sequence_dir + sequence_dir_lst = sorted(self.dataset_dir.glob("*/")) + + self.log_lookup = {e.name: e for e in sequence_dir_lst} + + # Intersect with log_subset + if log_subset is not None: + self.log_lookup = { + k: v + for k, v in sorted(self.log_lookup.items()) + if k in set(log_subset) + } + + def get_sequence_ids(self): + return sorted(self.log_lookup.keys()) + + def load_sequence(self, log_id: str) -> WaymoSupervisedFlowSequence: + sequence_folder = self.log_lookup[log_id] + return WaymoSupervisedFlowSequence(sequence_folder, + verbose=self.verbose) diff --git a/dataloaders/waymo_unsupervised_flow.py b/dataloaders/waymo_unsupervised_flow.py new file mode 100644 index 0000000..06f5127 --- /dev/null +++ b/dataloaders/waymo_unsupervised_flow.py @@ -0,0 +1,94 @@ +from pathlib import Path +from collections import defaultdict +from typing import List, Tuple, Dict, Optional, Any +import pandas as pd +from pointclouds import PointCloud, SE3, SE2 +import numpy as np +from loader_utils import load_npz + +from . import WaymoSupervisedFlowSequenceLoader, WaymoSupervisedFlowSequence + + +class WaymoUnsupervisedFlowSequence(WaymoSupervisedFlowSequence): + + def __init__(self, + sequence_folder: Path, + flow_folder: Path, + verbose: bool = False): + super().__init__(sequence_folder, verbose) + self.flow_folder = Path(flow_folder) + self.flow_files = sorted(self.flow_folder.glob('*.npz')) + # We expect the number of flow files to be one fewer than the number + # of sequence files, with the last sequence file missing a flow + assert len(self.flow_files) == len(self.sequence_files) - 1, \ + f'Number of flow files {len(self.flow_files)} does not match number of sequence files {len(self.sequence_files)}' + + def __repr__(self) -> str: + return f'WaymoUnsupervisedFlowSequence with {len(self)} frames' + + def __len__(self): + return len(self.flow_files) + + def _load_idx(self, idx: int): + pc, _, labels, pose = super()._load_idx(idx) + flow_path = self.flow_files[idx] + flow_dict = dict(load_npz(flow_path, verbose=False)) + flow = flow_dict['flow'] + valid_idxes = flow_dict['valid_idxes'] + return pc, flow, valid_idxes, labels, pose + + def load(self, idx: int, start_idx: int) -> Dict[str, Any]: + assert idx < len( + self + ), f'idx {idx} out of range, len {len(self)} for {self.dataset_dir}' + + idx_pc, idx_flow, idx_valid_idxes, idx_labels, idx_pose = self._load_idx( + idx) + _, _, _, _, start_pose = self._load_idx(start_idx) + + relative_pose = start_pose.inverse().compose(idx_pose) + relative_global_frame_pc = idx_pc.transform(relative_pose) + + return { + "relative_pc": relative_global_frame_pc, + "relative_pose": relative_pose, + "flow": idx_flow, + "valid_idxes": idx_valid_idxes, + "log_id": self.sequence_folder.name, + "log_idx": idx, + } + + +class WaymoUnsupervisedFlowSequenceLoader(WaymoSupervisedFlowSequenceLoader): + + def __init__(self, + raw_data_path: Path, + flow_data_path: Path, + verbose: bool = False): + super().__init__(raw_data_path, verbose=verbose) + self.flow_data_path = Path(flow_data_path) + assert self.flow_data_path.is_dir( + ), f"flow_data_path {self.flow_data_path} is not a directory" + # Load list of flows from flow_data_path + flow_dir_list = sorted(self.flow_data_path.glob("*/")) + assert len(flow_dir_list + ) > 0, f"flow_data_path {self.flow_data_path} is empty" + self.flow_lookup = {e.name: e for e in flow_dir_list} + + # Intersect the flow_lookup keys with log_lookup + self.log_lookup = { + k: v + for k, v in self.log_lookup.items() if k in self.flow_lookup + } + + # Set of keys that are in both log_lookup and flow_lookup should now be the same + assert set(self.log_lookup.keys()) == set( + self.flow_lookup.keys() + ), f"Log and flow lookup keys do not match. Missing keys: {set(self.log_lookup.keys()) - set(self.flow_lookup.keys())}" + + def load_sequence(self, log_id: str) -> WaymoSupervisedFlowSequence: + sequence_folder = self.log_lookup[log_id] + flow_folder = self.flow_lookup[log_id] + return WaymoUnsupervisedFlowSequence(sequence_folder, + flow_folder, + verbose=self.verbose) diff --git a/docker/Dockerfile b/docker/Dockerfile new file mode 100644 index 0000000..a282cec --- /dev/null +++ b/docker/Dockerfile @@ -0,0 +1,34 @@ +FROM nvidia/cudagl:11.3.0-devel-ubuntu20.04 +SHELL ["/bin/bash", "-c"] +# Set the timezone info because otherwise tzinfo blocks install +# flow and ignores the non-interactive frontend command 🤬🤬🤬 +RUN ln -snf /usr/share/zoneinfo/America/New_York /etc/localtime && echo "/usr/share/zoneinfo/America/New_York" > /etc/timezone + +# Core system packages +RUN apt-get update --fix-missing +RUN apt install -y software-properties-common wget curl gpg gcc git make + +# Install miniconda to /miniconda +RUN curl -LO http://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh +RUN bash Miniconda3-latest-Linux-x86_64.sh -p /miniconda -b +RUN rm Miniconda3-latest-Linux-x86_64.sh +ENV PATH=/miniconda/bin:${PATH} +RUN conda update -y conda + +RUN apt install -y apt-utils + +ENV TORCH_CUDA_ARCH_LIST="Ampere;Turing;Pascal" +ENV FORCE_CUDA="1" +RUN conda update -y conda +RUN conda install numpy python=3.10 pytorch==1.12.1 torchvision torchaudio cudatoolkit=11.3 -c pytorch -y +RUN conda install -c fvcore -c iopath -c conda-forge fvcore iopath -y +RUN pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu113_pyt1121/download.html +RUN pip install mmcv-full==1.7.1 -f https://download.openmmlab.com/mmcv/dist/cu113/torch1.12.0/index.html +RUN pip install pyarrow pyquaternion open3d +RUN pip install tensorboard +RUN pip install pytorch-lightning==1.9.4 +RUN pip install nntime + +ENV PYTHONPATH=/project:${PYTHONPATH} +COPY bashrc /root/.bashrc +WORKDIR /project \ No newline at end of file diff --git a/docker/Dockerfileav2 b/docker/Dockerfileav2 new file mode 100644 index 0000000..9e3c766 --- /dev/null +++ b/docker/Dockerfileav2 @@ -0,0 +1,43 @@ +FROM nvidia/cudagl:11.3.0-devel-ubuntu20.04 +SHELL ["/bin/bash", "-c"] +# Set the timezone info because otherwise tzinfo blocks install +# flow and ignores the non-interactive frontend command 🤬🤬🤬 +RUN ln -snf /usr/share/zoneinfo/America/New_York /etc/localtime && echo "/usr/share/zoneinfo/America/New_York" > /etc/timezone + +# Core system packages +RUN apt-get update --fix-missing +RUN apt install -y software-properties-common wget curl gpg gcc git make + +# Install miniconda to /miniconda +RUN curl -LO http://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh +RUN bash Miniconda3-latest-Linux-x86_64.sh -p /miniconda -b +RUN rm Miniconda3-latest-Linux-x86_64.sh +ENV PATH=/miniconda/bin:${PATH} +RUN conda update -y conda + +RUN apt install -y apt-utils + +ENV TORCH_CUDA_ARCH_LIST="Ampere;Turing;Pascal" +ENV FORCE_CUDA="1" +RUN conda update -y conda +RUN conda install numpy python=3.10 pytorch==1.12.1 torchvision torchaudio cudatoolkit=11.3 -c pytorch -y +RUN conda install -c fvcore -c iopath -c conda-forge fvcore iopath -y +RUN pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu113_pyt1121/download.html +RUN pip install mmcv-full==1.7.1 -f https://download.openmmlab.com/mmcv/dist/cu113/torch1.12.0/index.html +RUN pip install pyarrow pyquaternion open3d +RUN pip install tensorboard +RUN pip install pytorch-lightning==1.9.4 +RUN pip install nntime + +RUN git clone https://github.com/argoverse/av2-api.git /av2-api +WORKDIR /av2-api +RUN git checkout 74f949ea8f09df2e3145abf9bf4a14232804584e +RUN conda install -c conda-forge mamba -y +RUN mamba install -c conda-forge conda-forge/label/rust_dev::rust av click joblib kornia maturin nox numba polars pyarrow pyproj universal_pathlib -y +RUN pip install git+https://github.com/JonathonLuiten/TrackEval.git +ENV OPENSSL_DIR=/av2-api +RUN pip install -e . +RUN apt install -y nano +ENV PYTHONPATH=/project:${PYTHONPATH} +COPY bashrc /root/.bashrc +WORKDIR /project \ No newline at end of file diff --git a/docker/Dockerfilewaymo b/docker/Dockerfilewaymo new file mode 100644 index 0000000..f681520 --- /dev/null +++ b/docker/Dockerfilewaymo @@ -0,0 +1,36 @@ +FROM nvidia/cudagl:11.3.0-devel-ubuntu20.04 +SHELL ["/bin/bash", "-c"] +# Set the timezone info because otherwise tzinfo blocks install +# flow and ignores the non-interactive frontend command 🤬🤬🤬 +RUN ln -snf /usr/share/zoneinfo/America/New_York /etc/localtime && echo "/usr/share/zoneinfo/America/New_York" > /etc/timezone + +# Core system packages +RUN apt-get update --fix-missing +RUN apt install -y software-properties-common wget curl gpg gcc git make + +# Install miniconda to /miniconda +RUN curl -LO http://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh +RUN bash Miniconda3-latest-Linux-x86_64.sh -p /miniconda -b +RUN rm Miniconda3-latest-Linux-x86_64.sh +ENV PATH=/miniconda/bin:${PATH} +RUN conda update -y conda + +RUN apt install -y apt-utils + +ENV TORCH_CUDA_ARCH_LIST="Ampere;Turing;Pascal" +ENV FORCE_CUDA="1" +RUN conda update -y conda + +RUN conda install python=3.9 pip -y +RUN pip install --upgrade pip +# ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CONDA_PREFIX/lib/ + +RUN apt install -y openexr libopenexr-dev +RUN pip install waymo-open-dataset-tf-2.11.0==1.5.0 +RUN pip install pytorch-lightning +RUN pip install opencv-python +RUN pip install open3d + +ENV PYTHONPATH=/project/data_prep_scripts/waymo:/project/:${PYTHONPATH} +COPY bashrc /root/.bashrc +WORKDIR /project/data_prep_scripts/waymo diff --git a/docker/bashrc b/docker/bashrc new file mode 100644 index 0000000..dec163a --- /dev/null +++ b/docker/bashrc @@ -0,0 +1,127 @@ +# ~/.bashrc: executed by bash(1) for non-login shells. +# see /usr/share/doc/bash/examples/startup-files (in the package bash-doc) +# for examples +# If not running interactively, don't do anything +case $- in + *i*) ;; + *) return;; +esac +# don't put duplicate lines or lines starting with space in the history. +# See bash(1) for more options +HISTCONTROL=ignoreboth + +# append to the history file, don't overwrite it +shopt -s histappend + +# check the window size after each command and, if necessary, +# update the values of LINES and COLUMNS. +shopt -s checkwinsize + +# If set, the pattern "**" used in a pathname expansion context will +# match all files and zero or more directories and subdirectories. +#shopt -s globstar + +# make less more friendly for non-text input files, see lesspipe(1) +[ -x /usr/bin/lesspipe ] && eval "$(SHELL=/bin/sh lesspipe)" + +# set variable identifying the chroot you work in (used in the prompt below) +if [ -z "${debian_chroot:-}" ] && [ -r /etc/debian_chroot ]; then + debian_chroot=$(cat /etc/debian_chroot) +fi + +# set a fancy prompt (non-color, unless we know we "want" color) +case "$TERM" in + xterm-color|*-256color) color_prompt=yes;; +esac + +# uncomment for a colored prompt, if the terminal has the capability; turned +# off by default to not distract the user: the focus in a terminal window +# should be on the output of commands, not on the prompt +force_color_prompt=yes + +if [ -n "$force_color_prompt" ]; then + if [ -x /usr/bin/tput ] && tput setaf 1 >&/dev/null; then + # We have color support; assume it's compliant with Ecma-48 + # (ISO/IEC-6429). (Lack of such support is extremely rare, and such + # a case would tend to support setf rather than setaf.) + color_prompt=yes + else + color_prompt= + fi +fi + +INFO_STR="\$? " +if [ "$color_prompt" = yes ]; then + PS1='${debian_chroot:+($debian_chroot)}'"$INFO_STR"'\[\033[01;31m\]\u@\h\[\033[00m\]:\[\033[01;34m\]\w\[\033[00m\]\$ ' + #PS1='${debian_chroot:+($debian_chroot)}\[\033[38;5;46m\]\u@\h\[\033[00m\]:\[\033[38;5;33m\]\w\[\033[00m\]\$ ' +else + PS1='${debian_chroot:+($debian_chroot)}'"$INFO_STR"'\u@\h:\w\$ ' +fi +unset color_prompt force_color_prompt + +# If this is an xterm set the title to user@host:dir +case "$TERM" in +xterm*|rxvt*) + PS1="\[\e]0;${debian_chroot:+($debian_chroot)}\u@\h: \w\a\]$PS1" + ;; +*) + ;; +esac + +# enable color support of ls and also add handy aliases +if [ -x /usr/bin/dircolors ]; then + test -r ~/.dircolors && eval "$(dircolors -b ~/.dircolors)" || eval "$(dircolors -b)" + alias ls='ls --color=auto' + #alias dir='dir --color=auto' + #alias vdir='vdir --color=auto' + + alias grep='grep --color=auto' + alias fgrep='fgrep --color=auto' + alias egrep='egrep --color=auto' +fi + +# colored GCC warnings and errors +export GCC_COLORS='error=01;31:warning=01;35:note=01;36:caret=01;32:locus=01:quote=01' + +# some more ls aliases +alias ll='ls -alF' +alias la='ls -A' +alias l='ls -CF' + +# Add an "alert" alias for long running commands. Use like so: +# sleep 10; alert +alias alert='notify-send --urgency=low -i "$([ $? = 0 ] && echo terminal || echo error)" "$(history|tail -n1|sed -e '\''s/^\s*[0-9]\+\s*//;s/[;&|]\s*alert$//'\'')"' + +# Alias definitions. +# You may want to put all your additions into a separate file like +# ~/.bash_aliases, instead of adding them here directly. +# See /usr/share/doc/bash-doc/examples in the bash-doc package. + +if [ -f ~/.bash_aliases ]; then + . ~/.bash_aliases +fi + +# enable programmable completion features (you don't need to enable +# this, if it's already enabled in /etc/bash.bashrc and /etc/profile +# sources /etc/bash.bashrc). +if ! shopt -oq posix; then + if [ -f /usr/share/bash-completion/bash_completion ]; then + . /usr/share/bash-completion/bash_completion + elif [ -f /etc/bash_completion ]; then + . /etc/bash_completion + fi +fi + +# Eternal bash history. +# --------------------- +# Undocumented feature which sets the size to "unlimited". +# http://stackoverflow.com/questions/9457233/unlimited-bash-history +export HISTFILESIZE= +export HISTSIZE= +export HISTTIMEFORMAT="[%F %T] " +# Change the file location because certain bash sessions truncate .bash_history file upon close. +# http://superuser.com/questions/575479/bash-history-truncated-to-500-lines-on-each-login +export HISTFILE=/root/.bash_history +# Force prompt to write history after every command. +# http://superuser.com/questions/20900/bash-history-loss +PROMPT_COMMAND="history -a; $PROMPT_COMMAND" \ No newline at end of file diff --git a/launch.sh b/launch.sh new file mode 100755 index 0000000..643816c --- /dev/null +++ b/launch.sh @@ -0,0 +1,13 @@ +#!/bin/bash +xhost + +docker run --gpus=all --rm -it \ + --shm-size=16gb \ + -v `pwd`:/project \ + -v /efs:/efs \ + -v /bigdata:/bigdata \ + -v /tmp/.X11-unix:/tmp/.X11-unix \ + -v `pwd`/docker_history.txt:/root/.bash_history \ + -e DISPLAY=$DISPLAY \ + -h $HOSTNAME \ + --privileged \ + kylevedder/zeroflow:latest diff --git a/launch_av2.sh b/launch_av2.sh new file mode 100755 index 0000000..a2e0225 --- /dev/null +++ b/launch_av2.sh @@ -0,0 +1,13 @@ +#!/bin/bash +xhost + +docker run --gpus=all --rm -it \ + --shm-size=16gb \ + -v `pwd`:/project \ + -v /efs:/efs \ + -v /bigdata:/bigdata \ + -v /tmp/.X11-unix:/tmp/.X11-unix \ + -v `pwd`/docker_history.txt:/root/.bash_history \ + -e DISPLAY=$DISPLAY \ + -h $HOSTNAME \ + --privileged \ + kylevedder/zeroflow_av2:latest diff --git a/launch_waymo.sh b/launch_waymo.sh new file mode 100755 index 0000000..acc0a22 --- /dev/null +++ b/launch_waymo.sh @@ -0,0 +1,15 @@ +#!/bin/bash +xhost + +docker run --gpus=all --rm -it \ + --shm-size=16gb \ + -v `pwd`:/project \ + -v /efs:/Datasets \ + -v /efs:/efs \ + -v /bigdata:/bigdata \ + -v /scratch:/scratch \ + -v /tmp/.X11-unix:/tmp/.X11-unix \ + -v `pwd`/docker_history.txt:/root/.bash_history \ + -e DISPLAY=$DISPLAY \ + -h $HOSTNAME \ + --privileged \ + kylevedder/zeroflow_waymo:latest diff --git a/loader_utils/__init__.py b/loader_utils/__init__.py new file mode 100644 index 0000000..1d68adf --- /dev/null +++ b/loader_utils/__init__.py @@ -0,0 +1,13 @@ +from .loaders import load_pickle, save_pickle, load_json, save_json, load_csv, save_csv, run_cmd, load_npz, save_npz, load_npy, save_npy, load_txt, save_txt, save_by_extension, load_by_extension +from .results import ResultInfo, CATEGORY_ID_TO_NAME, CATEGORY_NAME_TO_IDX, SPEED_BUCKET_SPLITS_METERS_PER_SECOND, ENDPOINT_ERROR_SPLITS_METERS, BACKGROUND_CATEGORIES, PEDESTRIAN_CATEGORIES, SMALL_VEHICLE_CATEGORIES, VEHICLE_CATEGORIES, ANIMAL_CATEGORIES, METACATAGORIES, METACATAGORY_TO_SHORTNAME + +__all__ = [ + 'load_pickle', 'save_pickle', 'load_json', 'save_json', 'load_csv', + 'save_csv', 'load_npz', 'save_npz', 'load_npy', 'save_npy', 'run_cmd', + 'load_txt', 'save_txt', 'save_by_extension', 'ResultInfo', + 'CATEGORY_NAME_TO_IDX', 'SPEED_BUCKET_SPLITS_METERS_PER_SECOND', + 'ENDPOINT_ERROR_SPLITS_METERS', 'BACKGROUND_CATEGORIES', + 'PEDESTRIAN_CATEGORIES', 'SMALL_VEHICLE_CATEGORIES', 'VEHICLE_CATEGORIES', + 'ANIMAL_CATEGORIES', 'METACATAGORIES', 'METACATAGORY_TO_SHORTNAME', + 'CATEGORY_ID_TO_NAME' +] diff --git a/loader_utils/loaders.py b/loader_utils/loaders.py new file mode 100644 index 0000000..81f94ef --- /dev/null +++ b/loader_utils/loaders.py @@ -0,0 +1,209 @@ +from pathlib import Path +import json +import pickle +import subprocess +import numpy as np + + +def _compute_size_metric(filepath: Path): + size = filepath.stat().st_size + if size < 1024: + return f'{size} B' + elif size < 1024 * 1024: + return f'{size / 1024:.2f} KB' + elif size < 1024 * 1024 * 1024: + return f'{size / (1024 * 1024):.2f} MB' + else: + return f'{size / (1024 * 1024 * 1024):.2f} GB' + + +def load_txt(filepath: Path, verbose: bool = True): + filepath = Path(filepath) + assert filepath.exists(), f'{filepath} does not exist' + if verbose: + print(f'Loading {filepath} of size {_compute_size_metric(filepath)}') + with open(filepath, 'r') as f: + return f.read() + + +def save_txt(filepath: Path, txt, verbose: bool = True): + filepath = Path(filepath) + if verbose: + print(f'Saving {filepath}', end='') + if filepath.exists(): + filepath.unlink() + filepath.parent.mkdir(parents=True, exist_ok=True) + with open(filepath, 'w') as f: + f.write(txt) + if verbose: + print(f"\rSaved {filepath} of size {_compute_size_metric(filepath)}") + + +def load_npz(filepath: Path, verbose: bool = True): + filepath = Path(filepath) + assert filepath.exists(), f'{filepath} does not exist' + if verbose: + print(f'Loading {filepath} of size {_compute_size_metric(filepath)}') + return np.load(filepath, allow_pickle=True) + + +def save_npz(filepath: Path, npz, verbose: bool = True): + filepath = Path(filepath) + if verbose: + print(f'Saving {filepath}', end='') + if filepath.exists(): + filepath.unlink() + filepath.parent.mkdir(parents=True, exist_ok=True) + np.savez_compressed(filepath, **npz) + if verbose: + print(f"\rSaved {filepath} of size {_compute_size_metric(filepath)}") + + +def load_npy(filepath: Path, verbose: bool = True): + filepath = Path(filepath) + assert filepath.exists(), f'{filepath} does not exist' + if verbose: + print(f'Loading {filepath} of size {_compute_size_metric(filepath)}') + return np.load(filepath, allow_pickle=True) + + +def save_npy(filepath: Path, npy, verbose: bool = True): + filepath = Path(filepath) + if verbose: + print(f'Saving {filepath}', end='') + if filepath.exists(): + filepath.unlink() + filepath.parent.mkdir(parents=True, exist_ok=True) + np.save(filepath, npy) + if verbose: + print(f"\rSaved {filepath} of size {_compute_size_metric(filepath)}") + +def load_pickle(filepath: Path, verbose: bool = True): + filepath = Path(filepath) + assert filepath.exists(), f'{filepath} does not exist' + if verbose: + print(f'Loading {filepath} of size {_compute_size_metric(filepath)}') + with open(filepath, 'rb') as f: + return pickle.load(f) + + +def save_pickle(filepath: Path, pkl, verbose: bool = True): + filepath = Path(filepath) + if verbose: + print(f'Saving {filepath}', end='') + if filepath.exists(): + filepath.unlink() + filepath.parent.mkdir(parents=True, exist_ok=True) + with open(filepath, 'wb') as f: + pickle.dump(pkl, f) + if verbose: + print(f"\rSaved {filepath} of size {_compute_size_metric(filepath)}") + + +def load_json(filename: Path, verbose: bool = True): + filename = Path(filename) + assert filename.exists(), f'{filename} does not exist' + if verbose: + print(f'Loading {filename} of size {_compute_size_metric(filename)}') + with open(filename) as f: + return json.load(f) + + +def save_json(filename: Path, contents, indent=None, verbose: bool = True): + if verbose: + print(f'Saving {filename}', end='') + filename = Path(filename) + if filename.exists(): + filename.unlink() + filename.parent.mkdir(parents=True, exist_ok=True) + with open(filename, 'w') as f: + json.dump(contents, f, indent=indent) + if verbose: + print(f"\rSaved {filename} of size {_compute_size_metric(filename)}") + + +def load_csv(filename: Path, dtype=str): + print(f'Loading {filename} of size {_compute_size_metric(filename)}') + with open(filename) as f: + return [[dtype(e.strip()) for e in line.strip().split(',')] + for line in f.readlines()] + + +def save_csv(filename: Path, contents: list, verbose: bool = True): + filename = Path(filename) + if verbose: + print(f'Saving {filename}', end='') + if filename.exists(): + filename.unlink() + filename.parent.mkdir(parents=True, exist_ok=True) + with open(filename, 'w') as f: + for line in contents: + f.write(f'{",".join([str(e) for e in line])}\n') + if verbose: + print(f"\rSaved {filename} of size {_compute_size_metric(filename)}") + +def save_by_extension(filename: Path, contents, verbose : bool =True): + filename = Path(filename) + # Make parents if they don't exist + filename.parent.mkdir(parents=True, exist_ok=True) + if filename.suffix == '.txt': + save_txt(filename, contents, verbose=verbose) + elif filename.suffix == '.npz': + save_npz(filename, contents, verbose=verbose) + elif filename.suffix == '.npy': + save_npy(filename, contents, verbose=verbose) + elif filename.suffix == '.pkl': + save_pickle(filename, contents, verbose=verbose) + elif filename.suffix == '.json': + save_json(filename, contents, verbose=verbose) + elif filename.suffix == '.csv': + save_csv(filename, contents, verbose=verbose) + else: + raise ValueError(f'Unknown file extension: {filename.suffix}') + +def load_by_extension(filename: Path, verbose : bool =True): + filename = Path(filename) + if filename.suffix == '.txt': + return load_txt(filename, verbose=verbose) + elif filename.suffix == '.npz': + return load_npz(filename, verbose=verbose) + elif filename.suffix == '.npy': + return load_npy(filename, verbose=verbose) + elif filename.suffix == '.pkl': + return load_pickle(filename, verbose=verbose) + elif filename.suffix == '.json': + return load_json(filename, verbose=verbose) + elif filename.suffix == '.csv': + return load_csv(filename, verbose=verbose) + else: + raise ValueError(f'Unknown file extension: {filename.suffix}') + +def symlink_files(old_dir: Path, new_dir: Path, file_name_list: list): + old_dir = Path(old_dir) + assert old_dir.exists(), f'{old_dir} does not exist' + new_dir = Path(new_dir) + assert new_dir.exists(), f'{new_dir} does not exist' + for fn in file_name_list: + if (new_dir / fn).exists(): + (new_dir / fn).unlink() + (new_dir / fn).symlink_to(old_dir / fn) + + +def run_cmd(cmd: str, return_stdout: bool = False): + res = subprocess.run(cmd, + shell=True, + encoding='utf-8', + capture_output=True) + if res.returncode != 0: + print(f'Error running command: {cmd}') + print(res.stdout) + print(res.stderr) + else: + lines = res.stdout.strip().splitlines() + if len(lines) > 0: + print(lines[0].strip()) + print('...') + print(lines[-1].strip()) + assert res.returncode == 0, f'Command {cmd} failed with return code {res.returncode}' + if return_stdout: + return res.stdout \ No newline at end of file diff --git a/loader_utils/results.py b/loader_utils/results.py new file mode 100644 index 0000000..76ce82d --- /dev/null +++ b/loader_utils/results.py @@ -0,0 +1,319 @@ +import numpy as np +from pathlib import Path +from .loaders import load_pickle +from typing import List, Dict, Tuple, Union, Optional, Any +import copy +import enum + +CATEGORY_ID_TO_NAME = { + -1: 'BACKGROUND', + 0: 'ANIMAL', + 1: 'ARTICULATED_BUS', + 2: 'BICYCLE', + 3: 'BICYCLIST', + 4: 'BOLLARD', + 5: 'BOX_TRUCK', + 6: 'BUS', + 7: 'CONSTRUCTION_BARREL', + 8: 'CONSTRUCTION_CONE', + 9: 'DOG', + 10: 'LARGE_VEHICLE', + 11: 'MESSAGE_BOARD_TRAILER', + 12: 'MOBILE_PEDESTRIAN_CROSSING_SIGN', + 13: 'MOTORCYCLE', + 14: 'MOTORCYCLIST', + 15: 'OFFICIAL_SIGNALER', + 16: 'PEDESTRIAN', + 17: 'RAILED_VEHICLE', + 18: 'REGULAR_VEHICLE', + 19: 'SCHOOL_BUS', + 20: 'SIGN', + 21: 'STOP_SIGN', + 22: 'STROLLER', + 23: 'TRAFFIC_LIGHT_TRAILER', + 24: 'TRUCK', + 25: 'TRUCK_CAB', + 26: 'VEHICULAR_TRAILER', + 27: 'WHEELCHAIR', + 28: 'WHEELED_DEVICE', + 29: 'WHEELED_RIDER' +} + +CATEGORY_NAME_TO_IDX = { + v: idx + for idx, (_, v) in enumerate(sorted(CATEGORY_ID_TO_NAME.items())) +} + +SPEED_BUCKET_SPLITS_METERS_PER_SECOND = [0, 0.5, 2.0, np.inf] +ENDPOINT_ERROR_SPLITS_METERS = [0, 0.05, 0.1, np.inf] + +BACKGROUND_CATEGORIES = [ + 'BOLLARD', 'CONSTRUCTION_BARREL', 'CONSTRUCTION_CONE', + 'MOBILE_PEDESTRIAN_CROSSING_SIGN', 'SIGN', 'STOP_SIGN' +] +PEDESTRIAN_CATEGORIES = [ + 'PEDESTRIAN', 'STROLLER', 'WHEELCHAIR', 'OFFICIAL_SIGNALER' +] +SMALL_VEHICLE_CATEGORIES = [ + 'BICYCLE', 'BICYCLIST', 'MOTORCYCLE', 'MOTORCYCLIST', 'WHEELED_DEVICE', + 'WHEELED_RIDER' +] +VEHICLE_CATEGORIES = [ + 'ARTICULATED_BUS', 'BOX_TRUCK', 'BUS', 'LARGE_VEHICLE', 'RAILED_VEHICLE', + 'REGULAR_VEHICLE', 'SCHOOL_BUS', 'TRUCK', 'TRUCK_CAB', 'VEHICULAR_TRAILER', + 'TRAFFIC_LIGHT_TRAILER', 'MESSAGE_BOARD_TRAILER' +] +ANIMAL_CATEGORIES = ['ANIMAL', 'DOG'] + +METACATAGORIES = { + "BACKGROUND": BACKGROUND_CATEGORIES, + "PEDESTRIAN": PEDESTRIAN_CATEGORIES, + "SMALL_MOVERS": SMALL_VEHICLE_CATEGORIES, + "LARGE_MOVERS": VEHICLE_CATEGORIES +} + +METACATAGORY_TO_SHORTNAME = { + "BACKGROUND": "BG", + "PEDESTRIAN": "PED", + "SMALL_MOVERS": "SMALL", + "LARGE_MOVERS": "LARGE" +} + + +# Enum for close, far, and full +class DistanceThreshold(enum.Enum): + CLOSE = 0 + FAR = 1 + ALL = 2 + + +class ResultInfo(): + + def __init__(self, + name, + path: Union[Path, dict], + full_distance: DistanceThreshold = DistanceThreshold.ALL): + self.name = name + self.path = path + if isinstance(full_distance, str): + full_distance = DistanceThreshold[full_distance.upper()] + if isinstance(path, dict): + assert 'per_class_bucketed_error_sum' in path, f"Path {path} does not contain per_class_bucketed_error_sum" + assert 'per_class_bucketed_error_count' in path, f"Path {path} does not contain per_class_bucketed_error_count" + self.results = copy.deepcopy(path) + else: + self.results = load_pickle(path) + + # Shapes should be (distance threshold, num_classes, num_speed_buckets, num_endpoint_error_buckets) + assert self.results[ + 'per_class_bucketed_error_sum'].ndim == 4, f"per_class_bucketed_error_sum should have 4 dimensions but has {self.results['per_class_bucketed_error_sum'].ndim}; shape is {self.results['per_class_bucketed_error_sum'].shape}" + assert self.results[ + 'per_class_bucketed_error_count'].ndim == 4, f"per_class_bucketed_error_count should have 4 dimensions but has {self.results['per_class_bucketed_error_count'].ndim}; shape is {self.results['per_class_bucketed_error_count'].shape}" + + self.full_distance = full_distance + if full_distance == DistanceThreshold.ALL: + # We need to sum the first axis of the sum and count arrays + self.results['per_class_bucketed_error_sum'] = np.sum( + self.results['per_class_bucketed_error_sum'], axis=0) + self.results['per_class_bucketed_error_count'] = np.sum( + self.results['per_class_bucketed_error_count'], axis=0) + elif full_distance == DistanceThreshold.CLOSE: + # We need to select the first element of the sum and count arrays + self.results['per_class_bucketed_error_sum'] = self.results[ + 'per_class_bucketed_error_sum'][0] + self.results['per_class_bucketed_error_count'] = self.results[ + 'per_class_bucketed_error_count'][0] + elif full_distance == DistanceThreshold.FAR: + # We need to select the second element of the sum and count arrays + self.results['per_class_bucketed_error_sum'] = self.results[ + 'per_class_bucketed_error_sum'][1] + self.results['per_class_bucketed_error_count'] = self.results[ + 'per_class_bucketed_error_count'][1] + else: + raise ValueError(f"Unknown distance threshold {full_distance}") + + assert self.results['per_class_bucketed_error_sum'].shape == ( + len(CATEGORY_ID_TO_NAME), + len(SPEED_BUCKET_SPLITS_METERS_PER_SECOND) - 1, + len(ENDPOINT_ERROR_SPLITS_METERS) - 1 + ), f"Shape of per_class_bucketed_error_sum is {self.results['per_class_bucketed_error_sum'].shape} but should be {(len(CATEGORY_ID_TO_NAME), len(SPEED_BUCKET_SPLITS_METERS_PER_SECOND) - 1, len(ENDPOINT_ERROR_SPLITS_METERS) - 1)}" + assert self.results['per_class_bucketed_error_count'].shape == ( + len(CATEGORY_ID_TO_NAME), + len(SPEED_BUCKET_SPLITS_METERS_PER_SECOND) - 1, + len(ENDPOINT_ERROR_SPLITS_METERS) - 1 + ), f"Shape of per_class_bucketed_error_count is {self.results['per_class_bucketed_error_count'].shape} but should be {(len(CATEGORY_ID_TO_NAME), len(SPEED_BUCKET_SPLITS_METERS_PER_SECOND) - 1, len(ENDPOINT_ERROR_SPLITS_METERS) - 1)}" + + def __repr__(self): + return self.name + " @ " + str(self.path) + + def pretty_name(self): + name_dict = { + "fastflow3d_nsfp_distilatation": "DRLFS (Ours)", + "fastflow3d_nsfp_distilatation_2x": "DRLFS 2x (Ours)", + "fastflow3d_nsfp_distilatation_0.5x": "DRLFS 0.5x (Ours)", + "fastflow3d_nsfp_distilatation_0.1x": "DRLFS 0.1x (Ours)", + "fastflow3d_nsfp_distilatation_speed_scaled": "DRLFS (Ours, speed scaled)", + "fastflow3d_supervised": "FastFlow3D", + "fastflow3d_chodosh_distilatation": "DRLFS (Chodosh et al. pseudolabels)", + "fastflow3d_supervised_no_scale": "FastFlow3D (no scale)", + "nsfp_unsupervised_sensor_val_cached": "NSFP", + } + if self.name in name_dict: + return name_dict[self.name] + print("WARNING: No pretty name for", self.name) + return self.name + + def speed_bucket_categories(self): + return list( + zip(SPEED_BUCKET_SPLITS_METERS_PER_SECOND, + SPEED_BUCKET_SPLITS_METERS_PER_SECOND[1:])) + + def endpoint_error_categories(self): + return list( + zip(ENDPOINT_ERROR_SPLITS_METERS, + ENDPOINT_ERROR_SPLITS_METERS[1:])) + + def get_metacatagory_error_count_by_speed(self, metacatagory): + full_error_sum = self.results['per_class_bucketed_error_sum'] + full_error_count = self.results['per_class_bucketed_error_count'] + category_idxes = [ + CATEGORY_NAME_TO_IDX[cat] for cat in METACATAGORIES[metacatagory] + ] + error_sum = full_error_sum[category_idxes] + error_count = full_error_count[category_idxes] + # Sum up all the catagories into a single metacatagory + error_sum = np.sum(error_sum, axis=0) + error_count = np.sum(error_count, axis=0) + # Sum up all the EPEs + error_sum = np.sum(error_sum, axis=1) + error_count = np.sum(error_count, axis=1) + # Only remaining axis is speed + if metacatagory == "BACKGROUND": + return np.array([np.sum(error_sum) + ]), np.array([np.sum(error_count)]) + return error_sum, error_count + + def get_metacatagory_error_count_by_speed_and_epe(self, metacatagory): + full_error_sum = self.results['per_class_bucketed_error_sum'] + full_error_count = self.results['per_class_bucketed_error_count'] + category_idxes = [ + CATEGORY_NAME_TO_IDX[cat] for cat in METACATAGORIES[metacatagory] + ] + error_sum = full_error_sum[category_idxes] + error_count = full_error_count[category_idxes] + # Sum up all the catagories into a single metacatagory + error_sum = np.sum(error_sum, axis=0) + error_count = np.sum(error_count, axis=0) + # Only remaining axis is speed + if metacatagory == "BACKGROUND": + return np.array([np.sum(error_sum, axis=0) + ]), np.array([np.sum(error_count, axis=0)]) + return error_sum, error_count + + def get_metacatagory_epe_by_speed(self, metacatagory): + error_sum, error_count = self.get_metacatagory_error_count_by_speed( + metacatagory) + return error_sum / error_count + + def get_mover_point_all_epe(self): + mover_metacatagories = set(METACATAGORIES.keys()) - {"BACKGROUND"} + total_error = 0 + total_count = 0 + + # Sum up the moving components of the movers + for metacatagory in mover_metacatagories: + error_sum, error_count = self.get_metacatagory_error_count_by_speed( + metacatagory) + # Extract the non-background error + total_error += np.sum(error_sum) + total_count += np.sum(error_count) + + return total_error / total_count + + def get_mover_point_dynamic_epe(self): + mover_metacatagories = set(METACATAGORIES.keys()) - {"BACKGROUND"} + total_error = 0 + total_count = 0 + + # Sum up the moving components of the movers + for metacatagory in mover_metacatagories: + error_sum, error_count = self.get_metacatagory_error_count_by_speed( + metacatagory) + # Extract the non-background error + total_error += np.sum(error_sum[1:]) + total_count += np.sum(error_count[1:]) + + return total_error / total_count + + def get_mover_point_static_epe(self): + mover_metacatagories = set(METACATAGORIES.keys()) - {"BACKGROUND"} + total_error = 0 + total_count = 0 + + # Sum up the moving components of the movers + for metacatagory in mover_metacatagories: + error_sum, error_count = self.get_metacatagory_error_count_by_speed( + metacatagory) + # Extract the non-background error + total_error += error_sum[0] + total_count += error_count[0] + + return total_error / total_count + + def get_nonmover_point_epe(self): + mover_metacatagories = set(METACATAGORIES.keys()) - {"BACKGROUND"} + + total_error = 0 + total_count = 0 + + # Sum up the stationary components of the movers + for metacatagory in mover_metacatagories: + error_sum, error_count = self.get_metacatagory_error_count_by_speed( + metacatagory) + total_error += error_sum[0] + total_count += error_count[0] + + error_sum, error_count = self.get_metacatagory_error_count_by_speed( + "BACKGROUND") + total_error += np.sum(error_sum) + total_count += np.sum(error_count) + + return total_error / total_count + + def get_metacatagory_count_by_epe(self, metacatagory): + full_error_count = self.results['per_class_bucketed_error_count'] + category_idxes = [ + CATEGORY_NAME_TO_IDX[cat] for cat in METACATAGORIES[metacatagory] + ] + error_count = full_error_count[category_idxes] + # Sum up all the catagories into a single metacatagory + error_count = np.sum(error_count, axis=0) + # Sum up all the speeds + error_count = np.sum(error_count, axis=0) + # Only remaining axis is epe + return error_count + + def get_foreground_counts(self): + total_epe_count = np.zeros((len(ENDPOINT_ERROR_SPLITS_METERS) - 1)) + mover_metacatagories = set(METACATAGORIES.keys()) - {"BACKGROUND"} + for metacatagory in mover_metacatagories: + _, error_count = self.get_metacatagory_error_count_by_speed_and_epe( + metacatagory) + # Extract the non-background error for all EPE buckets + total_epe_count += np.sum(error_count[1:], axis=0) + return total_epe_count + + def get_metacatagory_count_by_epe_strict(self, metacatagory): + count_by_epe = self.get_metacatagory_count_by_epe(metacatagory) + return count_by_epe[0] + + def get_metacatagory_count_by_epe_relaxed(self, metacatagory): + count_by_epe = self.get_metacatagory_count_by_epe(metacatagory) + return np.sum(count_by_epe[:2]) + + def get_metacatagory_count_by_epe_all(self, metacatagory): + count_by_epe = self.get_metacatagory_count_by_epe(metacatagory) + return np.sum(count_by_epe) + + def get_latency(self): + return self.results['average_forward_time'] diff --git a/model_wrapper.py b/model_wrapper.py new file mode 100644 index 0000000..4515581 --- /dev/null +++ b/model_wrapper.py @@ -0,0 +1,497 @@ +import numpy as np +import torch +import torch.nn as nn +import torch.optim as optim +import models +from pathlib import Path +from pointclouds import from_fixed_array + +import pytorch_lightning as pl +import torchmetrics +from typing import Dict, List, Tuple +from loader_utils import * +import nntime +import time +import multiprocessing as mp +from functools import partial + + +class EndpointDistanceMetricRawTorch(): + + def __init__(self, + class_id_to_name_map: Dict[int, str], + speed_bucket_splits_meters_per_second: List[float], + endpoint_error_splits_meters: List[float], + close_object_threshold_meters: float = 35.0, + per_frame_to_per_second_scale_factor: float = 10.0): + self.class_index_to_name_map = {} + self.class_id_to_index_map = {} + for cls_index, (cls_id, + cls_name) in enumerate(class_id_to_name_map.items()): + self.class_index_to_name_map[cls_index] = cls_name + self.class_id_to_index_map[cls_id] = cls_index + + self.speed_bucket_splits_meters_per_second = speed_bucket_splits_meters_per_second + self.endpoint_error_splits_meters = endpoint_error_splits_meters + + speed_bucket_bounds = self.speed_bucket_bounds() + endpoint_error_bucket_bounds = self.epe_bucket_bounds() + + # Bucket by IS_CLOSE x CLASS x SPEED x EPE + + self.per_class_bucketed_error_sum = torch.zeros( + (2, len(class_id_to_name_map), len(speed_bucket_bounds), + len(endpoint_error_bucket_bounds)), + dtype=torch.float) + self.per_class_bucketed_error_count = torch.zeros( + (2, len(class_id_to_name_map), len(speed_bucket_bounds), + len(endpoint_error_bucket_bounds)), + dtype=torch.long) + + self.close_object_threshold_meters = close_object_threshold_meters + self.per_frame_to_per_second_scale_factor = per_frame_to_per_second_scale_factor + + self.total_forward_time = torch.tensor(0.0) + self.total_forward_count = torch.tensor(0, dtype=torch.long) + + def to(self, device): + self.per_class_bucketed_error_sum = self.per_class_bucketed_error_sum.to( + device) + self.per_class_bucketed_error_count = self.per_class_bucketed_error_count.to( + device) + self.total_forward_time = self.total_forward_time.to(device) + self.total_forward_count = self.total_forward_count.to(device) + + def gather(self, gather_fn): + per_class_bucketed_error_sum = torch.sum(gather_fn( + self.per_class_bucketed_error_sum), + dim=0) + per_class_bucketed_error_count = torch.sum(gather_fn( + self.per_class_bucketed_error_count), + dim=0) + total_forward_time = torch.sum(gather_fn(self.total_forward_time), + dim=0) + total_forward_count = torch.sum(gather_fn(self.total_forward_count), + dim=0) + return per_class_bucketed_error_sum, per_class_bucketed_error_count, total_forward_time, total_forward_count + + def speed_bucket_bounds(self) -> List[Tuple[float, float]]: + return list( + zip(self.speed_bucket_splits_meters_per_second, + self.speed_bucket_splits_meters_per_second[1:])) + + def epe_bucket_bounds(self) -> List[Tuple[float, float]]: + return list( + zip(self.endpoint_error_splits_meters, + self.endpoint_error_splits_meters[1:])) + + def update_class_error(self, pc: torch.Tensor, class_id: int, + regressed_flow: torch.Tensor, + gt_flow: torch.Tensor): + + assert regressed_flow.shape == gt_flow.shape, f"Shapes do not match: {regressed_flow.shape} vs {gt_flow.shape}" + assert regressed_flow.shape[0] == pc.shape[ + 0], f"Shapes do not match: {regressed_flow.shape[0]} vs {pc.shape[0]}" + assert pc.shape[ + 1] == 3, f"Shapes do not match: {regressed_flow.shape[1]} vs 3" + + # L_\infty norm the XY coordinates needs to be within the close object threshold. + xy_points = pc[:, :2] + point_xy_distances = torch.norm(xy_points, dim=1, p=np.inf) + is_close_mask = point_xy_distances <= self.close_object_threshold_meters + + class_index = self.class_id_to_index_map[class_id] + endpoint_errors = torch.norm(regressed_flow - gt_flow, dim=1, p=2) + + gt_speeds = torch.norm(gt_flow, dim=1, + p=2) * self.per_frame_to_per_second_scale_factor + + # IS CLOSE DISAGGREGATION + for close_mask_idx, close_mask in enumerate( + [is_close_mask, ~is_close_mask]): + # SPEED DISAGGREGATION + for speed_idx, (lower_speed_bound, upper_speed_bound) in enumerate( + self.speed_bucket_bounds()): + speed_mask = (gt_speeds >= lower_speed_bound) & ( + gt_speeds < upper_speed_bound) + + # ENDPOINT ERROR DISAGGREGATION + for epe_idx, (lower_epe_bound, upper_epe_bound) in enumerate( + self.epe_bucket_bounds()): + endpoint_error_mask = ( + endpoint_errors >= + lower_epe_bound) & (endpoint_errors < upper_epe_bound) + total_mask = close_mask & speed_mask & endpoint_error_mask + + self.per_class_bucketed_error_sum[ + close_mask_idx, class_index, speed_idx, + epe_idx] += torch.sum(endpoint_errors[total_mask]) + self.per_class_bucketed_error_count[ + close_mask_idx, class_index, speed_idx, + epe_idx] += torch.sum(total_mask) + + def update_runtime(self, run_time: float, run_count: int): + self.total_forward_time += run_time + self.total_forward_count += run_count + + def reset(self): + self.per_class_bucketed_error_sum.zero_() + self.per_class_bucketed_error_count.zero_() + self.total_forward_time.zero_() + self.total_forward_count.zero_() + + +class ModelWrapper(pl.LightningModule): + + def __init__(self, cfg): + super().__init__() + self.cfg = cfg + self.model = getattr(models, cfg.model.name)(**cfg.model.args) + + if not hasattr(cfg, "is_trainable") or cfg.is_trainable: + self.loss_fn = getattr(models, + cfg.loss_fn.name)(**cfg.loss_fn.args) + + self.lr = cfg.learning_rate + if hasattr(cfg, "train_forward_args"): + self.train_forward_args = cfg.train_forward_args + else: + self.train_forward_args = {} + + if hasattr(cfg, "val_forward_args"): + self.val_forward_args = cfg.val_forward_args + else: + self.val_forward_args = {} + + self.has_labels = True if not hasattr(cfg, + "has_labels") else cfg.has_labels + + self.save_output_folder = None if not hasattr( + cfg, "save_output_folder") else cfg.save_output_folder + self.save_pool = None if self.save_output_folder is None else mp.Pool( + mp.cpu_count()) + + self.metric = EndpointDistanceMetricRawTorch( + CATEGORY_ID_TO_NAME, SPEED_BUCKET_SPLITS_METERS_PER_SECOND, + ENDPOINT_ERROR_SPLITS_METERS) + + def on_load_checkpoint(self, checkpoint): + checkpoint_lrs = set() + + for optimizer_state_idx in range(len(checkpoint['optimizer_states'])): + for param_group_idx in range( + len(checkpoint['optimizer_states'][optimizer_state_idx] + ['param_groups'])): + checkpoint_lrs.add( + checkpoint['optimizer_states'][optimizer_state_idx] + ['param_groups'][param_group_idx]['lr']) + + # If there are multiple learning rates, or if the learning rate is not the same as the one in the config, reset the optimizer. + # This is to handle the case where we want to resume training with a different learning rate. + + reset_learning_rate = (len(set(checkpoint_lrs)) != + 1) or (self.lr != list(checkpoint_lrs)[0]) + + if reset_learning_rate: + print("Resetting learning rate to the one in the config.") + checkpoint.pop('optimizer_states') + checkpoint.pop('lr_schedulers') + + def configure_optimizers(self): + self.optimizer = optim.Adam(self.parameters(), lr=self.lr) + return self.optimizer + + def training_step(self, input_batch, batch_idx): + model_res = self.model(input_batch, **self.train_forward_args) + loss_res = self.loss_fn(input_batch, model_res) + loss = loss_res.pop("loss") + self.log("train/loss", loss, on_step=True) + for k, v in loss_res.items(): + self.log(f"train/{k}", v, on_step=True) + return {"loss": loss} + + def _visualize_regressed_ground_truth_pcs(self, pc0_pc, pc1_pc, + regressed_flowed_pc0_to_pc1, + ground_truth_flowed_pc0_to_pc1): + import open3d as o3d + import numpy as np + pc0_pc = pc0_pc.cpu().numpy() + pc1_pc = pc1_pc.cpu().numpy() + regressed_flowed_pc0_to_pc1 = regressed_flowed_pc0_to_pc1.cpu().numpy() + ground_truth_flowed_pc0_to_pc1 = ground_truth_flowed_pc0_to_pc1.cpu( + ).numpy() + # make open3d visualizer + vis = o3d.visualization.Visualizer() + vis.create_window() + vis.get_render_option().point_size = 1.5 + vis.get_render_option().background_color = (0, 0, 0) + vis.get_render_option().show_coordinate_frame = True + # set up vector + vis.get_view_control().set_up([0, 0, 1]) + + # Add input PC + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(pc0_pc) + pc_color = np.zeros_like(pc0_pc) + pc_color[:, 0] = 1 + pc_color[:, 1] = 1 + pcd.colors = o3d.utility.Vector3dVector(pc_color) + vis.add_geometry(pcd) + + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(pc1_pc) + pc_color = np.zeros_like(pc1_pc) + pc_color[:, 1] = 1 + pc_color[:, 2] = 1 + pcd.colors = o3d.utility.Vector3dVector(pc_color) + vis.add_geometry(pcd) + + # Add line set between pc0 and gt pc1 + line_set = o3d.geometry.LineSet() + assert len(pc0_pc) == len( + ground_truth_flowed_pc0_to_pc1 + ), f"{len(pc0_pc)} != {len(ground_truth_flowed_pc0_to_pc1)}" + line_set_points = np.concatenate( + [pc0_pc, ground_truth_flowed_pc0_to_pc1], axis=0) + + lines = np.array([[i, i + len(ground_truth_flowed_pc0_to_pc1)] + for i in range(len(pc0_pc))]) + line_set.points = o3d.utility.Vector3dVector(line_set_points) + line_set.lines = o3d.utility.Vector2iVector(lines) + line_set.colors = o3d.utility.Vector3dVector( + [[0, 1, 0] for _ in range(len(lines))]) + vis.add_geometry(line_set) + + # Add line set between pc0 and regressed pc1 + line_set = o3d.geometry.LineSet() + assert len(pc0_pc) == len( + regressed_flowed_pc0_to_pc1 + ), f"{len(pc0_pc)} != {len(regressed_flowed_pc0_to_pc1)}" + line_set_points = np.concatenate([pc0_pc, regressed_flowed_pc0_to_pc1], + axis=0) + + lines = np.array([[i, i + len(regressed_flowed_pc0_to_pc1)] + for i in range(len(pc0_pc))]) + line_set.points = o3d.utility.Vector3dVector(line_set_points) + line_set.lines = o3d.utility.Vector2iVector(lines) + line_set.colors = o3d.utility.Vector3dVector( + [[0, 0, 1] for _ in range(len(lines))]) + vis.add_geometry(line_set) + + vis.run() + + def _save_output(self, input_batch, output_batch, batch_idx, delta_time): + + def _to_numpy(x): + if isinstance(x, torch.Tensor): + return x.detach().cpu().numpy() + if isinstance(x, dict): + return {k: _to_numpy(v) for k, v in x.items()} + if isinstance(x, list): + return [_to_numpy(v) for v in x] + return x + + log_ids = np.transpose(_to_numpy(input_batch["log_ids"])) + log_idxes = np.transpose(_to_numpy(input_batch["log_idxes"])) + pc_array_stack = _to_numpy(input_batch['pc_array_stack']) + est_flows = output_batch['flow'] + est_pc1_flows_valid_idxes = output_batch['pc0_valid_point_idxes'] + est_pc2_flows_valid_idxes = output_batch['pc1_valid_point_idxes'] + + if self.has_labels: + gt_flows = input_batch['flowed_pc_array_stack'] - input_batch[ + 'pc_array_stack'] + else: + gt_flows = [None] * len(pc_array_stack) + + save_list = [] + + for pc_arrays, log_ids, log_idxs, gt_flows, est_flow, est_pc1_flows_valid_idx, est_pc2_flows_valid_idx in zip( + pc_array_stack, log_ids, log_idxes, gt_flows, est_flows, + est_pc1_flows_valid_idxes, est_pc2_flows_valid_idxes): + pc1 = _to_numpy(from_fixed_array(pc_arrays[0])) + pc2 = _to_numpy(from_fixed_array(pc_arrays[1])) + log_id = log_ids[0] + log_idx = log_idxs[0] + est_flow = _to_numpy(from_fixed_array(est_flow)) + est_pc1_flows_valid_idx = _to_numpy(est_pc1_flows_valid_idx) + est_pc2_flows_valid_idx = _to_numpy(est_pc2_flows_valid_idx) + + standard_save_file = Path( + self.save_output_folder) / log_id / f"{log_idx:06d}.npz" + expanded_save_file = Path( + self.save_output_folder) / log_id / f"{log_idx:06d}.npy" + + standard_data_dict = { + 'flow': est_flow, + 'valid_idxes': est_pc1_flows_valid_idx, + 'delta_time': delta_time + } + + save_list.append((standard_save_file, standard_data_dict)) + + if self.has_labels: + gt_flow = _to_numpy(from_fixed_array(gt_flows[0])) + + expanded_data_dict = { + "pc1": pc1[est_pc1_flows_valid_idx], + "pc2": pc2[est_pc2_flows_valid_idx], + "gt_flow": gt_flow[est_pc1_flows_valid_idx], + "est_flow": est_flow, + "pc1_flows_valid_idx": est_pc1_flows_valid_idx, + "pc2_flows_valid_idx": est_pc2_flows_valid_idx, + } + + save_list.append((expanded_save_file, expanded_data_dict)) + + self.save_pool.starmap(partial(save_by_extension, verbose=False), + save_list) + + def validation_step(self, input_batch, batch_idx): + nntime.timer_start(self, "validation_forward") + start_time = time.time() + model_res = self.model(input_batch, **self.val_forward_args) + end_time = time.time() + nntime.timer_end(self, "validation_forward") + output_batch = model_res["forward"] + self.metric.to(self.device) + + if self.save_output_folder is not None: + self._save_output(input_batch, output_batch, batch_idx, + end_time - start_time) + + if not self.has_labels: + return + + self.metric.update_runtime(output_batch["batch_delta_time"], + len(input_batch["pc_array_stack"])) + + # Decode the mini-batch. + for minibatch_idx, (pc_array, flowed_pc_array, regressed_flow, + pc0_valid_point_idxes, pc1_valid_point_idxes, + class_info) in enumerate( + zip(input_batch["pc_array_stack"], + input_batch["flowed_pc_array_stack"], + output_batch["flow"], + output_batch["pc0_valid_point_idxes"], + output_batch["pc1_valid_point_idxes"], + input_batch["pc_class_mask_stack"])): + # This is written to support an arbitrary sequence length, but we only want to compute a flow + # off of the last frame. + pc0_pc = pc_array[-2][pc0_valid_point_idxes] + pc1_pc = pc_array[-1][pc1_valid_point_idxes] + ground_truth_flowed_pc0_to_pc1 = flowed_pc_array[-2][ + pc0_valid_point_idxes] + pc0_pc_class_info = class_info[-2][pc0_valid_point_idxes] + + ground_truth_flow = ground_truth_flowed_pc0_to_pc1 - pc0_pc + + assert pc0_pc.shape == ground_truth_flowed_pc0_to_pc1.shape, f"The input and ground truth pointclouds are not the same shape. {pc0_pc.shape} != {ground_truth_flowed_pc0_to_pc1.shape}" + assert pc0_pc.shape == regressed_flow.shape, f"The input pc and output flow are not the same shape. {pc0_pc.shape} != {regressed_flow.shape}" + + assert regressed_flow.shape == ground_truth_flow.shape, f"The regressed and ground truth flowed pointclouds are not the same shape." + + # regressed_flowed_pc0_to_pc1 = pc0_pc + regressed_flow + # if batch_idx % 64 == 0 and minibatch_idx == 0: + # self._visualize_regressed_ground_truth_pcs( + # pc0_pc, pc1_pc, regressed_flowed_pc0_to_pc1, + # ground_truth_flowed_pc0_to_pc1) + + # ======================== Compute Metrics Split By Class ======================== + + for cls_id in torch.unique(pc0_pc_class_info): + cls_mask = (pc0_pc_class_info == cls_id) + self.metric.update_class_error(pc0_pc[cls_mask], cls_id.item(), + regressed_flow[cls_mask], + ground_truth_flow[cls_mask]) + + def _dict_vals_to_numpy(self, d): + for k, v in d.items(): + if isinstance(v, dict): + d[k] = self._dict_vals_to_numpy(v) + else: + d[k] = v.cpu().numpy() + return d + + def _save_validation_data(self, save_dict): + save_pickle(f"validation_results/{self.cfg.filename}.pkl", save_dict) + try: + timing_out = f"validation_results/{self.cfg.filename}_timing.csv" + nntime.export_timings(self, timing_out) + except AssertionError as e: + print(f"Could not export timings. Skipping.") + + def _log_validation_metrics(self, validation_result_dict, verbose=True): + result_full_info = ResultInfo(Path(self.cfg.filename).stem, + validation_result_dict, + full_distance='ALL') + result_close_info = ResultInfo(Path(self.cfg.filename).stem, + validation_result_dict, + full_distance='CLOSE') + self.log("val/full/nonmover_epe", + result_full_info.get_nonmover_point_epe(), + sync_dist=False, + rank_zero_only=True) + self.log("val/full/mover_epe", + result_full_info.get_mover_point_dynamic_epe(), + sync_dist=False, + rank_zero_only=True) + self.log("val/close/nonmover_epe", + result_close_info.get_nonmover_point_epe(), + sync_dist=False, + rank_zero_only=True) + self.log("val/close/mover_epe", + result_close_info.get_mover_point_dynamic_epe(), + sync_dist=False, + rank_zero_only=True) + + if verbose: + print("Validation Results:") + print( + f"Close Mover EPE: {result_close_info.get_mover_point_dynamic_epe()}" + ) + print( + f"Close Nonmover EPE: {result_close_info.get_nonmover_point_epe()}" + ) + print( + f"Full Mover EPE: {result_full_info.get_mover_point_dynamic_epe()}" + ) + print( + f"Full Nonmover EPE: {result_full_info.get_nonmover_point_epe()}" + ) + + def validation_epoch_end(self, batch_parts): + import time + before_gather = time.time() + + # These are copies of the metric values on each rank. + per_class_bucketed_error_sum, per_class_bucketed_error_count, total_forward_time, total_forward_count = self.metric.gather( + self.all_gather) + + after_gather = time.time() + + print( + f"Rank {self.global_rank} gathers done in {after_gather - before_gather}." + ) + + # Reset the metric for the next epoch. We have to do this on each rank, and because we are using + # copies of the metric values above, we don't have to worry about over-writing the values. + self.metric.reset() + + if self.global_rank != 0: + return {} + + validation_result_dict = { + "per_class_bucketed_error_sum": per_class_bucketed_error_sum, + "per_class_bucketed_error_count": per_class_bucketed_error_count, + "average_forward_time": total_forward_time / total_forward_count + } + + validation_result_dict = self._dict_vals_to_numpy( + validation_result_dict) + + self._log_validation_metrics(validation_result_dict) + + self._save_validation_data(validation_result_dict) + + return {} \ No newline at end of file diff --git a/models/__init__.py b/models/__init__.py new file mode 100644 index 0000000..e429b56 --- /dev/null +++ b/models/__init__.py @@ -0,0 +1,13 @@ +from .joint_flow import JointFlow, JointFlowLoss, JointFlowOptimizationLoss +from .pretrain import PretrainEmbedding +from .fast_flow_3d import FastFlow3D, FastFlow3DSelfSupervisedLoss, FastFlow3DSupervisedLoss, FastFlow3DDistillationLoss +from .nsfp import NSFP, NSFPCached +from .nearest_neighbor import NearestNeighborFlow +from .chodosh import ChodoshDirectSolver, ChodoshProblemDumper + +__all__ = [ + 'JointFlow', 'NSFP', 'JointFlowLoss', 'JointFlowOptimizationLoss', + 'PretrainEmbedding', 'FastFlow3D', 'FastFlow3DSelfSupervisedLoss', + 'FastFlow3DSupervisedLoss', 'FastFlow3DDistillationLoss', + 'NearestNeighborFlow', 'NSFPCached', 'ChodoshDirectSolver', 'ChodoshProblemDumper' +] \ No newline at end of file diff --git a/models/attention/__init__.py b/models/attention/__init__.py new file mode 100644 index 0000000..5b49800 --- /dev/null +++ b/models/attention/__init__.py @@ -0,0 +1,3 @@ +from .attention import JointConvAttention + +__all__ = ["JointConvAttention"] \ No newline at end of file diff --git a/models/attention/attention.py b/models/attention/attention.py new file mode 100644 index 0000000..0f27065 --- /dev/null +++ b/models/attention/attention.py @@ -0,0 +1,86 @@ +import torch +import torch.nn as nn +import numpy as np + + +class ConvWithNonLinearity(nn.Module): + + def __init__(self, in_num_channels: int, out_num_channels: int, + kernel_size: int, stride: int, padding: int): + super().__init__() + self.conv = nn.Conv2d(in_num_channels, out_num_channels, kernel_size, + stride, padding) + self.gelu = nn.GELU() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + conv_res = self.conv(x) + return self.gelu(conv_res) + + +class JointConvAttention(nn.Module): + + def __init__(self, + pseudoimage_dims: tuple, + sequence_length: int, + per_element_num_channels: int, + per_element_output_params: int, + latent_image_area_power: int = 2): + super().__init__() + assert len( + pseudoimage_dims + ) == 2, f"pseudoimage_dims must be a tuple of length 2, got {len(pseudoimage_dims)}" + assert sequence_length > 0, f"sequence_length must be > 0, got {sequence_length}" + assert per_element_num_channels > 0, f"per_element_num_channels must be > 0, got {per_element_num_channels}" + assert per_element_output_params > 0, f"per_element_output_params must be > 0, got {per_element_output_params}" + + self.pseudoimage_x = pseudoimage_dims[0] + self.pseudoimage_y = pseudoimage_dims[1] + + assert self.pseudoimage_x > 0, f"pseudoimage_x must be > 0, got {self.pseudoimage_x}" + assert self.pseudoimage_y > 0, f"pseudoimage_y must be > 0, got {self.pseudoimage_y}" + + self.pseudoimage_stepdowns = np.ceil( + np.max(np.log2(np.array([ + self.pseudoimage_x, self.pseudoimage_y + ])))).astype(int) - latent_image_area_power + print(f"pseudoimage_stepdowns: {self.pseudoimage_stepdowns}") + + self.sequence_length = sequence_length + self.total_num_channels = per_element_num_channels * sequence_length + + self.latent_image_area = (2**latent_image_area_power)**2 + self.total_output_channels = np.ceil( + per_element_output_params * sequence_length / + self.latent_image_area).astype(int) + print(f"output_channels: {self.total_output_channels}") + self.reshaped_out_size = per_element_output_params * sequence_length + + self.conv_layers = nn.ModuleList() + print( + f"Appending first conv layer of {self.total_num_channels} channels to {self.total_output_channels} channels" + ) + self.conv_layers.append( + ConvWithNonLinearity(self.total_num_channels, + self.total_output_channels, 3, 2, 1)) + for _ in range(self.pseudoimage_stepdowns - 2): + print( + f"Appending conv layer of {self.total_output_channels} channels to {self.total_output_channels} channels" + ) + self.conv_layers.append( + ConvWithNonLinearity(self.total_output_channels, + self.total_output_channels, 3, 2, 1)) + self.conv_layers.append( + nn.Conv2d(self.total_output_channels, self.total_output_channels, + 3, 2, 1)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + assert isinstance( + x, torch.Tensor), f"x must be a torch.Tensor, got {type(x)}" + assert x.shape[ + 1] == self.total_num_channels, f"x must have shape (batch_size, {self.total_num_channels}, pseudoimage_x, pseudoimage_y), got {x.shape}" + + for conv_layer in self.conv_layers: + x = conv_layer(x) + + x = x.reshape(x.shape[0], -1)[:, :self.reshaped_out_size] + return x diff --git a/models/backbones/__init__.py b/models/backbones/__init__.py new file mode 100644 index 0000000..4b810a8 --- /dev/null +++ b/models/backbones/__init__.py @@ -0,0 +1,4 @@ +from .feature_pyramid import FeaturePyramidNetwork +from .fast_flow_unet import FastFlowUNet + +__all__ = ["FeaturePyramidNetwork", "FastFlowUNet"] \ No newline at end of file diff --git a/models/backbones/fast_flow_unet.py b/models/backbones/fast_flow_unet.py new file mode 100644 index 0000000..13c30ec --- /dev/null +++ b/models/backbones/fast_flow_unet.py @@ -0,0 +1,147 @@ +import torch +import torch.nn as nn + +from typing import Tuple + + +class ConvWithNorms(nn.Module): + + def __init__(self, in_num_channels: int, out_num_channels: int, + kernel_size: int, stride: int, padding: int): + super().__init__() + self.conv = nn.Conv2d(in_num_channels, out_num_channels, kernel_size, + stride, padding) + self.batchnorm = nn.BatchNorm2d(out_num_channels) + self.nonlinearity = nn.GELU() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + conv_res = self.conv(x) + if conv_res.shape[2] == 1 and conv_res.shape[3] == 1: + # This is a hack to get around the fact that batchnorm doesn't support + # 1x1 convolutions + batchnorm_res = conv_res + else: + batchnorm_res = self.batchnorm(conv_res) + return self.nonlinearity(batchnorm_res) + + +class BilinearDecoder(nn.Module): + + def __init__(self, scale_factor: int): + super().__init__() + self.scale_factor = scale_factor + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return nn.functional.interpolate(x, + scale_factor=self.scale_factor, + mode="bilinear", + align_corners=False) + + +# class UpsampleSkip(nn.Module): + +# def __init__(self, in_num_channels: int, out_num_channels: int): +# super().__init__() +# self.before_concat = nn.Sequential( +# nn.Conv2d(in_num_channels, in_num_channels, 1, 1, 0), +# BilinearDecoder(2), +# nn.Conv2d(in_num_channels, in_num_channels, 1, 1, 0), +# ) +# self.after_concat = nn.Sequential( +# nn.Conv2d(in_num_channels * 2, out_num_channels, 3, 1, 1), +# nn.Conv2d(out_num_channels, out_num_channels, 3, 1, 1)) + +# def forward(self, x: torch.Tensor, skip: torch.Tensor) -> torch.Tensor: +# print("UpsampleSkip forward") +# upsample_before = self.before_concat(x) + +# cat_result = torch.cat([upsample_before, skip], dim=1) +# try: +# upsample_after = self.after_concat(cat_result) +# except RuntimeError: +# breakpoint() +# return upsample_after + + +class UpsampleSkip(nn.Module): + + def __init__(self, skip_channels: int, latent_channels: int, + out_channels: int): + super().__init__() + self.u1_u2 = nn.Sequential( + nn.Conv2d(skip_channels, latent_channels, 1, 1, 0), + BilinearDecoder(2)) + self.u3 = nn.Conv2d(latent_channels, latent_channels, 1, 1, 0) + self.u4_u5 = nn.Sequential( + nn.Conv2d(2 * latent_channels, out_channels, 3, 1, 1), + nn.Conv2d(out_channels, out_channels, 3, 1, 1)) + + def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + u2_res = self.u1_u2(a) + u3_res = self.u3(b) + u5_res = self.u4_u5(torch.cat([u2_res, u3_res], dim=1)) + return u5_res + + +class FastFlowUNet(nn.Module): + """ + Standard UNet with a few modifications: + - Uses Bilinear interpolation instead of transposed convolutions + """ + + def __init__(self) -> None: + super().__init__() + + self.encoder_step_1 = nn.Sequential(ConvWithNorms(32, 64, 3, 2, 1), + ConvWithNorms(64, 64, 3, 1, 1), + ConvWithNorms(64, 64, 3, 1, 1), + ConvWithNorms(64, 64, 3, 1, 1)) + self.encoder_step_2 = nn.Sequential(ConvWithNorms(64, 128, 3, 2, 1), + ConvWithNorms(128, 128, 3, 1, 1), + ConvWithNorms(128, 128, 3, 1, 1), + ConvWithNorms(128, 128, 3, 1, 1), + ConvWithNorms(128, 128, 3, 1, 1), + ConvWithNorms(128, 128, 3, 1, 1)) + self.encoder_step_3 = nn.Sequential(ConvWithNorms(128, 256, 3, 2, 1), + ConvWithNorms(256, 256, 3, 1, 1), + ConvWithNorms(256, 256, 3, 1, 1), + ConvWithNorms(256, 256, 3, 1, 1), + ConvWithNorms(256, 256, 3, 1, 1), + ConvWithNorms(256, 256, 3, 1, 1)) + self.decoder_step1 = UpsampleSkip(512, 256, 256) + self.decoder_step2 = UpsampleSkip(256, 128, 128) + self.decoder_step3 = UpsampleSkip(128, 64, 64) + self.decoder_step4 = nn.Conv2d(64, 64, 3, 1, 1) + + def forward(self, pc0_B: torch.Tensor, + pc1_B: torch.Tensor) -> torch.Tensor: + + expected_channels = 32 + assert pc0_B.shape[ + 1] == expected_channels, f"Expected {expected_channels} channels, got {pc0_B.shape[1]}" + assert pc1_B.shape[ + 1] == expected_channels, f"Expected {expected_channels} channels, got {pc1_B.shape[1]}" + + pc0_F = self.encoder_step_1(pc0_B) + pc0_L = self.encoder_step_2(pc0_F) + pc0_R = self.encoder_step_3(pc0_L) + + pc1_F = self.encoder_step_1(pc1_B) + pc1_L = self.encoder_step_2(pc1_F) + pc1_R = self.encoder_step_3(pc1_L) + + Rstar = torch.cat([pc0_R, pc1_R], + dim=1) # torch.Size([1, 512, 64, 64]) + Lstar = torch.cat([pc0_L, pc1_L], + dim=1) # torch.Size([1, 256, 128, 128]) + Fstar = torch.cat([pc0_F, pc1_F], + dim=1) # torch.Size([1, 128, 256, 256]) + Bstar = torch.cat([pc0_B, pc1_B], + dim=1) # torch.Size([1, 64, 512, 512]) + + S = self.decoder_step1(Rstar, Lstar) + T = self.decoder_step2(S, Fstar) + U = self.decoder_step3(T, Bstar) + V = self.decoder_step4(U) + + return V \ No newline at end of file diff --git a/models/backbones/feature_pyramid.py b/models/backbones/feature_pyramid.py new file mode 100644 index 0000000..484b6b8 --- /dev/null +++ b/models/backbones/feature_pyramid.py @@ -0,0 +1,157 @@ +import torch +import torch.nn as nn +import numpy as np + + +class ConvWithNorms(nn.Module): + + def __init__(self, num_channels: int, kernel_size: int, stride: int, + padding: int): + super().__init__() + self.conv = nn.Conv2d(num_channels, num_channels, kernel_size, stride, + padding) + self.batchnorm = nn.BatchNorm2d(num_channels) + self.gelu = nn.GELU() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + conv_res = self.conv(x) + if conv_res.shape[2] == 1 and conv_res.shape[3] == 1: + # This is a hack to get around the fact that batchnorm doesn't support + # 1x1 convolutions + batchnorm_res = conv_res + else: + batchnorm_res = self.batchnorm(conv_res) + return self.gelu(batchnorm_res) + + +class ConvTransposeWithNorms(nn.Module): + + def __init__(self, + num_channels: int, + kernel_size: int, + stride: int, + padding: int, + output_padding: int = 0): + super().__init__() + self.conv = nn.ConvTranspose2d(num_channels, + num_channels, + kernel_size, + stride, + padding, + output_padding=output_padding) + self.batchnorm = nn.BatchNorm2d(num_channels) + self.gelu = nn.GELU() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + conv_res = self.conv(x) + if conv_res.shape[2] == 1 and conv_res.shape[3] == 1: + # This is a hack to get around the fact that batchnorm doesn't support + # 1x1 convolutions + batchnorm_res = conv_res + else: + batchnorm_res = self.batchnorm(conv_res) + return self.gelu(batchnorm_res) + + +class StepDown(nn.Module): + + def __init__(self, num_channels: int, num_filters_per_block: int): + super().__init__() + # Single stride 2 conv layer, followed by num_filters stride 1 conv layers + # Each conv is followed by a batchnorm and a gelu + self.conv_layers = nn.ModuleList() + self.conv_layers.append(ConvWithNorms(num_channels, 3, 2, 1)) + for _ in range(num_filters_per_block): + self.conv_layers.append(ConvWithNorms(num_channels, 3, 1, 1)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for conv_layer in self.conv_layers: + x = conv_layer(x) + return x + + +class StepUp(nn.Module): + + def __init__(self, num_channels: int, num_filters_per_block: int): + super().__init__() + # Single stride 2 conv layer, followed by num_filters stride 1 conv layers + # Each conv is followed by a batchnorm and a gelu + self.conv_layers = nn.ModuleList() + self.conv_layers.append( + ConvTransposeWithNorms(num_channels, 3, 2, 1, 1)) + for _ in range(num_filters_per_block): + self.conv_layers.append( + ConvTransposeWithNorms(num_channels, 3, 1, 1)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for conv_layer in self.conv_layers: + x = conv_layer(x) + return x + + +class StepUpPyramid(nn.Module): + + def __init__(self, num_steps: int, num_channels: int, num_filters: int): + super().__init__() + assert num_steps > 0, f"num_steps must be > 0, got {num_steps}" + assert num_channels > 0, f"num_channels must be > 0, got {num_channels}" + assert num_filters > 0, f"num_filters must be > 0, got {num_filters}" + + self.step_up_layers = nn.ModuleList() + for _ in range(num_steps): + self.step_up_layers.append(StepUp(num_channels, num_filters)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for step_up_layer in self.step_up_layers: + x = step_up_layer(x) + return x + + +class FeaturePyramidNetwork(nn.Module): + + def __init__(self, pseudoimage_dims: tuple, input_num_channels: int, + num_filters_per_block: int, num_layers_of_pyramid: int): + super().__init__() + pseudoimage_x, pseudoimage_y, = pseudoimage_dims + assert pseudoimage_x > 0, f"pseudoimage_x must be > 0, got {pseudoimage_x}" + assert pseudoimage_y > 0, f"pseudoimage_y must be > 0, got {pseudoimage_y}" + assert num_filters_per_block > 0, f"num_filters must be > 0, got {num_filters_per_block}" + assert num_layers_of_pyramid > 0, f"num_layers must be > 0, got {num_layers_of_pyramid}" + + self.num_filters = num_filters_per_block + self.num_layers = num_layers_of_pyramid + + self.down_layers = nn.ModuleList() + # Create down layers + for _ in range(num_layers_of_pyramid): + self.down_layers.append( + StepDown(input_num_channels, num_filters_per_block)) + + # For each down layer, we need a separate up pyramid layer + self.up_pyramids = nn.ModuleList() + for idx in range(num_layers_of_pyramid): + self.up_pyramids.append( + StepUpPyramid(idx + 1, input_num_channels, + num_filters_per_block)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # x is of shape (batch_size, num_channels, height, width) + # We need to downsample x num_layers times + # Each downsampled layer will be used in the up pyramid + downsampled_layers = [] + + down_x = x + for down_layer in self.down_layers: + down_x = down_layer(down_x) + downsampled_layers.append(down_x) + + # Now we need to upsample x num_layers times + # Each upsampled layer will be concatenated with the corresponding + # downsampled layer + upsampled_layers = [x] + for up_x, up_pyramid in zip(downsampled_layers, self.up_pyramids): + upsampled_layers.append(up_pyramid(up_x)) + + # Concatenate the upsampled layers together into a single tensor + upsampled_latent = torch.cat(upsampled_layers, dim=1) + return upsampled_latent diff --git a/models/chodosh.py b/models/chodosh.py new file mode 100644 index 0000000..5d774a7 --- /dev/null +++ b/models/chodosh.py @@ -0,0 +1,243 @@ +import numpy as np +from sklearn.cluster import DBSCAN +import torch +import torch.nn as nn + +import numpy as np +from models.embedders import DynamicVoxelizer +from models.nsfp_baseline import NSFPProcessor +from models.nsfp import NSFPCached + +from typing import Dict, Any, Optional +from collections import defaultdict +from pathlib import Path +import time + + +def fit_rigid(A, B): + """ + Fit Rigid transformation A @ R.T + t = B + """ + assert A.shape == B.shape + num_rows, num_cols = A.shape + + # find mean column wise + centroid_A = np.mean(A, axis=0) + centroid_B = np.mean(B, axis=0) + + # ensure centroids are 1x3 + centroid_A = centroid_A.reshape(1, -1) + centroid_B = centroid_B.reshape(1, -1) + + # subtract mean + Am = A - centroid_A + Bm = B - centroid_B + + H = Am.T @ Bm + # find rotation + U, S, Vt = np.linalg.svd(H) + R = Vt.T @ U.T + + # special reflection case + if np.linalg.det(R) < 0: + Vt[2, :] *= -1 + R = Vt.T @ U.T + + t = -centroid_A @ R.T + centroid_B + + return R, t + + +def ransac_fit_rigid(pts1, pts2, inlier_thresh, ntrials): + best_R = np.eye(3) + best_t = np.zeros((3, )) + best_inliers = (np.linalg.norm(pts1 - pts2, axis=-1) < inlier_thresh) + best_inliers_sum = best_inliers.sum() + for i in range(ntrials): + choice = np.random.choice(len(pts1), 3) + R, t = fit_rigid(pts1[choice], pts2[choice]) + inliers = (np.linalg.norm(pts1 @ R.T + t - pts2, axis=-1) < + inlier_thresh) + if inliers.sum() > best_inliers_sum: + best_R = R + best_t = t + best_inliers = inliers + best_inliers_sum = best_inliers.sum() + if best_inliers_sum / len(pts1) > 0.5: + break + best_R, best_t = fit_rigid(pts1[best_inliers], pts2[best_inliers]) + return best_R, best_t + + +def refine_flow(pc0, + flow_pred, + eps=0.4, + min_points=10, + motion_threshold=0.05, + inlier_thresh=0.2, + ntrials=250): + labels = DBSCAN(eps=eps, min_samples=min_points).fit_predict(pc0) + max_label = labels.max() + refined_flow = np.zeros_like(flow_pred) + for l in range(max_label + 1): + label_mask = labels == l + cluster_pts = pc0[label_mask] + cluster_flows = flow_pred[label_mask] + R, t = ransac_fit_rigid(cluster_pts, cluster_pts + cluster_flows, + inlier_thresh, ntrials) + if np.linalg.norm(t) < motion_threshold: + R = np.eye(3) + t = np.zeros((3, )) + refined_flow[label_mask] = (cluster_pts @ R.T + t) - cluster_pts + + refined_flow[labels == -1] = flow_pred[labels == -1] + return refined_flow + + +class ChodoshDirectSolver(NSFPCached): + + def __init__(self, VOXEL_SIZE, POINT_CLOUD_RANGE, SEQUENCE_LENGTH, + nsfp_save_folder: Path, chodosh_save_folder: Path) -> None: + super().__init__(VOXEL_SIZE, POINT_CLOUD_RANGE, SEQUENCE_LENGTH, + nsfp_save_folder) + self.out_folder = Path(chodosh_save_folder) + + def forward(self, batched_sequence: Dict[str, torch.Tensor]): + pc0_points_lst, pc0_valid_point_idxes, pc1_points_lst, pc1_valid_point_idxes = self._voxelize_batched_sequence( + batched_sequence) + + log_ids_lst = self._transpose_list(batched_sequence['log_ids']) + log_idxes_lst = self._transpose_list(batched_sequence['log_idxes']) + + flows = [] + delta_times = [] + refine_times = [] + for minibatch_idx, (pc0_points, pc0_valid_points_idx, log_ids, + log_idxes) in enumerate( + zip(pc0_points_lst, pc0_valid_point_idxes, + log_ids_lst, log_idxes_lst)): + assert len( + log_ids) == 2, f"Expect 2 frames for NSFP, got {len(log_ids)}" + assert len( + log_idxes + ) == 2, f"Expect 2 frames for NSFP, got {len(log_idxes)}" + + flow, flow_valid_idxes, delta_time = self._load_result( + log_ids[0], log_idxes[0]) + + assert flow_valid_idxes.shape == pc0_valid_points_idx.shape, f"{flow_valid_idxes.shape} != {pc0_valid_points_idx.shape}" + idx_equality = flow_valid_idxes == pc0_valid_points_idx.detach( + ).cpu().numpy() + assert idx_equality.all( + ), f"Points that disagree: {(~idx_equality).astype(int).sum()}" + assert flow.shape[0] == 1, f"{flow.shape[0]} != 1" + flow = flow.squeeze(0) + assert flow.shape == pc0_points.shape, f"{flow.shape} != {pc0_points.shape}" + + valid_pc0_points = pc0_points.detach().cpu().numpy() + + before_time = time.time() + refined_flow = refine_flow(valid_pc0_points, + flow).astype(np.float32) + after_time = time.time() + refine_time = after_time - before_time + delta_time += refine_time + + # warped_pc0_points = valid_pc0_points + refined_flow + + # self._visualize_result(valid_pc0_points, warped_pc0_points) + + # breakpoint() + + flows.append( + torch.from_numpy(refined_flow).to( + pc0_points.device).unsqueeze(0)) + delta_times.append(delta_time) + refine_times.append(refine_time) + print(f"Average delta time: {np.mean(delta_times)}") + print(f"Average refine time: {np.mean(refine_times)}") + return { + "forward": { + "flow": flows, + "batch_delta_time": np.sum(delta_times), + "pc0_points_lst": pc0_points_lst, + "pc0_valid_point_idxes": pc0_valid_point_idxes, + "pc1_points_lst": pc1_points_lst, + "pc1_valid_point_idxes": pc1_valid_point_idxes + } + } + + +class ChodoshProblemDumper(ChodoshDirectSolver): + + def _dump_problem(self, log_id: str, batch_idx: int, minibatch_idx: int, + pc, flow, flow_valid_idxes, nsfp_delta_time: float): + self.out_folder.mkdir(parents=True, exist_ok=True) + save_file = self.out_folder / f"{log_id}_{batch_idx}_{minibatch_idx}.npz" + # Delete save_file if exists + if save_file.exists(): + save_file.unlink() + np.savez(save_file, + pc=pc, + flow=flow, + valid_idxes=flow_valid_idxes, + nsfp_delta_time=nsfp_delta_time) + + def forward(self, batched_sequence: Dict[str, torch.Tensor]): + pc0_points_lst, pc0_valid_point_idxes, pc1_points_lst, pc1_valid_point_idxes = self._voxelize_batched_sequence( + batched_sequence) + + log_ids_lst = self._transpose_list(batched_sequence['log_ids']) + log_idxes_lst = self._transpose_list(batched_sequence['log_idxes']) + + flows = [] + delta_times = [] + for minibatch_idx, (pc0_points, pc0_valid_points_idx, log_ids, + log_idxes) in enumerate( + zip(pc0_points_lst, pc0_valid_point_idxes, + log_ids_lst, log_idxes_lst)): + log_id = batched_sequence['log_ids'][minibatch_idx][0] + batch_idx = batched_sequence['data_index'].item() + + assert len( + log_ids) == 2, f"Expect 2 frames for NSFP, got {len(log_ids)}" + assert len( + log_idxes + ) == 2, f"Expect 2 frames for NSFP, got {len(log_idxes)}" + + flow, flow_valid_idxes, delta_time = self._load_result( + log_ids[0], log_idxes[0]) + + assert flow_valid_idxes.shape == pc0_valid_points_idx.shape, f"{flow_valid_idxes.shape} != {pc0_valid_points_idx.shape}" + idx_equality = flow_valid_idxes == pc0_valid_points_idx.detach( + ).cpu().numpy() + assert idx_equality.all( + ), f"Points that disagree: {(~idx_equality).astype(int).sum()}" + assert flow.shape[0] == 1, f"{flow.shape[0]} != 1" + flow = flow.squeeze(0) + assert flow.shape == pc0_points.shape, f"{flow.shape} != {pc0_points.shape}" + + valid_pc0_points = pc0_points.detach().cpu().numpy() + + self._dump_problem(log_id, batch_idx, minibatch_idx, + valid_pc0_points, flow, flow_valid_idxes, + delta_time) + + # warped_pc0_points = valid_pc0_points + refined_flow + + # self._visualize_result(valid_pc0_points, warped_pc0_points) + + # breakpoint() + + flows.append(torch.from_numpy(flow).to(pc0_points.device)) + delta_times.append(delta_time) + return { + "forward": { + "flow": flows, + "batch_delta_time": np.sum(delta_times), + "pc0_points_lst": pc0_points_lst, + "pc0_valid_point_idxes": pc0_valid_point_idxes, + "pc1_points_lst": pc1_points_lst, + "pc1_valid_point_idxes": pc1_valid_point_idxes + } + } \ No newline at end of file diff --git a/models/embedders/__init__.py b/models/embedders/__init__.py new file mode 100644 index 0000000..593319a --- /dev/null +++ b/models/embedders/__init__.py @@ -0,0 +1,11 @@ +from .make_voxels import HardVoxelizer, DynamicVoxelizer +from .process_voxels import HardSimpleVFE, PillarFeatureNet, DynamicPillarFeatureNet +from .scatter import PointPillarsScatter + +from .embedder_model import HardEmbedder, DynamicEmbedder + +__all__ = [ + 'HardEmbedder', 'DynamicEmbedder', 'HardVoxelizer', 'HardSimpleVFE', + 'PillarFeatureNet', 'PointPillarsScatter', 'DynamicVoxelizer', + 'DynamicPillarFeatureNet' +] diff --git a/models/embedders/embedder_model.py b/models/embedders/embedder_model.py new file mode 100644 index 0000000..ef3c01a --- /dev/null +++ b/models/embedders/embedder_model.py @@ -0,0 +1,71 @@ +import torch +import torch.nn as nn + +from . import HardVoxelizer, DynamicVoxelizer, PillarFeatureNet, DynamicPillarFeatureNet, PointPillarsScatter + + +class HardEmbedder(nn.Module): + + def __init__(self, + voxel_size=(0.2, 0.2, 4), + pseudo_image_dims=(350, 350), + point_cloud_range=(-35, -35, -3, 35, 35, 1), + max_points_per_voxel=128, + feat_channels=64) -> None: + super().__init__() + self.voxelizer = HardVoxelizer( + voxel_size=voxel_size, + point_cloud_range=point_cloud_range, + max_points_per_voxel=max_points_per_voxel) + self.feature_net = PillarFeatureNet( + in_channels=3, + feat_channels=(feat_channels, ), + point_cloud_range=point_cloud_range, + voxel_size=voxel_size) + self.scatter = PointPillarsScatter(in_channels=feat_channels, + output_shape=pseudo_image_dims) + + def forward(self, points: torch.Tensor) -> torch.Tensor: + assert isinstance( + points, + torch.Tensor), f"points must be a torch.Tensor, got {type(points)}" + + output_voxels, output_voxel_coords, points_per_voxel = self.voxelizer( + points) + output_features = self.feature_net(output_voxels, points_per_voxel, + output_voxel_coords) + pseudoimage = self.scatter(output_features, output_voxel_coords) + + return pseudoimage + + +class DynamicEmbedder(nn.Module): + + def __init__(self, voxel_size, pseudo_image_dims, point_cloud_range, + feat_channels: int) -> None: + super().__init__() + self.voxelizer = DynamicVoxelizer(voxel_size=voxel_size, + point_cloud_range=point_cloud_range) + self.feature_net = DynamicPillarFeatureNet( + in_channels=3, + feat_channels=(feat_channels, ), + point_cloud_range=point_cloud_range, + voxel_size=voxel_size, + mode='avg') + self.scatter = PointPillarsScatter(in_channels=feat_channels, + output_shape=pseudo_image_dims) + + def forward(self, points: torch.Tensor) -> torch.Tensor: + + # List of points and coordinates for each batch + voxel_info_list = self.voxelizer(points) + + pseudoimage_lst = [] + for voxel_info_dict in voxel_info_list: + points = voxel_info_dict['points'] + coordinates = voxel_info_dict['voxel_coords'] + voxel_feats, voxel_coors = self.feature_net(points, coordinates) + pseudoimage = self.scatter(voxel_feats, voxel_coors) + pseudoimage_lst.append(pseudoimage) + # Concatenate the pseudoimages along the batch dimension + return torch.cat(pseudoimage_lst, dim=0), voxel_info_list diff --git a/models/embedders/make_voxels.py b/models/embedders/make_voxels.py new file mode 100644 index 0000000..c3284e0 --- /dev/null +++ b/models/embedders/make_voxels.py @@ -0,0 +1,89 @@ +import torch +import torch.nn as nn +from mmcv.ops import Voxelization +from typing import List, Tuple + + +class HardVoxelizer(nn.Module): + + def __init__(self, voxel_size, point_cloud_range, + max_points_per_voxel: int): + super().__init__() + assert max_points_per_voxel > 0, f"max_points_per_voxel must be > 0, got {max_points_per_voxel}" + + self.voxelizer = Voxelization(voxel_size, + point_cloud_range, + max_points_per_voxel, + deterministic=False) + + def forward(self, points: torch.Tensor): + assert isinstance( + points, + torch.Tensor), f"points must be a torch.Tensor, got {type(points)}" + not_nan_mask = ~torch.isnan(points).any(dim=2) + return {"voxel_coords": self.voxelizer(points[not_nan_mask])} + + +class DynamicVoxelizer(nn.Module): + + def __init__(self, voxel_size, point_cloud_range): + super().__init__() + self.voxel_size = voxel_size + self.point_cloud_range = point_cloud_range + self.voxelizer = Voxelization(voxel_size, + point_cloud_range, + max_num_points=-1) + + def _get_point_offsets(self, points: torch.Tensor, + voxel_coords: torch.Tensor): + + point_cloud_range = torch.tensor(self.point_cloud_range, + dtype=points.dtype, + device=points.device) + min_point = point_cloud_range[:3] + voxel_size = torch.tensor(self.voxel_size, + dtype=points.dtype, + device=points.device) + + # Voxel coords are in the form Z, Y, X :eyeroll:, convert to X, Y, Z + voxel_coords = voxel_coords[:, [2, 1, 0]] + + # Offsets are computed relative to min point + voxel_centers = voxel_coords * voxel_size + min_point + voxel_size / 2 + + return points - voxel_centers + + def forward( + self, + points: torch.Tensor) -> List[Tuple[torch.Tensor, torch.Tensor]]: + + batch_results = [] + for batch_idx in range(len(points)): + batch_points = points[batch_idx] + valid_point_idxes = torch.arange(batch_points.shape[0], + device=batch_points.device) + not_nan_mask = ~torch.isnan(batch_points).any(dim=1) + batch_non_nan_points = batch_points[not_nan_mask] + valid_point_idxes = valid_point_idxes[not_nan_mask] + batch_voxel_coords = self.voxelizer(batch_non_nan_points) + # If any of the coords are -1, then the point is not in the voxel grid and should be discarded + batch_voxel_coords_mask = (batch_voxel_coords != -1).all(dim=1) + + valid_batch_voxel_coords = batch_voxel_coords[ + batch_voxel_coords_mask] + valid_batch_non_nan_points = batch_non_nan_points[ + batch_voxel_coords_mask] + valid_point_idxes = valid_point_idxes[batch_voxel_coords_mask] + + point_offsets = self._get_point_offsets(valid_batch_non_nan_points, + valid_batch_voxel_coords) + + result_dict = { + "points": valid_batch_non_nan_points, + "voxel_coords": valid_batch_voxel_coords, + "point_idxes": valid_point_idxes, + "point_offsets": point_offsets + } + + batch_results.append(result_dict) + return batch_results \ No newline at end of file diff --git a/models/embedders/process_voxels.py b/models/embedders/process_voxels.py new file mode 100644 index 0000000..3cc6b3b --- /dev/null +++ b/models/embedders/process_voxels.py @@ -0,0 +1,416 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +from mmcv.ops import DynamicScatter + + +class HardSimpleVFE(nn.Module): + """Simple voxel feature encoder used in SECOND. + + It simply averages the values of points in a voxel. + + Args: + num_features (int, optional): Number of features to use. Default: 4. + """ + + def __init__(self, num_features=4): + super(HardSimpleVFE, self).__init__() + self.num_features = num_features + + def forward(self, features, num_points, coors): + """Forward function. + + Args: + features (torch.Tensor): Point features in shape + (N, M, 3(4)). N is the number of voxels and M is the maximum + number of points inside a single voxel. + num_points (torch.Tensor): Number of points in each voxel, + shape (N, ). + coors (torch.Tensor): Coordinates of voxels. + + Returns: + torch.Tensor: Mean of points inside each voxel in shape (N, 3(4)) + """ + points_mean = features[:, :, :self.num_features].sum( + dim=1, keepdim=False) / num_points.type_as(features).view(-1, 1) + return points_mean.contiguous() + + +class PFNLayer(nn.Module): + """Pillar Feature Net Layer. + The Pillar Feature Net is composed of a series of these layers, but the + PointPillars paper results only used a single PFNLayer. + Args: + in_channels (int): Number of input channels. + out_channels (int): Number of output channels. + last_layer (bool, optional): If last_layer, there is no + concatenation of features. Defaults to False. + mode (str, optional): Pooling model to gather features inside voxels. + Defaults to 'max'. + """ + + def __init__(self, + in_channels, + out_channels, + last_layer=False, + mode='max'): + + super().__init__() + self.fp16_enabled = False + self.name = 'PFNLayer' + self.last_vfe = last_layer + if not self.last_vfe: + out_channels = out_channels // 2 + self.units = out_channels + + self.norm = nn.BatchNorm1d(self.units, eps=1e-3, momentum=0.01) + self.linear = nn.Linear(in_channels, self.units, bias=False) + + assert mode in ['max', 'avg'] + self.mode = mode + + def forward(self, inputs, num_voxels=None, aligned_distance=None): + """Forward function. + Args: + inputs (torch.Tensor): Pillar/Voxel inputs with shape (N, M, C). + N is the number of voxels, M is the number of points in + voxels, C is the number of channels of point features. + num_voxels (torch.Tensor, optional): Number of points in each + voxel. Defaults to None. + aligned_distance (torch.Tensor, optional): The distance of + each points to the voxel center. Defaults to None. + Returns: + torch.Tensor: Features of Pillars. + """ + x = self.linear(inputs) + x = self.norm(x.permute(0, 2, 1).contiguous()).permute(0, 2, + 1).contiguous() + x = F.gelu(x) + + if self.mode == 'max': + if aligned_distance is not None: + x = x.mul(aligned_distance.unsqueeze(-1)) + x_max = torch.max(x, dim=1, keepdim=True)[0] + elif self.mode == 'avg': + if aligned_distance is not None: + x = x.mul(aligned_distance.unsqueeze(-1)) + x_max = x.sum(dim=1, + keepdim=True) / num_voxels.type_as(inputs).view( + -1, 1, 1) + + if self.last_vfe: + return x_max + else: + x_repeat = x_max.repeat(1, inputs.shape[1], 1) + x_concatenated = torch.cat([x, x_repeat], dim=2) + return x_concatenated + + +def get_paddings_indicator(actual_num, max_num, axis=0): + """Create boolean mask by actually number of a padded tensor. + Args: + actual_num (torch.Tensor): Actual number of points in each voxel. + max_num (int): Max number of points in each voxel + Returns: + torch.Tensor: Mask indicates which points are valid inside a voxel. + """ + actual_num = torch.unsqueeze(actual_num, axis + 1) + # tiled_actual_num: [N, M, 1] + max_num_shape = [1] * len(actual_num.shape) + max_num_shape[axis + 1] = -1 + max_num = torch.arange(max_num, dtype=torch.int, + device=actual_num.device).view(max_num_shape) + # tiled_actual_num: [[3,3,3,3,3], [4,4,4,4,4], [2,2,2,2,2]] + # tiled_max_num: [[0,1,2,3,4], [0,1,2,3,4], [0,1,2,3,4]] + paddings_indicator = actual_num.int() > max_num + # paddings_indicator shape: [batch_size, max_num] + return paddings_indicator + + +class PillarFeatureNet(nn.Module): + """Pillar Feature Net. + The network prepares the pillar features and performs forward pass + through PFNLayers. + Args: + in_channels (int, optional): Number of input features, + either x, y, z or x, y, z, r. Defaults to 4. + feat_channels (tuple, optional): Number of features in each of the + N PFNLayers. Defaults to (64, ). + with_distance (bool, optional): Whether to include Euclidean distance + to points. Defaults to False. + with_cluster_center (bool, optional): [description]. Defaults to True. + with_voxel_center (bool, optional): [description]. Defaults to True. + voxel_size (tuple[float], optional): Size of voxels, only utilize x + and y size. Defaults to (0.2, 0.2, 4). + point_cloud_range (tuple[float], optional): Point cloud range, only + utilizes x and y min. Defaults to (0, -40, -3, 70.4, 40, 1). + mode (str, optional): The mode to gather point features. Options are + 'max' or 'avg'. Defaults to 'max'. + legacy (bool, optional): Whether to use the new behavior or + the original behavior. Defaults to True. + """ + + def __init__(self, + in_channels=4, + feat_channels=(64, ), + with_distance=False, + with_cluster_center=True, + with_voxel_center=True, + voxel_size=(0.2, 0.2, 4), + point_cloud_range=(0, -40, -3, 70.4, 40, 1), + mode='max'): + super(PillarFeatureNet, self).__init__() + assert len(feat_channels) > 0 + if with_cluster_center: + in_channels += 3 + if with_voxel_center: + in_channels += 3 + if with_distance: + in_channels += 1 + self._with_distance = with_distance + self._with_cluster_center = with_cluster_center + self._with_voxel_center = with_voxel_center + self.fp16_enabled = False + # Create PillarFeatureNet layers + self.in_channels = in_channels + feat_channels = [in_channels] + list(feat_channels) + pfn_layers = [] + for i in range(len(feat_channels) - 1): + in_filters = feat_channels[i] + out_filters = feat_channels[i + 1] + if i < len(feat_channels) - 2: + last_layer = False + else: + last_layer = True + pfn_layers.append( + PFNLayer(in_filters, + out_filters, + last_layer=last_layer, + mode=mode)) + self.pfn_layers = nn.ModuleList(pfn_layers) + + # Need pillar (voxel) size and x/y offset in order to calculate offset + self.vx = voxel_size[0] + self.vy = voxel_size[1] + self.vz = voxel_size[2] + self.x_offset = self.vx / 2 + point_cloud_range[0] + self.y_offset = self.vy / 2 + point_cloud_range[1] + self.z_offset = self.vz / 2 + point_cloud_range[2] + self.point_cloud_range = point_cloud_range + + def forward(self, features, num_points, coors): + """Forward function. + Args: + features (torch.Tensor): Point features or raw points in shape + (N, M, C). + num_points (torch.Tensor): Number of points in each pillar. + coors (torch.Tensor): Coordinates of each voxel. + Returns: + torch.Tensor: Features of pillars. + """ + features_ls = [features] + # Find distance of x, y, and z from cluster center + if self._with_cluster_center: + points_mean = features[:, :, :3].sum( + dim=1, keepdim=True) / num_points.type_as(features).view( + -1, 1, 1) + f_cluster = features[:, :, :3] - points_mean + features_ls.append(f_cluster) + + # Find distance of x, y, and z from pillar center + dtype = features.dtype + if self._with_voxel_center: + f_center = torch.zeros_like(features[:, :, :3]) + f_center[:, :, 0] = features[:, :, 0] - ( + coors[:, 2].to(dtype).unsqueeze(1) * self.vx + self.x_offset) + f_center[:, :, 1] = features[:, :, 1] - ( + coors[:, 1].to(dtype).unsqueeze(1) * self.vy + self.y_offset) + f_center[:, :, 2] = features[:, :, 2] - ( + coors[:, 0].to(dtype).unsqueeze(1) * self.vz + self.z_offset) + features_ls.append(f_center) + + if self._with_distance: + points_dist = torch.norm(features[:, :, :3], 2, 2, keepdim=True) + features_ls.append(points_dist) + + # Combine together feature decorations + features = torch.cat(features_ls, dim=-1) + # The feature decorations were calculated without regard to whether + # pillar was empty. Need to ensure that + # empty pillars remain set to zeros. + voxel_count = features.shape[1] + mask = get_paddings_indicator(num_points, voxel_count, axis=0) + mask = torch.unsqueeze(mask, -1).type_as(features) + features *= mask + + for pfn in self.pfn_layers: + features = pfn(features, num_points) + + return features.squeeze(1) + + +class DynamicPillarFeatureNet(PillarFeatureNet): + """Pillar Feature Net using dynamic voxelization. + The network prepares the pillar features and performs forward pass + through PFNLayers. The main difference is that it is used for + dynamic voxels, which contains different number of points inside a voxel + without limits. + Args: + in_channels (int, optional): Number of input features, + either x, y, z or x, y, z, r. Defaults to 4. + feat_channels (tuple, optional): Number of features in each of the + N PFNLayers. Defaults to (64, ). + with_distance (bool, optional): Whether to include Euclidean distance + to points. Defaults to False. + with_cluster_center (bool, optional): [description]. Defaults to True. + with_voxel_center (bool, optional): [description]. Defaults to True. + voxel_size (tuple[float], optional): Size of voxels, only utilize x + and y size. Defaults to (0.2, 0.2, 4). + point_cloud_range (tuple[float], optional): Point cloud range, only + utilizes x and y min. Defaults to (0, -40, -3, 70.4, 40, 1). + norm_cfg ([type], optional): [description]. + Defaults to dict(type='BN1d', eps=1e-3, momentum=0.01). + mode (str, optional): The mode to gather point features. Options are + 'max' or 'avg'. Defaults to 'max'. + legacy (bool, optional): Whether to use the new behavior or + the original behavior. Defaults to True. + """ + + def __init__(self, + in_channels, + voxel_size, + point_cloud_range, + feat_channels=(64, ), + with_distance=False, + with_cluster_center=True, + with_voxel_center=True, + mode='max'): + super(DynamicPillarFeatureNet, + self).__init__(in_channels, + feat_channels, + with_distance, + with_cluster_center=with_cluster_center, + with_voxel_center=with_voxel_center, + voxel_size=voxel_size, + point_cloud_range=point_cloud_range, + mode=mode) + self.fp16_enabled = False + feat_channels = [self.in_channels] + list(feat_channels) + pfn_layers = [] + # TODO: currently only support one PFNLayer + + for i in range(len(feat_channels) - 1): + in_filters = feat_channels[i] + out_filters = feat_channels[i + 1] + if i > 0: + in_filters *= 2 + pfn_layers.append( + nn.Sequential( + nn.Linear(in_filters, out_filters, bias=False), + nn.BatchNorm1d(out_filters, eps=1e-3, momentum=0.01), + nn.ReLU(inplace=True))) + self.num_pfn = len(pfn_layers) + self.pfn_layers = nn.ModuleList(pfn_layers) + self.pfn_scatter = DynamicScatter(voxel_size, point_cloud_range, + (mode != 'max')) + self.cluster_scatter = DynamicScatter(voxel_size, + point_cloud_range, + average_points=True) + + def map_voxel_center_to_point(self, pts_coors, voxel_mean, voxel_coors): + """Map the centers of voxels to its corresponding points. + Args: + pts_coors (torch.Tensor): The coordinates of each points, shape + (M, 3), where M is the number of points. + voxel_mean (torch.Tensor): The mean or aggregated features of a + voxel, shape (N, C), where N is the number of voxels. + voxel_coors (torch.Tensor): The coordinates of each voxel. + Returns: + torch.Tensor: Corresponding voxel centers of each points, shape + (M, C), where M is the number of points. + """ + if pts_coors.shape[0] == 0: + return torch.zeros((0, voxel_mean.shape[1]), + dtype=voxel_mean.dtype, + device=voxel_mean.device) + # Step 1: scatter voxel into canvas + # Calculate necessary things for canvas creation + assert voxel_mean.shape[0] == voxel_coors.shape[ + 0], f"voxel_mean.shape[0] {voxel_mean.shape[0]} != voxel_coors.shape[0] {voxel_coors.shape[0]}" + assert pts_coors.shape[ + 1] == 3, f"pts_coors.shape[1] {pts_coors.shape[1]} != 3" + assert voxel_coors.shape[ + 1] == 3, f"voxel_coors.shape[1] {voxel_coors.shape[1]} != 3" + + canvas_y = int( + (self.point_cloud_range[4] - self.point_cloud_range[1]) / self.vy) + canvas_x = int( + (self.point_cloud_range[3] - self.point_cloud_range[0]) / self.vx) + canvas_channel = voxel_mean.size(1) + batch_size = pts_coors[:, 0].max() + 1 + + canvas_len = canvas_y * canvas_x * batch_size + # Create the canvas for this sample + canvas = voxel_mean.new_zeros(canvas_channel, canvas_len) + # Only include non-empty pillars + indices = (voxel_coors[:, 0] * canvas_y * canvas_x + + voxel_coors[:, 2] * canvas_x + voxel_coors[:, 1]) + assert indices.long().max() < canvas_len, 'Index out of range' + assert indices.long().min() >= 0, 'Index out of range' + # Scatter the blob back to the canvas + canvas[:, indices.long()] = voxel_mean.t() + # Step 2: get voxel mean for each point + voxel_index = (pts_coors[:, 0] * canvas_y * canvas_x + + pts_coors[:, 2] * canvas_x + pts_coors[:, 1]) + assert voxel_index.long().max() < canvas_len, 'Index out of range' + assert voxel_index.long().min() >= 0, 'Index out of range' + center_per_point = canvas[:, voxel_index.long()].t() + return center_per_point + + def forward(self, features, coors): + """Forward function. + Args: + features (torch.Tensor): Point features or raw points in shape + (N, M, C). + coors (torch.Tensor): Coordinates of each voxel + Returns: + torch.Tensor: Features of pillars. + """ + features_ls = [features] + # Find distance of x, y, and z from cluster center + if self._with_cluster_center: + voxel_mean, mean_coors = self.cluster_scatter(features, coors) + points_mean = self.map_voxel_center_to_point( + coors, voxel_mean, mean_coors) + # TODO: maybe also do cluster for reflectivity + f_cluster = features[:, :3] - points_mean[:, :3] + features_ls.append(f_cluster) + + # Find distance of x, y, and z from pillar center + if self._with_voxel_center: + f_center = features.new_zeros(size=(features.size(0), 3)) + f_center[:, 0] = features[:, 0] - ( + coors[:, 2].type_as(features) * self.vx + self.x_offset) + f_center[:, 1] = features[:, 1] - ( + coors[:, 1].type_as(features) * self.vy + self.y_offset) + f_center[:, 2] = features[:, 2] - ( + coors[:, 0].type_as(features) * self.vz + self.z_offset) + features_ls.append(f_center) + + if self._with_distance: + points_dist = torch.norm(features[:, :3], 2, 1, keepdim=True) + features_ls.append(points_dist) + + # Combine together feature decorations + features = torch.cat(features_ls, dim=-1) + for i, pfn in enumerate(self.pfn_layers): + point_feats = pfn(features) + voxel_feats, voxel_coors = self.pfn_scatter(point_feats, coors) + if i != len(self.pfn_layers) - 1: + # need to concat voxel feats if it is not the last pfn + feat_per_point = self.map_voxel_center_to_point( + coors, voxel_feats, voxel_coors) + features = torch.cat([point_feats, feat_per_point], dim=1) + + return voxel_feats, voxel_coors diff --git a/models/embedders/scatter.py b/models/embedders/scatter.py new file mode 100644 index 0000000..aedd533 --- /dev/null +++ b/models/embedders/scatter.py @@ -0,0 +1,94 @@ +import torch +import torch.nn as nn + +class PointPillarsScatter(nn.Module): + """Point Pillar's Scatter. + + Converts learned features from dense tensor to sparse pseudo image. + + Args: + in_channels (int): Channels of input features. + output_shape (list[int]): Required output shape of features. + """ + + def __init__(self, in_channels, output_shape): + super().__init__() + self.output_shape = output_shape + self.ny = output_shape[0] + self.nx = output_shape[1] + self.in_channels = in_channels + self.fp16_enabled = False + + def forward(self, voxel_features, coors, batch_size=None): + """Foraward function to scatter features.""" + # TODO: rewrite the function in a batch manner + # no need to deal with different batch cases + if batch_size is not None: + return self.forward_batch(voxel_features, coors, batch_size) + else: + return self.forward_single(voxel_features, coors) + + def forward_single(self, voxel_features, coors): + """Scatter features of single sample. + + Args: + voxel_features (torch.Tensor): Voxel features in shape (N, M, C). + coors (torch.Tensor): Coordinates of each voxel. + """ + # Create the canvas for this sample + canvas = torch.zeros( + self.in_channels, + self.nx * self.ny, + dtype=voxel_features.dtype, + device=voxel_features.device) + + indices = coors[:, 1] * self.nx + coors[:, 2] + indices = indices.long() + voxels = voxel_features.t() + # Now scatter the blob back to the canvas. + canvas[:, indices] = voxels + # Undo the column stacking to final 4-dim tensor + canvas = canvas.view(1, self.in_channels, self.ny, self.nx) + return canvas + + def forward_batch(self, voxel_features, coors, batch_size): + """Scatter features of single sample. + + Args: + voxel_features (torch.Tensor): Voxel features in shape (N, M, C). + coors (torch.Tensor): Coordinates of each voxel in shape (N, 4). + The first column indicates the sample ID. + batch_size (int): Number of samples in the current batch. + """ + # batch_canvas will be the final output. + batch_canvas = [] + for batch_itt in range(batch_size): + # Create the canvas for this sample + canvas = torch.zeros( + self.in_channels, + self.nx * self.ny, + dtype=voxel_features.dtype, + device=voxel_features.device) + + # Only include non-empty pillars + batch_mask = coors[:, 0] == batch_itt + this_coors = coors[batch_mask, :] + indices = this_coors[:, 2] * self.nx + this_coors[:, 3] + indices = indices.type(torch.long) + voxels = voxel_features[batch_mask, :] + voxels = voxels.t() + + # Now scatter the blob back to the canvas. + canvas[:, indices] = voxels + + # Append to a list for later stacking. + batch_canvas.append(canvas) + + # Stack to 3-dim tensor (batch-size, in_channels, nrows*ncols) + batch_canvas = torch.stack(batch_canvas, 0) + + # Undo the column stacking to final 4-dim tensor + batch_canvas = batch_canvas.view(batch_size, self.in_channels, self.ny, + self.nx) + + return batch_canvas diff --git a/models/fast_flow_3d.py b/models/fast_flow_3d.py new file mode 100644 index 0000000..db352ae --- /dev/null +++ b/models/fast_flow_3d.py @@ -0,0 +1,405 @@ +import torch +import torch.nn as nn + +import numpy as np +from models.embedders import HardEmbedder, DynamicEmbedder +from models.backbones import FastFlowUNet +from models.heads import FastFlowDecoder, FastFlowDecoderStepDown +from pointclouds import from_fixed_array +from pointclouds.losses import warped_pc_loss + +from typing import Dict, Any, Optional +from collections import defaultdict +import time + + +class FastFlow3DSelfSupervisedLoss(): + + def __init__(self, warp_upscale: float, device: str = None): + super().__init__() + self.warp_upscale = warp_upscale + + def _warped_loss(self, model_res): + flows = model_res["flow"] + pc0_points_lst = model_res["pc0_points_lst"] + pc1_points_lst = model_res["pc1_points_lst"] + + warped_loss = 0 + for flow, pc0_points, pc1_points in zip(flows, pc0_points_lst, + pc1_points_lst): + pc0_warped_to_pc1_points = pc0_points + flow + warped_loss += warped_pc_loss(pc0_warped_to_pc1_points, + pc1_points) * self.warp_upscale + return warped_loss + + def _triangle_loss(self, forward_res, backward_res): + forward_flows = forward_res["flow"] + backward_flows = backward_res["flow"] + # These are the idxes of the forward flow present in the backwards flow + valid_forward_flow_idxes = backward_res["pc0_valid_point_idxes"] + + triangle_loss = 0 + for forward_flow, backward_flow, valid_forward_flow_idx in zip( + forward_flows, backward_flows, valid_forward_flow_idxes): + forward_flow_valid_in_backward_flow = forward_flow[ + valid_forward_flow_idx] + + triangle_error = torch.norm(forward_flow_valid_in_backward_flow + + backward_flow, + dim=1, + p=2) + triangle_loss += torch.mean(triangle_error) + return triangle_loss + + def _symmetry_loss(self, model_res, flipped_res, idx): + flows = model_res["flow"] + flipped_flows = flipped_res["flow"] + symmetry_loss = 0 + for flow, flipped_flow in zip(flows, flipped_flows): + # Other than the y axis, the flows should be the same + # Negate the y axis of the flipped flow, then subtract the two flows should produce a zero flow + flipped_flow[:, idx] = -flipped_flow[:, idx] + symmetry_loss += torch.mean( + torch.norm(flow - flipped_flow, dim=1, p=2)) + + return symmetry_loss + + def __call__(self, input_batch, model_res_dict: Dict[str, Dict[str, Any]]): + warped_loss = self._warped_loss(model_res_dict["forward"]) + # self._visualize_regressed_ground_truth_pcs(model_res_dict["forward"]) + + res_dict = { + "warped_loss": warped_loss, + } + + total_loss = warped_loss + + if "backward" in model_res_dict: + triangle_loss = self._triangle_loss(model_res_dict["forward"], + model_res_dict["backward"]) + total_loss += triangle_loss + res_dict["triangle_loss"] = triangle_loss + + if "symmetry_x" in model_res_dict: + symmetry_loss = self._symmetry_loss(model_res_dict["forward"], + model_res_dict["symmetry_x"], + 0) + total_loss += symmetry_loss + res_dict["symmetry_x_loss"] = symmetry_loss + + if "symmetry_y" in model_res_dict: + symmetry_loss = self._symmetry_loss(model_res_dict["forward"], + model_res_dict["symmetry_y"], + 1) + total_loss += symmetry_loss + res_dict["symmetry_y_loss"] = symmetry_loss + + res_dict["loss"] = total_loss + + return res_dict + + def _visualize_regressed_ground_truth_pcs(self, model_res): + regressed_flowed_pc0_to_pc1 = model_res["flow"] + pc0_pc = model_res["pc0_points_lst"] + pc1_pc = model_res["pc1_points_lst"] + import open3d as o3d + import numpy as np + pc0_pc = pc0_pc.detach().cpu().numpy() + pc1_pc = pc1_pc.detach().cpu().numpy() + regressed_flowed_pc0_to_pc1 = regressed_flowed_pc0_to_pc1.detach().cpu( + ).numpy() + # make open3d visualizer + vis = o3d.visualization.Visualizer() + vis.create_window() + vis.get_render_option().point_size = 1.5 + vis.get_render_option().background_color = (0, 0, 0) + vis.get_render_option().show_coordinate_frame = True + # set up vector + vis.get_view_control().set_up([0, 0, 1]) + + # Add input PC + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(pc0_pc) + pc_color = np.zeros_like(pc0_pc) + pc_color[:, 0] = 1 + pc_color[:, 1] = 1 + pcd.colors = o3d.utility.Vector3dVector(pc_color) + vis.add_geometry(pcd) + + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(pc1_pc) + pc_color = np.zeros_like(pc1_pc) + pc_color[:, 1] = 1 + pc_color[:, 2] = 1 + pcd.colors = o3d.utility.Vector3dVector(pc_color) + vis.add_geometry(pcd) + + # Add line set between pc0 and regressed pc1 + line_set = o3d.geometry.LineSet() + assert len(pc0_pc) == len( + regressed_flowed_pc0_to_pc1 + ), f"{len(pc0_pc)} != {len(regressed_flowed_pc0_to_pc1)}" + line_set_points = np.concatenate([pc0_pc, regressed_flowed_pc0_to_pc1], + axis=0) + + lines = np.array([[i, i + len(regressed_flowed_pc0_to_pc1)] + for i in range(len(pc0_pc))]) + line_set.points = o3d.utility.Vector3dVector(line_set_points) + line_set.lines = o3d.utility.Vector2iVector(lines) + line_set.colors = o3d.utility.Vector3dVector( + [[0, 0, 1] for _ in range(len(lines))]) + vis.add_geometry(line_set) + + vis.run() + + +class FastFlow3DDistillationLoss(): + + def __init__(self, device: str = None, fast_mover_scale: bool = False): + super().__init__() + self.fast_mover_scale = fast_mover_scale + + def __call__(self, input_batch, model_res_dict): + model_res = model_res_dict["forward"] + estimated_flows = model_res["flow"] + + gt_flow_array_stack = input_batch['flow_array_stack'] + # self._visualize_regressed_ground_truth_pcs(model_res, + # gt_flow_array_stack) + + total_loss = 0 + # Iterate through the batch + for est_flow, gt_flow_array in zip(estimated_flows, + gt_flow_array_stack): + gt_flow = from_fixed_array(gt_flow_array[0]) + assert est_flow.shape == gt_flow.shape, f"estimated_flow {est_flow.shape} != ground_truth_flow {gt_flow.shape}" + + importance_scale = torch.ones_like(gt_flow[:, 0]) + if self.fast_mover_scale: + # Compute the importance scale using m/s speed. + gt_speed = torch.norm(gt_flow, dim=1, p=2) * 10.0 + mins = torch.ones_like(gt_speed) * 0.1 + maxs = torch.ones_like(gt_speed) + # Plot \max\left(0.1,\min\left(1,1.8x-0.8\right)\right) in Desmos + importance_scale = torch.min( + mins, torch.max(1.8 * gt_speed - 0.8, maxs)) + + total_loss += (torch.norm(est_flow - gt_flow, dim=1, p=2) * + importance_scale).mean() + + return { + "loss": total_loss, + } + + def _visualize_regressed_ground_truth_pcs(self, model_res, + gt_flow_array_stack): + regressed_flowed_pc0_to_pc1_lst = model_res["flow"] + pc0_pc_lst = model_res["pc0_points_lst"] + pc1_pc_lst = model_res["pc1_points_lst"] + import open3d as o3d + import numpy as np + for regressed_flowed_pc0_to_pc1, pc0_pc, pc1_pc, gt_flow_array in zip( + regressed_flowed_pc0_to_pc1_lst, pc0_pc_lst, pc1_pc_lst, + gt_flow_array_stack): + gt_flow = from_fixed_array(gt_flow_array[0]) + pc0_pc = pc0_pc.detach().cpu().numpy() + pc1_pc = pc1_pc.detach().cpu().numpy() + regressed_flowed_pc0_to_pc1 = pc0_pc + gt_flow.detach().cpu( + ).numpy() + #regressed_flowed_pc0_to_pc1.detach().cpu().numpy() + # breakpoint() + + # make open3d visualizer + vis = o3d.visualization.Visualizer() + vis.create_window() + vis.get_render_option().point_size = 1.5 + vis.get_render_option().background_color = (0, 0, 0) + vis.get_render_option().show_coordinate_frame = True + # set up vector + vis.get_view_control().set_up([0, 0, 1]) + + # Add input PC + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(pc0_pc) + pc_color = np.zeros_like(pc0_pc) + pc_color[:, 0] = 1 + pc_color[:, 1] = 1 + pcd.colors = o3d.utility.Vector3dVector(pc_color) + vis.add_geometry(pcd) + + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(pc1_pc) + pc_color = np.zeros_like(pc1_pc) + pc_color[:, 1] = 1 + pc_color[:, 2] = 1 + pcd.colors = o3d.utility.Vector3dVector(pc_color) + vis.add_geometry(pcd) + + # Add line set between pc0 and regressed pc1 + line_set = o3d.geometry.LineSet() + assert len(pc0_pc) == len( + regressed_flowed_pc0_to_pc1 + ), f"{len(pc0_pc)} != {len(regressed_flowed_pc0_to_pc1)}" + line_set_points = np.concatenate( + [pc0_pc, regressed_flowed_pc0_to_pc1], axis=0) + + lines = np.array([[i, i + len(regressed_flowed_pc0_to_pc1)] + for i in range(len(pc0_pc))]) + line_set.points = o3d.utility.Vector3dVector(line_set_points) + line_set.lines = o3d.utility.Vector2iVector(lines) + line_set.colors = o3d.utility.Vector3dVector( + [[0, 0, 1] for _ in range(len(lines))]) + vis.add_geometry(line_set) + + vis.run() + + +class FastFlow3DSupervisedLoss(): + + def __init__(self, device: str = None, scale_background: bool = True): + super().__init__() + self._scale_background = scale_background + + def __call__(self, input_batch, model_res_dict): + model_res = model_res_dict["forward"] + estimated_flows = model_res["flow"] + + pc0_valid_point_idxes = model_res["pc0_valid_point_idxes"] + input_pcs = input_batch['pc_array_stack'] + flowed_input_pcs = input_batch['flowed_pc_array_stack'] + pc_classes = input_batch['pc_class_mask_stack'] + input_flows = flowed_input_pcs - input_pcs + + total_loss = 0 + # Iterate through the batch + for estimated_flow, input_flow, pc0_valid_point_idx, pc_class_arr in zip( + estimated_flows, input_flows, pc0_valid_point_idxes, + pc_classes): + ground_truth_flow = input_flow[0, pc0_valid_point_idx] + + error = torch.norm(estimated_flow - ground_truth_flow, dim=1, p=2) + + if self._scale_background: + classes = pc_class_arr[0, pc0_valid_point_idx] + # Background class is -1 + is_foreground_class = (classes >= 0) + # We want to build a scalar array where the background points are 0.1 and the foreground points are 1 + # This is because we want to downweight the background points by 10x + # We do this simply by converting the bool, where the foreground points are 1, to float, and then + # multiplying by 0.9 and adding 0.1 to get a scalar in {0.1, 1} + background_scalar = is_foreground_class.float() * 0.9 + 0.1 + error = error * background_scalar + + total_loss += error.mean() + + return {"loss": total_loss} + + +class FastFlow3D(nn.Module): + """ + FastFlow3D based on the paper: + https://arxiv.org/abs/2103.01306v5 + + Note that there are several small differences between this implementation and the paper: + - We use a different loss function (predict flow for P_-1 to P_0 instead of P_0 to and + unseen P_1); referred to as pc0 and pc1 in the code. + """ + + def __init__(self, + VOXEL_SIZE, + PSEUDO_IMAGE_DIMS, + POINT_CLOUD_RANGE, + FEATURE_CHANNELS, + SEQUENCE_LENGTH, + bottleneck_head=False) -> None: + super().__init__() + self.SEQUENCE_LENGTH = SEQUENCE_LENGTH + assert self.SEQUENCE_LENGTH == 2, "This implementation only supports a sequence length of 2." + self.embedder = DynamicEmbedder(voxel_size=VOXEL_SIZE, + pseudo_image_dims=PSEUDO_IMAGE_DIMS, + point_cloud_range=POINT_CLOUD_RANGE, + feat_channels=FEATURE_CHANNELS) + + self.backbone = FastFlowUNet() + if bottleneck_head: + self.head = FastFlowDecoderStepDown( + voxel_pillar_size=VOXEL_SIZE[:2], num_stepdowns=3) + else: + self.head = FastFlowDecoder() + + def _model_forward(self, pc0s, pc1s): + + before_forward = time.time() + pc0_before_pseudoimages, pc0_voxel_infos_lst = self.embedder(pc0s) + pc1_before_pseudoimages, pc1_voxel_infos_lst = self.embedder(pc1s) + + grid_flow_pseudoimage = self.backbone(pc0_before_pseudoimages, + pc1_before_pseudoimages) + flows = self.head( + torch.cat((pc0_before_pseudoimages, pc1_before_pseudoimages), + dim=1), grid_flow_pseudoimage, pc0_voxel_infos_lst) + after_forward = time.time() + + pc0_points_lst = [e["points"] for e in pc0_voxel_infos_lst] + pc0_valid_point_idxes = [e["point_idxes"] for e in pc0_voxel_infos_lst] + pc1_points_lst = [e["points"] for e in pc1_voxel_infos_lst] + pc1_valid_point_idxes = [e["point_idxes"] for e in pc1_voxel_infos_lst] + + pc0_warped_pc1_points_lst = [ + pc0_points + flow + for pc0_points, flow in zip(pc0_points_lst, flows) + ] + + return { + "flow": flows, + "batch_delta_time": after_forward - before_forward, + "pc0_points_lst": pc0_points_lst, + "pc0_warped_pc1_points_lst": pc0_warped_pc1_points_lst, + "pc0_valid_point_idxes": pc0_valid_point_idxes, + "pc1_points_lst": pc1_points_lst, + "pc1_valid_point_idxes": pc1_valid_point_idxes + } + + def forward(self, + batched_sequence: Dict[str, torch.Tensor], + compute_cycle=False, + compute_symmetry_x=False, + compute_symmetry_y=False): + pc_arrays = batched_sequence['pc_array_stack'] + pc0s = pc_arrays[:, 0] + pc1s = pc_arrays[:, 1] + model_res = self._model_forward(pc0s, pc1s) + + ret_dict = {"forward": model_res} + + if compute_cycle: + # The warped pointcloud, original pointcloud should be the input to the model + pc0_warped_pc1_points_lst = model_res["pc0_warped_pc1_points_lst"] + pc0_points_lst = model_res["pc0_points_lst"] + # Some of the warped points may be outside the pseudoimage range, causing them to be clipped. + # When we compute this reverse flow, we need to solve for the original points that were warped to the clipped points. + backward_model_res = self._model_forward(pc0_warped_pc1_points_lst, + pc0_points_lst) + # model_res["reverse_flow"] = backward_model_res["flow"] + # model_res[ + # "flow_valid_point_idxes_for_reverse_flow"] = backward_model_res[ + # "pc0_valid_point_idxes"] + ret_dict["backward"] = backward_model_res + + if compute_symmetry_x: + pc0s_sym = pc0s.clone() + pc0s_sym[:, :, 0] *= -1 + pc1s_sym = pc1s.clone() + pc1s_sym[:, :, 0] *= -1 + model_res_sym = self._model_forward(pc0s_sym, pc1s_sym) + ret_dict["symmetry_x"] = model_res_sym + + if compute_symmetry_y: + pc0s_sym = pc0s.clone() + pc0s_sym[:, :, 1] *= -1 + pc1s_sym = pc1s.clone() + pc1s_sym[:, :, 1] *= -1 + model_res_sym = self._model_forward(pc0s_sym, pc1s_sym) + ret_dict["symmetry_y"] = model_res_sym + + return ret_dict diff --git a/models/heads/__init__.py b/models/heads/__init__.py new file mode 100644 index 0000000..f571b38 --- /dev/null +++ b/models/heads/__init__.py @@ -0,0 +1,8 @@ +from .nsfp import NeuralSceneFlowPrior +from .nsfp_optimizable import NeuralSceneFlowPriorOptimizable +from .fast_flow_decoder import FastFlowDecoder, FastFlowDecoderStepDown + +__all__ = [ + "NeuralSceneFlowPrior", "NeuralSceneFlowPriorOptimizable", + "FastFlowDecoder", "FastFlowDecoderStepDown" +] diff --git a/models/heads/fast_flow_decoder.py b/models/heads/fast_flow_decoder.py new file mode 100644 index 0000000..bc0dfb7 --- /dev/null +++ b/models/heads/fast_flow_decoder.py @@ -0,0 +1,227 @@ +import torch +import torch.nn as nn +from typing import List, Tuple, Dict + + +class FastFlowDecoder(nn.Module): + + def __init__(self, pseudoimage_channels: int = 64): + super().__init__() + + self.decoder = nn.Sequential( + nn.Linear(pseudoimage_channels * 2 + 3, 32), nn.GELU(), + nn.Linear(32, 3)) + + def forward_single(self, before_pseudoimage: torch.Tensor, + after_pseudoimage: torch.Tensor, + point_offsets: torch.Tensor, + voxel_coords: torch.Tensor) -> torch.Tensor: + voxel_coords = voxel_coords.long() + assert (voxel_coords[:, 0] == 0).all(), "Z index must be 0" + + # Voxel coords are Z, Y, X, and the pseudoimage is Channel, Y, X + # I have confirmed via visualization that these coordinates are correct. + after_voxel_vectors = after_pseudoimage[:, voxel_coords[:, 1], + voxel_coords[:, 2]].T + before_voxel_vectors = before_pseudoimage[:, voxel_coords[:, 1], + voxel_coords[:, 2]].T + concatenated_vectors = torch.cat( + [before_voxel_vectors, after_voxel_vectors, point_offsets], dim=1) + + flow = self.decoder(concatenated_vectors) + return flow + + def forward( + self, before_pseudoimages: torch.Tensor, + after_pseudoimages: torch.Tensor, + voxelizer_infos: List[Dict[str, + torch.Tensor]]) -> List[torch.Tensor]: + + flow_results = [] + for before_pseudoimage, after_pseudoimage, voxelizer_info in zip( + before_pseudoimages, after_pseudoimages, voxelizer_infos): + point_offsets = voxelizer_info["point_offsets"] + voxel_coords = voxelizer_info["voxel_coords"] + flow = self.forward_single(before_pseudoimage, after_pseudoimage, + point_offsets, voxel_coords) + flow_results.append(flow) + return flow_results + + +class ConvWithNorms(nn.Module): + + def __init__(self, in_num_channels: int, out_num_channels: int, + kernel_size: int, stride: int, padding: int): + super().__init__() + self.conv = nn.Conv2d(in_num_channels, out_num_channels, kernel_size, + stride, padding) + self.batchnorm = nn.BatchNorm2d(out_num_channels) + self.nonlinearity = nn.GELU() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + conv_res = self.conv(x) + if conv_res.shape[2] == 1 and conv_res.shape[3] == 1: + # This is a hack to get around the fact that batchnorm doesn't support + # 1x1 convolutions + batchnorm_res = conv_res + else: + batchnorm_res = self.batchnorm(conv_res) + return self.nonlinearity(batchnorm_res) + + +class FastFlowDecoderStepDown(FastFlowDecoder): + + def __init__(self, voxel_pillar_size: Tuple[float, float], + num_stepdowns: int) -> None: + super().__init__(pseudoimage_channels=16) + assert num_stepdowns > 0, "stepdown_factor must be positive" + self.num_stepdowns = num_stepdowns + self.voxel_pillar_size = voxel_pillar_size + assert len(voxel_pillar_size) == 2, "voxel_pillar_size must be 2D" + + self.pseudoimage_stepdown_head = nn.ModuleList() + # Build convolutional pseudoimage stepdown head + for i in range(num_stepdowns): + if i == 0: + in_channels = 64 + else: + in_channels = 16 + out_channels = 16 + self.pseudoimage_stepdown_head.append( + ConvWithNorms(in_channels, out_channels, 3, 2, 1)) + + def _plot_pseudoimage(self, pseudoimage: torch.Tensor, + point_offsets: torch.Tensor, + voxel_coords: torch.Tensor, points: torch.Tensor, + pillar_size: Tuple[float, float]): + assert len(pillar_size) == 2, "pillar_size must be 2D" + + import matplotlib.pyplot as plt + import numpy as np + + plot_offset_bc_imshow_stupid = -0.5 + + full_res_before_pseudoimage_np = pseudoimage.detach().cpu().numpy( + ).sum(0) + voxel_coords_np = voxel_coords.detach().cpu().numpy() + offsets_np = point_offsets.detach().cpu().numpy() + points_np = points.detach().cpu().numpy() + plt.title(f"Full res: {full_res_before_pseudoimage_np.shape}") + + def _raw_points_to_scaled_points(raw_points: np.ndarray) -> np.ndarray: + scaled_points = raw_points.copy() + scaled_points[:, 0] /= pillar_size[0] + scaled_points[:, 1] /= pillar_size[1] + scaled_points[:, :2] += pseudoimage.shape[2] // 2 + return scaled_points + + scaled_points_np = _raw_points_to_scaled_points(points_np) + + # Plot pseudoimage and points from voxel plus offset info + + def _offset_points_to_scaled_points(offsets_np, voxel_coords_np): + offsets_np = offsets_np.copy() + # Scale voxel offsets to be in scale with the voxel size + offsets_np[:, 0] /= pillar_size[0] + offsets_np[:, 1] /= pillar_size[1] + + voxel_centers_np = voxel_coords_np[:, [2, 1]].astype( + np.float32) + 0.5 + + point_positions_np = voxel_centers_np + offsets_np[:, :2] + + voxel_centers_np = voxel_coords_np[:, [2, 1]].astype( + np.float32) + 0.5 + + point_positions_np = voxel_centers_np + offsets_np[:, :2] + + return voxel_centers_np, point_positions_np + + voxel_centers_np, point_positions_np = _offset_points_to_scaled_points( + offsets_np, voxel_coords_np) + + plt.scatter(scaled_points_np[:, 0] + plot_offset_bc_imshow_stupid, + scaled_points_np[:, 1] + plot_offset_bc_imshow_stupid, + s=1, + marker="x", + color='green') + + plt.scatter(point_positions_np[:, 0] + plot_offset_bc_imshow_stupid, + point_positions_np[:, 1] + plot_offset_bc_imshow_stupid, + s=1, + marker="x", + color='red') + plt.scatter(voxel_centers_np[:, 0] + plot_offset_bc_imshow_stupid, + voxel_centers_np[:, 1] + plot_offset_bc_imshow_stupid, + s=1, + marker="x", + color='orange') + + plt.imshow(full_res_before_pseudoimage_np, interpolation='nearest') + plt.show() + + def _step_down_pseudoimage(self, + pseudoimage: torch.Tensor) -> torch.Tensor: + for stepdown in self.pseudoimage_stepdown_head: + pseudoimage = stepdown(pseudoimage) + return pseudoimage + + def _stepdown_voxel_info( + self, stepped_down_pseudoimage: torch.Tensor, points: torch.Tensor, + voxel_coords: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + # We need to compute how the voxel coordinates change after the stepdown. + # We then recompute the point offsets. + + stepdown_scale = 2**self.num_stepdowns + assert stepdown_scale > 1, "stepdown_scale must be positive" + + scaled_voxel_pillar_size = torch.Tensor([ + self.voxel_pillar_size[0] * stepdown_scale, + self.voxel_pillar_size[1] * stepdown_scale + ]).to(points.device) + + # Coordinates are in ZYX order. + stepped_down_voxel_coords = torch.div(voxel_coords, + stepdown_scale, + rounding_mode='trunc') + + stepped_down_center_translation = torch.Tensor([ + stepped_down_pseudoimage.shape[1] / 2, + stepped_down_pseudoimage.shape[2] / 2 + ]).float().to(points.device) + stepped_down_voxel_centers = ( + stepped_down_voxel_coords[:, [2, 1]].float() - + stepped_down_center_translation[None, :] + + 0.5) * scaled_voxel_pillar_size[None, :] + + stepped_down_point_offsets = points.clone() + stepped_down_point_offsets[:, :2] -= stepped_down_voxel_centers + + return stepped_down_point_offsets, stepped_down_voxel_coords + + def forward( + self, full_res_before_pseudoimages: torch.Tensor, + full_res_after_pseudoimages: torch.Tensor, + voxelizer_infos: List[Dict[str, + torch.Tensor]]) -> List[torch.Tensor]: + + flow_results = [] + + # Step down the pseudoimages + reduced_res_before_pseudoimages = self._step_down_pseudoimage( + full_res_before_pseudoimages) + reduced_res_after_pseudoimages = self._step_down_pseudoimage( + full_res_after_pseudoimages) + + for before_pseudoimage, after_pseudoimage, voxelizer_info in zip( + reduced_res_before_pseudoimages, + reduced_res_after_pseudoimages, voxelizer_infos): + voxel_coords = voxelizer_info["voxel_coords"] + points = voxelizer_info["points"] + point_offsets, voxel_coords = self._stepdown_voxel_info( + before_pseudoimage, points, voxel_coords) + + flow = self.forward_single(before_pseudoimage, after_pseudoimage, + point_offsets, voxel_coords) + flow_results.append(flow) + return flow_results diff --git a/models/heads/nsfp.py b/models/heads/nsfp.py new file mode 100644 index 0000000..6bee459 --- /dev/null +++ b/models/heads/nsfp.py @@ -0,0 +1,94 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class NeuralSceneFlowPrior(nn.Module): + + def __init__(self, num_hidden_units=64, num_hidden_layers=4, nonlinearity=torch.sigmoid): + super().__init__() + assert num_hidden_units > 0, f"num_hidden_units must be > 0, got {num_hidden_units}" + assert num_hidden_layers > 0, f"num_hidden_layers must be > 0, got {num_hidden_layers}" + param_count = self._num_params_for_linear(3, num_hidden_units) + for _ in range(num_hidden_layers - 1): + param_count += self._num_params_for_linear(num_hidden_units, + num_hidden_units) + param_count += self._num_params_for_linear(num_hidden_units, 3) + self.param_count = param_count + self.num_hidden_layers = num_hidden_layers + self.num_hidden_units = num_hidden_units + self.nonlinearity = nonlinearity + + def _num_params_for_linear(self, input_dim, output_dim): + # Rectangular matrix plus bias + return input_dim * output_dim + output_dim + + def decode_linear_from_params(self, params: torch.Tensor, in_units: int, + out_units: int, idx_offset: int): + # Params are K, where K is the number of parameters in the network + assert isinstance( + params, + torch.Tensor), f"params must be a torch.Tensor, got {type(params)}" + assert params.shape == ( + self.param_count, + ), f"params must have shape ({self.param_count},), got {params.shape}" + + # The first idx_offset parameters are for the first linear layer + # The next in_units * out + # The next out_units parameters are for the bias + linear_size = in_units * out_units + linear = params[idx_offset:idx_offset + linear_size] + bias = params[idx_offset + linear_size:idx_offset + linear_size + + out_units] + linear = linear.reshape(in_units, out_units) + return linear, bias, idx_offset + linear_size + out_units + + def _evaluate_network(self, points, linears, biases): + # Evaluate the MLP on the points + points_offsets = points + for idx, (linear, bias) in enumerate(zip(linears, biases)): + points_offsets = torch.matmul(points_offsets, linear) + bias + if idx < len(linears) - 1: + points_offsets = self.nonlinearity(points_offsets) + return points_offsets + + + def forward(self, points: torch.Tensor, + params: torch.Tensor) -> torch.Tensor: + + # Points are N x 3 + # Params are K, where K is the number of parameters in the network + assert isinstance( + points, + torch.Tensor), f"x must be a torch.Tensor, got {type(points)}" + assert points.shape[ + 1] == 3, f"x must have shape (points, 3), got {points.shape}" + + assert isinstance( + params, + torch.Tensor), f"params must be a torch.Tensor, got {type(params)}" + assert params.shape == ( + self.param_count, + ), f"params must have shape ({self.param_count},), got {params.shape}" + + # Decode the network into linears and biases. + linears = [] + biases = [] + linear, bias, idx_offset = self.decode_linear_from_params( + params, 3, self.num_hidden_units, 0) + linears.append(linear) + biases.append(bias) + for _ in range(self.num_hidden_layers - 1): + linear, bias, idx_offset = self.decode_linear_from_params( + params, self.num_hidden_units, self.num_hidden_units, + idx_offset) + linears.append(linear) + biases.append(bias) + linear, bias, idx_offset = self.decode_linear_from_params( + params, self.num_hidden_units, 3, idx_offset) + linears.append(linear) + biases.append(bias) + assert idx_offset == self.param_count, f"idx_offset must be {self.param_count}, got {idx_offset}" + + # Evaluate the MLP on the points to get the points offsets + return points + self._evaluate_network(points, linears, biases) diff --git a/models/heads/nsfp_optimizable.py b/models/heads/nsfp_optimizable.py new file mode 100644 index 0000000..479c0db --- /dev/null +++ b/models/heads/nsfp_optimizable.py @@ -0,0 +1,81 @@ +import torch +import torch.nn as nn +# Import functional +import torch.nn.functional as F + + +class NeuralSceneFlowPriorOptimizable(nn.Module): + + def __init__(self, + params: torch.Tensor, + num_hidden_units: int = 64, + num_hidden_layers: int = 4, + nonlinearity=nn.Sigmoid()): + super().__init__() + + self.num_hidden_units = num_hidden_units + self.num_hidden_layers = num_hidden_layers + + # Decode the network into linears and biases. + self.module_list = nn.ModuleList() + linear_nn, idx_offset = self._decode_linear_from_params( + params, 3, self.num_hidden_units, 0) + self.module_list.append(linear_nn) + self.module_list.append(nonlinearity) + for _ in range(self.num_hidden_layers - 1): + linear_nn, idx_offset = self._decode_linear_from_params( + params, self.num_hidden_units, self.num_hidden_units, + idx_offset) + self.module_list.append(linear_nn) + self.module_list.append(nonlinearity) + linear_nn, idx_offset = self._decode_linear_from_params( + params, self.num_hidden_units, 3, idx_offset) + self.module_list.append(linear_nn) + self.module_list.append(nonlinearity) + + def _decode_linear_from_params(self, params: torch.Tensor, in_units: int, + out_units: int, idx_offset: int): + # Params are K, where K is the number of parameters in the network + assert isinstance( + params, + torch.Tensor), f"params must be a torch.Tensor, got {type(params)}" + linear_nn = nn.Linear(out_units, in_units, bias=True) + with torch.no_grad(): + # The first idx_offset parameters are for the first linear layer + # The next in_units * out + # The next out_units parameters are for the bias + linear_size = in_units * out_units + linear = params[idx_offset:idx_offset + linear_size] + bias = params[idx_offset + linear_size:idx_offset + linear_size + + out_units] + linear = linear.reshape(out_units, in_units) + + # assign the results to the linear module. + linear_nn.weight = nn.Parameter(linear) + linear_nn.bias = nn.Parameter(bias) + return linear_nn, idx_offset + linear_size + out_units + + def _decode_params_from_linear(self): + param_list = [] + for module in self.module_list: + if isinstance(module, nn.Linear): + with torch.no_grad(): + param_list.append(module.weight.reshape(-1).detach()) + param_list.append(module.bias.detach()) + return torch.cat(param_list, dim=0) + + def decode_params(self): + return self._decode_params_from_linear() + + def forward(self, points: torch.Tensor) -> torch.Tensor: + # Points are N x 3 + assert isinstance( + points, + torch.Tensor), f"points must be a torch.Tensor, got {type(points)}" + assert points.shape[ + 1] == 3, f"points must be N x 3, got {points.shape}" + points = torch.unsqueeze(points, dim=0) + for module in self.module_list: + points = module(points) + points = torch.squeeze(points, dim=0) + return points diff --git a/models/joint_flow.py b/models/joint_flow.py new file mode 100644 index 0000000..550172e --- /dev/null +++ b/models/joint_flow.py @@ -0,0 +1,319 @@ +import torch +import torch.nn as nn +from pytorch3d.ops.knn import knn_points +from pytorch3d.loss import chamfer_distance +import numpy as np +from tqdm import tqdm + +from models.embedders import HardEmbedder +from models.backbones import FeaturePyramidNetwork +from models.attention import JointConvAttention +from models.heads import NeuralSceneFlowPrior, NeuralSceneFlowPriorOptimizable + +from typing import List, Tuple, Dict +from pointclouds import PointCloud, SE3 +from pointclouds.losses import warped_pc_loss + + +def _torch_to_jagged_pc(torch_tensor: torch.Tensor) -> PointCloud: + assert isinstance( + torch_tensor, + torch.Tensor), f"Expected torch.Tensor, got {type(torch_tensor)}" + return PointCloud.from_fixed_array(torch_tensor) + + +def _pc_to_torch(pc: PointCloud, device) -> torch.Tensor: + assert isinstance(pc, PointCloud), f"Expected PointCloud, got {type(pc)}" + return pc.points.float().to(device) + + +def _pose_to_translation(transform_matrix: torch.Tensor) -> torch.Tensor: + assert isinstance( + transform_matrix, + torch.Tensor), f"Expected torch.Tensor, got {type(transform_matrix)}" + assert transform_matrix.shape == ( + 4, 4 + ), f"Expected transform_matrix.shape == (4, 4), got {transform_matrix.shape[1:]}" + return transform_matrix[:3, 3] + + +class JointFlowOptimizationLoss(): + + def __init__(self, + device: str, + hidden_units=64, + hidden_layers=4, + max_iters=50) -> None: + self.device = device + self.num_hidden_units = hidden_units + self.num_hidden_layers = hidden_layers + self.max_iters = max_iters + + def _optimize_nsfp(self, pc1, pc2, nsfp_params): + nsfp = NeuralSceneFlowPriorOptimizable( + nsfp_params, + num_hidden_units=self.num_hidden_units, + num_hidden_layers=self.num_hidden_layers).to(self.device) + optimizer = torch.optim.Adam(nsfp.parameters(), + lr=0.008, + weight_decay=0) + for _ in range(self.max_iters): + optimizer.zero_grad() + flow = nsfp(pc1) + warped_pc1 = pc1 + flow + loss = warped_pc_loss(warped_pc1, pc2, dist_threshold=None) + loss.backward() + optimizer.step() + + return nsfp.decode_params() + + def __call__(self, + batched_sequence: List[List[Tuple[torch.Tensor, torch.Tensor, + torch.Tensor]]], + visualize=False): + delta = 1 + assert delta > 0, f"delta must be positive, got {delta}" + loss = 0 + for sequence in batched_sequence: + sequence_length = len(sequence) + assert delta < sequence_length, f"delta must be less than sequence_length, got {delta}" + for i in tqdm(range(sequence_length - delta), + leave=False, + desc="Optimizing NSFP weight targets"): + pc_t0, _, nsfp_params = sequence[i] + pc_t1, _, _ = sequence[i + delta] + pc_t0 = _torch_to_jagged_pc(pc_t0) + pc_t1 = _torch_to_jagged_pc(pc_t1) + pc_t0 = _pc_to_torch(pc_t0, self.device) + pc_t1 = _pc_to_torch(pc_t1, self.device) + + target_params = self._optimize_nsfp( + pc_t0, pc_t1, + nsfp_params.clone().detach()) + assert target_params.shape == nsfp_params.shape, f"target_params.shape == nsfp_params.shape, got {target_params.shape} != {nsfp_params.shape}" + loss += torch.linalg.norm(target_params - nsfp_params, ord=1) + + return loss + + +class JointFlowLoss(): + + def __init__(self, + device: str, + NSFP_FILTER_SIZE=64, + NSFP_NUM_LAYERS=4, + zero_prior: bool = False) -> None: + self.device = device + self.nsfp = NeuralSceneFlowPrior(num_hidden_units=NSFP_FILTER_SIZE, + num_hidden_layers=NSFP_NUM_LAYERS).to( + self.device) + self.zero_prior = zero_prior + self.warp_idx = 0 + + def __call__(self, + batched_sequence: List[List[Tuple[torch.Tensor, torch.Tensor, + torch.Tensor]]], + delta: int, + visualize=False) -> torch.Tensor: + loss = 0 + zero_prior_loss = 0 + for sequence in batched_sequence: + sequence_length = len(sequence) + assert delta > 0, f"delta must be positive, got {delta}" + assert delta < sequence_length, f"delta must be less than sequence_length, got {delta}" + for i in range(sequence_length - delta): + pc_t0, _, _ = sequence[i] + pc_t1, _, _ = sequence[i + delta] + pc_t0 = _torch_to_jagged_pc(pc_t0) + pc_t1 = _torch_to_jagged_pc(pc_t1) + pc_t0 = _pc_to_torch(pc_t0, self.device) + pc_t1 = _pc_to_torch(pc_t1, self.device) + + pc_t0_warped_to_t1 = pc_t0 + for j in range(delta): + _, _, nsfp_params = sequence[i + j] + # nsfp_params = torch.zeros_like(nsfp_params) + pc_t0_warped_to_t1 = self.nsfp(pc_t0_warped_to_t1, + nsfp_params) + + if visualize: + self._visualize_bev_warp(pc_t0, pc_t0_warped_to_t1) + # self._visualize_o3d(pc_t0, pc_t1, pc_t0_warped_to_t1) + loss += warped_pc_loss(pc_t0_warped_to_t1, + pc_t1, + dist_threshold=None) + if self.zero_prior: + zero_prior_loss += warped_pc_loss(pc_t0, + pc_t1, + dist_threshold=None) + if self.zero_prior: + return loss, zero_prior_loss + return loss + + def _visualize_bev_warp(self, pc_t0, pc_t0_warped_to_t1): + + pc_xy = pc_t0[:, :2].detach().cpu().numpy() + flow = pc_t0_warped_to_t1 - pc_t0 + flow = flow.detach().cpu().numpy() + flow_xy = flow[:, :2] + flow_magnitudes = np.linalg.norm(flow_xy, axis=1) + flow_angles = np.arctan2(flow_xy[:, 1], flow_xy[:, 0]) + flow_colors = (flow_angles + np.pi) / (2 * np.pi) * 360 + + # Scatterplot with matplotlib, using flow colors to color the dots + import matplotlib.pyplot as plt + plt.scatter(pc_xy[:, 0], + pc_xy[:, 1], + c=flow_colors, + cmap='hsv', + vmin=0, + vmax=360) + plt.colorbar() + plt.savefig(f"/efs/bev_warps/bev_warp{self.warp_idx:09d}.png") + self.warp_idx += 1 + plt.clf() + + def _visualize_o3d(self, pc_t0, pc_t1, pc_t0_warped_to_t1): + import open3d as o3d + pc_t0_npy = pc_t0.cpu().numpy() + pc_t1_npy = pc_t1.cpu().numpy() + pc_t0_warped_to_t1_npy = pc_t0_warped_to_t1.detach().cpu().numpy() + + pc_t0_color = np.zeros_like(pc_t0_npy) + pc_t0_color[:, 0] = 1.0 + pc_t1_color = np.zeros_like(pc_t1_npy) + pc_t1_color[:, 1] = 1.0 + pc_t0_warped_to_t1_color = np.zeros_like(pc_t0_warped_to_t1_npy) + pc_t0_warped_to_t1_color[:, 2] = 1.0 + + pc_t0_o3d = o3d.geometry.PointCloud() + pc_t0_o3d.points = o3d.utility.Vector3dVector(pc_t0_npy) + pc_t0_o3d.colors = o3d.utility.Vector3dVector(pc_t0_color) + + pc_t1_o3d = o3d.geometry.PointCloud() + pc_t1_o3d.points = o3d.utility.Vector3dVector(pc_t1_npy) + pc_t1_o3d.colors = o3d.utility.Vector3dVector(pc_t1_color) + + pc_t0_warped_to_t1_o3d = o3d.geometry.PointCloud() + pc_t0_warped_to_t1_o3d.points = o3d.utility.Vector3dVector( + pc_t0_warped_to_t1_npy) + pc_t0_warped_to_t1_o3d.colors = o3d.utility.Vector3dVector( + pc_t0_warped_to_t1_color) + + vis = o3d.visualization.Visualizer() + vis.create_window() + vis.get_render_option().point_size = 2 + vis.get_render_option().background_color = (0, 0, 0) + vis.get_render_option().show_coordinate_frame = True + # set up vector + vis.get_view_control().set_up([0, 0, 1]) + # add geometry + vis.add_geometry(pc_t0_o3d) + vis.add_geometry(pc_t1_o3d) + vis.add_geometry(pc_t0_warped_to_t1_o3d) + # run + vis.run() + + +class JointFlow(nn.Module): + + def __init__(self, + batch_size: int, + device: str, + VOXEL_SIZE=(0.14, 0.14, 4), + PSEUDO_IMAGE_DIMS=(512, 512), + POINT_CLOUD_RANGE=(-33.28, -33.28, -3, 33.28, 33.28, 1), + MAX_POINTS_PER_VOXEL=128, + FEATURE_CHANNELS=16, + FILTERS_PER_BLOCK=3, + PYRAMID_LAYERS=1, + SEQUENCE_LENGTH=5, + NSFP_FILTER_SIZE=64, + NSFP_NUM_LAYERS=4) -> None: + super().__init__() + self.batch_size = batch_size + self.device = device + self.embedder = HardEmbedder(voxel_size=VOXEL_SIZE, + pseudo_image_dims=PSEUDO_IMAGE_DIMS, + point_cloud_range=POINT_CLOUD_RANGE, + max_points_per_voxel=MAX_POINTS_PER_VOXEL, + feat_channels=FEATURE_CHANNELS) + + self.pyramid = FeaturePyramidNetwork( + pseudoimage_dims=PSEUDO_IMAGE_DIMS, + input_num_channels=FEATURE_CHANNELS + 3, + num_filters_per_block=FILTERS_PER_BLOCK, + num_layers_of_pyramid=PYRAMID_LAYERS) + + self.nsfp = NeuralSceneFlowPrior(num_hidden_units=NSFP_FILTER_SIZE, + num_hidden_layers=NSFP_NUM_LAYERS) + self.attention = JointConvAttention( + pseudoimage_dims=PSEUDO_IMAGE_DIMS, + sequence_length=SEQUENCE_LENGTH, + per_element_num_channels=(FEATURE_CHANNELS + 3) * + (PYRAMID_LAYERS + 1), + per_element_output_params=self.nsfp.param_count) + self.SEQUENCE_LENGTH = SEQUENCE_LENGTH + self.PSEUDO_IMAGE_DIMS = PSEUDO_IMAGE_DIMS + + def forward( + self, batched_sequence: Dict[str, torch.Tensor] + ) -> List[Tuple[PointCloud, SE3, torch.Tensor]]: + # print(f"Sequence info: {sequence_info}") + batched_results = [] + for batch_idx in range(self.batch_size): + pc_array = batched_sequence['pc_array_stack'][batch_idx] + pose_array = batched_sequence['pose_array_stack'][batch_idx] + sequence = list(zip(pc_array, pose_array)) + pc_embedding_list = [ + self._embed_pc(pc, pose) for pc, pose in sequence + ] + + pc_embedding_stack = torch.cat(pc_embedding_list, dim=1) + nsfp_weights_tensor = torch.squeeze( + self.attention(pc_embedding_stack), dim=0) + nsfp_weights_list = torch.split(nsfp_weights_tensor, + self.nsfp.param_count, + dim=0) + res = [(pc, pose, nsfp_weights) + for (pc, + pose), nsfp_weights in zip(sequence, nsfp_weights_list) + ] + batched_results.append(res) + return batched_results + + def _get_device(self): + return self.device + + def _embed_pc(self, pc: torch.Tensor, pose: torch.Tensor) -> torch.Tensor: + assert isinstance( + pc, torch.Tensor), f"Expected torch.Tensor, got {type(pc)}" + assert isinstance( + pose, torch.Tensor), f"Expected torch.Tensor, got {type(pose)}" + + batch_size = 1 + + # The pc is in a global frame relative to the initial pose of the sequence. + # The pose is the pose of the current point cloud relative to the initial pose of the sequence. + # In order to voxelize the point cloud, we need to translate it to the origin. + translation = _pose_to_translation(pose) + translated_pc = (pc - translation).float() + pseudoimage = self.embedder(translated_pc) + + translation_x_latent = torch.ones( + (batch_size, 1, self.PSEUDO_IMAGE_DIMS[0], + self.PSEUDO_IMAGE_DIMS[1])).to( + self._get_device()) * translation[0] + translation_y_latent = torch.ones( + (batch_size, 1, self.PSEUDO_IMAGE_DIMS[0], + self.PSEUDO_IMAGE_DIMS[1])).to( + self._get_device()) * translation[1] + translation_z_latent = torch.ones( + (batch_size, 1, self.PSEUDO_IMAGE_DIMS[0], + self.PSEUDO_IMAGE_DIMS[1])).to( + self._get_device()) * translation[2] + pseudoimage = torch.cat((pseudoimage, translation_x_latent, + translation_y_latent, translation_z_latent), + dim=1) + latent = self.pyramid(pseudoimage) + return latent diff --git a/models/nearest_neighbor.py b/models/nearest_neighbor.py new file mode 100644 index 0000000..2282ca1 --- /dev/null +++ b/models/nearest_neighbor.py @@ -0,0 +1,156 @@ +import torch +import torch.nn as nn + +import numpy as np +from models.embedders import DynamicVoxelizer +from models.nsfp_baseline import NSFPProcessor + +from typing import Dict, Any, Optional +from collections import defaultdict +from pathlib import Path +import time +from pytorch3d.ops.knn import knn_points + + +class NearestNeighborFlow(nn.Module): + + def __init__(self, + VOXEL_SIZE, + POINT_CLOUD_RANGE, + SEQUENCE_LENGTH, + flow_save_folder: Path, + skip_existing: bool = False) -> None: + super().__init__() + self.SEQUENCE_LENGTH = SEQUENCE_LENGTH + assert self.SEQUENCE_LENGTH == 2, "This implementation only supports a sequence length of 2." + self.voxelizer = DynamicVoxelizer(voxel_size=VOXEL_SIZE, + point_cloud_range=POINT_CLOUD_RANGE) + self.nsfp_processor = NSFPProcessor() + self.flow_save_folder = Path(flow_save_folder) + self.flow_save_folder.mkdir(parents=True, exist_ok=True) + self.skip_existing_cache = dict() + self.skip_existing = skip_existing + if self.skip_existing: + self.skip_existing_cache = self._build_skip_existing_cache() + + def _build_skip_existing_cache(self): + skip_existing_cache = dict() + for log_id in self.flow_save_folder.iterdir(): + skip_existing_cache[log_id.name] = set() + for npz_idx, npz_file in enumerate(sorted((log_id).glob("*.npz"))): + npz_filename_idx = int(npz_file.stem.split('_')[0]) + skip_existing_cache[log_id.name].add(npz_idx) + skip_existing_cache[log_id.name].add(npz_filename_idx) + return skip_existing_cache + + def _save_result(self, log_id: str, batch_idx: int, minibatch_idx: int, + delta_time: float, flow: torch.Tensor, + valid_idxes: torch.Tensor): + flow = flow.cpu().numpy() + valid_idxes = valid_idxes.cpu().numpy() + data = { + 'delta_time': delta_time, + 'flow': flow, + 'valid_idxes': valid_idxes, + } + seq_save_folder = self.flow_save_folder / log_id + seq_save_folder.mkdir(parents=True, exist_ok=True) + np.savez(seq_save_folder / f'{batch_idx:010d}_{minibatch_idx:03d}.npz', + **data) + + def _voxelize_batched_sequence(self, batched_sequence: Dict[str, + torch.Tensor]): + pc_arrays = batched_sequence['pc_array_stack'] + pc0s = pc_arrays[:, 0] + pc1s = pc_arrays[:, 1] + + pc0_voxel_infos_lst = self.voxelizer(pc0s) + pc1_voxel_infos_lst = self.voxelizer(pc1s) + + pc0_points_lst = [e["points"] for e in pc0_voxel_infos_lst] + pc0_valid_point_idxes = [e["point_idxes"] for e in pc0_voxel_infos_lst] + pc1_points_lst = [e["points"] for e in pc1_voxel_infos_lst] + pc1_valid_point_idxes = [e["point_idxes"] for e in pc1_voxel_infos_lst] + + return pc0_points_lst, pc0_valid_point_idxes, pc1_points_lst, pc1_valid_point_idxes + + def _visualize_result(self, pc0_points: torch.Tensor, + warped_pc0_points: torch.Tensor): + + pc0_points = pc0_points.cpu().numpy()[0] + warped_pc0_points = warped_pc0_points.cpu().numpy()[0] + + import open3d as o3d + + line_set = o3d.geometry.LineSet() + assert len(pc0_points) == len( + warped_pc0_points + ), f'pc and flowed_pc must have same length, but got {len(pc0_pcd)} and {len(warped_pc0_points)}' + line_set_points = np.concatenate([pc0_points, warped_pc0_points], + axis=0) + + pc0_pcd = o3d.geometry.PointCloud() + pc0_pcd.points = o3d.utility.Vector3dVector(pc0_points) + warped_pc0_pcd = o3d.geometry.PointCloud() + warped_pc0_pcd.points = o3d.utility.Vector3dVector(warped_pc0_points) + + line_set = o3d.geometry.LineSet() + line_set.points = o3d.utility.Vector3dVector(line_set_points) + lines = np.array([[i, i + len(pc0_points)] + for i in range(len(pc0_points))]) + line_set.lines = o3d.utility.Vector2iVector(lines) + + o3d.visualization.draw_geometries([pc0_pcd, warped_pc0_pcd, line_set]) + + def forward(self, batched_sequence: Dict[str, torch.Tensor]): + pc0_points_lst, pc0_valid_point_idxes, pc1_points_lst, pc1_valid_point_idxes = self._voxelize_batched_sequence( + batched_sequence) + + # process minibatch + flows = [] + delta_times = [] + for minibatch_idx, (pc0_points, pc1_points, + pc0_valid_point_idx) in enumerate( + zip(pc0_points_lst, pc1_points_lst, + pc0_valid_point_idxes)): + log_id = batched_sequence['log_ids'][minibatch_idx][0] + batch_idx = batched_sequence['data_index'].item() + if self.skip_existing: + if log_id in self.skip_existing_cache: + existing_set = self.skip_existing_cache[log_id] + if batch_idx in existing_set: + print(f'Skip existing flow for {log_id} {batch_idx}') + continue + + pc0_points = torch.unsqueeze(pc0_points, 0) + pc1_points = torch.unsqueeze(pc1_points, 0) + + before_time = time.time() + pc0_to_pc1_knn = knn_points(p1=pc0_points, p2=pc1_points, K=1) + + paired_indices = pc0_to_pc1_knn.idx[0, :, 0] + warped_pc0_points = pc1_points[0, paired_indices, :] + after_time = time.time() + + delta_time = after_time - before_time + # self._visualize_result(pc0_points, warped_pc0_points) + flow = warped_pc0_points - pc0_points + self._save_result(log_id=log_id, + batch_idx=batch_idx, + minibatch_idx=minibatch_idx, + delta_time=delta_time, + flow=flow, + valid_idxes=pc0_valid_point_idx) + flows.append(flow.squeeze(0)) + delta_times.append(delta_time) + + return { + "forward": { + "flow": flows, + "batch_delta_time": np.sum(delta_times), + "pc0_points_lst": pc0_points_lst, + "pc0_valid_point_idxes": pc0_valid_point_idxes, + "pc1_points_lst": pc1_points_lst, + "pc1_valid_point_idxes": pc1_valid_point_idxes + } + } diff --git a/models/nsfp.py b/models/nsfp.py new file mode 100644 index 0000000..4c0035c --- /dev/null +++ b/models/nsfp.py @@ -0,0 +1,251 @@ +import torch +import torch.nn as nn + +import numpy as np +from models.embedders import DynamicVoxelizer +from models.nsfp_baseline import NSFPProcessor + +from typing import Dict, Any, Optional +from collections import defaultdict +from pathlib import Path +import time + + +class NSFP(nn.Module): + + def __init__(self, + VOXEL_SIZE, + POINT_CLOUD_RANGE, + SEQUENCE_LENGTH, + flow_save_folder: Path, + skip_existing: bool = False) -> None: + super().__init__() + self.SEQUENCE_LENGTH = SEQUENCE_LENGTH + assert self.SEQUENCE_LENGTH == 2, "This implementation only supports a sequence length of 2." + self.voxelizer = DynamicVoxelizer(voxel_size=VOXEL_SIZE, + point_cloud_range=POINT_CLOUD_RANGE) + self.nsfp_processor = NSFPProcessor() + self.flow_save_folder = Path(flow_save_folder) + self.flow_save_folder.mkdir(parents=True, exist_ok=True) + self.skip_existing_cache = dict() + self.skip_existing = skip_existing + if self.skip_existing: + self.skip_existing_cache = self._build_skip_existing_cache() + + def _build_skip_existing_cache(self): + skip_existing_cache = dict() + for log_id in self.flow_save_folder.iterdir(): + skip_existing_cache[log_id.name] = set() + for npz_idx, npz_file in enumerate(sorted((log_id).glob("*.npz"))): + npz_filename_idx = int(npz_file.stem.split('_')[0]) + skip_existing_cache[log_id.name].add(npz_idx) + skip_existing_cache[log_id.name].add(npz_filename_idx) + return skip_existing_cache + + def _save_result(self, log_id: str, batch_idx: int, minibatch_idx: int, + delta_time: float, flow: torch.Tensor, + valid_idxes: torch.Tensor): + flow = flow.cpu().numpy() + valid_idxes = valid_idxes.cpu().numpy() + data = { + 'delta_time': delta_time, + 'flow': flow, + 'valid_idxes': valid_idxes, + } + seq_save_folder = self.flow_save_folder / log_id + seq_save_folder.mkdir(parents=True, exist_ok=True) + np.savez(seq_save_folder / f'{batch_idx:010d}_{minibatch_idx:03d}.npz', + **data) + + def _voxelize_batched_sequence(self, batched_sequence: Dict[str, + torch.Tensor]): + pc_arrays = batched_sequence['pc_array_stack'] + pc0s = pc_arrays[:, 0] + pc1s = pc_arrays[:, 1] + + pc0_voxel_infos_lst = self.voxelizer(pc0s) + pc1_voxel_infos_lst = self.voxelizer(pc1s) + + pc0_points_lst = [e["points"] for e in pc0_voxel_infos_lst] + pc0_valid_point_idxes = [e["point_idxes"] for e in pc0_voxel_infos_lst] + pc1_points_lst = [e["points"] for e in pc1_voxel_infos_lst] + pc1_valid_point_idxes = [e["point_idxes"] for e in pc1_voxel_infos_lst] + + return pc0_points_lst, pc0_valid_point_idxes, pc1_points_lst, pc1_valid_point_idxes + + def _visualize_result(self, pc0_points: torch.Tensor, + warped_pc0_points: torch.Tensor): + + # if pc0_points is torch tensor, convert to numpy + if isinstance(pc0_points, torch.Tensor): + pc0_points = pc0_points.cpu().numpy()[0] + if isinstance(warped_pc0_points, torch.Tensor): + warped_pc0_points = warped_pc0_points.cpu().numpy()[0] + + import open3d as o3d + + line_set = o3d.geometry.LineSet() + assert len(pc0_points) == len( + warped_pc0_points + ), f'pc and flowed_pc must have same length, but got {len(pc0_pcd)} and {len(warped_pc0_points)}' + line_set_points = np.concatenate([pc0_points, warped_pc0_points], + axis=0) + + pc0_pcd = o3d.geometry.PointCloud() + pc0_pcd.points = o3d.utility.Vector3dVector(pc0_points) + warped_pc0_pcd = o3d.geometry.PointCloud() + warped_pc0_pcd.points = o3d.utility.Vector3dVector(warped_pc0_points) + + line_set = o3d.geometry.LineSet() + line_set.points = o3d.utility.Vector3dVector(line_set_points) + lines = np.array([[i, i + len(pc0_points)] + for i in range(len(pc0_points))]) + line_set.lines = o3d.utility.Vector2iVector(lines) + + o3d.visualization.draw_geometries([pc0_pcd, warped_pc0_pcd, line_set]) + + def forward(self, batched_sequence: Dict[str, torch.Tensor]): + pc0_points_lst, pc0_valid_point_idxes, pc1_points_lst, pc1_valid_point_idxes = self._voxelize_batched_sequence( + batched_sequence) + + # process minibatch + flows = [] + delta_times = [] + for minibatch_idx, (pc0_points, pc1_points, + pc0_valid_point_idx) in enumerate( + zip(pc0_points_lst, pc1_points_lst, + pc0_valid_point_idxes)): + log_id = batched_sequence['log_ids'][minibatch_idx][0] + batch_idx = batched_sequence['data_index'].item() + if self.skip_existing: + if log_id in self.skip_existing_cache: + existing_set = self.skip_existing_cache[log_id] + if batch_idx in existing_set: + print(f'Skip existing flow for {log_id} {batch_idx}') + continue + else: + print("batch_idx", batch_idx, "not in", existing_set) + else: + print("log_id", log_id, "not in", self.skip_existing_cache.keys()) + # and ( + # log_id in self.skip_existing_cache) and ( + # batch_idx in self.skip_existing_cache[log_id]): + # print(f'Skip existing flow for {log_id} {batch_idx}') + # continue + # else: + # print(self.skip_existing_cache.keys()) + + pc0_points = torch.unsqueeze(pc0_points, 0) + pc1_points = torch.unsqueeze(pc1_points, 0) + + with torch.inference_mode(False): + with torch.enable_grad(): + pc0_points_new = pc0_points.clone().detach( + ).requires_grad_(True) + pc1_points_new = pc1_points.clone().detach( + ).requires_grad_(True) + + self.nsfp_processor.train() + + before_time = time.time() + warped_pc0_points, _ = self.nsfp_processor( + pc0_points_new, pc1_points_new, pc1_points_new.device) + after_time = time.time() + + delta_time = after_time - before_time + # self._visualize_result(pc0_points, warped_pc0_points) + flow = warped_pc0_points - pc0_points + self._save_result(log_id=log_id, + batch_idx=batch_idx, + minibatch_idx=minibatch_idx, + delta_time=delta_time, + flow=flow, + valid_idxes=pc0_valid_point_idx) + flows.append(flow.squeeze(0)) + delta_times.append(delta_time) + + return { + "forward": { + "flow": flows, + "batch_delta_time": np.sum(delta_times), + "pc0_points_lst": pc0_points_lst, + "pc0_valid_point_idxes": pc0_valid_point_idxes, + "pc1_points_lst": pc1_points_lst, + "pc1_valid_point_idxes": pc1_valid_point_idxes + } + } + + +class NSFPCached(NSFP): + + def __init__(self, VOXEL_SIZE, POINT_CLOUD_RANGE, SEQUENCE_LENGTH, + flow_save_folder: Path) -> None: + super().__init__(VOXEL_SIZE, POINT_CLOUD_RANGE, SEQUENCE_LENGTH, + flow_save_folder) + # Implement basic caching to avoid repeated folder reads for the same log. + self.flow_folder_id = None + self.flow_folder_list = None + + def _load_result(self, log_id: str, log_idx: int): + if self.flow_folder_id != log_id: + self.flow_folder_id = log_id + flow_folder = self.flow_save_folder / log_id + assert flow_folder.is_dir(), f"{flow_folder} does not exist" + self.flow_folder_list = sorted(flow_folder.glob('*.npz')) + + assert log_idx < len( + self.flow_folder_list), f"Log index {log_idx} is out of range" + flow_file = self.flow_folder_list[log_idx] + print(f"Loading flow from {flow_file}") + data = dict(np.load(flow_file, allow_pickle=True)) + flow = data['flow'] + valid_idxes = data['valid_idxes'] + delta_time = data['delta_time'] + return flow, valid_idxes, delta_time + + def _transpose_list(self, lst): + if isinstance(lst[0], torch.Tensor): + return torch.stack(lst).T.tolist() + return np.array(lst).T.tolist() + + def forward(self, batched_sequence: Dict[str, torch.Tensor]): + pc0_points_lst, pc0_valid_point_idxes, pc1_points_lst, pc1_valid_point_idxes = self._voxelize_batched_sequence( + batched_sequence) + + log_ids_lst = self._transpose_list(batched_sequence['log_ids']) + log_idxes_lst = self._transpose_list(batched_sequence['log_idxes']) + + flows = [] + delta_times = [] + for minibatch_idx, (pc0_points, pc0_valid_points_idx, log_ids, + log_idxes) in enumerate( + zip(pc0_points_lst, pc0_valid_point_idxes, + log_ids_lst, log_idxes_lst)): + assert len( + log_ids) == 2, f"Expect 2 frames for NSFP, got {len(log_ids)}" + assert len( + log_idxes + ) == 2, f"Expect 2 frames for NSFP, got {len(log_idxes)}" + flow, flow_valid_idxes, delta_time = self._load_result( + log_ids[0], log_idxes[0]) + assert flow_valid_idxes.shape == pc0_valid_points_idx.shape, f"{flow_valid_idxes.shape} != {pc0_valid_points_idx.shape}" + idx_equality = flow_valid_idxes == pc0_valid_points_idx.detach( + ).cpu().numpy() + assert idx_equality.all( + ), f"Points that disagree: {(~idx_equality).astype(int).sum()}" + if flow.ndim == 3 and flow.shape[0] == 1: + flow = flow.squeeze(0) + assert flow.shape == pc0_points.shape, f"{flow.shape} != {pc0_points.shape}" + flows.append(torch.from_numpy(flow).to(pc0_points.device)) + delta_times.append(delta_time) + + return { + "forward": { + "flow": flows, + "batch_delta_time": np.sum(delta_times), + "pc0_points_lst": pc0_points_lst, + "pc0_valid_point_idxes": pc0_valid_point_idxes, + "pc1_points_lst": pc1_points_lst, + "pc1_valid_point_idxes": pc1_valid_point_idxes + } + } diff --git a/models/nsfp_baseline/__init__.py b/models/nsfp_baseline/__init__.py new file mode 100644 index 0000000..cad486b --- /dev/null +++ b/models/nsfp_baseline/__init__.py @@ -0,0 +1,6 @@ +from .optimize import NSFPProcessor, NSFPLoss + +__all__ = [ + 'NSFPProcessor', + 'NSFPLoss', +] \ No newline at end of file diff --git a/models/nsfp_baseline/model.py b/models/nsfp_baseline/model.py new file mode 100644 index 0000000..df18871 --- /dev/null +++ b/models/nsfp_baseline/model.py @@ -0,0 +1,36 @@ +import torch + + +class Neural_Prior(torch.nn.Module): + def __init__(self, dim_x=3, filter_size=128, act_fn='relu', layer_size=8): + super().__init__() + self.layer_size = layer_size + + self.nn_layers = torch.nn.ModuleList([]) + # input layer (default: xyz -> 128) + if layer_size >= 1: + self.nn_layers.append(torch.nn.Sequential(torch.nn.Linear(dim_x, filter_size))) + if act_fn == 'relu': + self.nn_layers.append(torch.nn.ReLU()) + elif act_fn == 'sigmoid': + self.nn_layers.append(torch.nn.Sigmoid()) + for _ in range(layer_size-1): + self.nn_layers.append(torch.nn.Sequential(torch.nn.Linear(filter_size, filter_size))) + if act_fn == 'relu': + self.nn_layers.append(torch.nn.ReLU()) + elif act_fn == 'sigmoid': + self.nn_layers.append(torch.nn.Sigmoid()) + self.nn_layers.append(torch.nn.Linear(filter_size, dim_x)) + else: + self.nn_layers.append(torch.nn.Sequential(torch.nn.Linear(dim_x, dim_x))) + + def forward(self, in_x): + """ points -> features + [B, N, 3] -> [B, K] + """ + x = in_x + for layer in self.nn_layers: + x = layer(x) + + return x + \ No newline at end of file diff --git a/models/nsfp_baseline/optimize.py b/models/nsfp_baseline/optimize.py new file mode 100644 index 0000000..927daf8 --- /dev/null +++ b/models/nsfp_baseline/optimize.py @@ -0,0 +1,331 @@ +import copy +import torch +import torch.nn.functional as F +from pytorch3d.ops.knn import knn_gather, knn_points +from pytorch3d.structures.pointclouds import Pointclouds +import time + +from pointclouds.losses import warped_pc_loss + +from .utils import EarlyStopping +from .model import Neural_Prior + +from typing import Union + +import torch +import torch.nn as nn +import tqdm + + +def _validate_chamfer_reduction_inputs(batch_reduction: Union[str, None], + point_reduction: str): + """Check the requested reductions are valid. + Args: + batch_reduction: Reduction operation to apply for the loss across the + batch, can be one of ["mean", "sum"] or None. + point_reduction: Reduction operation to apply for the loss across the + points, can be one of ["mean", "sum"]. + """ + if batch_reduction is not None and batch_reduction not in ["mean", "sum"]: + raise ValueError( + 'batch_reduction must be one of ["mean", "sum"] or None') + if point_reduction not in ["mean", "sum"]: + raise ValueError('point_reduction must be one of ["mean", "sum"]') + + +def _handle_pointcloud_input( + points: Union[torch.Tensor, Pointclouds], + lengths: Union[torch.Tensor, None], +): + """ + If points is an instance of Pointclouds, retrieve the padded points tensor + along with the number of points per batch and the padded normals. + Otherwise, return the input points (and normals) with the number of points per cloud + set to the size of the second dimension of `points`. + """ + if isinstance(points, Pointclouds): + X = points.points_padded() + lengths = points.num_points_per_cloud() + elif torch.is_tensor(points): + if points.ndim != 3: + raise ValueError( + f"Expected points to be of shape (N, P, D), got {points.shape}" + ) + X = points + if lengths is not None and (lengths.ndim != 1 + or lengths.shape[0] != X.shape[0]): + raise ValueError("Expected lengths to be of shape (N,)") + if lengths is None: + lengths = torch.full((X.shape[0], ), + X.shape[1], + dtype=torch.int64, + device=points.device) + else: + raise ValueError("The input pointclouds should be either " + + "Pointclouds objects or torch.Tensor of shape " + + "(minibatch, num_points, 3).") + return X, lengths + + +def my_chamfer_fn( + x, + y, + x_lengths=None, + y_lengths=None, + x_normals=None, + y_normals=None, + weights=None, + batch_reduction: Union[str, None] = "mean", + point_reduction: str = "mean", +): + """ + Chamfer distance between two pointclouds x and y. + + Args: + x: FloatTensor of shape (N, P1, D) or a Pointclouds object representing + a batch of point clouds with at most P1 points in each batch element, + batch size N and feature dimension D. + y: FloatTensor of shape (N, P2, D) or a Pointclouds object representing + a batch of point clouds with at most P2 points in each batch element, + batch size N and feature dimension D. + x_lengths: Optional LongTensor of shape (N,) giving the number of points in each + cloud in x. + y_lengths: Optional LongTensor of shape (N,) giving the number of points in each + cloud in x. + x_normals: Optional FloatTensor of shape (N, P1, D). + y_normals: Optional FloatTensor of shape (N, P2, D). + weights: Optional FloatTensor of shape (N,) giving weights for + batch elements for reduction operation. + batch_reduction: Reduction operation to apply for the loss across the + batch, can be one of ["mean", "sum"] or None. + point_reduction: Reduction operation to apply for the loss across the + points, can be one of ["mean", "sum"]. + + Returns: + 2-element tuple containing + + - **loss**: Tensor giving the reduced distance between the pointclouds + in x and the pointclouds in y. + - **loss_normals**: Tensor giving the reduced cosine distance of normals + between pointclouds in x and pointclouds in y. Returns None if + x_normals and y_normals are None. + """ + _validate_chamfer_reduction_inputs(batch_reduction, point_reduction) + + x, x_lengths = _handle_pointcloud_input(x, None) + y, y_lengths = _handle_pointcloud_input(y, None) + + assert x_lengths.item() > 0, f"x_lengths is {x_lengths.item()}" + assert y_lengths.item() > 0, f"y_lengths is {y_lengths.item()}" + + return_normals = x_normals is not None and y_normals is not None + + N, P1, D = x.shape + P2 = y.shape[1] + + # Check if inputs are heterogeneous and create a lengths mask. + is_x_heterogeneous = (x_lengths != P1).any() + is_y_heterogeneous = (y_lengths != P2).any() + x_mask = (torch.arange(P1, device=x.device)[None] >= x_lengths[:, None] + ) # shape [N, P1] + y_mask = (torch.arange(P2, device=y.device)[None] >= y_lengths[:, None] + ) # shape [N, P2] + + if y.shape[0] != N or y.shape[2] != D: + raise ValueError("y does not have the correct shape.") + if weights is not None: + if weights.size(0) != N: + raise ValueError("weights must be of shape (N,).") + if not (weights >= 0).all(): + raise ValueError("weights cannot be negative.") + if weights.sum() == 0.0: + weights = weights.view(N, 1) + if batch_reduction in ["mean", "sum"]: + return ( + (x.sum((1, 2)) * weights).sum() * 0.0, + (x.sum((1, 2)) * weights).sum() * 0.0, + ) + return ((x.sum((1, 2)) * weights) * 0.0, (x.sum( + (1, 2)) * weights) * 0.0) + + cham_norm_x = x.new_zeros(()) + cham_norm_y = x.new_zeros(()) + + x_nn = knn_points(x, y, lengths1=x_lengths, lengths2=y_lengths, K=1) + y_nn = knn_points(y, x, lengths1=y_lengths, lengths2=x_lengths, K=1) + + cham_x = x_nn.dists[..., 0] # (N, P1) + cham_y = y_nn.dists[..., 0] # (N, P2) + + # NOTE: truncated Chamfer distance. + dist_thd = 2 + x_mask[cham_x >= dist_thd] = True + y_mask[cham_y >= dist_thd] = True + cham_x[x_mask] = 0.0 + cham_y[y_mask] = 0.0 + + if is_x_heterogeneous: + cham_x[x_mask] = 0.0 + if is_y_heterogeneous: + cham_y[y_mask] = 0.0 + + if weights is not None: + cham_x *= weights.view(N, 1) + cham_y *= weights.view(N, 1) + + if return_normals: + # Gather the normals using the indices and keep only value for k=0 + x_normals_near = knn_gather(y_normals, x_nn.idx, y_lengths)[..., 0, :] + y_normals_near = knn_gather(x_normals, y_nn.idx, x_lengths)[..., 0, :] + + cham_norm_x = 1 - torch.abs( + F.cosine_similarity(x_normals, x_normals_near, dim=2, eps=1e-6)) + cham_norm_y = 1 - torch.abs( + F.cosine_similarity(y_normals, y_normals_near, dim=2, eps=1e-6)) + + if is_x_heterogeneous: + # pyre-fixme[16]: `int` has no attribute `__setitem__`. + cham_norm_x[x_mask] = 0.0 + if is_y_heterogeneous: + cham_norm_y[y_mask] = 0.0 + + if weights is not None: + cham_norm_x *= weights.view(N, 1) + cham_norm_y *= weights.view(N, 1) + + # Apply point reduction + cham_x = cham_x.sum(1) # (N,) + cham_y = cham_y.sum(1) # (N,) + + if point_reduction == "mean": + cham_x /= x_lengths + cham_y /= y_lengths + + if batch_reduction is not None: + # batch_reduction == "sum" + cham_x = cham_x.sum() + cham_y = cham_y.sum() + if batch_reduction == "mean": + cham_x /= N + cham_y /= N + + cham_dist = cham_x + cham_y + cham_normals = cham_norm_x + cham_norm_y if return_normals else None + + return cham_dist, cham_normals + + +def optimize(pc1, pc2, device, iterations=5000, lr=8e-3, min_delta=0.00005): + net = Neural_Prior(filter_size=128, act_fn='relu', layer_size=8) + # if torch.__version__.split('.')[0] == '2': + # net = torch.compile(net) + net = net.to(device) + net.train() + early_stopping = EarlyStopping(patience=100, min_delta=min_delta) + + net_inv = copy.deepcopy(net) + + params = [{ + 'params': net.parameters(), + 'lr': lr, + 'weight_decay': 0 + }, { + 'params': net_inv.parameters(), + 'lr': lr, + 'weight_decay': 0 + }] + best_forward = {'loss': torch.inf} + optimizer = torch.optim.Adam(params, lr=0.008, weight_decay=0) + position = 1 + if isinstance(device, torch.device): + position += device.index + + net1_forward_time = 0 + net1_loss_forward_time = 0 + net2_forward_time = 0 + net2_loss_forward_time = 0 + + loss_backwards_time = 0 + optimizer_step_time = 0 + + for epoch in tqdm.tqdm(range(iterations), + position=position, + leave=False, + desc='Optimizing NSFP'): + + optimizer.zero_grad() + + net1_forward_before = time.time() + forward_flow = net(pc1) + net1_forward_after = time.time() + pc1_warped_to_pc2 = pc1 + forward_flow + + net2_forward_before = time.time() + reverse_flow = net_inv(pc1_warped_to_pc2) + net2_forward_after = time.time() + est_pc2_warped_to_pc1 = pc1_warped_to_pc2 - reverse_flow + + net1_loss_forward_before = time.time() + forward_loss, _ = my_chamfer_fn(pc1_warped_to_pc2, pc2, None, None) + net1_loss_forward_after = time.time() + + net2_loss_forward_before = time.time() + reversed_loss, _ = my_chamfer_fn(est_pc2_warped_to_pc1, pc1, None, + None) + net2_loss_forward_after = time.time() + + loss = forward_loss + reversed_loss + + # ANCHOR: get best metrics + if forward_loss <= best_forward['loss']: + best_forward['loss'] = forward_loss.item() + best_forward['warped_pc'] = pc1_warped_to_pc2 + best_forward['target_pc'] = pc2 + + if early_stopping.step(loss): + break + + loss_backwards_before = time.time() + loss.backward() + loss_backwards_after = time.time() + optimizer.step() + optimizer_step_afterwards = time.time() + + net1_forward_time += net1_forward_after - net1_forward_before + net1_loss_forward_time += net1_loss_forward_after - net1_loss_forward_before + net2_forward_time += net2_forward_after - net2_forward_before + net2_loss_forward_time += net2_loss_forward_after - net2_loss_forward_before + loss_backwards_time += loss_backwards_after - loss_backwards_before + optimizer_step_time += optimizer_step_afterwards - loss_backwards_after + + # print(f'net1_forward_time: {net1_forward_time}') + # print(f'net1_loss_forward_time: {net1_loss_forward_time}') + # print(f'net2_forward_time: {net2_forward_time}') + # print(f'net2_loss_forward_time: {net2_loss_forward_time}') + # print(f'loss_backwards_time: {loss_backwards_time}') + # print(f'optimizer_step_time: {optimizer_step_time}') + + # print( + # f"Total time: {net1_forward_time + net1_loss_forward_time + net2_forward_time + net2_loss_forward_time + loss_backwards_time + optimizer_step_time}" + # ) + + return best_forward + + +class NSFPProcessor(nn.Module): + + def __init__(self): + super().__init__() + + def forward(self, pc1, pc2, device): + res = optimize(pc1, pc2, device) + return res['warped_pc'], res['target_pc'] + + +class NSFPLoss(): + + def __init__(self): + pass + + def __call__(self, warped_pc, target_pc): + return warped_pc_loss(warped_pc, target_pc) diff --git a/models/nsfp_baseline/utils.py b/models/nsfp_baseline/utils.py new file mode 100644 index 0000000..7cd7ef4 --- /dev/null +++ b/models/nsfp_baseline/utils.py @@ -0,0 +1,150 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function +from __future__ import unicode_literals + +import torch +import time +from collections import defaultdict + + +# ANCHOR: metrics computation, follow FlowNet3D metrics.... +def scene_flow_metrics(pred, labels): + l2_norm = torch.sqrt(torch.sum((pred - labels) ** 2, 2)).cpu() # Absolute distance error. + labels_norm = torch.sqrt(torch.sum(labels * labels, 2)).cpu() + relative_err = l2_norm / (labels_norm + 1e-20) + + EPE3D = torch.mean(l2_norm).item() # Mean absolute distance error + + # NOTE: Acc_5 + error_lt_5 = torch.BoolTensor((l2_norm < 0.05)) + relative_err_lt_5 = torch.BoolTensor((relative_err < 0.05)) + acc3d_strict = torch.mean((error_lt_5 | relative_err_lt_5).float()).item() + + # NOTE: Acc_10 + error_lt_10 = torch.BoolTensor((l2_norm < 0.1)) + relative_err_lt_10 = torch.BoolTensor((relative_err < 0.1)) + acc3d_relax = torch.mean((error_lt_10 | relative_err_lt_10).float()).item() + + # NOTE: outliers + l2_norm_gt_3 = torch.BoolTensor(l2_norm > 0.3) + relative_err_gt_10 = torch.BoolTensor(relative_err > 0.1) + outlier = torch.mean((l2_norm_gt_3 | relative_err_gt_10).float()).item() + + # NOTE: angle error + unit_label = labels / labels.norm(dim=2, keepdim=True) + unit_pred = pred / pred.norm(dim=2, keepdim=True) + eps = 1e-7 + dot_product = (unit_label * unit_pred).sum(2).clamp(min=-1+eps, max=1-eps) + dot_product[dot_product != dot_product] = 0 # Remove NaNs + angle_error = torch.acos(dot_product).mean().item() + + return EPE3D, acc3d_strict, acc3d_relax, outlier, angle_error + + +# ANCHOR: timer! +class Timers(object): + def __init__(self): + self.timers = defaultdict(Timer) + + def tic(self, key): + self.timers[key].tic() + + def toc(self, key): + self.timers[key].toc() + + def print(self, key=None): + if key is None: + for k, v in self.timers.items(): + print("Average time for {:}: {:}".format(k, v.avg())) + else: + print("Average time for {:}: {:}".format(key, self.timers[key].avg())) + + def get_avg(self, key): + return self.timers[key].avg() + + +class Timer(object): + def __init__(self): + self.reset() + + def tic(self): + self.start_time = time.time() + + def toc(self, average=True): + self.diff = time.time() - self.start_time + self.total_time += self.diff + self.calls += 1 + + def total(self): + return self.total_time + + def avg(self): + return self.total_time / float(self.calls) + + def reset(self): + self.total_time = 0. + self.calls = 0 + self.start_time = 0. + self.diff = 0. + + +# ANCHOR: generator +class GeneratorWrap: + def __init__(self, gen): + self.gen = gen + + def __iter__(self): + self.value = yield from self.gen + + +# ANCHOR: early stopping strategy +class EarlyStopping(object): + def __init__(self, mode='min', min_delta=0, patience=10, percentage=False): + self.mode = mode + self.min_delta = min_delta + self.patience = patience + self.best = None + self.num_bad_epochs = 0 + self.is_better = None + self._init_is_better(mode, min_delta, percentage) + + if patience == 0: + self.is_better = lambda a, b: True + self.step = lambda a: False + + def step(self, metrics): + if self.best is None: + self.best = metrics + return False + + if torch.isnan(metrics): + return True + + if self.is_better(metrics, self.best): + self.num_bad_epochs = 0 + self.best = metrics + else: + self.num_bad_epochs += 1 + + if self.num_bad_epochs >= self.patience: + return True + + return False + + def _init_is_better(self, mode, min_delta, percentage): + if mode not in {'min', 'max'}: + raise ValueError('mode ' + mode + ' is unknown!') + if not percentage: + if mode == 'min': + self.is_better = lambda a, best: a < best - min_delta + if mode == 'max': + self.is_better = lambda a, best: a > best + min_delta + else: + if mode == 'min': + self.is_better = lambda a, best: a < best - ( + best * min_delta / 100) + if mode == 'max': + self.is_better = lambda a, best: a > best + ( + best * min_delta / 100) + \ No newline at end of file diff --git a/models/pretrain.py b/models/pretrain.py new file mode 100644 index 0000000..ce5aa20 --- /dev/null +++ b/models/pretrain.py @@ -0,0 +1,120 @@ +import torch +import torch.nn as nn +import numpy as np +from typing import Tuple + +from models.embedders import HardEmbedder +from models.backbones import FeaturePyramidNetwork + + +class ConvWithNorms(nn.Module): + + def __init__(self, in_num_channels: int, out_num_channels: int, + kernel_size: int, stride: int, padding: int): + super().__init__() + self.conv = nn.Conv2d(in_num_channels, out_num_channels, kernel_size, + stride, padding) + self.batchnorm = nn.BatchNorm2d(out_num_channels) + self.gelu = nn.GELU() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + conv_res = self.conv(x) + if conv_res.shape[2] == 1 and conv_res.shape[3] == 1: + # This is a hack to get around the fact that batchnorm doesn't support + # 1x1 convolutions + batchnorm_res = conv_res + else: + batchnorm_res = self.batchnorm(conv_res) + return self.gelu(batchnorm_res) + + +class ShearMatrixRegression(nn.Module): + + def __init__(self, pseudoimage_dims: Tuple[int, int], + pseudoimage_channels: int): + super().__init__() + self.pseudoimage_x = pseudoimage_dims[0] + self.pseudoimage_y = pseudoimage_dims[1] + self.num_channels = pseudoimage_channels + + # Scale, shear value, binary classification of upper or lower off diag shear + out_channels = 3 + + assert self.pseudoimage_x > 0, f"pseudoimage_x must be > 0, got {self.pseudoimage_x}" + assert self.pseudoimage_y > 0, f"pseudoimage_y must be > 0, got {self.pseudoimage_y}" + + pseudoimage_stepdowns = np.ceil( + np.max(np.log2(np.array([self.pseudoimage_x, + self.pseudoimage_y])))).astype(int) + print(f"pseudoimage_stepdowns: {pseudoimage_stepdowns}") + + conv_list = [ + ConvWithNorms(pseudoimage_channels, out_channels, 3, 1, 1) + ] + for _ in range(pseudoimage_stepdowns - 1): + conv_list.append(ConvWithNorms(out_channels, out_channels, 3, 2, + 1)) + # Standard convolution + conv_list.append(nn.Conv2d(out_channels, 3, 3, 2, 1)) + + self.ops = nn.Sequential(*conv_list) + + def forward(self, pseudoimage: torch.Tensor) -> torch.Tensor: + return self.ops(pseudoimage) + + +class PretrainEmbedding(nn.Module): + + def __init__(self, batch_size: int, device: str, VOXEL_SIZE, + PSEUDO_IMAGE_DIMS, POINT_CLOUD_RANGE, MAX_POINTS_PER_VOXEL, + FEATURE_CHANNELS, FILTERS_PER_BLOCK, PYRAMID_LAYERS) -> None: + super().__init__() + self.PSEUDO_IMAGE_DIMS = PSEUDO_IMAGE_DIMS + self.batch_size = batch_size + self.device = device + self.embedder = HardEmbedder(voxel_size=VOXEL_SIZE, + pseudo_image_dims=PSEUDO_IMAGE_DIMS, + point_cloud_range=POINT_CLOUD_RANGE, + max_points_per_voxel=MAX_POINTS_PER_VOXEL, + feat_channels=FEATURE_CHANNELS) + + self.pyramid = FeaturePyramidNetwork( + pseudoimage_dims=PSEUDO_IMAGE_DIMS, + input_num_channels=FEATURE_CHANNELS + 3, + num_filters_per_block=FILTERS_PER_BLOCK, + num_layers_of_pyramid=PYRAMID_LAYERS) + + self.head = ShearMatrixRegression( + PSEUDO_IMAGE_DIMS, (FEATURE_CHANNELS + 3) * (PYRAMID_LAYERS + 1)) + + def forward(self, points, translations): + shears = [] + for batch_idx in range(self.batch_size): + embeddings = self.embedder(points[batch_idx]) + + translation = translations[batch_idx] + + translation_x_latent = torch.ones( + (1, 1, self.PSEUDO_IMAGE_DIMS[0], + self.PSEUDO_IMAGE_DIMS[1])).to( + self._get_device()) * translation[0] + translation_y_latent = torch.ones( + (1, 1, self.PSEUDO_IMAGE_DIMS[0], + self.PSEUDO_IMAGE_DIMS[1])).to( + self._get_device()) * translation[1] + translation_z_latent = torch.ones( + (1, 1, self.PSEUDO_IMAGE_DIMS[0], + self.PSEUDO_IMAGE_DIMS[1])).to( + self._get_device()) * translation[2] + embeddings = torch.cat( + (embeddings, translation_x_latent, translation_y_latent, + translation_z_latent), + dim=1) + + pseudoimage = self.pyramid(embeddings) + shear = self.head(pseudoimage) + shears.append(shear) + return torch.vstack(shears) + + def _get_device(self): + return self.device \ No newline at end of file diff --git a/plot_performance.py b/plot_performance.py new file mode 100644 index 0000000..a9ba0ba --- /dev/null +++ b/plot_performance.py @@ -0,0 +1,1244 @@ +import numpy as np +from pathlib import Path +from loader_utils import * +import matplotlib +import matplotlib.pyplot as plt +import argparse +from typing import List, Tuple, Dict, Any, Optional +import shutil + +# Get path to methods from command line +parser = argparse.ArgumentParser() +parser.add_argument('results_folder', type=Path) +parser.add_argument('--dataset', + type=str, + default="argo", + choices=["argo", "waymo"]) +args = parser.parse_args() + +assert args.results_folder.exists( +), f"Results folder {args.results_folder} does not exist" + +save_folder = args.results_folder / f"plots_{args.dataset}" +save_folder.mkdir(exist_ok=True, parents=True) + + +def set_font(size): + matplotlib.rcParams.update({# Use mathtext, not LaTeX + 'text.usetex': False, + # Use the Computer modern font + 'font.family': 'serif', + 'font.serif': ['cmr10'], + 'font.size' : size, + 'mathtext.fontset': 'cm', + # Use ASCII minus + 'axes.unicode_minus': False, + }) + + +def color_map(rev: bool = False): + # return 'gist_earth' + if rev: + return 'magma_r' + return 'magma' + + +def color(count, total_elements, intensity=1.3): + start = 0.2 + stop = 0.7 + + colormap = matplotlib.cm.get_cmap(color_map()) + cm_subsection = np.linspace(start, stop, total_elements) + #color = [matplotlib.cm.gist_earth(x) for x in cm_subsection][count] + color = [colormap(x) for x in cm_subsection][count] + # Scale the color by intensity while leaving the 4th channel (alpha) unchanged + return [min(x * intensity, 1) for x in color[:3]] + [color[3]] + + +def color2d(count_x, count_y, total_elements_x, total_elements_y): + + # Select the actual color, then scale along the intensity axis + start = 1.7 + stop = 1 + intensity_scale = np.linspace(start, stop, total_elements_y) + intensity = intensity_scale[count_y] + return color(count_x, total_elements_x, intensity) + + +linewidth = 0.5 +minor_tick_color = (0.9, 0.9, 0.9) + + +def grid(minor=True, axis='both'): + plt.grid(linewidth=linewidth / 2, axis=axis) + if minor: + plt.grid(which='minor', + color=minor_tick_color, + linestyle='--', + alpha=0.7, + clip_on=True, + linewidth=linewidth / 4, + zorder=0) + + +def savefig(name, pad: float = 0): + for ext in ['pdf', 'png']: + outfile = save_folder / f"{name}.{ext}" + print("Saving", outfile) + plt.savefig(outfile, bbox_inches='tight', pad_inches=pad) + plt.clf() + + +def savetable(name, content: List[List[Any]]): + outfile = save_folder / f"{name}.txt" + + def fmt(e): + if type(e) == float or type(e) == np.float64 or type( + e) == np.float32 or type(e) == np.float16: + return f"{e:.3f}" + return str(e) + + print("Saving", outfile) + with open(outfile, 'w') as f: + + assert type(content) == list, "Table must be a list of rows" + for row in content: + assert type(row) == list, "Table rows must be lists" + f.write(" & ".join([fmt(e) for e in row]) + "\\\\\n") + + +def load_results(validation_folder: Path, + dataset: str = args.dataset, + full_distance: str = "ALL"): + print("Loading results from", validation_folder) + config_folder = validation_folder / "configs" + print() + assert config_folder.exists( + ), f"Config folder {config_folder} does not exist" + result_lst = [] + for architecture_folder in sorted(config_folder.glob("*/")): + if "bak" in architecture_folder.name: + continue + for result_file in (architecture_folder / dataset).glob("*.pkl"): + result_lst.append( + ResultInfo(architecture_folder.name + "_" + + ".".join(result_file.stem.split(".")[:-1]), + result_file, + full_distance=full_distance)) + + return sorted(result_lst, key=lambda x: x.pretty_name()) + + +print("Loading results...") +results = load_results(args.results_folder) +results_close = load_results(args.results_folder, full_distance="CLOSE") +results_far = load_results(args.results_folder, full_distance="FAR") +print("Done loading results.") +print(results) +assert len( + results) > 0, f"No results found in {args.results_folder.absolute()}" + + +def plot_scaling(x_label: bool = True, log_scale: bool = False): + + ours_results = [ + (1.0, 0.087), + (0.5, 0.101), + (0.2, 0.116), + (0.1, 0.139), + (0.01, 0.216), + ] + + fastflow_results = [ + (1.0, 0.076), + (0.5, 0.091), + (0.2, 0.104), + (0.1, 0.125), + (0.01, 0.202), + ] + + plt.plot(*zip(*ours_results), label="Ours", color='red', marker='o') + plt.plot(*zip(*fastflow_results), + label="FastFlow3D", + color='black', + marker='o') + if x_label: + plt.xlabel("Dataset Fraction") + plt.ylabel("Threeway EPE (m)") + + # Set x ticks + plt.xticks([0.01, 0.1, 0.2, 0.5, 1.0], ["1%", "10%", "20%", "50%", "100%"]) + # Horizontal line + plt.axhline(y=0.087, + color='black', + linestyle='--', + linewidth=linewidth / 2) + + plt.legend() + + # Draw arrow pointing from top right to bottom left + plt.annotate('', + xy=(0.01, 0.19), + xytext=(0.01, 0.10), + color='gray', + arrowprops=dict(arrowstyle="<-", color='gray')) + + + if log_scale: + plt.yscale('log') + plt.xscale('log') + # plt.yticks([ + # 0.08, 0.09, 0.1, 0.11, 0.12, 0.13, 0.14, 0.15, 0.16, 0.17, 0.18, + # 0.19, 0.20 + # ], [ + # "0.08", "", "0.10", "", "0.12", "", "0.14", "", "0.16", "", "0.18", + # "", "0.20" + # ]) + plt.yticks( + [0.08, 0.1, 0.12, 0.14, 0.16, 0.18, 0.20, 0.22], + ["0.08", "0.10", "0.12", "0.14", "0.16", "0.18", "0.20", "0.22"]) + plt.xticks([0.01, 0.1, 0.2, 0.5, 1.0], + ["1%", "10%", "20%", "50%", "100%"]) + # Draw text along arrow + plt.text(0.0096, + 0.14, + 'Better', + color='gray', + ha='center', + rotation=90, + rotation_mode='anchor') + else: + # Draw text along arrow + plt.text(0, + 0.15, + 'Better', + color='gray', + ha='center', + rotation=90, + rotation_mode='anchor') + + +def plot_speed_vs_performance_tradeoff(perf_error_bar: bool = True, + runtime_error_bar: bool = True, + gradient_bg: bool = True): + runtimes = { + 'Ours': 29.33, + 'FastFlow3D': 29.33, + 'NSFP': 26285.0, + 'Chodosh': 26285.0 + 8996.4, + 'Gojcic': 6087.87, + 'Sim2Real': 99.3477, + 'EgoFlow': 2116.34, + 'PPWC': 79.4275, + 'FlowStep3D': 687.536, + 'ICP': 523.11, + } + + runtimes_error_bar = { + 'Ours': 2.38, + 'FastFlow3D': 2.38, + 'NSFP': 18139.3, + 'Chodosh': 20247.7, + 'Gojcic': 1690.56, + 'Sim2Real': 13.88, + 'EgoFlow': 292.32, + 'PPWC': 2.20, + 'FlowStep3D': 3.13, + 'ICP': 169.34, + } + + perf_error_bar_dims = { + 'Ours': (0.00, 0.002), + 'FastFlow3D': (0.002, 0.003), + } + + points_processed = { + 'Ours': 52871.6, + 'FastFlow3D': 52871.6, + 'NSFP': 52871.6, + 'Chodosh': 52871.6, + 'Gojcic': 20000, + 'Sim2Real': 8192, + 'EgoFlow': 8192, + 'PPWC': 8192, + 'FlowStep3D': 8192, + 'ICP': 8192, + } + + performance = { + 'Ours': 0.087, + 'NSFP': 0.068, + 'Chodosh': 0.061, + 'FastFlow3D': 0.076, + 'PPWC': 0.130, + 'FlowStep3D': 0.161, + 'Sim2Real': 0.157, + 'EgoFlow': 0.205, + 'Gojcic': 0.083, + 'ICP': 0.204, + } + + uses_labels = { + 'Ours': False, + 'NSFP': False, + 'FastFlow3D': True, + 'PPWC': False, + 'FlowStep3D': False, + 'Sim2Real': False, + 'EgoFlow': False, + 'Gojcic': True, + 'Chodosh': False, + 'ICP': False, + } + + horiz_offset = 2.5 + vert_offset = 6.0 + + label_offset = { + 'ICP': (-horiz_offset, -vert_offset), + 'Sim2Real': (-horiz_offset, -vert_offset), + 'NSFP': (horiz_offset, vert_offset), + 'Gojcic': (-horiz_offset, vert_offset), + 'Chodosh': (horiz_offset, -vert_offset), + } + + keys = runtimes.keys() + runtimes = [runtimes[k] for k in keys] + performance = [performance[k] for k in keys] + shapes = ['x' if uses_labels[k] else 'o' for k in keys] + alphas = [1.0 if uses_labels[k] else 1.0 for k in keys] + gliph_colors = [ + 'red' if k == 'Ours' else ('black' if uses_labels[k] else 'black') + for k in keys + ] + text_colors = ['red' if k == 'Ours' else 'black' for k in keys] + points = [points_processed[k] for k in keys] + + worst_runtime = max(runtimes) + best_runtime = min(runtimes) + worst_perf = max(performance) + best_perf = min(performance) + print(worst_runtime, best_runtime, worst_perf, best_perf) + # breakpoint() + + plt.gca().axvspan(0, 100, alpha=0.08, facecolor='blue') + + def gradient_image(ax, direction=0.3, cmap_range=(0, 1), **kwargs): + """ + Draw a gradient image based on a colormap. + + Parameters + ---------- + ax : Axes + The axes to draw on. + direction : float + The direction of the gradient. This is a number in + range 0 (=vertical) to 1 (=horizontal). + cmap_range : float, float + The fraction (cmin, cmax) of the colormap that should be + used for the gradient, where the complete colormap is (0, 1). + **kwargs + Other parameters are passed on to `.Axes.imshow()`. + In particular, *cmap*, *extent*, and *transform* may be useful. + """ + phi = direction * np.pi / 2 + v = np.array([np.cos(phi), np.sin(phi)]) + X = np.array([[v @ [1, 0], v @ [1, 1]], [v @ [0, 0], v @ [0, 1]]]) + a, b = cmap_range + X = a + (b - a) / X.max() * X + im = ax.imshow(X, + interpolation='bicubic', + clim=(0, 1), + aspect='auto', + **kwargs) + return im + + if gradient_bg: + gradient_image(plt.gca(), + direction=0.65, + extent=(0.0, 1, 0, 1), + transform=plt.gca().transAxes, + cmap=plt.get_cmap('Greys'), + alpha=0.35) + + plt.ylim(bottom=best_perf - 0.014, top=worst_perf + 0.008) + plt.xlim(left=18, right=worst_runtime * 1.7) + + plt.text(43, + worst_perf, + 'Real Time at 10Hz', + color='blue', + ha='center', + rotation_mode='anchor') + + for key, runtime, perf, shape, alpha, point_count, gliph_color, text_color in zip( + keys, runtimes, performance, shapes, alphas, points, gliph_colors, + text_colors): + + fontweight = 'bold' if key == 'Ours' else 'normal' + dot_scale = 35 + alignment = 'left' if runtime < 20000 else 'right' + x_offset_sign = 1 if alignment == 'left' else -1 + plt.scatter(runtime, + perf, + marker=shape, + label=key, + color=gliph_color, + zorder=10, + s=dot_scale, + alpha=alpha, + edgecolors='none') + if perf_error_bar: + if key in perf_error_bar_dims: + plt.errorbar(runtime, + perf, + yerr=np.array([perf_error_bar_dims[key]]).T, + color=gliph_color, + capsize=1.5, + zorder=0, + alpha=0.2) + if runtime_error_bar: + plt.errorbar(runtime, + perf, + xerr=runtimes_error_bar[key], + color=gliph_color, + capsize=1.5, + zorder=0, + alpha=0.2) + # Annotate with name + plt.annotate(f"{key} ({point_count / 1000.0:0.1f}k points)", + (runtime, perf), + xytext=(6 * x_offset_sign + label_offset.get(key, + (0, 0))[0], + -2.75 + label_offset.get(key, (0, 0))[1]), + textcoords='offset points', + color=text_color, + ha=alignment, + fontweight=fontweight) + # Set x axis log scale + plt.xscale('log') + # Make vertical bar at 100 on x axis + # plt.axvline(100, color='blue', linestyle='--', linewidth=2, zorder=0) + plt.xlabel("Runtime (ms)") + plt.ylabel("Threeway EPE (m)") + + plt.gca().spines['right'].set_visible(False) + plt.gca().spines['top'].set_visible(False) + + # Draw arrow pointing from top right to bottom left + plt.annotate('', + xy=(worst_runtime, worst_perf), + xytext=(worst_runtime / 4, worst_perf - 0.04), + color='gray', + arrowprops=dict(arrowstyle="<-", color='gray')) + # Draw text along arrow + plt.text(worst_runtime / 2, + worst_perf - 0.017, + 'Better', + color='gray', + ha='center', + rotation=35, + rotation_mode='anchor') + + # Set plt y axis min to 0.05 + + # plt.annotate('Worse', xy=(worst_runtime, worst_perf), color='gray') + + # grid(minor=False, axis='x') + + +def process_metacategory_counts(result): + full_error_count = result['per_class_bucketed_error_count'] + metacatagory_results = {} + # How do we do on vehicles by speed? + for metacatagory in METACATAGORIES: + category_names = METACATAGORIES[metacatagory] + category_idxes = [CATEGORY_NAME_TO_IDX[cat] for cat in category_names] + + metacatagory_counts = full_error_count[category_idxes] + # Sum up other axes + metacatagory_counts = np.sum(metacatagory_counts, axis=(1, 2)) + metacatagory_results[metacatagory] = {} + for category_result, category_name in zip(metacatagory_counts, + category_names): + metacatagory_results[metacatagory][category_name] = category_result + + return metacatagory_results + + +def process_category_speed_counts(result): + category_speed_counts_raw = result['per_class_bucketed_error_count'].sum( + axis=2) + category_speed_counts_normalized = category_speed_counts_raw / category_speed_counts_raw.sum( + axis=1, keepdims=True) + return category_speed_counts_normalized, category_speed_counts_raw + + +def bar_offset(pos_idx, num_pos, position_offset=0.2): + """ + Compute the X offset for a bar plot given the number of bars and the + index of the bar. + """ + return -(num_pos + 1) / 2 * position_offset + position_offset * (pos_idx + + 1) + + +BAR_WIDTH = 0.07 + +num_metacatagories = len(METACATAGORIES) + + +def merge_dict_list(dict_list): + result_dict = {} + for d in dict_list: + for k, v in d.items(): + if k not in result_dict: + result_dict[k] = [] + result_dict[k].append(v) + return result_dict + + +def plot_mover_nonmover_vs_error_by_category(results: List[ResultInfo], + metacatagory, vmax): + for result_idx, result in enumerate(results): + metacatagory_epe_by_speed = result.get_metacatagory_epe_by_speed( + metacatagory) + xs = np.arange(len(metacatagory_epe_by_speed)) + bar_offset( + result_idx, len(results), BAR_WIDTH) + plt.bar(xs, + metacatagory_epe_by_speed, + label=result.pretty_name(), + width=BAR_WIDTH, + color=color(result_idx, len(results)), + zorder=3) + if metacatagory == "BACKGROUND": + plt.xticks([-1, 0, 1], ["", "All", ""]) + plt.xlabel(" ") + plt.legend() + else: + speed_buckets = result.speed_bucket_categories() + plt.xticks(np.arange(len(speed_buckets)), [ + f"{l:0.1f}-" + (f"{u:0.1f}" if u != np.inf else "$\infty$") + for l, u in speed_buckets + ], + rotation=0) + plt.xlabel("Speed Bucket (m/s)") + plt.ylabel("Average EPE (m)") + plt.ylim(0, vmax) + grid() + + +def table_speed_vs_error(results: List[ResultInfo]): + # for metacatagory in METACATAGORIES: + rows = [] + + for result in results: + row = [result.pretty_name()] + for metacatagory in METACATAGORIES: + + metacatagory_epe_by_speed = result.get_metacatagory_epe_by_speed( + metacatagory) + print(metacatagory, metacatagory_epe_by_speed) + row.extend(metacatagory_epe_by_speed.tolist()) + rows.append(row) + + return rows + + +def plot_nonmover_epe_overall(results: List[ResultInfo], vmax): + for result_idx, result in enumerate(results): + nonmover_epe = result.get_nonmover_point_epe() + xs = bar_offset(result_idx, len(results), BAR_WIDTH) + plt.bar(xs, + nonmover_epe, + label=result.pretty_name(), + width=BAR_WIDTH, + color=color(result_idx, len(results)), + zorder=3) + plt.ylabel("Average EPE (m)") + plt.xticks([], []) + plt.ylim(0, vmax) + plt.legend() + grid() + + +def plot_mover_epe_overall(results: List[ResultInfo], vmax): + for result_idx, result in enumerate(results): + mover_epe = result.get_mover_point_all_epe() + xs = bar_offset(result_idx, len(results), BAR_WIDTH) + plt.bar(xs, + mover_epe, + label=result.pretty_name(), + width=BAR_WIDTH, + color=color(result_idx, len(results)), + zorder=3) + plt.ylabel("Average EPE (m)") + plt.xticks([], []) + plt.ylim(0, vmax) + plt.legend() + grid() + + +def table_mover_nonmover_epe_overall(results: List[ResultInfo]): + rows = [] + for result in results: + row = [result.pretty_name()] + row.append(result.get_nonmover_point_epe()) + row.append(result.get_mover_point_all_epe()) + rows.append(row) + return rows + + +def plot_metacatagory_epe_counts(results: List[ResultInfo]): + + for meta_idx, metacatagory in enumerate(sorted(METACATAGORIES.keys())): + for result_idx, result in enumerate(results): + # Each error bucket is a single bar in the stacked bar plot. + metacatagory_epe_counts = result.get_metacatagory_count_by_epe( + metacatagory) + x_pos = meta_idx + bar_offset( + len(results) - result_idx - 1, len(results), BAR_WIDTH) + y_height = metacatagory_epe_counts.sum( + ) / metacatagory_epe_counts.sum() + bottom = 0 + for epe_idx, (epe_count, + (bucket_lower, bucket_upper)) in enumerate( + zip(metacatagory_epe_counts, + result.speed_bucket_categories())): + y_height = epe_count / metacatagory_epe_counts.sum() + epe_color = color2d(result_idx, epe_idx, len(results), + len(metacatagory_epe_counts)) + + label = None + if epe_idx == len( + metacatagory_epe_counts) - 1 and meta_idx == 0: + label = result.pretty_name() + rect = plt.barh([x_pos], [y_height], + label=label, + height=BAR_WIDTH, + color=epe_color, + left=bottom) + bottom += y_height + # Draw text in middle of bar + plt.text(bottom - y_height / 2, + x_pos, + f"{y_height * 100:0.1f}%", + ha="center", + va="center", + color="white", + fontsize=4) + + # xlabels to be the metacatagories + plt.yticks(np.arange(len(METACATAGORIES)), + [METACATAGORY_TO_SHORTNAME[e] for e in METACATAGORIES.keys()], + rotation=0) + plt.xticks(np.linspace(0, 1, 5), + [f"{e}%" for e in np.linspace(0, 100, 5).astype(int)]) + plt.xlabel("Percentage of Endpoints Within Error Threshold") + legend = plt.legend(loc="lower left", fancybox=False) + # set the boarder of the legend artist to be transparent + # legend.get_frame().set_edgecolor('none') + plt.tight_layout() + # plt.legend() + + +def plot_metacatagory_epe_counts_v15(results: List[ResultInfo]): + + for meta_idx, metacatagory in enumerate(sorted(METACATAGORIES.keys())): + for result_idx, result in enumerate(results): + # Each error bucket is a single bar in the stacked bar plot. + metacatagory_epe_counts = result.get_metacatagory_count_by_epe( + metacatagory) + x_pos = meta_idx + bar_offset( + len(results) - result_idx - 1, len(results), BAR_WIDTH) + y_height = metacatagory_epe_counts.sum( + ) / metacatagory_epe_counts.sum() + bottom = 0 + metacatagory_epe_counts_subset = metacatagory_epe_counts[:2] + for epe_idx, epe_count in enumerate( + metacatagory_epe_counts_subset): + y_height = epe_count / metacatagory_epe_counts.sum() + y_sum = metacatagory_epe_counts[:epe_idx + 1].sum( + ) / metacatagory_epe_counts.sum() + epe_color = color2d(result_idx, epe_idx, len(results), + len(metacatagory_epe_counts_subset)) + + label = None + if epe_idx == 0 and meta_idx == 0: + label = result.pretty_name() + rect = plt.barh([x_pos], [y_height], + label=label, + height=BAR_WIDTH, + color=epe_color, + left=bottom) + bottom += y_height + # Draw text in middle of bar + # plt.text(bottom - y_height / 2, + # x_pos, + # f"{y_sum * 100:0.1f}%", + # ha="center", + # va="center", + # color="white", + # fontsize=4) + + # xlabels to be the metacatagories + plt.yticks(np.arange(len(METACATAGORIES)), + [METACATAGORY_TO_SHORTNAME[e] for e in METACATAGORIES.keys()], + rotation=0) + plt.xticks(np.linspace(0, 1, 5), + [f"{e}%" for e in np.linspace(0, 100, 5).astype(int)]) + plt.xlabel("Endpoints Within Error Threshold") + legend = plt.legend(loc="lower left", fancybox=False) + # set the boarder of the legend artist to be transparent + # legend.get_frame().set_edgecolor('none') + plt.tight_layout() + # plt.legend() + + +def plot_metacatagory_epe_counts_v2(results: List[ResultInfo]): + + for meta_idx, metacatagory in enumerate(sorted(METACATAGORIES.keys())): + for result_idx, result in enumerate(results): + # Each error bucket is a single bar in the stacked bar plot. + metacatagory_epe_counts = result.get_metacatagory_count_by_epe( + metacatagory) + x_pos_center = meta_idx + bar_offset( + len(results) - result_idx - 1, len(results), BAR_WIDTH) + y_height = metacatagory_epe_counts.sum( + ) / metacatagory_epe_counts.sum() + for epe_idx, epe_count in enumerate(metacatagory_epe_counts[:2]): + x_offset = bar_offset(epe_idx, 2, BAR_WIDTH / 2) + y_height = metacatagory_epe_counts[:epe_idx + 1].sum( + ) / metacatagory_epe_counts.sum() + epe_color = color2d(result_idx, epe_idx, len(results), 2) + + label = None + if meta_idx == 0: + label = result.pretty_name() + if epe_idx == 0: + label += " (Strict)" + else: + label += " (Relaxed)" + + # if meta_idx == 0 and epe_idx == 0: + # label = result.pretty_name() + + plt.barh([x_pos_center + x_offset], [y_height], + label=label, + height=BAR_WIDTH / 2, + color=epe_color) + + # xlabels to be the metacatagories + plt.yticks(np.arange(len(METACATAGORIES)), + [METACATAGORY_TO_SHORTNAME[e] for e in METACATAGORIES.keys()], + rotation=0) + plt.xticks(np.linspace(0, 1, 5), + [f"{e}%" for e in np.linspace(0, 100, 5).astype(int)]) + plt.xlabel("Percentage of Endpoints Within Error Threshold") + legend = plt.legend(loc="lower left", fancybox=False) + # set the boarder of the legend artist to be transparent + # legend.get_frame().set_edgecolor('none') + plt.tight_layout() + # plt.legend() + + +def plot_metacatagory_epe_counts_v3(results: List[ResultInfo], + epe_bucket_idx: int): + + for meta_idx, metacatagory in enumerate(sorted(METACATAGORIES.keys())): + for result_idx, result in enumerate(results): + # Each error bucket is a single bar in the stacked bar plot. + metacatagory_epe_counts = result.get_metacatagory_count_by_epe( + metacatagory) + x_pos = meta_idx + bar_offset( + len(results) - result_idx - 1, len(results), BAR_WIDTH) + y_height = metacatagory_epe_counts.sum( + ) / metacatagory_epe_counts.sum() + + y_height = metacatagory_epe_counts[:epe_bucket_idx + 1].sum( + ) / metacatagory_epe_counts.sum() + epe_color = color(result_idx, len(results)) + + label = None + if meta_idx == 0: + label = result.pretty_name() + plt.barh([x_pos], [y_height], + label=label, + height=BAR_WIDTH, + color=epe_color) + + # xlabels to be the metacatagories + plt.yticks(np.arange(len(METACATAGORIES)), + [METACATAGORY_TO_SHORTNAME[e] for e in METACATAGORIES.keys()], + rotation=0) + plt.xticks(np.linspace(0, 1, 5), + [f"{e}%" for e in np.linspace(0, 100, 5).astype(int)]) + plt.xlabel("Endpoints Within Error Threshold") + legend = plt.legend(loc="lower left", fancybox=False) + # set the boarder of the legend artist to be transparent + # legend.get_frame().set_edgecolor('none') + plt.tight_layout() + # plt.legend() + + +def table_latency(results: List[ResultInfo]): + table = [] + for result in results: + table_row = [ + result.pretty_name(), + f"{result.get_latency():0.4f}", + ] + table.append(table_row) + return table + + +def table_epe(results: List[ResultInfo]): + table_rows = [] + for result in results: + foreground_counts = result.get_foreground_counts() + relaxed_percentage = foreground_counts[:2].sum( + ) / foreground_counts.sum() + strict_percentage = foreground_counts[0] / foreground_counts.sum() + + def fmt(val): + return f"${val:0.3f}$" + + epe_entries = [ + result.get_mover_point_dynamic_epe(), + result.get_mover_point_static_epe(), + result.get_nonmover_point_epe() + ] + average_epe = np.average(epe_entries) + + entries = [ + result.pretty_name(), + fmt(average_epe), *[fmt(e) for e in epe_entries], + fmt(relaxed_percentage), + fmt(strict_percentage) + ] + table_rows.append(entries) + return table_rows + + +def plot_validation_pointcloud_size(dataset: str, max_x=160000): + validation_data_counts_path = args.results_folder / f"{dataset}_validation_pointcloud_point_count.pkl" + assert validation_data_counts_path.exists( + ), f"Could not find {validation_data_counts_path}" + point_cloud_counts = load_pickle(validation_data_counts_path) + point_cloud_counts = np.array(point_cloud_counts) + point_cloud_counts = np.sort(point_cloud_counts) + + print(f"Lowest 20 point cloud counts: {point_cloud_counts[:20]}") + + mean = np.mean(point_cloud_counts) + std = np.std(point_cloud_counts) + print("VVVVVVVVVVVVVVVVVVVV") + print(f"Mean point cloud count {dataset}: {mean}, std: {std}") + print("^^^^^^^^^^^^^^^^^^^^") + if dataset == "argo": + point_cloud_counts = point_cloud_counts[point_cloud_counts < 100000] + elif dataset == "waymo": + point_cloud_counts = point_cloud_counts[point_cloud_counts < 200000] + else: + raise ValueError(f"Unknown dataset {dataset}") + + # Make histogram of point cloud counts + plt.hist(point_cloud_counts, bins=100, zorder=3, color=color(0, 1)) + plt.xlabel("Number of points") + plt.ylabel("Number of point clouds") + plt.xlim(left=0, right=max_x) + plt.gca().xaxis.set_major_formatter( + matplotlib.ticker.StrMethodFormatter('{x:,.0f}')) + plt.gca().yaxis.set_major_formatter( + matplotlib.ticker.StrMethodFormatter('{x:,.0f}')) + plt.tight_layout() + + +def plot_val_endpoint_error_distribution(source: str, + use_log_scale: bool = False, + unrotated: bool = False, + ticks: bool = False): + + base_path = Path( + f"/efs/argoverse2/val_{source}_unsupervised_vs_supervised_flow/") + glob_str = "*_error_distribution.npy" if not unrotated else "*_error_distribution_unrotated.npy" + error_distribution_files = sorted(base_path.glob(glob_str)) + npy_file_arr = np.array([ + load_npy(error_distribution_file, verbose=False) + for error_distribution_file in error_distribution_files + ]) + distribution = np.sum(npy_file_arr, axis=0) + + distribution_x = distribution.sum(axis=0) + distribution_y = distribution.sum(axis=1) + + average_x = np.average(np.arange(distribution_x.shape[0]), + weights=distribution_x) + average_y = np.average(np.arange(distribution_y.shape[0]), + weights=distribution_y) + + if (not unrotated) and (not use_log_scale): + print("VVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVVV") + print( + f"{source} Average x: {average_x}, average y: {average_y} is log scale: {use_log_scale} is rotated {not unrotated}" + ) + print("^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^") + + plot_mat = np.flip(distribution, axis=1).T + if use_log_scale: + plot_mat = np.log(plot_mat) + print("++++++++++++++++++++++++++++") + + print("Num finite cells:", np.isfinite(plot_mat).sum()) + print(plot_mat[np.isfinite(plot_mat)].max(), + plot_mat[np.isfinite(plot_mat)].min()) + plot_mat[~np.isfinite(plot_mat)] = -2 + print("++++++++++++++++++++++++++++") + plt.imshow(plot_mat, cmap=color_map()) + grid_radius_meters = 1.5 + cells_per_meter = 50 + plt.xticks( + np.linspace(0, grid_radius_meters * 2 * cells_per_meter, + int(grid_radius_meters * 2 + 1)), + [ + f"{e}m" if ticks else "" + for e in np.linspace(-grid_radius_meters, grid_radius_meters, + int(grid_radius_meters * 2 + 1)) + ]) + plt.yticks( + np.linspace(0, grid_radius_meters * 2 * cells_per_meter, + int(grid_radius_meters * 2 + 1)), + [ + f"{-e}m" if ticks else "" + for e in np.linspace(-grid_radius_meters, grid_radius_meters, + int(grid_radius_meters * 2 + 1)) + ]) + + +################################################################################ + +set_font(8) + +vmax = 0.4 + +for metacatagory in METACATAGORIES: + plt.gcf().set_size_inches(5.5 / 2, 5.5 / 1.6 / 2) + plot_mover_nonmover_vs_error_by_category(results, metacatagory, vmax=vmax) + print("saving", f"speed_vs_error_{metacatagory}") + savefig(f"speed_vs_error_{metacatagory}") + plt.clf() + +for metacatagory in METACATAGORIES: + plt.gcf().set_size_inches(5.5 / 2, 5.5 / 1.6 / 2) + plot_mover_nonmover_vs_error_by_category(results_close, + metacatagory, + vmax=vmax) + print("saving", f"speed_vs_error_{metacatagory}") + savefig(f"speed_vs_error_{metacatagory}_close") + plt.clf() + +for metacatagory in METACATAGORIES: + plt.gcf().set_size_inches(5.5 / 2, 5.5 / 1.6 / 2) + plot_mover_nonmover_vs_error_by_category(results_far, + metacatagory, + vmax=vmax) + print("saving", f"speed_vs_error_{metacatagory}") + savefig(f"speed_vs_error_{metacatagory}_far") + plt.clf() + +################################################################################ + +savetable("epe_table", table_epe(results_close)) + +################################################################################ + +plt.gcf().set_size_inches(6.5, 6.5 / 2.2) +plot_speed_vs_performance_tradeoff(gradient_bg=False) +savefig(f"speed_vs_performance_tradeoff") + +plt.gcf().set_size_inches(6.5, 6.5 / 2.2) +plot_speed_vs_performance_tradeoff(gradient_bg=True) +savefig(f"speed_vs_performance_tradeoff_gradient") + +################################################################################ + +plt.gcf().set_size_inches(6.5, 6.5 / 2.2) +plot_scaling() +savefig(f"scaling") + +plt.gcf().set_size_inches(6.5, 6.5 / 2.2) +plot_scaling(log_scale=True) +savefig(f"scaling_log") + +################################################################################ + +savetable("speed_vs_error", table_speed_vs_error(results)) + +################################################################################ + +error_dist_scalar = 0.6 + +# # nsfp + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('nsfp', use_log_scale=True) +# savefig(f"val_endpoint_error_distribution_log_nsfp") + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('nsfp', use_log_scale=False) +# savefig(f"val_endpoint_error_distribution_absolute_nsfp") + +# # Odom + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('odom', use_log_scale=True) +# savefig(f"val_endpoint_error_distribution_log_odom") + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('odom', use_log_scale=False) +# savefig(f"val_endpoint_error_distribution_absolute_odom") + +# # Nearest neighbor + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('nearest_neighbor', use_log_scale=True) +# savefig(f"val_endpoint_error_distribution_log_nearest_neighbor") + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('nearest_neighbor', use_log_scale=False) +# savefig(f"val_endpoint_error_distribution_absolute_nearest_neighbor") + +# # Chodosh + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('chodosh', use_log_scale=True) +# savefig(f"val_endpoint_error_distribution_log_chodosh") + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('chodosh', use_log_scale=False) +# savefig(f"val_endpoint_error_distribution_absolute_chodosh") + +# # Supervised + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('supervised', use_log_scale=True) +# savefig(f"val_endpoint_error_distribution_log_supervised") + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('supervised', use_log_scale=False) +# savefig(f"val_endpoint_error_distribution_absolute_supervised") + +# # Distilation + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('distilation', use_log_scale=True) +# savefig(f"val_endpoint_error_distribution_log_distilation") + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('distilation', use_log_scale=False) +# savefig(f"val_endpoint_error_distribution_absolute_distilation") + +# # Chodosh distilation + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('chodosh_distilation', use_log_scale=True) +# savefig(f"val_endpoint_error_distribution_log_chodosh_distilation") + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('chodosh_distilation', +# use_log_scale=False) +# savefig(f"val_endpoint_error_distribution_absolute_chodosh_distilation") + +# ##### Unrotated + +# # nsfp + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('nsfp', +# use_log_scale=True, +# unrotated=True) +# savefig(f"val_endpoint_error_distribution_log_nsfp_unrotated") + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('nsfp', +# use_log_scale=False, +# unrotated=True) +# savefig(f"val_endpoint_error_distribution_absolute_nsfp_unrotated") + +# # Odom + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('odom', +# use_log_scale=True, +# unrotated=True) +# savefig(f"val_endpoint_error_distribution_log_odom_unrotated") + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('odom', +# use_log_scale=False, +# unrotated=True) +# savefig(f"val_endpoint_error_distribution_absolute_odom_unrotated") + +# # Nearest neighbor + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('nearest_neighbor', +# use_log_scale=True, +# unrotated=True) +# savefig(f"val_endpoint_error_distribution_log_nearest_neighbor_unrotated") + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('nearest_neighbor', +# use_log_scale=False, +# unrotated=True) +# savefig(f"val_endpoint_error_distribution_absolute_nearest_neighbor_unrotated") + +# # Chodosh + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('chodosh', +# use_log_scale=True, +# unrotated=True) +# savefig(f"val_endpoint_error_distribution_log_chodosh_unrotated") + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('chodosh', +# use_log_scale=False, +# unrotated=True) +# savefig(f"val_endpoint_error_distribution_absolute_chodosh_unrotated") + +# # Supervised + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('supervised', +# use_log_scale=True, +# unrotated=True) +# savefig(f"val_endpoint_error_distribution_log_supervised_unrotated") + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('supervised', +# use_log_scale=False, +# unrotated=True) +# savefig(f"val_endpoint_error_distribution_absolute_supervised_unrotated") + +# # Distilation + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('distilation', +# use_log_scale=True, +# unrotated=True) +# savefig(f"val_endpoint_error_distribution_log_distilation_unrotated") + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('distilation', +# use_log_scale=False, +# unrotated=True) +# savefig(f"val_endpoint_error_distribution_absolute_distilation_unrotated") + +# # Chodosh distilation + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('chodosh_distilation', +# use_log_scale=True, +# unrotated=True) +# savefig(f"val_endpoint_error_distribution_log_chodosh_distilation_unrotated") + +# plt.gcf().set_size_inches(2.5 * error_dist_scalar, 2.5 * error_dist_scalar) +# plot_val_endpoint_error_distribution('chodosh_distilation', +# use_log_scale=False, +# unrotated=True) +# savefig( +# f"val_endpoint_error_distribution_absolute_chodosh_distilation_unrotated") + +################################################################################ + +plt.gcf().set_size_inches(5.5 / 2, 5.5 / 1.6) +plot_nonmover_epe_overall(results, 0.3) +savefig(f"nonmover_epe_overall") + +plt.gcf().set_size_inches(5.5 / 2, 5.5 / 1.6) +plot_mover_epe_overall(results, 0.3) +savefig(f"mover_epe_overall") + +################################################################################ + +savetable("mover_nonmover_epe_overall", + table_mover_nonmover_epe_overall(results)) + +################################################################################ + +plt.gcf().set_size_inches(5.5 / 2, 5.5 / 1.6) +plot_nonmover_epe_overall(results_close, 0.3) +savefig(f"nonmover_epe_overall_close") + +plt.gcf().set_size_inches(5.5 / 2, 5.5 / 1.6) +plot_mover_epe_overall(results_close, 0.3) +savefig(f"mover_epe_overall_close") + +plt.gcf().set_size_inches(5.5 / 2, 5.5 / 1.6) +plot_nonmover_epe_overall(results_far, 0.4) +savefig(f"nonmover_epe_overall_far") + +plt.gcf().set_size_inches(5.5 / 2, 5.5 / 1.6) +plot_mover_epe_overall(results_far, 0.4) +savefig(f"mover_epe_overall_far") + +################################################################################ + +plt.gcf().set_size_inches(5.5, 2.5) +plot_metacatagory_epe_counts(results) +savefig(f"epe_counts") + +plt.gcf().set_size_inches(5.5, 2.5) +plot_metacatagory_epe_counts(results_close) +savefig(f"epe_counts_close") + +plt.gcf().set_size_inches(5.5, 2.5) +plot_metacatagory_epe_counts(results_far) +savefig(f"epe_counts_close_far") + +################################################################################ + +plt.gcf().set_size_inches(5.5, 2.5) +plot_metacatagory_epe_counts_v15(results) +savefig(f"epe_counts_v15") + +plt.gcf().set_size_inches(5.5, 2.5) +plot_metacatagory_epe_counts_v15(results_close) +savefig(f"epe_counts_v15_close") + +plt.gcf().set_size_inches(5.5, 2.5) +plot_metacatagory_epe_counts_v15(results_far) +savefig(f"epe_counts_v15_far") + +################################################################################ + +plt.gcf().set_size_inches(5.5, 2.5) +plot_metacatagory_epe_counts_v2(results) +savefig(f"epe_counts_v2") + +################################################################################ + +plt.gcf().set_size_inches(5.5 / 2, 2.5) +plot_metacatagory_epe_counts_v3(results, 0) +savefig(f"epe_counts_v3_strict") + +plt.gcf().set_size_inches(5.5 / 2, 2.5) +plot_metacatagory_epe_counts_v3(results, 1) +savefig(f"epe_counts_v3_loose") + +################################################################################ + +plt.gcf().set_size_inches(5.5 / 2, 5.5 / 1.6 / 2) +plot_validation_pointcloud_size("argo", max_x=90000) +grid() +savefig(f"validation_pointcloud_size_argo", pad=0.02) + +plt.gcf().set_size_inches(5.5 / 2, 5.5 / 1.6 / 2) +plot_validation_pointcloud_size("waymo") +grid() +savefig(f"validation_pointcloud_size_waymo", pad=0.02) + +################################################################################ + +savetable("latency_table", table_latency(results)) diff --git a/pointclouds/__init__.py b/pointclouds/__init__.py new file mode 100644 index 0000000..4372dcb --- /dev/null +++ b/pointclouds/__init__.py @@ -0,0 +1,5 @@ +from .pointcloud import PointCloud, to_fixed_array, from_fixed_array +from .se3 import SE3 +from .se2 import SE2 + +__all__ = ['PointCloud', 'SE3', 'SE2'] diff --git a/pointclouds/losses/__init__.py b/pointclouds/losses/__init__.py new file mode 100644 index 0000000..496b303 --- /dev/null +++ b/pointclouds/losses/__init__.py @@ -0,0 +1,3 @@ +from .losses import warped_pc_loss, pc0_to_pc1_distance + +__all__ = ["warped_pc_loss", "pc0_to_pc1_distance"] diff --git a/pointclouds/losses/losses.py b/pointclouds/losses/losses.py new file mode 100644 index 0000000..5d57d03 --- /dev/null +++ b/pointclouds/losses/losses.py @@ -0,0 +1,63 @@ +import torch +from pytorch3d.ops.knn import knn_points +from pytorch3d.loss import chamfer_distance +from pointclouds import PointCloud +from typing import Union + + +def warped_pc_loss(warped_pc: torch.Tensor, + target_pc: torch.Tensor, + dist_threshold=2.0): + if warped_pc.ndim == 2: + warped_pc = warped_pc.unsqueeze(0) + if target_pc.ndim == 2: + target_pc = target_pc.unsqueeze(0) + + assert warped_pc.ndim == 3, f"warped_pc.ndim = {warped_pc.ndim}, not 3; shape = {warped_pc.shape}" + assert target_pc.ndim == 3, f"target_pc.ndim = {target_pc.ndim}, not 3; shape = {target_pc.shape}" + + loss = 0 + + if dist_threshold is None: + loss += chamfer_distance(warped_pc, target_pc, + point_reduction="mean")[0].sum() + loss += chamfer_distance(target_pc, warped_pc, + point_reduction="mean")[0].sum() + return loss + + # Compute min distance between warped point cloud and point cloud at t+1. + warped_to_target_knn = knn_points(p1=warped_pc, p2=target_pc, K=1) + warped_to_target_distances = warped_to_target_knn.dists[0] + target_to_warped_knn = knn_points(p1=target_pc, p2=warped_pc, K=1) + target_to_warped_distances = target_to_warped_knn.dists[0] + # Throw out distances that are too large (beyond the dist threshold). + loss += warped_to_target_distances[ + warped_to_target_distances < dist_threshold].mean() + + loss += target_to_warped_distances[ + target_to_warped_distances < dist_threshold].mean() + + return loss + + +def pc0_to_pc1_distance(pc0: Union[PointCloud, torch.Tensor], + pc1: Union[PointCloud, torch.Tensor]): + if isinstance(pc0, PointCloud): + pc0 = pc0.points + if isinstance(pc1, PointCloud): + pc1 = pc1.points + + if pc0.ndim == 2: + pc0 = pc0.unsqueeze(0) + if pc1.ndim == 2: + pc1 = pc1.unsqueeze(0) + + pc0_shape_tensor = torch.LongTensor([pc0.shape[0]]).to(pc0.device) + pc1_shape_tensor = torch.LongTensor([pc1.shape[0]]).to(pc1.device) + pc0_to_pc1_knn = knn_points(p1=pc0, + p2=pc1, + lengths1=pc0_shape_tensor, + lengths2=pc1_shape_tensor, + K=1) + pc0_to_pc1_distances = pc0_to_pc1_knn.dists[0] + return pc0_to_pc1_distances \ No newline at end of file diff --git a/pointclouds/pointcloud.py b/pointclouds/pointcloud.py new file mode 100644 index 0000000..c808271 --- /dev/null +++ b/pointclouds/pointcloud.py @@ -0,0 +1,123 @@ +import numpy as np + +from .se3 import SE3 + + +def to_fixed_array(array: np.ndarray, + max_len: int, + pad_val=np.nan) -> np.ndarray: + if len(array) > max_len: + np.random.RandomState(len(array)).shuffle(array) + sliced_pts = array[:max_len] + return sliced_pts + else: + pad_tuples = [(0, max_len - len(array))] + for _ in range(array.ndim - 1): + pad_tuples.append((0, 0)) + return np.pad(array, pad_tuples, constant_values=pad_val) + + +def from_fixed_array(array: np.ndarray) -> np.ndarray: + if isinstance(array, np.ndarray): + if len(array.shape) == 2: + check_array = array[:, 0] + elif len(array.shape) == 1: + check_array = array + else: + raise ValueError(f'unknown array shape {array.shape}') + are_valid_points = np.logical_not(np.isnan(check_array)) + are_valid_points = are_valid_points.astype(bool) + else: + import torch + if len(array.shape) == 2: + check_array = array[:, 0] + elif len(array.shape) == 1: + check_array = array + else: + raise ValueError(f'unknown array shape {array.shape}') + are_valid_points = torch.logical_not(torch.isnan(check_array)) + are_valid_points = are_valid_points.bool() + return array[are_valid_points] + + +class PointCloud(): + + def __init__(self, points: np.ndarray) -> None: + assert points.ndim == 2, f'points must be a 2D array, got {points.ndim}' + assert points.shape[1] == 3, f'points must be a Nx3 array, got {points.shape}' + self.points = points + + def __eq__(self, o: object) -> bool: + if not isinstance(o, PointCloud): + return False + return np.allclose(self.points, o.points) + + def __len__(self): + return self.points.shape[0] + + def __repr__(self) -> str: + return f'PointCloud with {len(self)} points' + + def __getitem__(self, idx): + return self.points[idx] + + def transform(self, se3: SE3) -> 'PointCloud': + assert isinstance(se3, SE3) + return PointCloud(se3.transform_points(self.points)) + + def translate(self, translation: np.ndarray) -> 'PointCloud': + assert translation.shape == (3, ) + return PointCloud(self.points + translation) + + def flow(self, flow: np.ndarray) -> 'PointCloud': + assert flow.shape == self.points.shape, f"flow shape {flow.shape} must match point cloud shape {self.points.shape}" + return PointCloud(self.points + flow) + + def to_fixed_array(self, max_points: int) -> np.ndarray: + return to_fixed_array(self.points, max_points) + + def matched_point_diffs(self, other: 'PointCloud') -> np.ndarray: + assert len(self) == len(other) + return self.points - other.points + + def matched_point_distance(self, other: 'PointCloud') -> np.ndarray: + assert len(self) == len(other) + return np.linalg.norm(self.points - other.points, axis=1) + + @staticmethod + def from_fixed_array(points) -> 'PointCloud': + return PointCloud(from_fixed_array(points)) + + def to_array(self) -> np.ndarray: + return self.points + + @staticmethod + def from_array(points: np.ndarray) -> 'PointCloud': + return PointCloud(points) + + def mask_points(self, mask: np.ndarray) -> 'PointCloud': + assert isinstance(mask, np.ndarray) + assert mask.ndim == 1 + if mask.dtype == np.bool: + assert mask.shape[0] == len(self) + else: + in_bounds = mask < len(self) + assert np.all( + in_bounds + ), f"mask values must be in bounds, got {(~in_bounds).sum()} indices not in bounds out of {len(self)} points" + + return PointCloud(self.points[mask]) + + def within_region(self, x_min, x_max, y_min, y_max, z_min, + z_max) -> 'PointCloud': + mask = np.logical_and(self.points[:, 0] < x_max, + self.points[:, 0] > x_min) + mask = np.logical_and(mask, self.points[:, 1] < y_max) + mask = np.logical_and(mask, self.points[:, 1] > y_min) + mask = np.logical_and(mask, self.points[:, 2] < z_max) + mask = np.logical_and(mask, self.points[:, 2] > z_min) + return self.mask_points(mask) + + @property + def shape(self) -> tuple: + return self.points.shape diff --git a/pointclouds/se2.py b/pointclouds/se2.py new file mode 100644 index 0000000..e22dee7 --- /dev/null +++ b/pointclouds/se2.py @@ -0,0 +1,70 @@ +import numpy as np + + +class SE2: + + def __init__(self, rotation: np.ndarray, translation: np.ndarray) -> None: + """Initialize. + Args: + rotation: np.ndarray of shape (2,2). + translation: np.ndarray of shape (2,1). + Raises: + ValueError: if rotation or translation do not have the required shapes. + """ + assert rotation.shape == (2, 2) + assert translation.shape == (2, ) + self.rotation = rotation + self.translation = translation + self.transform_matrix = np.eye(3) + self.transform_matrix[:2, :2] = self.rotation + self.transform_matrix[:2, 2] = self.translation + + def transform_point_cloud(self, point_cloud: np.ndarray) -> np.ndarray: + """Apply the SE(2) transformation to point_cloud. + Args: + point_cloud: np.ndarray of shape (N, 2). + Returns: + transformed_point_cloud: np.ndarray of shape (N, 2). + Raises: + ValueError: if point_cloud does not have the required shape. + """ + assert point_cloud.ndim == 2 + assert point_cloud.shape[1] == 2 + num_points = point_cloud.shape[0] + homogeneous_pts = np.hstack([point_cloud, np.ones((num_points, 1))]) + transformed_point_cloud = homogeneous_pts.dot(self.transform_matrix.T) + return transformed_point_cloud[:, :2] + + def inverse(self) -> "SE2": + """Return the inverse of the current SE2 transformation. + For example, if the current object represents target_SE2_src, we will return instead src_SE2_target. + Returns: + inverse of this SE2 transformation. + """ + return SE2(rotation=self.rotation.T, + translation=self.rotation.T.dot(-self.translation)) + + def inverse_transform_point_cloud(self, + point_cloud: np.ndarray) -> np.ndarray: + """Transform the point_cloud by the inverse of this SE2. + Args: + point_cloud: Numpy array of shape (N,2). + Returns: + point_cloud transformed by the inverse of this SE2. + """ + return self.inverse().transform_point_cloud(point_cloud) + + def compose(self, right_se2: "SE2") -> "SE2": + """Multiply this SE2 from right by right_se2 and return the composed transformation. + Args: + right_se2: SE2 object to multiply this object by from right. + Returns: + The composed transformation. + """ + chained_transform_matrix = self.transform_matrix.dot( + right_se2.transform_matrix) + chained_se2 = SE2( + rotation=chained_transform_matrix[:2, :2], + translation=chained_transform_matrix[:2, 2], + ) + return chained_se2 \ No newline at end of file diff --git a/pointclouds/se3.py b/pointclouds/se3.py new file mode 100644 index 0000000..a5bbe5a --- /dev/null +++ b/pointclouds/se3.py @@ -0,0 +1,86 @@ +import numpy as np +from pyquaternion import Quaternion + + +class SE3: + """An SE3 class allows point cloud rotation and translation operations.""" + + def __init__(self, rotation_matrix: np.ndarray, + translation: np.ndarray) -> None: + """Initialize an SE3 instance with its rotation and translation matrices. + Args: + rotation: Array of shape (3, 3) + translation: Array of shape (3,) + """ + assert rotation_matrix.shape == (3, 3) + assert translation.shape == (3, ) + self.rotation_matrix = rotation_matrix + self.translation = translation + + if isinstance(self.rotation_matrix, np.ndarray): + self.transform_matrix = np.eye(4) + else: + import torch + self.transform_matrix = torch.eye(4) + self.transform_matrix[:3, :3] = self.rotation_matrix + self.transform_matrix[:3, 3] = self.translation + + @staticmethod + def from_rot_w_x_y_z_translation_x_y_z(rw, rx, ry, rz, tx, ty, tz) -> None: + rotation_matrix = Quaternion(w=rw, x=rx, y=ry, z=rz).rotation_matrix + translation = np.array([tx, ty, tz]) + return SE3(rotation_matrix, translation) + + def transform_points(self, point_cloud: np.ndarray) -> np.ndarray: + """Apply the SE(3) transformation to this point cloud. + Args: + point_cloud: Array of shape (N, 3). If the transform represents dst_SE3_src, + then point_cloud should consist of points in frame `src` + Returns: + Array of shape (N, 3) representing the transformed point cloud, i.e. points in frame `dst` + """ + return point_cloud @ self.rotation_matrix.T + self.translation + + def inverse_transform_points(self, point_cloud: np.ndarray) -> np.ndarray: + """Undo the translation and then the rotation (Inverse SE(3) transformation).""" + return (point_cloud.copy() - self.translation) @ self.rotation_matrix + + def inverse(self) -> "SE3": + """Return the inverse of the current SE3 transformation. + For example, if the current object represents target_SE3_src, we will return instead src_SE3_target. + Returns: + src_SE3_target: instance of SE3 class, representing + inverse of SE3 transformation target_SE3_src + """ + return SE3(rotation_matrix=self.rotation_matrix.T, + translation=self.rotation_matrix.T.dot(-self.translation)) + + def compose(self, right_se3: "SE3") -> "SE3": + """Compose (right multiply) this class' transformation matrix T with another SE3 instance. + Algebraic representation: chained_se3 = T * right_se3 + Args: + right_se3: another instance of SE3 class + Returns: + chained_se3: new instance of SE3 class + """ + chained_transform_matrix = self.transform_matrix @ right_se3.transform_matrix + chained_se3 = SE3( + rotation_matrix=chained_transform_matrix[:3, :3], + translation=chained_transform_matrix[:3, 3], + ) + return chained_se3 + + def to_array(self) -> np.ndarray: + """Return the SE3 transformation matrix as a numpy array.""" + return self.transform_matrix + + @staticmethod + def from_array(transform_matrix: np.ndarray) -> "SE3": + """Initialize an SE3 instance from a numpy array.""" + return SE3( + rotation_matrix=transform_matrix[:3, :3], + translation=transform_matrix[:3, 3], + ) + + def __repr__(self) -> str: + return "SE3 transform" \ No newline at end of file diff --git a/test_pl.py b/test_pl.py new file mode 100644 index 0000000..67552ba --- /dev/null +++ b/test_pl.py @@ -0,0 +1,79 @@ +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.optim as optim +from pathlib import Path +import argparse +import pytorch_lightning as pl +from pytorch_lightning.strategies import DDPStrategy +import torchmetrics +import warnings + +from tqdm import tqdm +import dataloaders +from dataloaders import ArgoverseRawSequenceLoader, ArgoverseRawSequence, SubsequenceRawDataset, OriginMode + +from pointclouds import PointCloud, SE3 + +import models +from model_wrapper import ModelWrapper + +from pathlib import Path + +from mmcv import Config + +# Get config file from command line +parser = argparse.ArgumentParser() +parser.add_argument('config', type=Path) +parser.add_argument('--checkpoint', type=Path, default=None) +parser.add_argument('--gpus', type=int, default=1) +args = parser.parse_args() + +assert args.config.exists(), f"Config file {args.config} does not exist" +cfg = Config.fromfile(args.config) + + +def make_test_dataloader(cfg): + # Setup val infra + test_sequence_loader = getattr( + dataloaders, cfg.test_loader.name)(**cfg.test_loader.args) + test_dataset = getattr(dataloaders, cfg.test_dataset.name)( + sequence_loader=test_sequence_loader, **cfg.test_dataset.args) + test_dataloader = torch.utils.data.DataLoader(test_dataset, + **cfg.test_dataloader.args) + + return test_dataloader + + +def setup_model(cfg): + if hasattr(cfg, "is_trainable") and not cfg.is_trainable: + model = ModelWrapper(cfg) + else: + assert args.checkpoint is not None, "Must provide checkpoint for validation" + assert args.checkpoint.exists( + ), f"Checkpoint file {args.checkpoint} does not exist" + model = ModelWrapper.load_from_checkpoint(args.checkpoint, cfg=cfg) + return model + + +pl.seed_everything(42069) + +test_dataloader = make_test_dataloader(cfg) + +print("Val dataloader length:", len(test_dataloader)) + +model = setup_model(cfg) +trainer = pl.Trainer(devices=args.gpus, + accelerator="gpu", + strategy=DDPStrategy(find_unused_parameters=False), + move_metrics_to_cpu=False, + num_sanity_val_steps=2, + log_every_n_steps=2, + val_check_interval=cfg.validate_every, + max_epochs=cfg.epochs, + resume_from_checkpoint=args.checkpoint, + accumulate_grad_batches=cfg.accumulate_grad_batches + if hasattr(cfg, "accumulate_grad_batches") else 1, + gradient_clip_val=cfg.gradient_clip_val if hasattr( + cfg, "gradient_clip_val") else 0.0) +trainer.validate(model, dataloaders=test_dataloader) diff --git a/train_pl.py b/train_pl.py new file mode 100644 index 0000000..2e70644 --- /dev/null +++ b/train_pl.py @@ -0,0 +1,189 @@ +import os +# PL_FAULT_TOLERANT_TRAINING=1 +# to enable fault tolerant training +#os.environ['PL_FAULT_TOLERANT_TRAINING'] = '1' + +import datetime +import torch +from pathlib import Path +import argparse +import pytorch_lightning as pl +from pytorch_lightning.strategies import DDPStrategy +import torchmetrics +from pytorch_lightning.loggers import TensorBoardLogger +from pytorch_lightning.callbacks import ModelCheckpoint + +from tqdm import tqdm +import dataloaders +from dataloaders import ConcatDataset + +from pointclouds import PointCloud, SE3 + +import models +from model_wrapper import ModelWrapper + +from pathlib import Path + +from mmcv import Config + + +def get_rank() -> int: + # SLURM_PROCID can be set even if SLURM is not managing the multiprocessing, + # therefore LOCAL_RANK needs to be checked first + rank_keys = ("RANK", "LOCAL_RANK", "SLURM_PROCID", "JSM_NAMESPACE_RANK") + for key in rank_keys: + rank = os.environ.get(key) + if rank is not None: + return int(rank) + return 0 + + +def get_checkpoint_path(cfg, checkpoint_dir_name: str): + cfg_filename = Path(cfg.filename) + config_name = cfg_filename.stem + parent_name = cfg_filename.parent.name + parent_path = Path(f"model_checkpoints/{parent_name}/{config_name}/") + rank = get_rank() + if rank == 0: + # Since we're rank 0, we can create the directory + return parent_path / checkpoint_dir_name, checkpoint_dir_name + else: + # Since we're not rank 0, we shoulds grab the most recent directory + checkpoint_path = sorted(parent_path.glob("*"))[-1] + return checkpoint_path, checkpoint_path.name + + +def make_train_dataloader(cfg): + # Handle single loader case + if not isinstance(cfg.loader, list) and not isinstance(cfg.dataset, list): + try: + train_sequence_loader = getattr(dataloaders, + cfg.loader.name)(**cfg.loader.args) + except Exception as e: + print("Error loading loader:", cfg.loader.name) + print("Config:", cfg.loader) + raise e + train_dataset = getattr(dataloaders, cfg.dataset.name)( + sequence_loader=train_sequence_loader, **cfg.dataset.args) + return torch.utils.data.DataLoader(train_dataset, + **cfg.dataloader.args) + + # Handle multiple loader case + assert isinstance(cfg.loader, list) and isinstance(cfg.dataset, list), \ + f"Either both loader and dataset should be lists, or neither should be. Got loader: {type(cfg.loader)} and dataset: {type(cfg.dataset)}" + + assert len(cfg.loader) == len(cfg.dataset), \ + f"Length of loader list {len(cfg.loader)} should be equal to length of dataset list {len(cfg.dataset)}" + + print("Using multiple loaders of length:", len(cfg.loader)) + train_sequence_loader_lst = [] + for loader in cfg.loader: + train_sequence_loader_lst.append( + getattr(dataloaders, loader.name)(**loader.args)) + + print("Using multiple datasets of length:", len(cfg.dataset)) + train_dataloader_lst = [] + for dataset, train_sequence_loader in zip(cfg.dataset, + train_sequence_loader_lst): + train_dataset = getattr(dataloaders, dataset.name)( + sequence_loader=train_sequence_loader, **dataset.args) + train_dataloader_lst.append(train_dataset) + + # Use the concat dataloader to combine the multiple dataloaders + concat_dataset = dataloaders.ConcatDataset(train_dataloader_lst) + + return torch.utils.data.DataLoader(concat_dataset, **cfg.dataloader.args) + + +def make_val_dataloader(cfg): + # Setup val infra + val_sequence_loader = getattr(dataloaders, + cfg.test_loader.name)(**cfg.test_loader.args) + val_dataset = getattr(dataloaders, cfg.test_dataset.name)( + sequence_loader=val_sequence_loader, **cfg.test_dataset.args) + val_dataloader = torch.utils.data.DataLoader(val_dataset, + **cfg.test_dataloader.args) + + return val_dataloader + + +def main(): + + # Get config file from command line + parser = argparse.ArgumentParser() + parser.add_argument('config', type=Path) + parser.add_argument('--gpus', type=int, default=torch.cuda.device_count()) + parser.add_argument('--resume_from_checkpoint', type=Path, default=None) + parser.add_argument( + '--checkpoint_dir_name', + type=str, + default=datetime.datetime.now().strftime("%Y_%m_%d-%I_%M_%S_%p")) + parser.add_argument('--dry_run', action='store_true') + args = parser.parse_args() + + assert args.config.exists(), f"Config file {args.config} does not exist" + cfg = Config.fromfile(args.config) + + if hasattr(cfg, "is_trainable") and not cfg.is_trainable: + raise ValueError("Config file indicates this model is not trainable.") + + if hasattr(cfg, "seed_everything"): + pl.seed_everything(cfg.seed_everything) + + resume_from_checkpoint = args.resume_from_checkpoint + + checkpoint_path, checkpoint_dir_name = get_checkpoint_path( + cfg, args.checkpoint_dir_name) + if not args.dry_run: + checkpoint_path.mkdir(parents=True, exist_ok=True) + # Save config file to checkpoint directory + cfg.dump(str(checkpoint_path / "config.py")) + + tbl = TensorBoardLogger("tb_logs", + name=cfg.filename, + version=checkpoint_dir_name) + + train_dataloader = make_train_dataloader(cfg) + val_dataloader = make_val_dataloader(cfg) + + print("Train dataloader length:", len(train_dataloader)) + print("Val dataloader length:", len(val_dataloader)) + + model = ModelWrapper(cfg) + + checkpoint_callback = ModelCheckpoint( + dirpath=checkpoint_path, + filename="checkpoint_{epoch:03d}_{step:010d}", + save_top_k=-1, + # every_n_train_steps=cfg.save_every, + every_n_epochs=1, + save_on_train_epoch_end=True) + + trainer = pl.Trainer(devices=args.gpus, + accelerator="gpu", + logger=tbl, + strategy=DDPStrategy(find_unused_parameters=False), + move_metrics_to_cpu=False, + num_sanity_val_steps=2, + log_every_n_steps=2, + val_check_interval=cfg.validate_every, + check_val_every_n_epoch=cfg.check_val_every_n_epoch if hasattr( + cfg, "check_val_every_n_epoch") else 1, + max_epochs=cfg.epochs, + resume_from_checkpoint=resume_from_checkpoint, + accumulate_grad_batches=cfg.accumulate_grad_batches + if hasattr(cfg, "accumulate_grad_batches") else 1, + gradient_clip_val=cfg.gradient_clip_val if hasattr( + cfg, "gradient_clip_val") else 0.0, + callbacks=[checkpoint_callback]) + if args.dry_run: + print("Dry run, exiting") + exit(0) + print("Starting training") + print("Length of train dataloader:", len(train_dataloader)) + print("Length of val dataloader:", len(val_dataloader)) + trainer.fit(model, train_dataloader, val_dataloader) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/util_scripts/jbackup b/util_scripts/jbackup new file mode 100755 index 0000000..1567557 --- /dev/null +++ b/util_scripts/jbackup @@ -0,0 +1,49 @@ +#!/usr/bin/env python3 +import sys +sys.path.insert(0, '.') +from pathlib import Path +import argparse +import shutil +import datetime +from loader_utils import load_json, save_json + +# Get path to weights file and name of model +parser = argparse.ArgumentParser() +parser.add_argument('weights', type=Path) +parser.add_argument('--save_folder', + type=Path, + default=Path('/efs/sceneflow_golden_weights/')) +args = parser.parse_args() + +# Check that weights file exists +assert args.weights.exists(), f"Weights file {args.weights} does not exist" + +dataset_name = args.weights.parent.parent.parent.name +model_name = args.weights.parent.parent.name +print(f"Model name: {model_name}") + +save_folder = args.save_folder / dataset_name / model_name +# Make sure save folder exists +save_folder.mkdir(parents=True, exist_ok=True) + +metadata_file = args.save_folder / 'metadata.json' +if metadata_file.exists(): + print(f"Loading metadata from {metadata_file}") + metadata = load_json(metadata_file) +else: + metadata = {} + +# Copy weights file to save folder +save_path = save_folder / (model_name + '.ckpt') +print(f"Saving weights to {save_path}") +shutil.copy(args.weights, save_path) + +# Update metadata +metadata[model_name] = { + 'weight_origin': str(args.weights), + 'weight_dest': str(save_path), + 'timestamp': datetime.datetime.now().strftime("%Y_%m_%d-%I_%M_%S_%p") +} + +# Save metadata +save_json(metadata_file, metadata, indent=4) diff --git a/util_scripts/jhist b/util_scripts/jhist new file mode 100755 index 0000000..f0bb3cf --- /dev/null +++ b/util_scripts/jhist @@ -0,0 +1,48 @@ +#!/usr/bin/env python3 +from pathlib import Path + +job_dir = Path() / 'job_dir' +assert job_dir.exists(), f"Job dir {job_dir} does not exist" +# Grab all the job logs +job_logs = sorted(job_dir.glob('*/')) + + +def load_job(job_folder: Path): + folder_name = job_folder.name + command_file = job_folder / 'command.sh' + if not command_file.exists(): + return None + assert command_file.exists(), f"Command file {command_file} does not exist" + # read the command file lines + with open(command_file, 'r') as f: + lines = f.readlines() + # Grab the command + command = lines[1].strip() + return {'folder_name': folder_name, 'command': command} + + +def get_max_len_keys(dicts): + max_len = {} + for d in dicts: + for k in d.keys(): + if k not in max_len: + max_len[k] = 0 + max_len[k] = max(max_len[k], len(str(d[k]))) + return max_len + + +# Load all the jobs +jobs = [load_job(job) for job in job_logs] +# Filter out the jobs that failed to load +jobs = [job for job in jobs if job is not None] +# Get the max length of each key +max_len = get_max_len_keys(jobs) +# Print the jobs in a table format + +for job in jobs: + total_len = sum(max_len.values()) + 2 * len(max_len) + 5 + print(f"{'-' * total_len}") + print("| ", end='') + for k, v in job.items(): + print(v.ljust(max_len[k] + 2), end='| ') + print() diff --git a/util_scripts/jlaunch b/util_scripts/jlaunch new file mode 100755 index 0000000..207e07a --- /dev/null +++ b/util_scripts/jlaunch @@ -0,0 +1,158 @@ +#!/usr/bin/env python3 +import sys + +sys.path.insert(0, '.') + +import argparse +from pathlib import Path +import time +from loader_utils import run_cmd + +parser = argparse.ArgumentParser() +parser.add_argument('command', type=str) +parser.add_argument('--job_dir', type=Path, default="./job_dir/") +parser.add_argument('--num_gpus', type=int, default=1) +parser.add_argument('--cpus_per_gpu', type=int, default=2) +parser.add_argument('--mem_per_gpu', type=int, default=12) +parser.add_argument('--runtime_mins', type=int, default=180) +parser.add_argument('--runtime_hours', type=int, default=None) +parser.add_argument('--job_name', type=str, default='ff3d') +parser.add_argument('--qos', type=str, default='ee-med') +parser.add_argument('--partition', type=str, default='eaton-compute') +parser.add_argument('--dry_run', action='store_true') +parser.add_argument('--use_srun', action='store_true') +parser.add_argument('--blacklist_substring', type=str, default=None) +args = parser.parse_args() + +num_prior_jobs = len(list(args.job_dir.glob("*"))) +jobdir_path = args.job_dir / f"{num_prior_jobs:06d}" +jobdir_path.mkdir(exist_ok=True, parents=True) +job_runtime_mins = args.runtime_mins if args.runtime_hours is None else args.runtime_hours * 60 + + +def load_available_nodes(): + res = run_cmd("sinfo --Node | awk '{print $1}' | tail +2", + return_stdout=True) + available_nodes = res.split('\n') + return [e.strip() for e in available_nodes] + + +def get_node_blacklist(): + if args.blacklist_substring is None: + return [] + + node_list = load_available_nodes() + node_blacklist_list = [] + if "," in args.blacklist_substring: + node_blacklist_list = [ + e.strip() for e in args.blacklist_substring.split(",") + ] + else: + node_blacklist_list = [args.blacklist_substring] + + print(f"Blacklisting nodes with substrings {node_blacklist_list}") + print(f"Available nodes: {node_list}") + + def is_blacklisted(node): + for blacklist_str in node_blacklist_list: + if blacklist_str in node: + return True + return False + + node_blacklist = [node for node in node_list if is_blacklisted(node)] + print(f"Blacklisted nodes: {node_blacklist}") + return node_blacklist + + +def get_runtime_format(runtime_mins): + hours = runtime_mins // 60 + minutes = runtime_mins % 60 + return f"{hours:02d}:{minutes:02d}:00" + + +def make_command_file(command): + command_path = jobdir_path / f"command.sh" + command_file_content = f"""#!/bin/bash +{command} +""" + with open(command_path, "w") as f: + f.write(command_file_content) + + +def make_srun(): + srun_path = jobdir_path / f"srun.sh" + docker_image_path = Path( + "kylevedder_offline_sceneflow_latest.sqsh").absolute() + assert docker_image_path.is_file( + ), f"Docker image {docker_image_path} squash file does not exist" + srun_file_content = f"""#!/bin/bash +srun --gpus={args.num_gpus} --nodes=1 --mem-per-gpu={args.mem_per_gpu}G --cpus-per-gpu={args.cpus_per_gpu} --time={get_runtime_format(job_runtime_mins)} --exclude={','.join(node_blacklist)} --job-name={args.job_name} --qos={args.qos} --partition={args.partition} --container-mounts=../../datasets/:/efs/,`pwd`:/project --container-image={docker_image_path} bash {jobdir_path}/command.sh +""" + with open(srun_path, "w") as f: + f.write(srun_file_content) + + +def make_sbatch(): + current_working_dir = Path.cwd().absolute() + sbatch_path = jobdir_path / f"sbatch.bash" + docker_image_path = Path( + "kylevedder_offline_sceneflow_latest.sqsh").absolute() + assert docker_image_path.is_file( + ), f"Docker image {docker_image_path} squash file does not exist" + sbatch_file_content = f"""#!/bin/bash +#SBATCH --job-name={args.job_name} +#SBATCH --qos={args.qos} +#SBATCH --partition={args.partition} +#SBATCH --nodes=1 +#SBATCH --output={jobdir_path}/job.out +#SBATCH --error={jobdir_path}/job.err +#SBATCH --time={get_runtime_format(job_runtime_mins)} +#SBATCH --gpus={args.num_gpus} +#SBATCH --mem-per-gpu={args.mem_per_gpu}G +#SBATCH --cpus-per-gpu={args.cpus_per_gpu} +#SBATCH --exclude={','.join(node_blacklist)} +#SBATCH --container-mounts=../../datasets/:/efs/,{current_working_dir}:/project +#SBATCH --container-image={docker_image_path} + +bash {jobdir_path}/command.sh && echo 'done' > {jobdir_path}/job.done +""" + with open(sbatch_path, "w") as f: + f.write(sbatch_file_content) + + +def make_screen(): + screen_path = jobdir_path / f"screen.sh" + screen_file_content = f"""#!/bin/bash +screen -L -Logfile {jobdir_path}/stdout.log -dmS {args.job_name} bash {jobdir_path}/srun.sh +""" + with open(screen_path, "w") as f: + f.write(screen_file_content) + + +node_blacklist = get_node_blacklist() +make_command_file(args.command) +if args.use_srun: + make_srun() + make_screen() +else: + make_sbatch() +if not args.dry_run: + if args.use_srun: + run_cmd(f"bash {jobdir_path}/screen.sh") + else: + print("RUN THIS COMMAND TO SUBMIT THE JOB:") + print("|") + print("|") + print("|") + print("|") + print("|") + print("V") + print(f"sbatch {jobdir_path}/sbatch.bash") + print("^") + print("|") + print("|") + print("|") + print("|") + print("|") + +print(f"Config files written to {jobdir_path.absolute()}") diff --git a/util_scripts/jlog b/util_scripts/jlog new file mode 100755 index 0000000..2374320 --- /dev/null +++ b/util_scripts/jlog @@ -0,0 +1,42 @@ +#!/usr/bin/env python3 +from pathlib import Path +import argparse + +job_dir = Path() / 'job_dir' +assert job_dir.exists(), f"Job dir {job_dir} does not exist" + +job_dir_folders = {int(f.name.replace('.', '')): f for f in job_dir.glob('*/')} + +# Grab job log id + +parser = argparse.ArgumentParser() +parser.add_argument('job_id', type=int, nargs='?', default=None) +args = parser.parse_args() + +if args.job_id is None: + print("Job ids:") + print(sorted(job_dir_folders.keys())) + exit(0) + +job_id = args.job_id +assert job_id in job_dir_folders, f"Job id {job_id} does not exist" + +job_folder = job_dir_folders[job_id] + +# load stdout and stderr +stdout_file = job_folder / 'job.out' +stderr_file = job_folder / 'job.err' + +assert stdout_file.exists(), f"Stdout file {stdout_file} does not exist" +assert stderr_file.exists(), f"Stderr file {stderr_file} does not exist" + +with open(stdout_file, 'r') as f: + stdout = f.readlines() + +with open(stderr_file, 'r') as f: + stderr = f.readlines() + +print(f"Stdout for job {job_id}:") +print(''.join(stdout)) +print(f"Stderr for job {job_id}:") +print(''.join(stderr)) diff --git a/visualization/dataset_count_compute.py b/visualization/dataset_count_compute.py new file mode 100644 index 0000000..48b50a9 --- /dev/null +++ b/visualization/dataset_count_compute.py @@ -0,0 +1,112 @@ +import torch +import pandas as pd +import open3d as o3d +from dataloaders import ArgoverseSupervisedFlowSequenceLoader, ArgoverseSupervisedFlowSequence +from pointclouds import PointCloud, SE3 +import numpy as np +import tqdm +import joblib +import multiprocessing +from loader_utils import save_pickle + +CATEGORY_NAME_TO_ID = { + "ANIMAL": 0, + "ARTICULATED_BUS": 1, + "BICYCLE": 2, + "BICYCLIST": 3, + "BOLLARD": 4, + "BOX_TRUCK": 5, + "BUS": 6, + "CONSTRUCTION_BARREL": 7, + "CONSTRUCTION_CONE": 8, + "DOG": 9, + "LARGE_VEHICLE": 10, + "MESSAGE_BOARD_TRAILER": 11, + "MOBILE_PEDESTRIAN_CROSSING_SIGN": 12, + "MOTORCYCLE": 13, + "MOTORCYCLIST": 14, + "OFFICIAL_SIGNALER": 15, + "PEDESTRIAN": 16, + "RAILED_VEHICLE": 17, + "REGULAR_VEHICLE": 18, + "SCHOOL_BUS": 19, + "SIGN": 20, + "STOP_SIGN": 21, + "STROLLER": 22, + "TRAFFIC_LIGHT_TRAILER": 23, + "TRUCK": 24, + "TRUCK_CAB": 25, + "VEHICULAR_TRAILER": 26, + "WHEELCHAIR": 27, + "WHEELED_DEVICE": 28, + "WHEELED_RIDER": 29, + "NONE": -1 +} + +CATEGORY_ID_TO_NAME = {v: k for k, v in CATEGORY_NAME_TO_ID.items()} + +CATEGORY_ID_TO_IDX = { + v: idx + for idx, v in enumerate(sorted(CATEGORY_NAME_TO_ID.values())) +} +CATEGORY_IDX_TO_ID = {v: k for k, v in CATEGORY_ID_TO_IDX.items()} + +speed_bucket_ticks = [ + 0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1, 1.1, 1.2, 1.3, 1.4, 1.5, + 1.6, 1.7, 1.8, 1.9, 2.0, np.inf +] + + +def get_speed_bucket_ranges(): + return list(zip(speed_bucket_ticks, speed_bucket_ticks[1:])) + + +def process_sequence(sequence: ArgoverseSupervisedFlowSequence): + + speed_bucket_ranges = get_speed_bucket_ranges() + count_array = np.zeros((len(CATEGORY_ID_TO_IDX), len(speed_bucket_ranges)), + dtype=np.int64) + + for idx in range(len(sequence)): + frame = sequence.load(idx, 0) + pc = frame['relative_pc'] + flowed_pc = frame['relative_flowed_pc'] + if flowed_pc is None: + continue + category_ids = frame['pc_classes'] + flow = flowed_pc.points - pc.points + for category_id in np.unique(category_ids): + category_idx = CATEGORY_ID_TO_IDX[category_id] + category_mask = (category_ids == category_id) + category_flow = flow[category_mask] + # convert to m/s + category_speeds = np.linalg.norm(category_flow, axis=1) * 10.0 + category_speed_buckets = np.digitize(category_speeds, + speed_bucket_ticks) - 1 + for speed_bucket_idx in range(len(speed_bucket_ranges)): + bucket_mask = (category_speed_buckets == speed_bucket_idx) + num_points = np.sum(bucket_mask) + count_array[category_idx, speed_bucket_idx] += num_points + return count_array + + +sequence_loader = ArgoverseSupervisedFlowSequenceLoader( + '/efs/argoverse2/val/', '/efs/argoverse2/val_sceneflow/') + +sequences = [ + sequence_loader.load_sequence(sequence_id) + for sequence_id in sequence_loader.get_sequence_ids() +] + +# counts_array_lst = [ +# process_sequence(sequence) for sequence in tqdm.tqdm(sequences) +# ] + +# process sequences in parallel with joblib +counts_array_lst = joblib.Parallel(n_jobs=multiprocessing.cpu_count())( + joblib.delayed(process_sequence)(sequence) + for sequence in tqdm.tqdm(sequences)) +total_count_array = sum(counts_array_lst) + +# save the counts array +save_pickle('total_count_array.pkl', total_count_array) diff --git a/visualization/dataset_count_vis.py b/visualization/dataset_count_vis.py new file mode 100644 index 0000000..70c6bff --- /dev/null +++ b/visualization/dataset_count_vis.py @@ -0,0 +1,92 @@ +from pathlib import Path +from loader_utils import load_pickle +import numpy as np + +import matplotlib.pyplot as plt + +CATEGORY_NAME_TO_ID = { + "ANIMAL": 0, + "ARTICULATED_BUS": 1, + "BICYCLE": 2, + "BICYCLIST": 3, + "BOLLARD": 4, + "BOX_TRUCK": 5, + "BUS": 6, + "CONSTRUCTION_BARREL": 7, + "CONSTRUCTION_CONE": 8, + "DOG": 9, + "LARGE_VEHICLE": 10, + "MESSAGE_BOARD_TRAILER": 11, + "MOBILE_PEDESTRIAN_CROSSING_SIGN": 12, + "MOTORCYCLE": 13, + "MOTORCYCLIST": 14, + "OFFICIAL_SIGNALER": 15, + "PEDESTRIAN": 16, + "RAILED_VEHICLE": 17, + "REGULAR_VEHICLE": 18, + "SCHOOL_BUS": 19, + "SIGN": 20, + "STOP_SIGN": 21, + "STROLLER": 22, + "TRAFFIC_LIGHT_TRAILER": 23, + "TRUCK": 24, + "TRUCK_CAB": 25, + "VEHICULAR_TRAILER": 26, + "WHEELCHAIR": 27, + "WHEELED_DEVICE": 28, + "WHEELED_RIDER": 29, + "NONE": -1 +} + +CATEGORY_ID_TO_NAME = {v: k for k, v in CATEGORY_NAME_TO_ID.items()} + +CATEGORY_ID_TO_IDX = { + v: idx + for idx, v in enumerate(sorted(CATEGORY_NAME_TO_ID.values())) +} +CATEGORY_IDX_TO_ID = {v: k for k, v in CATEGORY_ID_TO_IDX.items()} + +speed_bucket_ticks = [ + 0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1, 1.1, 1.2, 1.3, 1.4, 1.5, + 1.6, 1.7, 1.8, 1.9, 2.0, np.inf +] + + +def get_speed_bucket_ranges(): + return list(zip(speed_bucket_ticks, speed_bucket_ticks[1:])) + + +# Load the saved count data +count_array = load_pickle('total_count_array.pkl') + +category_normalized_count_array = count_array / count_array.sum(axis=1, + keepdims=True) + +ax = plt.gca() +fig = plt.gcf() +ax.matshow(category_normalized_count_array) +for (i, j), z in np.ndenumerate(count_array): + ax.text(j, + i, + f'{z}', + ha='center', + va='center', + bbox=dict(boxstyle='round', facecolor='white', edgecolor='0.3')) + +# Set X tick labels to be the speed bucket ranges, rotated 90 degrees +ax.set_xticks(range(len(get_speed_bucket_ranges())), + [f"{l}-{u} m/s" for l, u in get_speed_bucket_ranges()], + rotation=90) +# Set Y tick labels to be the category names +ax.set_yticks(range(len(CATEGORY_ID_TO_IDX)), [ + CATEGORY_ID_TO_NAME[CATEGORY_IDX_TO_ID[i]] + for i in range(len(CATEGORY_ID_TO_IDX)) +]) + +ax.set_xlabel('Speed Bucket Ranges') +ax.set_ylabel('Categories') + +# Set figure to be 30x30 +fig.set_size_inches(30, 30) +# Save the figure +fig.savefig('total_count_array.png') diff --git a/visualization/generate_flow_error_plot_data.py b/visualization/generate_flow_error_plot_data.py new file mode 100644 index 0000000..e752171 --- /dev/null +++ b/visualization/generate_flow_error_plot_data.py @@ -0,0 +1,263 @@ +import os + +# Set Open MP's num threads to 1 +os.environ["OMP_NUM_THREADS"] = "1" + +import torch +import pandas as pd +from dataloaders import ArgoverseUnsupervisedFlowSequenceLoader, ArgoverseSupervisedFlowSequenceLoader +from pointclouds import PointCloud, SE3 +import numpy as np +import tqdm +import argparse +from pathlib import Path +import joblib +import multiprocessing + +from models.embedders import DynamicVoxelizer + +from configs.pseudoimage import POINT_CLOUD_RANGE, VOXEL_SIZE + +parser = argparse.ArgumentParser() +# dataset type (train, val), default to val +parser.add_argument('--dataset_type', + type=str, + default='val', + choices=['train', 'val']) +parser.add_argument('--dataset_source', + type=str, + default='nsfp', + choices=['nsfp', 'odom', 'nearest_neighbor', 'chodosh', 'distilation', 'supervised', 'chodosh_distilation']) +parser.add_argument('--step_size', + type=int, + default=5, + help='select every nth frame') +parser.add_argument('--num_workers', + type=int, + default=multiprocessing.cpu_count(), + help='number of workers') +args = parser.parse_args() + +assert args.step_size > 0, 'step_size must be positive' + +voxelizer = DynamicVoxelizer(voxel_size=VOXEL_SIZE, + point_cloud_range=POINT_CLOUD_RANGE) + + +def voxel_restrict_pointcloud(pc: PointCloud): + pc0s = pc.points + assert pc0s.ndim == 2, "pc0s must be a batch of 3D points" + assert pc0s.shape[1] == 3, "pc0s must be a batch of 3D points" + # Convert to torch tensor + pc0s = pc0s.reshape(1, -1, 3) + pc0s = torch.from_numpy(pc0s).float() + + pc0_voxel_infos_lst = voxelizer(pc0s) + + pc0_points_lst = [e["points"] for e in pc0_voxel_infos_lst] + pc0_valid_point_idxes = [e["point_idxes"] for e in pc0_voxel_infos_lst] + + # Convert to numpy + pc0_points_lst = [e.numpy() for e in pc0_points_lst][0] + pc0_valid_point_idxes = [e.numpy() for e in pc0_valid_point_idxes][0] + + return PointCloud.from_array(pc0_points_lst), pc0_valid_point_idxes + + +def get_pcs_flows_pose(unsupervised_frame_dict, supervised_frame_dict): + + # Load unsupervised data + unrestricted_pc = unsupervised_frame_dict['relative_pc'] + + # pc, valid_points_pc = voxel_restrict_pointcloud(unrestricted_pc) + pose = unsupervised_frame_dict['relative_pose'] + unsupervised_flow = unsupervised_frame_dict['flow'] + if unsupervised_flow.ndim == 3 and unsupervised_flow.shape[0] == 1: + unsupervised_flow = unsupervised_flow[0] + else: + assert unsupervised_flow.ndim == 2, f"unsupervised_flow must be a batch of 3D points, but got {unsupervised_flow.ndim} dimensions" + valid_idxes = unsupervised_frame_dict['valid_idxes'] + pc = unrestricted_pc.mask_points(valid_idxes) + assert unsupervised_flow is not None, 'flow must not be None' + assert pc.shape == unsupervised_flow.shape, f"pc and flow must have same shape, but got {pc.shape} and {unsupervised_flow.shape}" + unsupervised_flowed_pc = pc.flow(unsupervised_flow) + + # Load supervised data + full_supervised_flowed_pc = supervised_frame_dict['relative_flowed_pc'] + full_pc_classes = supervised_frame_dict['pc_classes'] + # Extract the full PC flow. + supervised_flow = full_supervised_flowed_pc.points - unrestricted_pc.points + supervised_flow = supervised_flow[valid_idxes] + pc_classes = full_pc_classes[valid_idxes] + supervised_flowed_pc = pc.flow(supervised_flow) + + assert unsupervised_flowed_pc.points.shape == supervised_flowed_pc.points.shape, f"unsupervised_flowed_pc and supervised_flowed_pc must have same shape, but got {unsupervised_flowed_pc.points.shape} and {supervised_flowed_pc.points.shape}" + + return unsupervised_flowed_pc, unsupervised_flow, supervised_flowed_pc, supervised_flow, pc_classes + + +def accumulate_flow_scatter(supervised_flow, + unsupervised_flow, + classes, + normalize_rotation=True, + grid_radius_meters=1.5, + cells_per_meter=50): + # import matplotlib.pyplot as plt + not_background_mask = (classes >= 0) + moving_points_mask = np.linalg.norm(supervised_flow, axis=1) > 0.05 + + valid_mask = not_background_mask & moving_points_mask + + supervised_flow = supervised_flow[valid_mask][:, :2] + unsupervised_flow = unsupervised_flow[valid_mask][:, :2] + + if normalize_rotation: + # Canonocalize the supervised flows so that ground truth is all pointing in the positive X direction + supervised_norm = np.linalg.norm(supervised_flow, axis=1) + non_zero_mask = (supervised_norm > 0.001) + + rotatable_supervised_flow = supervised_flow[non_zero_mask] + rotatable_unsupervised_flow = unsupervised_flow[non_zero_mask] + + # Compute the angle between the supervised flow and the positive X axis + supervised_angles = np.arctan2(rotatable_supervised_flow[:, 0], + rotatable_supervised_flow[:, 1]) + + # Compute the rotation matrix + rotation_matrix = np.array( + [[np.cos(supervised_angles), -np.sin(supervised_angles)], + [np.sin(supervised_angles), + np.cos(supervised_angles)]]).transpose(2, 0, 1) + + # Rotate the supervised flow + supervised_flow[non_zero_mask] = np.einsum('ijk,ik->ij', + rotation_matrix, + rotatable_supervised_flow) + + # Rotate the unsupervised flow + unsupervised_flow[non_zero_mask] = np.einsum( + 'ijk,ik->ij', rotation_matrix, rotatable_unsupervised_flow) + + # Compute the flow difference + flow_error_xy = (unsupervised_flow - supervised_flow) + + # Create the grid + accumulation_grid = np.zeros( + (int(grid_radius_meters * 2 * cells_per_meter + 1), + int(grid_radius_meters * 2 * cells_per_meter + 1))) + + # Compute the grid indices for flow_error_xy using np digitize + grid_indices_x = np.digitize( + flow_error_xy[:, 0], + bins=np.linspace( + -grid_radius_meters, grid_radius_meters, + int(grid_radius_meters * 2 * cells_per_meter) + 1)) - 1 + grid_indices_y = np.digitize( + flow_error_xy[:, 1], + bins=np.linspace( + -grid_radius_meters, grid_radius_meters, + int(grid_radius_meters * 2 * cells_per_meter) + 1)) - 1 + + grid_indices = np.array([grid_indices_x, grid_indices_y]).T + unique_grid_indices, unique_counts = np.unique(grid_indices, + axis=0, + return_counts=True) + + accumulation_grid[unique_grid_indices[:, 0], + unique_grid_indices[:, 1]] = unique_counts + + # import matplotlib.pyplot as plt + # plt.matshow(accumulation_grid) + # plt.xticks( + # np.linspace(0, grid_radius_meters * 2 * cells_per_meter, + # grid_radius_meters * 2 + 1), + # np.linspace(-grid_radius_meters, grid_radius_meters, + # grid_radius_meters * 2 + 1)) + # plt.yticks( + # np.linspace(0, grid_radius_meters * 2 * cells_per_meter, + # grid_radius_meters * 2 + 1), + # np.linspace(-grid_radius_meters, grid_radius_meters, + # grid_radius_meters * 2 + 1)) + + # plt.colorbar() + # plt.show() + + return accumulation_grid + + +def process_sequence(sequence_id): + + unsupervised_sequence = unsupervised_sequence_loader.load_sequence( + sequence_id) + supervised_sequence = supervised_sequence_loader.load_sequence(sequence_id) + + unsupervised_frame_list = unsupervised_sequence.load_frame_list(None) + supervised_frame_list = supervised_sequence.load_frame_list(None) + + zipped_frame_list = list( + zip(unsupervised_frame_list, + supervised_frame_list))[:-1][::args.step_size] + + accumulation_grid_lst_rotated = [] + accumulation_grid_lst_unrotated = [] + for idx, (unsupervised_frame_dict, + supervised_frame_dict) in enumerate(zipped_frame_list): + + step_idx = idx * args.step_size + try: + unsupervised_flowed_pc, unsupervised_flow, supervised_flowed_pc, supervised_flow, pc_classes = get_pcs_flows_pose( + unsupervised_frame_dict, supervised_frame_dict) + except AssertionError as e: + print("Error processing sequence: ", sequence_id, " step: ", + step_idx) + print(e) + # raise e + + accumulation_grid_lst_rotated.append( + accumulate_flow_scatter(supervised_flow, unsupervised_flow, + pc_classes)) + accumulation_grid_lst_unrotated.append( + accumulate_flow_scatter(supervised_flow, + unsupervised_flow, + pc_classes, + normalize_rotation=False)) + + accumulation_grid = np.sum(accumulation_grid_lst_rotated, axis=0) + accumulation_grid_unrotated = np.sum(accumulation_grid_lst_unrotated, + axis=0) + save_dir.mkdir(parents=True, exist_ok=True) + np.save(save_dir / f'{sequence_id}_error_distribution.npy', + accumulation_grid) + np.save(save_dir / f'{sequence_id}_error_distribution_unrotated.npy', + accumulation_grid_unrotated) + + +data_root = Path('/efs/argoverse2/') + +unsupervised_sequence_loader = ArgoverseUnsupervisedFlowSequenceLoader( + data_root / args.dataset_type, + data_root / f'{args.dataset_type}_{args.dataset_source}_flow/') + +supervised_sequence_loader = ArgoverseSupervisedFlowSequenceLoader( + data_root / args.dataset_type, + data_root / f'{args.dataset_type}_sceneflow/') + +save_dir = data_root / f'{args.dataset_type}_{args.dataset_source}_unsupervised_vs_supervised_flow' + +unsupervised_sequences = set(unsupervised_sequence_loader.get_sequence_ids()) +supervised_sequences = set(supervised_sequence_loader.get_sequence_ids()) + +overlapping_sequences = unsupervised_sequences.intersection( + supervised_sequences) + +print("Number of overlapping sequences: ", len(overlapping_sequences)) + +# Use joblib to parallelize the processing of each sequence + +if args.num_workers <= 1: + for sequence_id in tqdm.tqdm(sorted(overlapping_sequences)): + process_sequence(sequence_id) +else: + joblib.Parallel(n_jobs=args.num_workers, verbose=10)( + joblib.delayed(process_sequence)(sequence_id) + for sequence_id in tqdm.tqdm(sorted(overlapping_sequences))) diff --git a/visualization/visualize_flow_for_figure.py b/visualization/visualize_flow_for_figure.py new file mode 100644 index 0000000..c0799b2 --- /dev/null +++ b/visualization/visualize_flow_for_figure.py @@ -0,0 +1,307 @@ +import torch +import pandas as pd +import open3d as o3d +from dataloaders import ArgoverseRawSequenceLoader, ArgoverseSupervisedFlowSequenceLoader, ArgoverseUnsupervisedFlowSequenceLoader, WaymoSupervisedFlowSequence, WaymoSupervisedFlowSequenceLoader, WaymoUnsupervisedFlowSequenceLoader +from pointclouds import PointCloud, SE3 +import numpy as np +import tqdm +from pathlib import Path +from loader_utils import load_pickle, save_pickle +import time + +# sequence_loader = ArgoverseSequenceLoader('/bigdata/argoverse_lidar/train/') +# sequence_loader = ArgoverseSupervisedFlowSequenceLoader( +# '/efs/argoverse2/val/', '/efs/argoverse2/val_sceneflow/') +supervised_sequence_loader = WaymoSupervisedFlowSequenceLoader( + '/efs/waymo_open_processed_flow/validation/') +our_sequence_loader = WaymoUnsupervisedFlowSequenceLoader( + '/efs/waymo_open_processed_flow/validation/', + '/bigdata/waymo_open_processed_flow/val_distillation_out/') +sup_sequence_loader = WaymoUnsupervisedFlowSequenceLoader( + '/efs/waymo_open_processed_flow/validation/', + '/bigdata/waymo_open_processed_flow/val_supervised_out/') + +# supervised_sequence_loader = ArgoverseSupervisedFlowSequenceLoader( +# '/efs/argoverse2/val/', '/efs/argoverse2/val_sceneflow/') +# unsupervised_sequence_loader = ArgoverseUnsupervisedFlowSequenceLoader( +# '/efs/argoverse2/val/', '/efs/argoverse2/val_distilation_flow/') + +camera_params_path = Path() / 'camera_params.pkl' +screenshot_path = Path() / 'screenshots' +screenshot_path.mkdir(exist_ok=True) + + +def load_sequence(sequence_idx: int): + sequence_id = supervised_sequence_loader.get_sequence_ids()[sequence_idx] + assert sequence_id in set( + our_sequence_loader.get_sequence_ids() + ), f"sequence_id {sequence_id} not in unsupervised_sequence_loader.get_sequence_ids() {our_sequence_loader.get_sequence_ids()}" + print("Sequence ID: ", sequence_id, "Sequence idx: ", sequence_idx) + supervised_sequence = supervised_sequence_loader.load_sequence(sequence_id) + our_sequence = our_sequence_loader.load_sequence(sequence_id) + sup_sequence = sup_sequence_loader.load_sequence(sequence_id) + + before_human_load = time.time() + supervised_frame_list = supervised_sequence.load_frame_list(0) + before_our_load = time.time() + print(f"Time to load human: {before_our_load - before_human_load}") + our_frame_list = our_sequence.load_frame_list(0) + before_sup_load = time.time() + print(f"Time to load our: {before_sup_load - before_our_load}") + sup_frame_list = sup_sequence.load_frame_list(0) + after_sup_load = time.time() + print(f"Time to load sup: {after_sup_load - before_sup_load}") + + return supervised_frame_list, our_frame_list, sup_frame_list + + +# make open3d visualizer +vis = o3d.visualization.VisualizerWithKeyCallback() +vis.create_window() +vis.get_render_option().point_size = 2 +vis.get_render_option().background_color = (1, 1, 1) +vis.get_render_option().show_coordinate_frame = False +# Set camera parameters + + +def save_camera_params(camera_params: o3d.camera.PinholeCameraParameters): + param_dict = {} + param_dict['intrinsics'] = camera_params.intrinsic.intrinsic_matrix + param_dict['extrinsics'] = camera_params.extrinsic + save_pickle(camera_params_path, param_dict) + + +ctr = vis.get_view_control() + +sequence_idx = 1 +# sequence_idx = 7 + +draw_flow_lines_color = None + +flow_type = 'gt' + +starter_idx = 122 + +supervised_frame_list, our_frame_list, sup_frame_list = load_sequence( + sequence_idx) + + +def toggle_flow_lines(vis): + global draw_flow_lines_color + + if draw_flow_lines_color is None: + draw_flow_lines_color = 'red' + elif draw_flow_lines_color == 'red': + draw_flow_lines_color = 'blue' + elif draw_flow_lines_color == 'blue': + draw_flow_lines_color = 'green' + elif draw_flow_lines_color == 'green': + draw_flow_lines_color = 'purple' + elif draw_flow_lines_color == 'purple': + draw_flow_lines_color = 'yellow' + elif draw_flow_lines_color == 'yellow': + draw_flow_lines_color = 'orange' + elif draw_flow_lines_color == 'orange': + draw_flow_lines_color = None + else: + raise ValueError( + f'Invalid draw_flow_lines_color: {draw_flow_lines_color}') + vis.clear_geometries() + draw_frames(reset_view=False) + + +def toggle_flow_type(vis): + global flow_type + if flow_type == 'gt': + flow_type = 'ours' + elif flow_type == 'ours': + flow_type = 'sup' + elif flow_type == 'sup': + flow_type = 'gt' + else: + raise ValueError(f'Invalid flow_type: {flow_type}') + vis.clear_geometries() + draw_frames(reset_view=False) + + +def increase_starter_idx(vis): + global starter_idx + starter_idx += 1 + if starter_idx >= min(len(supervised_frame_list), len(our_frame_list)) - 1: + starter_idx = 0 + print(starter_idx) + vis.clear_geometries() + draw_frames(reset_view=False) + + +def decrease_starter_idx(vis): + global starter_idx + starter_idx -= 1 + if starter_idx < 0: + starter_idx = min(len(supervised_frame_list), len(our_frame_list)) - 2 + print(starter_idx) + vis.clear_geometries() + draw_frames(reset_view=False) + + +def starter_idx_beginning(vis): + global starter_idx + starter_idx = 0 + print(starter_idx) + vis.clear_geometries() + draw_frames(reset_view=False) + + +def starter_idx_end(vis): + global starter_idx + starter_idx = min(len(supervised_frame_list), len(our_frame_list)) - 2 + print(starter_idx) + vis.clear_geometries() + draw_frames(reset_view=False) + + +def save_screenshot(vis): + save_name = screenshot_path / f'{supervised_sequence_loader.get_sequence_ids()[sequence_idx]}_{starter_idx}_{time.time():03f}.png' + vis.capture_screen_image(str(save_name)) + + +def decrease_sequence_idx(vis): + global sequence_idx + global supervised_frame_list + global our_frame_list + global sup_frame_list + sequence_idx -= 1 + supervised_frame_list, our_frame_list, sup_frame_list = load_sequence( + sequence_idx) + vis.clear_geometries() + draw_frames(reset_view=True) + + +def increase_sequence_idx(vis): + global sequence_idx + global supervised_frame_list + global our_frame_list + global sup_frame_list + sequence_idx += 1 + supervised_frame_list, our_frame_list, sup_frame_list = load_sequence( + sequence_idx) + vis.clear_geometries() + draw_frames(reset_view=True) + + +vis.register_key_callback(ord("F"), toggle_flow_lines) +vis.register_key_callback(ord("G"), toggle_flow_type) +# left arrow decrease starter_idx +vis.register_key_callback(263, decrease_starter_idx) +# # right arrow increase starter_idx +vis.register_key_callback(262, increase_starter_idx) +# Up arrow for next sequence +vis.register_key_callback(265, increase_sequence_idx) +# Down arrow for prior sequence +vis.register_key_callback(264, decrease_sequence_idx) +# Home key for beginning of sequence +vis.register_key_callback(268, starter_idx_beginning) +# End key for end of sequence +vis.register_key_callback(269, starter_idx_end) +vis.register_key_callback(ord("S"), save_screenshot) + + +def make_lineset(pc, flowed_pc, draw_color): + line_set = o3d.geometry.LineSet() + assert len(pc) == len( + flowed_pc + ), f'pc and flowed_pc must have same length, but got {len(pc)} and {len(flowed_pc)}' + line_set_points = np.concatenate([pc, flowed_pc], axis=0) + + lines = np.array([[i, i + len(pc)] for i in range(len(pc))]) + line_set.points = o3d.utility.Vector3dVector(line_set_points) + line_set.lines = o3d.utility.Vector2iVector(lines) + draw_color = { + 'green': [0, 1, 0], + 'blue': [0, 0, 1], + 'red': [1, 0, 0], + 'yellow': [1, 1, 0], + 'orange': [1, 0.5, 0], + 'purple': [0.5, 0, 1], + 'pink': [1, 0, 1], + 'cyan': [0, 1, 1], + 'white': [1, 1, 1], + }[draw_flow_lines_color] + line_set.colors = o3d.utility.Vector3dVector( + [draw_color for _ in range(len(lines))]) + return line_set + + +def limit_frames(frame_lst): + return frame_lst[starter_idx:starter_idx + 2] + + +def draw_frames(reset_view=False): + limited_supervised_frame_list = limit_frames(supervised_frame_list) + limited_our_frame_list = limit_frames(our_frame_list) + limited_sup_frame_list = limit_frames(sup_frame_list) + #color_scalar = np.linspace(0.5, 1, len(limited_supervised_frame_list)) + # Green then blue + color_scalar = [[0, 1, 0], [0, 0, 1]] + frame_dict_list = list( + zip(limited_supervised_frame_list, limited_our_frame_list, + limited_sup_frame_list)) + for idx, (supervised_frame_dict, our_frame_dict, + sup_frame_dict) in enumerate(tqdm.tqdm(frame_dict_list)): + supervised_pc = supervised_frame_dict['relative_pc'] + pose = supervised_frame_dict['relative_pose'] + supervised_flowed_pc = supervised_frame_dict['relative_flowed_pc'] + classes = supervised_frame_dict['pc_classes'] + + our_pc_valid_idxes = our_frame_dict['valid_idxes'] + our_pc = our_frame_dict['relative_pc'][our_pc_valid_idxes] + + sup_pc_valid_idxes = sup_frame_dict['valid_idxes'] + sup_pc = sup_frame_dict['relative_pc'][sup_pc_valid_idxes] + + # Add base point cloud + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(supervised_pc.points) + pc_color = np.ones_like(supervised_pc.points) * color_scalar[idx] + pcd.colors = o3d.utility.Vector3dVector(pc_color) + vis.add_geometry(pcd, reset_bounding_box=reset_view) + + our_flow = our_frame_dict['flow'] + if our_flow is None: + continue + if our_flow.ndim == 3: + our_flow = our_flow.squeeze(0) + our_flowed_pc = our_pc + our_flow + + sup_flow = sup_frame_dict['flow'] + if sup_flow is None: + continue + if sup_flow.ndim == 3: + sup_flow = sup_flow.squeeze(0) + sup_flowed_pc = sup_pc + sup_flow + + # Add flowed point cloud + if supervised_flowed_pc is not None and idx < len(frame_dict_list) - 1: + print("sequence_idx", sequence_idx, "starter_idx:", starter_idx, + "Flow type: ", flow_type) + + if draw_flow_lines_color is not None: + + if flow_type == 'gt': + line_set = make_lineset(supervised_pc.points, + supervised_flowed_pc.points, + draw_flow_lines_color) + elif flow_type == 'ours': + line_set = make_lineset(our_pc, our_flowed_pc, + draw_flow_lines_color) + elif flow_type == 'sup': + line_set = make_lineset(sup_pc, sup_flowed_pc, + draw_flow_lines_color) + else: + raise ValueError(f'Unknown flow_type {flow_type}') + vis.add_geometry(line_set, reset_bounding_box=reset_view) + + +draw_frames(reset_view=True) + +vis.run() diff --git a/visualization/visualize_flow_sequence_stats.py b/visualization/visualize_flow_sequence_stats.py new file mode 100644 index 0000000..d2515fc --- /dev/null +++ b/visualization/visualize_flow_sequence_stats.py @@ -0,0 +1,154 @@ +import torch +import pandas as pd +import open3d as o3d +from dataloaders import ArgoverseRawSequenceLoader, ArgoverseSupervisedFlowSequenceLoader +from pointclouds import PointCloud, SE3 +import numpy as np +import tqdm +from pytorch3d.ops.knn import knn_points +from pytorch3d.loss import chamfer_distance + +# sequence_loader = ArgoverseSequenceLoader('/bigdata/argoverse_lidar/train/') +sequence_loader = ArgoverseSupervisedFlowSequenceLoader( + '/efs/argoverse2/val/', '/efs/argoverse2/val_sceneflow/') + +sequence_id = sequence_loader.get_sequence_ids()[29] +print("Sequence ID: ", sequence_id) +sequence = sequence_loader.load_sequence(sequence_id) + +sequence_length = len(sequence) + +frame_list = sequence.load_frame_list(0) + + +def sequence_idx_to_color(idx): + return [1 - idx / sequence_length, idx / sequence_length, 0] + + +MAX_K = 10 +current_k = 1 +geometry_idx = 0 +knn_result_cache = {} + + +def next_geometry(_): + global geometry_idx + geometry_idx = (geometry_idx + 1) % len(frame_list) + draw_geometry(geometry_idx) + + +def previous_geometry(_): + global geometry_idx + geometry_idx = (geometry_idx - 1) % len(frame_list) + draw_geometry(geometry_idx) + + +def cycle_k(_): + global current_k + current_k = (current_k + 1) % (MAX_K + 1) + if current_k == 0: + current_k = 1 + print("K: ", current_k) + draw_geometry(geometry_idx) + + +def pc_knn(warped_pc: torch.Tensor, + target_pc: torch.Tensor, + K: int = 1) -> np.ndarray: + if warped_pc.ndim == 2: + warped_pc = warped_pc.unsqueeze(0) + if target_pc.ndim == 2: + target_pc = target_pc.unsqueeze(0) + + # Compute min distance between warped point cloud and point cloud at t+1. + print("Computing knn...") + warped_to_target_knn = knn_points(p1=warped_pc, + p2=target_pc, + K=K, + return_nn=True) + print("Done computing knn.") + + nearest_points = warped_to_target_knn.knn.squeeze(0) + return nearest_points.cpu().numpy() + + +# make open3d visualizer +vis = o3d.visualization.VisualizerWithKeyCallback() +vis.create_window("Point Cloud Nearest Neighbor Visualization") +vis.get_render_option().point_size = 1.5 +vis.get_render_option().background_color = (0, 0, 0) +vis.get_render_option().show_coordinate_frame = True +vis.register_key_callback(ord('K'), cycle_k) +# set up vector +vis.get_view_control().set_up([0, 0, 1]) + + +def draw_geometry(idx: int, reset_boxes=False): + frame_dict = frame_list[idx] + pc = frame_dict['relative_pc'] + pose = frame_dict['relative_pose'] + flowed_pc = frame_dict['relative_flowed_pc'] + classes = frame_dict['pc_classes'] + is_grounds = frame_dict['pc_is_ground'] + + vis.clear_geometries() + + # Add center of mass + sphere = o3d.geometry.TriangleMesh.create_sphere(radius=1) + sphere.translate(pose.translation) + sphere.paint_uniform_color(sequence_idx_to_color(idx)) + vis.add_geometry(sphere, reset_bounding_box=reset_boxes) + + # Add base point cloud + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(pc.points) + pc_color = np.zeros_like(pc.points) + # Base point cloud is red + pc_color[:, 0] = 1.0 + pcd.colors = o3d.utility.Vector3dVector(pc_color) + vis.add_geometry(pcd, reset_bounding_box=reset_boxes) + + if flowed_pc is None: + return + + # Add gt pc from next frame + next_pc = frame_list[idx + 1]['relative_pc'] + next_pcd = o3d.geometry.PointCloud() + next_pcd.points = o3d.utility.Vector3dVector(next_pc.points) + next_pc_color = np.zeros_like(next_pc.points) + # Next point cloud is green + next_pc_color[:, 1] = 1.0 + next_pcd.colors = o3d.utility.Vector3dVector(next_pc_color) + vis.add_geometry(next_pcd, reset_bounding_box=reset_boxes) + + # Add line set between pc0 and pc1 using KNN + pc_torch = torch.from_numpy(pc.points).float() + next_pc_torch = torch.from_numpy(next_pc.points).float() + + if idx not in knn_result_cache: + knn_result_cache[idx] = pc_knn(pc_torch, next_pc_torch, K=MAX_K) + nearest_points_arr = knn_result_cache[idx] + + def get_k_nearest_points(nearest_points_arr, k): + return nearest_points_arr[:, :k, :].mean(axis=1) + + nearest_points = get_k_nearest_points(nearest_points_arr, k=current_k) + + assert pc.shape == nearest_points.shape, f'pc and nearest_points must have same shape, but got {pc.shape} and {nearest_points.shape}' + + line_set = o3d.geometry.LineSet() + line_set_points = np.concatenate([pc, nearest_points], axis=0) + + lines = np.array([[i, i + len(pc)] for i in range(len(pc))]) + line_set.points = o3d.utility.Vector3dVector(line_set_points) + line_set.lines = o3d.utility.Vector2iVector(lines) + # Line set is blue + line_set.colors = o3d.utility.Vector3dVector([[0, 1, 1] + for _ in range(len(lines))]) + vis.add_geometry(line_set, reset_bounding_box=reset_boxes) + + +draw_geometry(geometry_idx, reset_boxes=True) + +vis.run() +vis.destroy_window() \ No newline at end of file diff --git a/visualization/visualize_sequence.py b/visualization/visualize_sequence.py new file mode 100644 index 0000000..b965da7 --- /dev/null +++ b/visualization/visualize_sequence.py @@ -0,0 +1,51 @@ +import torch +import pandas as pd +import open3d as o3d +from dataloaders import ArgoverseRawSequenceLoader, WaymoSupervisedFlowSequence, WaymoSupervisedFlowSequenceLoader, WaymoUnsupervisedFlowSequenceLoader +from pointclouds import PointCloud, SE3 +import numpy as np +import tqdm + +sequence_loader = ArgoverseRawSequenceLoader('/efs/argoverse2/val/') + +sequence_idx = 1 +sequence_id = sequence_loader.get_sequence_ids()[sequence_idx] +sequence = sequence_loader.load_sequence(sequence_id) + +print("Sequence idx", sequence_idx, "Sequence ID: ", sequence_id) + +# make open3d visualizer +vis = o3d.visualization.Visualizer() +vis.create_window() +vis.get_render_option().point_size = 1 +vis.get_render_option().background_color = (0, 0, 0) +vis.get_render_option().show_coordinate_frame = True +# set up vector +vis.get_view_control().set_up([0, 0, 1]) + +sequence_length = len(sequence) + + +def sequence_idx_to_color(idx): + if idx == 50: + return [1, 1, 1] + return [1 - idx / sequence_length, idx / sequence_length, 0] + + +frame_lst = sequence.load_frame_list(0) + +print("Frames: ", len(frame_lst)) +for idx, entry_dict in enumerate(frame_lst): + timestamp = sequence.timestamp_list[idx] + print("idx:", idx, "Timestamp: ", timestamp) + idx_pc = entry_dict['relative_pc'] + pose = entry_dict['relative_pose'] + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(idx_pc.points) + vis.add_geometry(pcd) + sphere = o3d.geometry.TriangleMesh.create_sphere(radius=1) + sphere.translate(pose.translation) + sphere.paint_uniform_color(sequence_idx_to_color(idx)) + vis.add_geometry(sphere) + +vis.run() diff --git a/visualization/visualize_supervised_flow.py b/visualization/visualize_supervised_flow.py new file mode 100644 index 0000000..ad00371 --- /dev/null +++ b/visualization/visualize_supervised_flow.py @@ -0,0 +1,97 @@ +import torch +import pandas as pd +import open3d as o3d +from dataloaders import ArgoverseRawSequenceLoader, ArgoverseSupervisedFlowSequenceLoader, WaymoSupervisedFlowSequence, WaymoSupervisedFlowSequenceLoader, WaymoUnsupervisedFlowSequenceLoader +from pointclouds import PointCloud, SE3 +import numpy as np +import tqdm + +# sequence_loader = ArgoverseSequenceLoader('/bigdata/argoverse_lidar/train/') +sequence_loader = ArgoverseSupervisedFlowSequenceLoader( + '/efs/argoverse2/val/', '/efs/argoverse2/val_sceneflow/') +# sequence_loader = WaymoSupervisedFlowSequenceLoader( +# '/efs/waymo_open_processed_flow/training/') + +sequence_id = sequence_loader.get_sequence_ids()[1] +print("Sequence ID: ", sequence_id) +sequence = sequence_loader.load_sequence(sequence_id) + +# make open3d visualizer +vis = o3d.visualization.VisualizerWithKeyCallback() +vis.create_window() +vis.get_render_option().point_size = 1.5 +vis.get_render_option().background_color = (0, 0, 0) +vis.get_render_option().show_coordinate_frame = True +# set up vector +vis.get_view_control().set_up([0, 0, 1]) + +sequence_length = len(sequence) + + +def sequence_idx_to_color(idx): + return [1 - idx / sequence_length, idx / sequence_length, 0] + + +frame_list = sequence.load_frame_list(0)[125:127] + + +draw_flow_lines_color = None + + +def toggle_flow_lines(vis ): + global draw_flow_lines_color + + if draw_flow_lines_color is None: + draw_flow_lines_color = 'green' + elif draw_flow_lines_color == 'green': + draw_flow_lines_color = 'blue' + elif draw_flow_lines_color == 'blue': + draw_flow_lines_color = None + else: + raise ValueError(f'Invalid draw_flow_lines_color: {draw_flow_lines_color}') + vis.clear_geometries() + draw_frames(reset_view=False) + + +vis.register_key_callback(ord("F"), toggle_flow_lines) + + + +def draw_frames(reset_view=False): + color_scalar = np.linspace(0.5, 1, len(frame_list)) + for idx, frame_dict in enumerate(tqdm.tqdm(frame_list)): + pc = frame_dict['relative_pc'] + pose = frame_dict['relative_pose'] + flowed_pc = frame_dict['relative_flowed_pc'] + classes = frame_dict['pc_classes'] + + # Add base point cloud + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(pc.points) + pc_color = np.ones_like(pc.points) * color_scalar[idx] + pcd.colors = o3d.utility.Vector3dVector(pc_color) + vis.add_geometry(pcd, reset_bounding_box=reset_view) + + # Add flowed point cloud + if flowed_pc is not None and idx < len(frame_list) - 1: + print("Max flow magnitude:", + pc.matched_point_distance(flowed_pc).max()) + + if draw_flow_lines_color is not None: + line_set = o3d.geometry.LineSet() + assert len(pc) == len( + flowed_pc + ), f'pc and flowed_pc must have same length, but got {len(pc)} and {len(flowed_pc)}' + line_set_points = np.concatenate([pc, flowed_pc], axis=0) + + lines = np.array([[i, i + len(pc)] for i in range(len(pc))]) + line_set.points = o3d.utility.Vector3dVector(line_set_points) + line_set.lines = o3d.utility.Vector2iVector(lines) + draw_color = [0, 1, 0] if draw_flow_lines_color == 'green' else [0, 0, 1] + line_set.colors = o3d.utility.Vector3dVector( + [draw_color for _ in range(len(lines))]) + vis.add_geometry(line_set, reset_bounding_box=reset_view) + +draw_frames(reset_view=True) + +vis.run() diff --git a/visualization/visualize_unsupervised_flow.py b/visualization/visualize_unsupervised_flow.py new file mode 100644 index 0000000..489a815 --- /dev/null +++ b/visualization/visualize_unsupervised_flow.py @@ -0,0 +1,114 @@ +import torch +import pandas as pd +import open3d as o3d +from dataloaders import ArgoverseUnsupervisedFlowSequenceLoader +from pointclouds import PointCloud, SE3 +import numpy as np +import tqdm + +from models.embedders import DynamicVoxelizer + +from configs.pseudoimage import POINT_CLOUD_RANGE, VOXEL_SIZE + +voxelizer = DynamicVoxelizer(voxel_size=VOXEL_SIZE, + point_cloud_range=POINT_CLOUD_RANGE) + + +def voxel_restrict_pointcloud(pc: PointCloud): + pc0s = pc.points + assert pc0s.ndim == 2, "pc0s must be a batch of 3D points" + assert pc0s.shape[1] == 3, "pc0s must be a batch of 3D points" + # Convert to torch tensor + pc0s = pc0s.reshape(1, -1, 3) + pc0s = torch.from_numpy(pc0s).float() + + pc0_voxel_infos_lst = voxelizer(pc0s) + + pc0_points_lst = [e["points"] for e in pc0_voxel_infos_lst] + pc0_valid_point_idxes = [e["point_idxes"] for e in pc0_voxel_infos_lst] + + # Convert to numpy + pc0_points_lst = [e.numpy() for e in pc0_points_lst][0] + pc0_valid_point_idxes = [e.numpy() for e in pc0_valid_point_idxes][0] + + return PointCloud.from_array(pc0_points_lst) + + +sequence_loader = ArgoverseUnsupervisedFlowSequenceLoader( + '/efs/argoverse2/val/', '/efs/argoverse2/val_nsfp_flow/') + +sequence_id = sequence_loader.get_sequence_ids()[1] +print("Sequence ID: ", sequence_id) +sequence = sequence_loader.load_sequence(sequence_id) + +# make open3d visualizer +vis = o3d.visualization.Visualizer() +vis.create_window() +vis.get_render_option().point_size = 1.5 +vis.get_render_option().background_color = (0, 0, 0) +vis.get_render_option().show_coordinate_frame = True +# set up vector +vis.get_view_control().set_up([0, 0, 1]) + +sequence_length = len(sequence) + + +def sequence_idx_to_color(idx): + return [1 - idx / sequence_length, idx / sequence_length, 0] + + +frame_list = sequence.load_frame_list(None) +for idx, frame_dict in enumerate(tqdm.tqdm(frame_list[1:13])): + pc = frame_dict['relative_pc'] + pc = voxel_restrict_pointcloud(pc) + pose = frame_dict['relative_pose'] + flow = frame_dict['flow'][0] + assert flow is not None, 'flow must not be None' + flowed_pc = pc.flow(flow) + + # Add base point cloud + pcd = o3d.geometry.PointCloud() + pcd.points = o3d.utility.Vector3dVector(pc.points) + vis.add_geometry(pcd) + + # Add flowed point cloud + if flowed_pc is not None: + # flowed_pcd = o3d.geometry.PointCloud() + # flowed_pcd.points = o3d.utility.Vector3dVector(flowed_pc.points) + # flowed_pc_color = np.zeros_like(flowed_pc.points) + # # Flow point cloud is blue + # flowed_pc_color[:, 2] = 1.0 + # flowed_pcd.colors = o3d.utility.Vector3dVector(flowed_pc_color) + # vis.add_geometry(flowed_pcd) + + # Add line set between pc0 and regressed pc1 + line_set = o3d.geometry.LineSet() + assert len(pc) == len( + flowed_pc + ), f'pc and flowed_pc must have same length, but got {len(pc)} and {len(flowed_pc)}' + line_set_points = np.concatenate([pc, flowed_pc], axis=0) + + lines = np.array([[i, i + len(pc)] for i in range(len(pc))]) + line_set.points = o3d.utility.Vector3dVector(line_set_points) + line_set.lines = o3d.utility.Vector2iVector(lines) + # Line set is blue + line_set.colors = o3d.utility.Vector3dVector( + [[1, 0, 0] for _ in range(len(lines))]) + vis.add_geometry(line_set) + + # # Add gt pc from next frame + # next_pc = frame_list[idx + 1]['relative_pc'] + # next_pcd = o3d.geometry.PointCloud() + # next_pcd.points = o3d.utility.Vector3dVector(next_pc.points) + # next_pc_color = np.zeros_like(next_pc.points) + + # next_pcd.colors = o3d.utility.Vector3dVector(next_pc_color) + # vis.add_geometry(next_pcd) + + # Add center of mass + sphere = o3d.geometry.TriangleMesh.create_sphere(radius=1) + sphere.translate(pose.translation) + sphere.paint_uniform_color(sequence_idx_to_color(idx)) + vis.add_geometry(sphere) + +vis.run()