diff --git a/Dockerfile b/Dockerfile index 0a50ff6..55d295e 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,8 +1,15 @@ -FROM pytorch/pytorch:2.1.0-cuda11.8-cudnn8-devel +# check more: https://hub.docker.com/r/nvidia/cuda +FROM nvidia/cuda:11.7.1-devel-ubuntu20.04 ENV DEBIAN_FRONTEND noninteractive LABEL maintainer="Qingwen Zhang " -RUN apt update && apt install -y git tmux curl vim rsync libgl1 libglib2.0-0 ca-certificates +RUN apt update && apt install -y git curl vim rsync htop + +RUN curl -o ~/miniforge3.sh -LO https://github.com/conda-forge/miniforge/releases/latest/download/miniforge3-Linux-x86_64.sh && \ + chmod +x ~/miniforge3.sh && \ + ~/miniforge3.sh -b -p /opt/conda && \ + rm ~/miniforge3.sh && \ + /opt/conda/bin/conda clean -ya && /opt/conda/bin/conda init bash && /opt/conda/bin/conda init zsh # install zsh and oh-my-zsh RUN apt update && apt install -y wget git zsh tmux vim g++ @@ -14,23 +21,28 @@ RUN sh -c "$(wget -O- https://github.com/deluan/zsh-in-docker/releases/download/ -p https://github.com/zsh-users/zsh-syntax-highlighting RUN printf "y\ny\ny\n\n" | bash -c "$(curl -fsSL https://raw.githubusercontent.com/Kin-Zhang/Kin-Zhang/main/scripts/setup_ohmyzsh.sh)" -RUN /opt/conda/bin/conda init zsh # change to conda env ENV PATH /opt/conda/bin:$PATH -RUN /opt/conda/bin/conda config --set solver libmamba -RUN mkdir -p /home/kin/workspace && cd /home/kin/workspace && git clone https://github.com/Kin-Zhang/OpenSceneFlow +RUN mkdir -p /home/kin/workspace && cd /home/kin/workspace && git clone https://github.com/KTH-RPL/OpenSceneFlow.git WORKDIR /home/kin/workspace/OpenSceneFlow +RUN apt-get update && apt-get install libgl1 -y # need read the gpu device info to compile the cuda extension -RUN /opt/conda/bin/pip install -r /home/kin/workspace/OpenSceneFlow/requirements.txt -RUN /opt/conda/bin/pip install FastGeodis --no-build-isolation -RUN /opt/conda/bin/pip install --no-cache-dir -e ./assets/cuda/chamfer3D && /opt/conda/bin/pip install --no-cache-dir -e ./assets/cuda/mmcv +RUN cd /home/kin/workspace/OpenSceneFlow && /opt/conda/bin/conda env create -f environment.yaml +# To make images can run all methods in the codebase +RUN /opt/conda/envs/opensf/bin/pip install torch-scatter -f https://data.pyg.org/whl/torch-2.0.0+cu117.html +RUN /opt/conda/envs/opensf/bin/pip install FastGeodis --no-build-isolation --no-cache-dir +RUN /opt/conda/envs/opensf/bin/pip install mmengine-lite && \ + /opt/conda/bin/conda install -n opensf -y pytorch3d -c pytorch3d + +# custom cuda library +RUN cd /home/kin/workspace/OpenSceneFlow/assets/cuda/mmcv && /opt/conda/envs/opensf/bin/python ./setup.py install +RUN cd /home/kin/workspace/OpenSceneFlow/assets/cuda/chamfer3D && /opt/conda/envs/opensf/bin/python ./setup.py install -# environment for dataprocessing includes data-api -RUN /opt/conda/bin/conda env create -f envsftool.yaml +RUN cd /home/kin/workspace/OpenSceneFlow && /opt/conda/bin/conda env create -f envsftool.yaml RUN /opt/conda/envs/sftool/bin/pip install numpy==1.22 # clean up apt cache -RUN rm -rf /var/lib/apt/lists/* && rm -rf /root/.cache/pip +RUN rm -rf /var/lib/apt/lists/* && rm -rf /root/.cache/pip && /opt/conda/bin/conda clean -ya diff --git a/README.md b/README.md index a1aa2ae..ad357a0 100644 --- a/README.md +++ b/README.md @@ -11,15 +11,20 @@ OpenSceneFlow is a codebase for point cloud scene flow estimation. It is also an official implementation of the following papers (sorted by the time of publication): +- **DeltaFlow: An Efficient Multi-frame Scene Flow Estimation Method** +*Qingwen Zhang, Xiaomeng Zhu, Yushan Zhang, Yixi Cai, Olov Andersson, Patric Jensfelt* +Preprint; Under review; 2025 +[ Backbone ] [ Supervised ] - [ [arXiv](https://arxiv.org/abs/2508.17054) ] [ [Project](https://github.com/Kin-Zhang/DeltaFlow) ] + - **HiMo: High-Speed Objects Motion Compensation in Point Clouds** (SeFlow++) *Qingwen Zhang, Ajinkya Khoche, Yi Yang, Li Ling, Sina Sharif Mansouri, Olov Andersson, Patric Jensfelt* -Preprint; Under review; 2025 -[ Strategy ] [ Self-Supervised ] - [ [arXiv](https://arxiv.org/abs/2503.00803) ] [ [Project](https://kin-zhang.github.io/HiMo/) ] +IEEE Transactions on Robotics (**T-RO**) 2025 +[ Strategy ] [ Self-Supervised ] - [ [arXiv](https://arxiv.org/abs/2503.00803) ] [ [Project](https://kin-zhang.github.io/HiMo/) ] → [here](#seflow-1) - **VoteFlow: Enforcing Local Rigidity in Self-Supervised Scene Flow** *Yancong Lin\*, Shiming Wang\*, Liangliang Nan, Julian Kooij, Holger Caesar* Conference on Computer Vision and Pattern Recognition (**CVPR**) 2025 -[ Backbone ] [ Self-Supervised ] - [ [arXiv](https://arxiv.org/abs/2503.22328) ] [ [Project](https://github.com/tudelft-iv/VoteFlow/)] → [here](#VoteFLow) +[ Backbone ] [ Self-Supervised ] - [ [arXiv](https://arxiv.org/abs/2503.22328) ] [ [Project](https://github.com/tudelft-iv/VoteFlow/) ] → [here](#VoteFLow) - **Flow4D: Leveraging 4D Voxel Network for LiDAR Scene Flow Estimation** *Jaeyeul Kim, Jungwan Woo, Ukcheol Shin, Jean Oh, Sunghoon Im* @@ -57,16 +62,6 @@ Additionally, *OpenSceneFlow* integrates following excellent works: [ICLR'24 Zer 💡: Want to learn how to add your own network in this structure? Check [Contribute section](CONTRIBUTING.md#adding-a-new-method) and know more about the code. Fee free to pull request and your bibtex [here](#cite-us). ---- - - - ## 0. Installation There are two ways to install the codebase: directly on your [local machine](#environment-setup) or in a [Docker container](#docker-recommended-for-isolation). @@ -77,7 +72,7 @@ We use conda to manage the environment, you can install it follow [here](assets/ ```bash git clone --recursive https://github.com/KTH-RPL/OpenSceneFlow.git -cd OpenSceneFlow && mamba env create -f environment.yaml +cd OpenSceneFlow && conda env create -f environment.yaml # You may need export your LD_LIBRARY_PATH with env lib # export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/kin/mambaforge/lib @@ -100,14 +95,14 @@ cd /home/kin/workspace/OpenSceneFlow && git pull cd /home/kin/workspace/OpenSceneFlow/assets/cuda/mmcv && /opt/conda/envs/opensf/bin/python ./setup.py install cd /home/kin/workspace/OpenSceneFlow/assets/cuda/chamfer3D && /opt/conda/envs/opensf/bin/python ./setup.py install cd /home/kin/workspace/OpenSceneFlow -mamba activate opensf +conda activate opensf ``` If you prefer to build the Docker image by yourself, Check [build-docker-image](assets/README.md#build-docker-image) section for more details. ## 1. Data Preparation -Refer to [dataprocess/README.md](dataprocess/README.md) for dataset download instructions. Currently, we support **Argoverse 2**, **Waymo**, **nuScenes** and **custom datasets** (more datasets will be added in the future). +Refer to [dataprocess/README.md](dataprocess/README.md) for dataset download instructions. Currently, we support **Argoverse 2**, **Waymo**, **nuScenes**, **ZOD** and **custom datasets** (more datasets will be added in the future). After downloading, convert the raw data to `.h5` format for easy training, evaluation, and visualization. Follow the steps in [dataprocess/README.md#process](dataprocess/README.md#process). @@ -115,8 +110,9 @@ For a quick start, use our **mini processed dataset**, which includes one scene ```bash -wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/demo_data.zip -unzip demo_data.zip -d /home/kin/data/av2/h5py +# around 1.3G +wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/demo-data-v2.zip +unzip demo-data-v2.zip -d /home/kin/data/av2/h5py ``` Once extracted, you can directly use this dataset to run the [training script](#2-quick-start) without further processing. @@ -131,56 +127,34 @@ Some tips before running the code: And free yourself from trainning, you can download the pretrained weight from [HuggingFace](https://huggingface.co/kin-zhang/OpenSceneFlow) and we provided the detail `wget` command in each model section. For optimization-based method, it's train-free so you can directly run with [3. Evaluation](#3-evaluation) (check more in the evaluation section). ```bash -mamba activate opensf +conda activate opensf ``` -### VoteFLow -Extra pakcges needed for VoteFlow, [pytorch3d](https://pytorch3d.org/) (prefer 0.7.7) and [torch-scatter](https://github.com/rusty1s/pytorch_scatter?tab=readme-ov-file) (prefer 2.1.2): - -```bash -# Install Pytorch3d -conda install pytorch3d -c pytorch3d +### Supervised Training -# Install torch-scatter -pip install torch-scatter -f https://data.pyg.org/whl/torch-2.0.0+cu117.html -``` - -Train VoteFlow with the leaderboard submit config. [Runtime: Around 32 hours in 4 x V100 GPUs.] -```bash -python train.py model=voteflow lr=2e-4 lr_scheduler=step epochs=12 batch_size=4 model.target.m=8 model.target.n=128 loss_fn=seflowLoss "add_seloss={chamfer_dis: 1.0, static_flow_loss: 1.0, dynamic_chamfer_dis: 1.0, cluster_based_pc0pc1: 1.0}" -``` - -Pretrained weight can be downloaded through: -```bash -wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/voteflow_best.ckpt -``` - -### Flow4D +#### Flow4D Train Flow4D with the leaderboard submit config. [Runtime: Around 18 hours in 4x RTX 3090 GPUs.] ```bash -python train.py model=flow4d lr=1e-3 epochs=15 batch_size=8 num_frames=5 loss_fn=deflowLoss "voxel_size=[0.2, 0.2, 0.2]" "point_cloud_range=[-51.2, -51.2, -3.2, 51.2, 51.2, 3.2]" -``` +python train.py model=flow4d optimizer.lr=1e-3 epochs=15 batch_size=8 num_frames=5 loss_fn=deflowLoss "voxel_size=[0.2, 0.2, 0.2]" "point_cloud_range=[-51.2, -51.2, -3.2, 51.2, 51.2, 3.2]" -Pretrained weight can be downloaded through: -```bash +# Pretrained weight can be downloaded through: wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/flow4d_best.ckpt ``` -### SSF +#### SSF Extra pakcges needed for SSF model: ```bash -pip install mmengine-lite torch-scatter -# torch-scatter might not working, then reinstall by: -pip install https://data.pyg.org/whl/torch-2.0.0%2Bcu118/torch_scatter-2.1.2%2Bpt20cu118-cp310-cp310-linux_x86_64.whl +pip install mmengine-lite +pip install torch-scatter -f https://data.pyg.org/whl/torch-2.0.0+cu117.html ``` Train SSF with the leaderboard submit config. [Runtime: Around 6 hours in 8x A100 GPUs.] ```bash -python train.py model=ssf lr=8e-3 epochs=25 batch_size=64 loss_fn=deflowLoss "voxel_size=[0.2, 0.2, 6]" "point_cloud_range=[-51.2, -51.2, -3, 51.2, 51.2, 3]" +python train.py model=ssf optimizer.lr=8e-3 epochs=25 batch_size=64 loss_fn=deflowLoss "voxel_size=[0.2, 0.2, 6]" "point_cloud_range=[-51.2, -51.2, -3, 51.2, 51.2, 3]" ``` Pretrained weight can be downloaded through: @@ -192,31 +166,61 @@ wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/ssf_best.ckpt wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/ssf_long.ckpt ``` +#### DeFlow -### SeFlow - -Train SeFlow needed to specify the loss function, we set the config of our best model in the leaderboard. [Runtime: Around 11 hours in 4x A100 GPUs.] +Train DeFlow with the leaderboard submit config. [Runtime: Around 6-8 hours in 4x A100 GPUs.] Please change `batch_size&lr` accoordingly if you don't have enough GPU memory. (e.g. `batch_size=6` for 24GB GPU) ```bash -python train.py model=deflow lr=2e-4 epochs=9 batch_size=16 loss_fn=seflowLoss "add_seloss={chamfer_dis: 1.0, static_flow_loss: 1.0, dynamic_chamfer_dis: 1.0, cluster_based_pc0pc1: 1.0}" "model.target.num_iters=2" +python train.py model=deflow optimizer.lr=2e-4 epochs=15 batch_size=16 loss_fn=deflowLoss + +# Pretrained weight can be downloaded through: +wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/deflow_best.ckpt ``` -Pretrained weight can be downloaded through: +### Feed-Forward Self-Supervised Model Training + +Train Feed-forward SSL methods (e.g. SeFlow/SeFlow++/VoteFlow etc), we needed to: +1) process auto-label process. +2) specify the loss function, we set the config here for our best model in the leaderboard. + +#### SeFlow + ```bash +# [Runtime: Around 11 hours in 4x A100 GPUs.] +python train.py model=deflow optimizer.lr=2e-4 epochs=9 batch_size=16 loss_fn=seflowLoss +ssl_label=seflow_auto "+add_seloss={chamfer_dis: 1.0, static_flow_loss: 1.0, dynamic_chamfer_dis: 1.0, cluster_based_pc0pc1: 1.0}" "model.target.num_iters=2" + +# Pretrained weight can be downloaded through: wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/seflow_best.ckpt ``` -### DeFlow +#### VoteFLow +Extra pakcges needed for VoteFlow, [pytorch3d](https://pytorch3d.org/) (prefer 0.7.7) and [torch-scatter](https://github.com/rusty1s/pytorch_scatter?tab=readme-ov-file) (prefer 2.1.2): -Train DeFlow with the leaderboard submit config. [Runtime: Around 6-8 hours in 4x A100 GPUs.] Please change `batch_size&lr` accoordingly if you don't have enough GPU memory. (e.g. `batch_size=6` for 24GB GPU) +```bash +# Install Pytorch3d +conda install pytorch3d -c pytorch3d +# Install torch-scatter +pip install torch-scatter -f https://data.pyg.org/whl/torch-2.0.0+cu117.html +``` + +Train VoteFlow with the leaderboard submit config. [Runtime: Around 32 hours in 4 x V100 GPUs.] ```bash -python train.py model=deflow lr=2e-4 epochs=15 batch_size=16 loss_fn=deflowLoss +python train.py model=voteflow optimizer.lr=2e-4 +optimizer.scheduler.name=StepLR +optimizer.scheduler.step_size=6 epochs=12 batch_size=4 model.target.m=8 model.target.n=128 loss_fn=seflowLoss "+add_seloss={chamfer_dis: 1.0, static_flow_loss: 1.0, dynamic_chamfer_dis: 1.0, cluster_based_pc0pc1: 1.0}" +ssl_label=seflow_auto + +# Pretrained weight can be downloaded through: +wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/voteflow_best.ckpt ``` -Pretrained weight can be downloaded through: + +#### SeFlow++ + ```bash -wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/deflow_best.ckpt +# [Runtime: Around 10 hours in 4x A100 GPUs.] for Argoverse 2 +python train.py model=deflowpp save_top_model=3 val_every=3 voxel_size="[0.2, 0.2, 6]" point_cloud_range="[-51.2, -51.2, -3, 51.2, 51.2, 3]" num_workers=16 epochs=9 optimizer.lr=2e-4 +optimizer.scheduler.name=StepLR "+add_seloss={chamfer_dis: 1.0, static_flow_loss: 1.0, dynamic_chamfer_dis: 1.0, cluster_based_pc0pc1: 1.0}" +ssl_label=seflowpp_auto loss_fn=seflowppLoss num_frames=3 batch_size=4 + +# Pretrained weight can be downloaded through: +wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/seflowpp_best.ckpt ``` ## 3. Evaluation @@ -227,14 +231,14 @@ Since in training, we save all hyper-parameters and model checkpoints, the only ```bash # (feed-forward): load ckpt and run it, it will directly prints all metric -python eval.py checkpoint=/home/kin/seflow_best.ckpt av2_mode=val +python eval.py checkpoint=/home/kin/seflow_best.ckpt data_mode=val # (optimization-based): it might need take really long time, maybe tmux for run it. python eval.py model=nsfp # it will output the av2_submit.zip or av2_submit_v2.zip for you to submit to leaderboard -python eval.py checkpoint=/home/kin/seflow_best.ckpt av2_mode=test leaderboard_version=1 -python eval.py checkpoint=/home/kin/seflow_best.ckpt av2_mode=test leaderboard_version=2 +python eval.py checkpoint=/home/kin/seflow_best.ckpt data_mode=test leaderboard_version=1 +python eval.py checkpoint=/home/kin/seflow_best.ckpt data_mode=test leaderboard_version=2 ``` ### **📊 Range-Wise Metric (New!)** @@ -251,13 +255,13 @@ In [SSF paper](https://arxiv.org/abs/2501.17821), we introduce a new distance-ba ### Submit result to public leaderboard -To submit your result to the public Leaderboard, if you select `av2_mode=test`, it should be a zip file for you to submit to the leaderboard. +To submit your result to the public Leaderboard, if you select `data_mode=test`, it should be a zip file for you to submit to the leaderboard. Note: The leaderboard result in DeFlow&SeFlow main paper is [version 1](https://eval.ai/web/challenges/challenge-page/2010/evaluation), as [version 2](https://eval.ai/web/challenges/challenge-page/2210/overview) is updated after DeFlow&SeFlow. ```bash # since the env may conflict we set new on deflow, we directly create new one: -mamba create -n py37 python=3.7 -mamba activate py37 +conda create -n py37 python=3.7 +conda activate py37 pip install "evalai" # Step 2: login in eval and register your team diff --git a/assets/README.md b/assets/README.md index 837b01a..d49ebe8 100644 --- a/assets/README.md +++ b/assets/README.md @@ -80,6 +80,7 @@ python -c "import torch; print(torch.__version__); print(torch.cuda.is_available python -c "import lightning.pytorch as pl; print(pl.__version__)" python -c "from assets.cuda.mmcv import Voxelization, DynamicScatter;print('successfully import on our lite mmcv package')" python -c "from assets.cuda.chamfer3D import nnChamferDis;print('successfully import on our chamfer3D package')" +python -c "from av2.utils.io import read_feather; print('av2 package ok')" ``` @@ -100,3 +101,6 @@ python -c "from assets.cuda.chamfer3D import nnChamferDis;print('successfully im 3. torch_scatter problem: `OSError: /home/kin/mambaforge/envs/opensf-v2/lib/python3.10/site-packages/torch_scatter/_version_cpu.so: undefined symbol: _ZN5torch3jit17parseSchemaOrNameERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE` Solved by install the torch-cuda version: `pip install https://data.pyg.org/whl/torch-2.0.0%2Bcu118/torch_scatter-2.1.2%2Bpt20cu118-cp310-cp310-linux_x86_64.whl` + +4. cuda package problem: `ValueError(f"Unknown CUDA arch ({arch}) or GPU not supported")` + Solved by [checking GPU compute](https://developer.nvidia.cn/cuda-gpus#compute) then manually assign: `export TORCH_CUDA_ARCH_LIST=8.6` \ No newline at end of file diff --git a/assets/cuda/chamfer3D/__init__.py b/assets/cuda/chamfer3D/__init__.py index 3aac3a5..fc5020d 100644 --- a/assets/cuda/chamfer3D/__init__.py +++ b/assets/cuda/chamfer3D/__init__.py @@ -74,11 +74,11 @@ def dis_res(self, input0, input1): dist0, dist1, _, _ = ChamferDis.apply(input0, input1) return dist0, dist1 - def truncated_dis(self, input0, input1): + def truncated_dis(self, input0, input1, truncate_dist=2): # nsfp: truncated distance way is set >= 2 to 0 but not nanmean cham_x, cham_y = self.dis_res(input0, input1) - cham_x[cham_x >= 2] = 0.0 - cham_y[cham_y >= 2] = 0.0 + cham_x[cham_x >= truncate_dist] = 0.0 + cham_y[cham_y >= truncate_dist] = 0.0 return torch.mean(cham_x) + torch.mean(cham_y) def disid_res(self, input0, input1): diff --git a/assets/cuda/chamfer3D/setup.py b/assets/cuda/chamfer3D/setup.py index 2c16070..ed79970 100755 --- a/assets/cuda/chamfer3D/setup.py +++ b/assets/cuda/chamfer3D/setup.py @@ -1,15 +1,25 @@ from setuptools import setup from torch.utils.cpp_extension import BuildExtension, CUDAExtension +extra_compile_args = { + 'cxx': ['-DCCCL_IGNORE_DEPRECATED_CUDA_BELOW_12'], + 'nvcc': ['-DCCCL_IGNORE_DEPRECATED_CUDA_BELOW_12'], +} + setup( name='chamfer3D', ext_modules=[ - CUDAExtension('chamfer3D', [ - "/".join(__file__.split('/')[:-1] + ['chamfer3D_cuda.cpp']), # must named as xxx_cuda.cpp - "/".join(__file__.split('/')[:-1] + ['chamfer3D.cu']), - ]), + CUDAExtension( + name='chamfer3D', + sources=[ + "/".join(__file__.split('/')[:-1] + ['chamfer3D_cuda.cpp']), # must named as xxx_cuda.cpp + "/".join(__file__.split('/')[:-1] + ['chamfer3D.cu']), + ], + extra_compile_args=extra_compile_args + ), ], cmdclass={ 'build_ext': BuildExtension }, - version='1.0.1') + version='1.0.2' +) \ No newline at end of file diff --git a/assets/cuda/mmcv/setup.py b/assets/cuda/mmcv/setup.py index ba9140a..70473c5 100644 --- a/assets/cuda/mmcv/setup.py +++ b/assets/cuda/mmcv/setup.py @@ -3,9 +3,15 @@ from setuptools import find_packages, setup from torch.utils.cpp_extension import BuildExtension, CUDAExtension +extra_compile_args = { + 'cxx': ['-DCCCL_IGNORE_DEPRECATED_CUDA_BELOW_12'], + 'nvcc': ['-DCCCL_IGNORE_DEPRECATED_CUDA_BELOW_12'], +} + + setup( name='mmcv', - version='1.0.1', + version='1.0.2', ext_modules=[ CUDAExtension( name='mmcv', @@ -18,14 +24,8 @@ "/".join(__file__.split("/")[:-1] + ["pybind.cpp"]), ], - # extra_compile_args={ - # 'cxx': ['-std=c++17'], - # 'nvcc': ['-std=c++17', - # '-D__CUDA_NO_HALF_OPERATORS__', - # '-D__CUDA_NO_HALF_CONVERSIONS__', - # '-D__CUDA_NO_HALF2_OPERATORS__', - # ],} - ), + extra_compile_args=extra_compile_args + ), ], cmdclass={'build_ext': BuildExtension}, diff --git a/assets/slurm/0_process.sh b/assets/slurm/0_process.sh index 2b0392b..44481f2 100644 --- a/assets/slurm/0_process.sh +++ b/assets/slurm/0_process.sh @@ -7,29 +7,31 @@ #SBATCH -t 1-00:00:00 #SBATCH --mail-type=END,FAIL #SBATCH --mail-user=qingwen@kth.se -#SBATCH --output /proj/berzelius-2023-154/users/x_qinzh/workspace/SeFlow/logs/slurm/%J_data.out -#SBATCH --error /proj/berzelius-2023-154/users/x_qinzh/workspace/SeFlow/logs/slurm/%J_data.err +#SBATCH --output /proj/berzelius-2023-154/users/x_qinzh/workspace/OpenSceneFlow/logs/slurm/%J_data.out +#SBATCH --error /proj/berzelius-2023-154/users/x_qinzh/workspace/OpenSceneFlow/logs/slurm/%J_data.err -cd /proj/berzelius-2023-154/users/x_qinzh/workspace/SeFlow +PYTHON=/proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/sftool/bin/python export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/proj/berzelius-2023-154/users/x_qinzh/mambaforge/lib +cd /proj/berzelius-2023-364/users/x_qinzh/workspace/OpenSceneFlow # export HYDRA_FULL_ERROR=1 -/proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/dataprocess/bin/python dataprocess/extract_av2.py --nproc 64 \ + +$PYTHON dataprocess/extract_av2.py --nproc 64 \ --av2_type sensor \ --data_mode train \ --argo_dir /proj/berzelius-2023-154/users/x_qinzh/av2 \ - --output_dir /proj/berzelius-2023-364/users/x_qinzh/data/av2/preprocess_v2 + --output_dir /proj/berzelius-2023-364/users/x_qinzh/data/av2/h5py -/proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/dataprocess/bin/python dataprocess/extract_av2.py --nproc 64 \ +$PYTHON dataprocess/extract_av2.py --nproc 64 \ --av2_type sensor \ --data_mode val \ --argo_dir /proj/berzelius-2023-154/users/x_qinzh/av2 \ - --output_dir /proj/berzelius-2023-364/users/x_qinzh/data/av2/preprocess_v2 \ + --output_dir /proj/berzelius-2023-364/users/x_qinzh/data/av2/h5py \ --mask_dir /proj/berzelius-2023-154/users/x_qinzh/av2/3d_scene_flow -/proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/dataprocess/bin/python dataprocess/extract_av2.py --nproc 64 \ +$PYTHON dataprocess/extract_av2.py --nproc 64 \ --av2_type sensor \ --data_mode test \ --argo_dir /proj/berzelius-2023-154/users/x_qinzh/av2 \ - --output_dir /proj/berzelius-2023-364/users/x_qinzh/data/av2/preprocess_v2 \ + --output_dir /proj/berzelius-2023-364/users/x_qinzh/data/av2/h5py \ --mask_dir /proj/berzelius-2023-154/users/x_qinzh/av2/3d_scene_flow \ No newline at end of file diff --git a/assets/slurm/1_train.sh b/assets/slurm/1_train.sh index 72c0d92..dd99a8d 100644 --- a/assets/slurm/1_train.sh +++ b/assets/slurm/1_train.sh @@ -7,9 +7,13 @@ #SBATCH --output /proj/berzelius-2023-154/users/x_qinzh/seflow/logs/slurm/%J_seflow.out #SBATCH --error /proj/berzelius-2023-154/users/x_qinzh/seflow/logs/slurm/%J_seflow.err -cd /proj/berzelius-2023-154/users/x_qinzh/seflow +PYTHON=/proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/opensf/bin/python +export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/proj/berzelius-2023-154/users/x_qinzh/mambaforge/lib +cd /proj/berzelius-2023-364/users/x_qinzh/workspace/OpenSceneFlow -SOURCE="/proj/berzelius-2023-154/users/x_qinzh/data/av2/preprocess_v2" + +# ===> to transfer data into local node disk, it can be ignored. <=== +SOURCE="/proj/berzelius-2023-364/users/x_qinzh/data/av2/autolabel" DEST="/scratch/local/av2" SUBDIRS=("sensor/train" "sensor/val") @@ -24,14 +28,7 @@ elapsed=$((end_time - start_time)) echo "Copy ${SOURCE} to ${DEST} Total time: ${elapsed} seconds" echo "Start training..." -# ====> paper model = seflow_official -# /proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/seflow/bin/python train.py \ -# slurm_id=$SLURM_JOB_ID wandb_mode=online train_data=/scratch/local/av2/sensor/train val_data=/scratch/local/av2/sensor/val \ -# num_workers=16 model=deflow lr=2e-6 epochs=50 batch_size=20 "model.target.num_iters=2" "model.val_monitor=val/Dynamic/Mean" \ -# loss_fn=seflowLoss "add_seloss={chamfer_dis: 1.0, static_flow_loss: 1.0, dynamic_chamfer_dis: 1.0, cluster_based_pc0pc1: 1.0}" - # ====> leaderboard model = seflow_best -/proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/seflow/bin/python train.py \ - slurm_id=$SLURM_JOB_ID wandb_mode=online train_data=/scratch/local/av2/sensor/train val_data=/scratch/local/av2/sensor/val \ +$PYTHON train.py slurm_id=$SLURM_JOB_ID wandb_mode=online train_data=/scratch/local/av2/sensor/train val_data=/scratch/local/av2/sensor/val \ num_workers=16 model=deflow lr=2e-4 epochs=9 batch_size=16 "model.target.num_iters=2" "model.val_monitor=val/Dynamic/Mean" \ loss_fn=seflowLoss "add_seloss={chamfer_dis: 1.0, static_flow_loss: 1.0, dynamic_chamfer_dis: 1.0, cluster_based_pc0pc1: 1.0}" diff --git a/assets/slurm/2_eval.sh b/assets/slurm/2_eval.sh index 30684cd..1a57440 100644 --- a/assets/slurm/2_eval.sh +++ b/assets/slurm/2_eval.sh @@ -5,29 +5,16 @@ #SBATCH --output /proj/berzelius-2023-154/users/x_qinzh/seflow/logs/slurm/%J_eval.out #SBATCH --error /proj/berzelius-2023-154/users/x_qinzh/seflow/logs/slurm/%J_eval.err -cd /proj/berzelius-2023-154/users/x_qinzh/seflow -SOURCE="/proj/berzelius-2023-154/users/x_qinzh/av2/preprocess_v2" -DEST="/scratch/local/av2" -SUBDIRS=("sensor/val") +PYTHON=/proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/opensf/bin/python +export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/proj/berzelius-2023-154/users/x_qinzh/mambaforge/lib +cd /proj/berzelius-2023-364/users/x_qinzh/workspace/OpenSceneFlow -start_time=$(date +%s) -for dir in "${SUBDIRS[@]}"; do - mkdir -p "${DEST}/${dir}" - find "${SOURCE}/${dir}" -type f -print0 | xargs -0 -n1 -P16 cp -t "${DEST}/${dir}" & -done -wait -end_time=$(date +%s) -elapsed=$((end_time - start_time)) -echo "Copy ${SOURCE} to ${DEST} Total time: ${elapsed} seconds" -echo "Start training..." # ====> leaderboard model -# /proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/seflow/bin/python eval.py \ -# wandb_mode=online dataset_path=/scratch/local/av2/sensor \ +# $PYTHON eval.py wandb_mode=online dataset_path=/proj/berzelius-2023-364/users/x_qinzh/data/av2/autolabel data_mode=test \ # checkpoint=/proj/berzelius-2023-154/users/x_qinzh/seflow/logs/wandb/seflow-10086990/checkpoints/epoch_19_seflow.ckpt \ -# av2_mode=test save_res=True +# save_res=True -/proj/berzelius-2023-154/users/x_qinzh/mambaforge/envs/seflow/bin/python eval.py \ - wandb_mode=online dataset_path=/scratch/local/av2/sensor av2_mode=val \ +$PYTHON eval.py wandb_mode=online dataset_path=/proj/berzelius-2023-364/users/x_qinzh/data/av2/autolabel data_mode=val \ checkpoint=/proj/berzelius-2023-154/users/x_qinzh/seflow/logs/wandb/seflow-10086990/checkpoints/epoch_19_seflow.ckpt \ No newline at end of file diff --git a/conf/config.yaml b/conf/config.yaml index 4a80895..a9e68f0 100644 --- a/conf/config.yaml +++ b/conf/config.yaml @@ -26,10 +26,12 @@ epochs: 3 gradient_clip_val: 5.0 # optimizer ==> Adam -lr: 2e-6 -lr_scheduler: None # choices: [cosine, step, linear] -loss_fn: deflowLoss # choices: [ff3dLoss, zeroflowLoss, deflowLoss, seflowLoss] -add_seloss: # {chamfer_dis: 1.0, static_flow_loss: 1.0, dynamic_chamfer_dis: 1.0, cluster_based_pc0pc1: 1.0} +optimizer: + name: Adam # [Adam, AdamW] + lr: 1e-4 +loss_fn: seflowLoss # choices: [ff3dLoss, zeroflowLoss, deflowLoss, seflowLoss] +# add_seloss: {chamfer_dis: 1.0, static_flow_loss: 1.0, dynamic_chamfer_dis: 1.0, cluster_based_pc0pc1: 1.0} +# ssl_label: # log settings seed: 42069 diff --git a/conf/eval.yaml b/conf/eval.yaml index 2c65dbe..a6b35db 100644 --- a/conf/eval.yaml +++ b/conf/eval.yaml @@ -1,7 +1,7 @@ dataset_path: /home/kin/data/av2/h5py/sensor checkpoint: /home/kin/model_zoo/deflow.ckpt -av2_mode: val # [val, test] +data_mode: val # [val, test] save_res: False # [True, False] leaderboard_version: 1 # [1, 2] diff --git a/conf/ground/zod.toml b/conf/ground/zod.toml new file mode 100644 index 0000000..ab5e85b --- /dev/null +++ b/conf/ground/zod.toml @@ -0,0 +1,31 @@ + +[important] +height = 2.1 # sensor height. 雷达传感器高度，主要是+offset + +# 整个雷点点云以自己为中心 分为多少个segment，每个segment又分成多少个bin +[segments] +r_min = 0.1 # minimum point distance. 感兴趣的区域 +r_max = 100 # maximum point distance. +n_segments = 360 # number of radial segments. +n_bins = 160 # number of radial bins. + + +[ground] +# 直线的 slope 斜率，一个seg里拟合的线斜率 +min_slope = 0.0 # minimum slope of a ground line. [T_m_small] +max_slope = 0.1 # maximum slope of a ground line. [T_m] +# 如果拟合点超过这个 long threshold，判断其 height 变化是否超过 max_long_height +long_threshold = 2.0 # Distance at which points are considered far from each other. [T_d_prev] +max_long_height = 0.2 # maximum height change to previous point in long line. [T_b] +# 如果没有超过 long threshold，判断现在拟合的点与fitline高度变化是否超过 max_start_height +max_start_height = 0.2 # Maximum heigh of starting line to be labelled ground. +# 最大的拟合误差，超过这个误差的点不会被加入 line 拟合 +max_fit_error = 0.1 # maximum error of a point during line fit. [T_RMSE: will sqaure inside code.] + +# 判断是否是地面点的时候，需要满足的条件： +max_dist_to_line = 0.2 # maximum vertical distance of point to line to be considered ground. [T_d_ground] +line_search_angle = 2.0 # How far to search for a line in angular direction [rad]. + +[general] +n_threads = 8 # number of threads for parallel processing. +verbose = false # if you don't want to see every output, set this to false. \ No newline at end of file diff --git a/conf/model/deflowpp.yaml b/conf/model/deflowpp.yaml new file mode 100644 index 0000000..c896f18 --- /dev/null +++ b/conf/model/deflowpp.yaml @@ -0,0 +1,11 @@ +name: deflowpp + +target: + _target_: src.models.DeFlowPP + decoder_option: gru # choices: [linear, gru] + num_iters: 2 + voxel_size: ${voxel_size} + point_cloud_range: ${point_cloud_range} + num_frames: ${num_frames} + +val_monitor: val/Dynamic/Mean \ No newline at end of file diff --git a/conf/save.yaml b/conf/save.yaml index 0332482..5b2d5ad 100644 --- a/conf/save.yaml +++ b/conf/save.yaml @@ -1,4 +1,4 @@ -dataset_path: /home/kin/data/av2/h5py/demo/sensor/val +dataset_path: /home/kin/data/av2/h5py/demo/val checkpoint: /home/kin/model_zoo/seflow_best.ckpt res_name: # if None will directly be the `model_name.ckpt` in checkpoint path diff --git a/dataprocess/README.md b/dataprocess/README.md index 893af64..2651afa 100644 --- a/dataprocess/README.md +++ b/dataprocess/README.md @@ -11,8 +11,8 @@ We've updated the process dataset for: - [x] Argoverse 2.0: check [here](#argoverse-20). The process script Involved from [DeFlow](https://github.com/KTH-RPL/DeFlow). - [x] Waymo: check [here](#waymo-dataset). The process script was involved from [SeFlow](https://github.com/KTH-RPL/SeFlow). - [x] nuScenes: check [here](#nuscenes), The process script was involved from [DeltaFlow](https://github.com/Kin-Zhang/DeltaFlow). +- [x] ZOD (w/o gt): check [here](#zod-dataset). The process script was involved from [HiMo](https://kin-zhang.github.io/HiMo). (It could be a good first reference for users to extract other datasets in the future.) - [ ] TruckScene: done coding, public after review. Will be involved later by another paper. -- [ ] ZOD (w/o gt): done coding, public after review. Will be involved later by another paper. If you want to **use all datasets above**, there is a **specific environment** in [envsftool.yaml](../envsftool.yaml) to install all the necessary packages. As Waymo package have different configuration and conflict with the main environment. Setup through the following command: @@ -123,20 +123,35 @@ tar -xvf waymo_map.tar.gz -C /home/kin/data/waymo/flowlabel | train | 799 | 155687 | | val | 203 | 39381 | +### ZOD Dataset + +Although ZOD have the most dense LiDAR sensor (128-channel), the dataset itself **does not include ground truth flow**. +We provide the extraction script for Self-Supervised Learning (SSL) to train and visualize the results etc, again **no evaluation available** here. + +To download the ZOD dataset, you need follow [the instruction here](https://zod.zenseact.com/download/): send email and ask for the download from the team. + +For HiMo, we only downloaded [drives-set](https://zod.zenseact.com/drives/) for test purpose etc. The total drives-set includes 29 sequences (Total size: 303G). Here are [quick video play](https://www.bilibili.com/video/BV1Sh4y1z7v2) for each scene in the drives. + +Please check the scripts: [dataprocess/extract_zod.py](./extract_zod.py) in detail, current we only process one scene while feel free to comment out for all scene etc. + + + ## Process -This directory contains the scripts to preprocess the datasets. + +This directory contains the scripts to preprocess the datasets into `.h5` files. - `extract_av2.py`: Process the datasets in Argoverse 2.0. - `extract_nus.py`: Process the datasets in nuScenes. - `extract_waymo.py`: Process the datasets in Waymo. +- `extract_zod.py`: Process the datasets in ZOD. Example Running command: ```bash # av2: -python dataprocess/extract_av2.py --av2_type sensor --data_mode train --argo_dir /home/kin/data/av2 --output_dir /home/kin/data/av2/preprocess +python dataprocess/extract_av2.py --av2_type sensor --data_mode train --argo_dir /home/kin/data/av2 --output_dir /home/kin/data/av2/h5py # waymo: -python dataprocess/extract_waymo.py --mode train --flow_data_dir /home/kin/data/waymo/flowlabel --map_dir /home/kin/data/waymo/flowlabel/map --output_dir /home/kin/data/waymo/preprocess --nproc 48 +python dataprocess/extract_waymo.py --mode train --flow_data_dir /home/kin/data/waymo/flowlabel --map_dir /home/kin/data/waymo/flowlabel/map --output_dir /home/kin/data/waymo/h5py --nproc 48 # nus: python dataprocess/extract_nus.py --mode v1.0-trainval --output_dir /home/kin/data/nus/h5py/full --nproc 24 @@ -146,7 +161,7 @@ python dataprocess/extract_nus.py --mode v1.0-trainval --output_dir /home/kin/da Extract all Argoverse 2.0 data to unified `.h5` format. [Runtime: Normally need 45 mins finished run following commands totally in setup mentioned in our paper] ```bash -python dataprocess/extract_av2.py --av2_type sensor --data_mode train --argo_dir /home/kin/data/av2 --output_dir /home/kin/data/av2/preprocess_v2 +python dataprocess/extract_av2.py --av2_type sensor --data_mode train --argo_dir /home/kin/data/av2 --output_dir /home/kin/data/av2/h5py python dataprocess/extract_av2.py --av2_type sensor --data_mode val --mask_dir /home/kin/data/av2/3d_scene_flow python dataprocess/extract_av2.py --av2_type sensor --data_mode test --mask_dir /home/kin/data/av2/3d_scene_flow ``` @@ -175,9 +190,9 @@ Or you can run testing file to visualize the data. ```bash # view gt flow -python tools/visualization.py --data_dir /home/kin/data/av2/preprocess/sensor/mini --res_name flow +python tools/visualization.py --data_dir /home/kin/data/av2/h5py/sensor/mini --res_name flow -python tools/visualization.py --data_dir /home/kin/data/waymo/preprocess/val --res_name flow +python tools/visualization.py --data_dir /home/kin/data/waymo/h5py/val --res_name flow ``` ### Self-Supervised Process @@ -186,5 +201,5 @@ Process train data for self-supervised learning. Only training data needs this s [Runtime: Normally need 15 hours for my desktop, 3 hours for the cluster with five available nodes parallel running.] ```bash -python process.py --data_dir /home/kin/data/av2/preprocess_v2/sensor/train --scene_range 0,701 +python process.py --data_dir /home/kin/data/av2/h5py/sensor/train --scene_range 0,701 ``` diff --git a/dataprocess/extract_av2.py b/dataprocess/extract_av2.py index 188b03c..de45acb 100644 --- a/dataprocess/extract_av2.py +++ b/dataprocess/extract_av2.py @@ -15,6 +15,8 @@ import os os.environ["OMP_NUM_THREADS"] = "1" +import warnings +warnings.filterwarnings("ignore", category=UserWarning) from av2.datasets.sensor.av2_sensor_dataloader import convert_pose_dataframe_to_SE3 from av2.structures.sweep import Sweep @@ -40,7 +42,7 @@ import os, sys BASE_DIR = os.path.abspath(os.path.join( os.path.dirname( __file__ ), '..' )) sys.path.append(BASE_DIR) -from dataprocess.misc_data import create_reading_index +from dataprocess.misc_data import create_reading_index, check_h5py_file_exists from src.utils.av2_eval import read_ego_SE3_sensor BOUNDING_BOX_EXPANSION: Final = 0.2 @@ -92,15 +94,23 @@ def read_pose_pc_ground(data_dir: Path, log_id: str, timestamp: int, avm: Argove ego2sensor_pose = read_ego_SE3_sensor((data_dir / log_id))['up_lidar'] filtered_log_poses_df = log_poses_df[log_poses_df["timestamp_ns"].isin([timestamp])] pose = convert_pose_dataframe_to_SE3(filtered_log_poses_df.loc[filtered_log_poses_df["timestamp_ns"] == timestamp]) - pc = Sweep.from_feather(data_dir / log_id / "sensors" / "lidar" / f"{timestamp}.feather").xyz + av2_sweep = Sweep.from_feather(data_dir / log_id / "sensors" / "lidar" / f"{timestamp}.feather") + pc = av2_sweep.xyz + + # ref: https://github.com/argoverse/av2-api/issues/77 + lidar_id = np.zeros(len(pc), dtype=np.uint8) + lidar_id[av2_sweep.laser_number < 32] = 1 + lidar_id[av2_sweep.laser_number >= 32] = 2 + lidar_dt = av2_sweep.offset_ns / 1e9 # to second + # transform to city coordinate since sweeps[0].xyz is in ego coordinate to get ground mask is_ground = avm.get_ground_points_boolean(pose.transform_point_cloud(pc)) # NOTE(SeFlow): transform to sensor coordinate, since some ray-casting based methods need sensor coordinate pc = ego2sensor_pose.inverse().transform_point_cloud(pc) - return pc, pose, is_ground + return pc, lidar_id, lidar_dt, pose, is_ground -def compute_sceneflow(data_dir: Path, log_id: str, timestamps: Tuple[int, int]) -> Dict[str, Union[np.ndarray, SE3]]: +def compute_sceneflow(data_dir: Path, log_id: str, timestamps: Tuple[int, int], dclass) -> Dict[str, Union[np.ndarray, SE3]]: """Compute sceneflow between the sweeps at the given timestamps. Args: data_dir: Argoverse 2.0 directory, e.g. /home/kin/data/av2/sensor/train @@ -127,12 +137,11 @@ def compute_flow(sweeps, cuboids, poses): ego1_SE3_ego0.translation = ego1_SE3_ego0.translation.astype(np.float32) flow = ego1_SE3_ego0.transform_point_cloud(sweeps[0].xyz) - sweeps[0].xyz - # Convert to float32s flow = flow.astype(np.float32) valid = np.ones(len(sweeps[0].xyz), dtype=np.bool_) - # classes = -np.ones(len(sweeps[0].xyz), dtype=np.int8) classes = np.zeros(len(sweeps[0].xyz), dtype=np.uint8) + instances = np.zeros(len(sweeps[0].xyz), dtype=np.int16) # # old version # for id in cuboids[0]: @@ -173,9 +182,10 @@ def compute_flow(sweeps, cuboids, poses): obj_flow = c1_SE3_c0.transform_point_cloud(obj_pts) - obj_pts classes[obj_mask] = CATEGORY_TO_INDEX[str(c0.category)] flow[obj_mask] = obj_flow.astype(np.float32) + instances[obj_mask] = dclass[id]+1 else: valid[obj_mask] = 0 - return flow, classes, valid, ego1_SE3_ego0 + return flow, classes, valid, ego1_SE3_ego0, instances sweeps = [Sweep.from_feather(data_dir / log_id / "sensors" / "lidar" / f"{ts}.feather") for ts in timestamps] # ================== Load annotations ================== @@ -203,24 +213,28 @@ def compute_flow(sweeps, cuboids, poses): filtered_log_poses_df = log_poses_df[log_poses_df["timestamp_ns"].isin(timestamps)] poses = [convert_pose_dataframe_to_SE3(filtered_log_poses_df.loc[filtered_log_poses_df["timestamp_ns"] == ts]) for ts in timestamps] - flow_0_1, classes_0, valid_0, ego_motion = compute_flow(sweeps, cuboids, poses) + flow_0_1, classes_0, valid_0, ego_motion, instances = compute_flow(sweeps, cuboids, poses) return {'pcl_0': sweeps[0].xyz, 'pcl_1' :sweeps[1].xyz, 'flow_0_1': flow_0_1, 'valid_0': valid_0, 'classes_0': classes_0, - 'pose_0': poses[0], 'pose_1': poses[1], + 'pose_0': poses[0], 'pose_1': poses[1], 'instances': instances, 'ego_motion': ego_motion} def process_log(data_dir: Path, log_id: str, output_dir: Path, n: Optional[int] = None) : - def create_group_data(group, pc, gm, pose, flow_0to1=None, flow_valid=None, flow_category=None, ego_motion=None): + def create_group_data(group, pc, pc_id, pc_dt, gm, pose, flow_0to1=None, flow_valid=None, flow_category=None, flow_instance=None, ego_motion=None): group.create_dataset('lidar', data=pc.astype(np.float32)) group.create_dataset('ground_mask', data=gm.astype(bool)) group.create_dataset('pose', data=pose.astype(np.float32)) + # lidar_id for visualization and lidar_dt for HiMo mainly: + group.create_dataset('lidar_id', data=pc_id.astype(np.uint8)) # sensor id + group.create_dataset('lidar_dt', data=pc_dt.astype(np.float32)) # deltaT if flow_0to1 is not None: # ground truth flow information group.create_dataset('flow', data=flow_0to1.astype(np.float32)) group.create_dataset('flow_is_valid', data=flow_valid.astype(bool)) group.create_dataset('flow_category_indices', data=flow_category.astype(np.uint8)) + group.create_dataset('flow_instance_id', data=flow_instance.astype(np.int16)) group.create_dataset('ego_motion', data=ego_motion.astype(np.float32)) log_map_dirpath = data_dir / log_id / "map" @@ -235,29 +249,31 @@ def create_group_data(group, pc, gm, pose, flow_0to1=None, flow_valid=None, flow for file in os.listdir(data_dir / log_id / "sensors/lidar") if file.endswith('.feather')]) + gt_flow_flag = False if not (data_dir / log_id / "annotations.feather").exists() else True - + if check_h5py_file_exists(output_dir/f'{log_id}.h5', timestamps): + return # if n is not None: # iter_bar = tqdm(zip(timestamps, timestamps[1:]), leave=False, # total=len(timestamps) - 1, position=n, # desc=f'Log {log_id}') # else: # iter_bar = zip(timestamps, timestamps[1:]) - + dclass = defaultdict(lambda: len(dclass)) with h5py.File(output_dir/f'{log_id}.h5', 'a') as f: for cnt, ts0 in enumerate(timestamps): group = f.create_group(str(ts0)) - pc0, pose0, is_ground_0 = read_pose_pc_ground(data_dir, log_id, ts0, avm) + pc0, lidar_id0, lidar_dt0, pose0, is_ground_0 = read_pose_pc_ground(data_dir, log_id, ts0, avm) if pc0.shape[0] < 256: print(f'{log_id}/{ts0} has less than 256 points, skip this scenarios. Please check the data if needed.') break if cnt == len(timestamps) - 1 or not gt_flow_flag: - create_group_data(group, pc0, is_ground_0.astype(np.bool_), pose0.transform_matrix.astype(np.float32)) + create_group_data(group, pc0, lidar_id0, lidar_dt0, is_ground_0.astype(np.bool_), pose0.transform_matrix.astype(np.float32)) else: ts1 = timestamps[cnt + 1] - scene_flow = compute_sceneflow(data_dir, log_id, (ts0, ts1)) - create_group_data(group, pc0, is_ground_0.astype(np.bool_), pose0.transform_matrix.astype(np.float32), - scene_flow['flow_0_1'], scene_flow['valid_0'], scene_flow['classes_0'], + scene_flow = compute_sceneflow(data_dir, log_id, (ts0, ts1), dclass) + create_group_data(group, pc0, lidar_id0, lidar_dt0, is_ground_0.astype(np.bool_), pose0.transform_matrix.astype(np.float32), + scene_flow['flow_0_1'], scene_flow['valid_0'], scene_flow['classes_0'], scene_flow['instances'], scene_flow['ego_motion'].transform_matrix.astype(np.float32)) def proc(x, ignore_current_process=False): diff --git a/dataprocess/extract_zod.py b/dataprocess/extract_zod.py new file mode 100644 index 0000000..d60ecdd --- /dev/null +++ b/dataprocess/extract_zod.py @@ -0,0 +1,140 @@ +""" +# Created: 2024-07-07 22:18 +# Copyright (C) 2024-now, Scania Sverige EEARP Group +# Author: Qingwen Zhang (https://kin-zhang.github.io/) +# +# License: GPLv2, allow it free only for academic use. +# Description: Preprocess Data, save as h5df format for faster loading +""" + +import fire, time, os, sys, json, h5py +from pathlib import Path +from multiprocessing import current_process +from tqdm import tqdm +import multiprocessing +import numpy as np +from typing import Optional + +# import zod +from zod.constants import Camera, Lidar +from zod.data_classes.sequence import ZodSequence +from zod._zod_dataset import _create_frame + +BASE_DIR = os.path.abspath(os.path.join( os.path.dirname( __file__ ), '..')) +sys.path.append(BASE_DIR) +from dataprocess.misc_data import create_reading_index + +from linefit import ground_seg +GROUNDSEG_config = f"{BASE_DIR}/conf/ground/zod.toml" + +def process_log(data_dir: Path, log_id: str, output_dir: Path, n: Optional[int] = None) : + + def create_group_data(group, pc, pc_dt, pose, pc_id, gm=None, flow_0to1=None, flow_valid=None, flow_category=None, ego_motion=None): + group.create_dataset('lidar', data=pc.astype(np.float32)) + group.create_dataset('lidar_id', data=pc_id.astype(np.uint8)) # sensor id + group.create_dataset('lidar_dt', data=pc_dt.astype(np.float32)) # deltaT + group.create_dataset('pose', data=pose.astype(np.float32)) + if gm is not None: + group.create_dataset('ground_mask', data=gm.astype(np.bool_)) + if flow_0to1 is not None: + # ground truth flow information + group.create_dataset('flow', data=flow_0to1.astype(np.float32)) + group.create_dataset('flow_is_valid', data=flow_valid.astype(bool)) + group.create_dataset('flow_category_indices', data=flow_category.astype(np.uint8)) + group.create_dataset('ego_motion', data=ego_motion.astype(np.float32)) + + with open(data_dir / log_id / 'info.json', "r") as f: + frames = _create_frame(json.load(f), data_dir.as_posix().replace('drives', '')) + seq = ZodSequence(frames) + mygroundseg = ground_seg(GROUNDSEG_config) + with h5py.File(output_dir/f'{log_id}.h5', 'a') as f: + for cnt, frame in enumerate(seq.info.get_camera_lidar_map()): + + # NOTE(Qingwen): for a quick check downsampled data.. otherwise one h5py might have really long frames. + # FIXME(Qingwen): maybe split too long scene to mini-scene etc later? + if cnt % 4 != 0: + continue + + camera_frame, lidar_frame = frame + # img = camera_frame.read() + pcd = seq.get_compensated_lidar(camera_frame.time) + pose = seq.ego_motion.get_poses(camera_frame.time.timestamp()) + + diode_idx = pcd.diode_idx + lidar_id = np.zeros(diode_idx.shape, dtype=np.uint8) + + lidar_id[(diode_idx >= 0) & (diode_idx < 128)] = 1 + lidar_id[(diode_idx >= 128) & (diode_idx < 144)] = 2 + lidar_id[(diode_idx >= 144) & (diode_idx < 160)] = 3 + ego2sensor = seq._calibration.lidars[Lidar.VELODYNE].extrinsics + pose = pose @ ego2sensor.transform + lidar_timestamp = pcd.core_timestamp + lidar_dt = pcd.timestamps - lidar_timestamp + group = f.create_group(str(int(camera_frame.time.timestamp() * 10e6))) + + # NOTE(Qingwen): only need 128-channel long-range lidar, while if you want feel free to comment these. + # points = pcd.points # for all three lidars (1 long, 2 really short) + points = pcd.points[lidar_id == 1] + lidar_dt = lidar_dt[lidar_id == 1] + lidar_id = lidar_id[lidar_id == 1] + ground_mask = mygroundseg.run(points[:, :3]) + + create_group_data(group, points, lidar_dt, pose.astype(np.float32), lidar_id, gm=np.array(ground_mask, dtype=np.bool_)) + +def proc(x, ignore_current_process=False): + if not ignore_current_process: + current=current_process() + pos = current._identity[0] + else: + pos = 1 + process_log(*x, n=pos) + +def process_logs(data_dir: Path, output_dir: Path, nproc: int): + """Compute sceneflow for all logs in the dataset. Logs are processed in parallel. + Args: + data_dir: Argoverse 2.0 directory + output_dir: Output directory. + """ + + if not data_dir.exists(): + print(f'{data_dir} not found') + return + + # NOTE(Qingwen): if you don't want to all data_dir, then change here: logs = logs[:10] only 10 scene. + logs = os.listdir(data_dir) + # like here only 000018 scene. + logs = ['000018'] + args = sorted([(data_dir, log, output_dir) for log in logs]) + print(f'Using {nproc} processes to process data: {data_dir} to .h5 format. (#scenes: {len(args)})') + # for debug + for x in tqdm(args): + print(x) + proc(x, ignore_current_process=True) + break + + # comment out if you want to process all scene. + # if nproc <= 1: + # for x in tqdm(args, ncols=120): + # proc(x, ignore_current_process=True) + # else: + # with Pool(processes=nproc) as p: + # res = list(tqdm(p.imap_unordered(proc, args), total=len(logs), ncols=120)) + +def main( + dataset_root: str = "/home/kin/DATA_HDD/public_data/zod/drives", + output_dir: str ="/home/kin/data/zod/h5py/himo", + nproc: int = (multiprocessing.cpu_count() - 1), + only_index: bool = False, +): + output_dir_ = Path(output_dir) + if only_index: + create_reading_index(output_dir_) + return + output_dir_.mkdir(exist_ok=True, parents=True) + process_logs(Path(dataset_root), output_dir_, nproc) + create_reading_index(output_dir_) + +if __name__ == '__main__': + start_time = time.time() + fire.Fire(main) + print(f"\nRunning {__file__} used: {(time.time() - start_time)/60:.2f} mins") \ No newline at end of file diff --git a/environment.yaml b/environment.yaml index d520844..32212d4 100644 --- a/environment.yaml +++ b/environment.yaml @@ -3,13 +3,12 @@ channels: - conda-forge - pytorch dependencies: - - python=3.10 + - python=3.8 - pytorch::pytorch=2.0.0 - pytorch::torchvision - - pytorch::pytorch-cuda=11.8 - - conda-forge::lightning=2.0.1 - - nvidia/label/cuda-11.8.0::cuda - - nvidia/label/cuda-11.8.0::cuda-toolkit + - pytorch::pytorch-cuda=11.7 + - nvidia/label/cuda-11.7.0::cuda + - lightning==2.0.1 - mkl==2024.0.0 - tensorboard - numba @@ -24,7 +23,7 @@ dependencies: - hydra-core - fire - tabulate - - scikit-learn + - scikit-learn==1.3.2 - hdbscan - setuptools==69.5.1 - gxx_linux-64==11.4.0 @@ -34,18 +33,20 @@ dependencies: - assets/cuda/chamfer3D - assets/cuda/mmcv - open3d==0.18.0 - - av2==0.3.1 - - spconv-cu118==2.3.6 - - numpy==1.26.4 + - av2==0.2.1 + - spconv-cu117==2.3.6 + # Qingwen's public pkg: - dztimer - dufomap==1.1.0 - linefit==1.1.0 - + # Reason about the version fixed: # setuptools==68.5.1: https://github.com/aws-neuron/aws-neuron-sdk/issues/893 # mkl==2024.0.0: https://github.com/pytorch/pytorch/issues/123097#issue-2218541307 # av2==0.2.1: in case other version deleted some functions. # lightning==2.0.1: https://stackoverflow.com/questions/76647518/how-to-fix-error-cannot-import-name-modelmetaclass-from-pydantic-main # open3d==0.18.0: because 0.17.0 have bug on set the view json file -# dufomap==1.1.0: in case later updating may not compatible with the code. +# dufomap==1.0.0: in case later updating may not compatible with the code. +# spconv-cu117==2.3.6: avoid error: KeyError: ((16, 8, 8), float, float) +# torch-scatter==2.1.2: in case later updating may not compatible with the code. \ No newline at end of file diff --git a/envsftool.yaml b/envsftool.yaml index 8ac1925..887457a 100644 --- a/envsftool.yaml +++ b/envsftool.yaml @@ -6,6 +6,8 @@ dependencies: - python=3.8 - pytorch::pytorch=2.0.0 - pytorch::torchvision + - pytorch::pytorch-cuda=11.7 + - nvidia/label/cuda-11.7.0::cuda - mkl==2024.0.0 - numba - numpy==1.22 @@ -13,9 +15,9 @@ dependencies: - pip - scipy - tqdm - - scikit-learn + - scikit-learn=1.2.2 - fire - - hdbscan + - hdbscan=0.8.29 - s5cmd - pip: - nuscenes-devkit diff --git a/eval.py b/eval.py index bfe7bdd..bdceb08 100644 --- a/eval.py +++ b/eval.py @@ -21,8 +21,8 @@ from src.trainer import ModelWrapper def precheck_cfg_valid(cfg): - if os.path.exists(cfg.dataset_path + f"/{cfg.av2_mode}") is False: - raise ValueError(f"Dataset {cfg.dataset_path}/{cfg.av2_mode} does not exist. Please check the path.") + if os.path.exists(cfg.dataset_path + f"/{cfg.data_mode}") is False: + raise ValueError(f"Dataset {cfg.dataset_path}/{cfg.data_mode} does not exist. Please check the path.") if cfg.supervised_flag not in [True, False]: raise ValueError(f"Supervised flag {cfg.supervised_flag} is not valid. Please set it to True or False.") if cfg.leaderboard_version not in [1, 2]: @@ -37,7 +37,7 @@ def main(cfg): if 'iter_only' in cfg.model and cfg.model.iter_only: from src.runner import launch_runner print(f"---LOG[eval]: Run optmization-based method: {cfg.model.name}") - launch_runner(cfg, cfg.av2_mode) + launch_runner(cfg, cfg.data_mode) return if not os.path.exists(cfg.checkpoint): @@ -46,8 +46,9 @@ def main(cfg): torch_load_ckpt = torch.load(cfg.checkpoint) checkpoint_params = DictConfig(torch_load_ckpt["hyper_parameters"]) - cfg.output = checkpoint_params.cfg.output + f"-e{torch_load_ckpt['epoch']}-{cfg.av2_mode}-v{cfg.leaderboard_version}" + cfg.output = checkpoint_params.cfg.output + f"-e{torch_load_ckpt['epoch']}-{cfg.data_mode}-v{cfg.leaderboard_version}" cfg.model.update(checkpoint_params.cfg.model) + cfg.num_frames = cfg.model.target.get('num_frames', checkpoint_params.cfg.get('num_frames', cfg.get('num_frames', 2))) mymodel = ModelWrapper.load_from_checkpoint(cfg.checkpoint, cfg=cfg, eval=True) print(f"\n---LOG[eval]: Loaded model from {cfg.checkpoint}. The backbone network is {checkpoint_params.cfg.model.name}.\n") @@ -62,8 +63,8 @@ def main(cfg): # NOTE(Qingwen): search & check: def eval_only_step_(self, batch, res_dict) trainer.validate(model = mymodel, \ dataloaders = DataLoader( \ - HDF5Dataset(cfg.dataset_path + f"/{cfg.av2_mode}", \ - n_frames=checkpoint_params.cfg.num_frames if 'num_frames' in checkpoint_params.cfg else 2, \ + HDF5Dataset(cfg.dataset_path + f"/{cfg.data_mode}", \ + n_frames=cfg.num_frames, \ eval=True, leaderboard_version=cfg.leaderboard_version), \ batch_size=1, shuffle=False)) wandb.finish() diff --git a/process.py b/process.py index 076747d..90a6851 100644 --- a/process.py +++ b/process.py @@ -3,35 +3,60 @@ # Copyright (C) 2023-now, RPL, KTH Royal Institute of Technology # Author: Qingwen Zhang (https://kin-zhang.github.io/) # -# This file is part of SeFlow (https://github.com/KTH-RPL/SeFlow). +# This file is part of +# * HiMo (https://kin-zhang.github.io/HiMo). # If you find this repo helpful, please cite the respective publication as # listed on the above website. # -# Description: run dufomap on the dataset we preprocessed for afterward ssl training. -# it's only needed for ssl train but not inference. -# Goal to segment dynamic and static point roughly. +# Description: follow seflow & seflow++ idea but more to run: +# (a) dufomap; (b) linefit; (c) hdbscan; (d) nnd. +# """ from pathlib import Path from tqdm import tqdm import numpy as np -import fire, time, h5py, os +import fire, time, h5py, os, sys from hdbscan import HDBSCAN +from src.utils import npcal_pose0to1 from src.utils.mics import HDF5Data, transform_to_array from dufomap import dufomap +from linefit import ground_seg MIN_AXIS_RANGE = 2 # HARD CODED: remove ego vehicle points MAX_AXIS_RANGE = 50 # HARD CODED: remove far away points +BASE_DIR = os.path.abspath(os.path.join( os.path.dirname( __file__ ))) + +def check_data_key(filekey, keyname, scene_id, ts): + """ + Checks if a specified key exists in the given file-like object and deletes it if present. + + This is typically used to remove old or outdated data before overwriting with new data. + + Args: + filekey: The file-like object (e.g., HDF5 group or dictionary) to check for the key. + keyname: The name of the key to check and potentially delete. + scene_id: Identifier for the scene (used for logging or warning purposes). + ts: Timestamp or frame identifier (used for logging or warning purposes). + """ + if keyname in filekey: + # print(f"Warning: {scene_id} {ts} has {keyname}, old data will be removed and overwritten.") + del filekey[keyname] -def run_cluster( +def run_dufocluster( data_dir: str ="/home/kin/data/av2/preprocess/sensor/train", scene_range: list = [0, 1], - interval: int = 1, # useless here, just for the same interface args - overwrite: bool = False, + interval: int = 1, # interval frames to run dufomap only + overwrite: bool = True, + + # NOTE (Qingwen): following is for ground segmentation only... at least for code now. + tag: str = "av2", # [nus, zod, man, sca] etc + run_gm: bool = False, # run ground segmentation + min_nnd: float = 0.14, # min nnd distance 1.4m/s pedestrain speed; For Scania data, we set 0.32 here. ): data_path = Path(data_dir) - dataset = HDF5Data(data_path) + dataset = HDF5Data(data_path) # single frame reading. all_scene_ids = list(dataset.scene_id_bounds.keys()) for scene_in_data_index, scene_id in enumerate(all_scene_ids): start_time = time.time() @@ -43,7 +68,7 @@ def run_cluster( with h5py.File(os.path.join(data_path, f'{scene_id}.h5'), 'r+') as f: for ii in range(bounds["min_index"], bounds["max_index"]+1): key = str(dataset[ii]['timestamp']) - if 'label' not in f[key]: + if 'dufocluster' not in f[key]: flag_exist_label = False break if flag_exist_label and not overwrite: @@ -56,35 +81,54 @@ def run_cluster( pc0 = data['pc0'][:,:3] cluster_label = np.zeros(pc0.shape[0], dtype= np.int16) - if "dufo_label" not in data: - print(f"Warning: {scene_id} {data['timestamp']} has no dufo_label, will be skipped. Better to rerun dufomap again in this scene.") + if "dufo" not in data: + print(f"Warning: {scene_id} {data['timestamp']} has no dufo, will be skipped. Better to rerun dufomap again in this scene.") continue - elif data["dufo_label"].sum() < 20: + elif data["dufo"].sum() < 20: print(f"Warning: {scene_id} {data['timestamp']} has no dynamic points, will be skipped. Better to check this scene.") else: - hdb.fit(pc0[data["dufo_label"]==1]) + hdb.fit(pc0[data["dufo"]==1]) # NOTE(Qingwen): since -1 will be assigned if no cluster. We set it to 0. - cluster_label[data["dufo_label"]==1] = hdb.labels_ + 1 + cluster_label[data["dufo"]==1] = hdb.labels_ + 1 # save labels timestamp = data['timestamp'] key = str(timestamp) with h5py.File(os.path.join(data_path, f'{scene_id}.h5'), 'r+') as f: - if 'label' in f[key]: + if 'dufocluster' in f[key]: # print(f"Warning: {scene_id} {timestamp} has label, will be overwritten.") - del f[key]['label'] - f[key].create_dataset('label', data=np.array(cluster_label).astype(np.int16)) + del f[key]['dufocluster'] + f[key].create_dataset('dufocluster', data=np.array(cluster_label).astype(np.int16)) print(f"==> Scene {scene_id} finished, used: {(time.time() - start_time)/60:.2f} mins") - print(f"Data inside {str(data_path)} finished. Check the result with tools/visulization.py if you want to visualize them.") + print(f"Data inside {str(data_path)} finished. Check the result with tools/visulization.py if you want to visualize them.\n") -def run_dufo( +# since it's the only one need cuda to do. in case you want to run two single jobs. +def run_nnd( data_dir: str ="/home/kin/data/av2/preprocess/sensor/train", scene_range: list = [0, 1], - interval: int = 1, # interval frames to run dufomap - overwrite: bool = False, + interval: int = 1, # interval frames to run dufomap only + overwrite: bool = True, + + # NOTE (Qingwen): following is for ground segmentation only... at least for code now. + tag: str = "av2", # [nus, zod, man, sca] etc + run_gm: bool = False, # run ground segmentation + min_nnd: float = 0.14, # min nnd distance 1.4m/s pedestrain speed; For Scania data, we set 0.32 here. ): + # nnd function + from assets.cuda.chamfer3D import nnChamferDis + MyCUDAChamferDis = nnChamferDis() + import torch + # exit() + def cuda_nnd(pc0: torch.tensor, pc1: torch.tensor, moving_threshold=0.14, truncated=4.4): # 4.4 ~= 160km/h * 0.1s + # pc0: (N,3), already ego motion transformed; pc1: (M,3) + pow2_dist0, _= MyCUDAChamferDis.dis_res(pc0, pc1) + pow2_dist0 = pow2_dist0.cpu().numpy() + label = np.zeros(pc0.shape[0], dtype=np.uint8) + label[(pow2_dist0>=pow(moving_threshold,2)) & (pow2_dist0= scene_range[1]): continue bounds = dataset.scene_id_bounds[scene_id] - flag_has_dufo_label = True - with h5py.File(os.path.join(data_path, f'{scene_id}.h5'), 'r+') as f: + exist_dict = {"nnd": True} + with h5py.File(os.path.join(data_path, f'{scene_id}.h5'), 'r') as f: for ii in range(bounds["min_index"], bounds["max_index"]+1): key = str(dataset[ii]['timestamp']) - if "dufo_label" not in f[key]: - flag_has_dufo_label = False + for datakey in exist_dict.keys(): + if datakey not in f[key]: + exist_dict[datakey] = False + if not all(exist_dict.values()): break - if flag_has_dufo_label and not overwrite: - print(f"==> Scene {scene_id} has dufo_label, skip.") + + if all(exist_dict.values()) and not overwrite: + print(f"==> Scene {scene_id} already processed, skip.") continue - mydufo = dufomap(0.2, 0.2, 1, num_threads=12) # resolution, d_s, d_p, hit_extension - mydufo.setCluster(0, 20, 0.2) # depth=0, min_points=20, max_dist=0.2 + for data_id in tqdm(range(bounds["min_index"], bounds["max_index"]+1), desc=f"CUDA nnd run: {scene_in_data_index+1}/{len(all_scene_ids)}", ncols=80): + data = dataset[data_id] + pc0 = data['pc0'][:,:3] + timestamp = data['timestamp'] + key = str(timestamp) - print(f"==> Scene {scene_id} start, data path: {data_path}") - for i in tqdm(range(bounds["min_index"], bounds["max_index"]+1), desc=f"Dufo run: {scene_in_data_index}/{len(all_scene_ids)}", ncols=80): - if interval != 1 and i % interval != 0 and (i + interval//2 < bounds["max_index"] or i - interval//2 > bounds["min_index"]): - continue - data = dataset[i] - assert data['scene_id'] == scene_id, f"Check the data, scene_id {scene_id} is not consistent in {i}th data in {scene_in_data_index}th scene." - # HARD CODED: remove points outside the range - norm_pc0 = np.linalg.norm(data['pc0'][:, :3], axis=1) - range_mask = ( - (norm_pc0>MIN_AXIS_RANGE) & - (norm_pc0= scene_range[1]): + continue + bounds = dataset.scene_id_bounds[scene_id] + # If you don't want to seflowpp label, then remove cluster: True here. It won't process then. + exist_dict = {"dufo": True, "ground_mask": True, "cluster": True} + with h5py.File(os.path.join(data_path, f'{scene_id}.h5'), 'r') as f: + for ii in range(bounds["min_index"], bounds["max_index"]+1): + key = str(dataset[ii]['timestamp']) + for datakey in exist_dict.keys(): + if datakey not in f[key]: + exist_dict[datakey] = False + if not all(exist_dict.values()): + break + + if all(exist_dict.values()) and not overwrite: + print(f"==> Scene {scene_id} already processed, skip.") + continue + + # double check + if not exist_dict["ground_mask"] or run_gm: + mygroundseg = ground_seg(gm_config_path) + elif not exist_dict["ground_mask"] and not run_gm: + raise ValueError("You set run_gm=False, but ground segmentation is not done. Please check the code.") + exist_dict["ground_mask"] = exist_dict["ground_mask"] and not run_gm # and not overwrite # this overwrite is for debug mainly. + + # if overwrite: + # exist_dict["nnd"] = False + # assign all exist_dict to False, so we can run all the process again. + + if 'cluster' in exist_dict and not exist_dict["cluster"]: + hdbscan_cluster = HDBSCAN(min_cluster_size=20, cluster_selection_epsilon=0.7, alpha=1.1) + + if not exist_dict["dufo"]: + mydufo = dufomap(0.2, 0.2, 1, num_threads=12) # resolution, d_s, d_p, hit_extension + mydufo.setCluster(0, 20, 0.2) # depth=0, min_points=20, max_dist=0.2 + + print(f"==> Scene {scene_id} start, data path: {data_path}") + for i in tqdm(range(bounds["min_index"], bounds["max_index"]+1), desc=f"Dufo run: {scene_in_data_index+1}/{len(all_scene_ids)}", ncols=80): + if interval != 1 and i % interval != 0 and (i + interval//2 < bounds["max_index"] or i - interval//2 > bounds["min_index"]): + continue + data = dataset[i] + assert data['scene_id'] == scene_id, f"Check the data, scene_id {scene_id} is not consistent in {i}th data in {scene_in_data_index}th scene." + # HARD CODED: remove points outside the range + norm_pc0 = np.linalg.norm(data['pc0'][:, :3], axis=1) + range_mask = ( + (norm_pc0>MIN_AXIS_RANGE) & + (norm_pc0 Scene {scene_id} finished, used: {(time.time() - start_time)/60:.2f} mins") - print(f"Data inside {str(data_path)} finished. Check the result with vis() function if you want to visualize them.") - + if __name__ == '__main__': start_time = time.time() - # step 1: run dufomap - fire.Fire(run_dufo) - # step 2: run cluster on dufolabel - fire.Fire(run_cluster) - - print(f"\nTime used: {(time.time() - start_time)/60:.2f} mins") \ No newline at end of file + fire.Fire(main) + fire.Fire(run_dufocluster) + print("\nAlready Finished the main labels: dufo, cluster, ground mask etc.\n") + # fire.Fire(run_nnd) + print(f"\nScript Time used: {(time.time() - start_time)/60:.2f} mins") \ No newline at end of file diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index cbce091..0000000 --- a/requirements.txt +++ /dev/null @@ -1,32 +0,0 @@ -# A minimal env for training env only. -# --index-url https://download.pytorch.org/whl/cu118 -# torch==2.0.0 -# torchvision==0.15.1 -# torchaudio==2.0.1 - -lightning -tensorboard -wandb -tqdm -fire -omegaconf -hydra-core -tabulate -pathtools -numpy==1.26.0 -scipy -numba -pandas -scikit-learn -h5py -hdbscan -av2==0.3.4 -spconv-cu118==2.3.8 -flash-attn==2.6.3 -mmengine-lite -torch-scatter==2.1.2 - -# Qingwen's public pkg -dufomap==1.1.0 -linefit==1.1.0 -dztimer diff --git a/save.py b/save.py index b9195e6..5746eed 100644 --- a/save.py +++ b/save.py @@ -14,7 +14,7 @@ import torch from torch.utils.data import DataLoader import lightning.pytorch as pl -from lightning.pytorch.loggers import WandbLogger +from lightning.pytorch.loggers import TensorBoardLogger from omegaconf import DictConfig, OmegaConf import hydra, wandb, os, sys from hydra.core.hydra_config import HydraConfig @@ -45,21 +45,17 @@ def main(cfg): checkpoint_params = DictConfig(torch.load(cfg.checkpoint)["hyper_parameters"]) cfg.output = checkpoint_params.cfg.output cfg.model.update(checkpoint_params.cfg.model) + cfg.num_frames = cfg.model.target.get('num_frames', checkpoint_params.cfg.get('num_frames', cfg.get('num_frames', 2))) mymodel = ModelWrapper.load_from_checkpoint(cfg.checkpoint, cfg=cfg, eval=True) - wandb_logger = WandbLogger(save_dir=output_dir, - entity="kth-rpl", - project=f"deflow-eval", - name=f"{cfg.output}", - offline=True) - - trainer = pl.Trainer(logger=wandb_logger, devices=1) + logger = TensorBoardLogger(save_dir=output_dir, name="logs") + + trainer = pl.Trainer(logger=logger, devices=1) # NOTE(Qingwen): search & check in pl_model.py : def test_step(self, batch, res_dict) trainer.test(model = mymodel, \ dataloaders = DataLoader(\ - HDF5Dataset(cfg.dataset_path, n_frames=checkpoint_params.cfg.num_frames if 'num_frames' in checkpoint_params.cfg else 2), \ + HDF5Dataset(cfg.dataset_path, n_frames=cfg.num_frames), \ batch_size=1, shuffle=False)) - wandb.finish() if __name__ == "__main__": main() \ No newline at end of file diff --git a/src/autolabel.py b/src/autolabel.py new file mode 100644 index 0000000..736bbb6 --- /dev/null +++ b/src/autolabel.py @@ -0,0 +1,62 @@ +""" +# Created: 2024-02-22 22:05 +# Updated: 2025-08-20 21:45 +# Copyright (C) 2024-now, Scania Sverige EEARP Group, KTH Royal Institute of Technology +# Author: Qingwen ZHANG (https://kin-zhang.github.io/) +# License: GPLv2, allow it free only for academic use. +# +# Change Logs: +# 2024-02-22: clean up from HiMo project. +# +# Description: Label strategy for self-supervised learning (SSL) scene flow estimation. +# +""" + +import numpy as np +from copy import deepcopy + +def shiftClusterid(cluster): + """ + Shift the cluster labels by 1: + 0: background, no cluster id + 1: save this label for no cluster_id but dynamic + 2+: cluster id + """ + shifted_cluster = np.zeros_like(cluster) + mask = cluster > 0 + shifted_cluster[mask] = cluster[mask] + 1 + return shifted_cluster + +# mainly based on dufo with label inside for cluster-loss. Check HiMo Fig. 6 Top +def seflow_auto(input_data): + dufo = input_data['dufo'][:].astype(np.uint8) + cluster = shiftClusterid(input_data['dufocluster'][:].astype(np.int16)) + cluster[dufo == 0] = 0 + return cluster + +# based on dufo and nnd for dynamic then cluster-wise checking with reassign. Check HiMo Fig. 6 Bottom +def seflowpp_auto(input_data, tau1=0.05, tau2=0.30): + """ + check HiMo paper (Eq. 5) to know more about paramter setting here. + We didn't explore this parameter too much feel free to adjust as you want after reading the paper. + * For highway Scania data we set tau1=0.01, tau2=0.05 + * For urban Argoverse 2 data we set tau1=0.05 tau2=0.3 + """ + dufo = input_data['dufo'][:].astype(np.uint8) + cluster = shiftClusterid(input_data['cluster'][:].astype(np.int16)) + nnd = input_data['nnd'][:].astype(np.uint8) + + dynamic = np.zeros_like(dufo) + for cluster_id in np.unique(cluster): + if cluster_id in [0, 1]: + continue + all_pts = np.sum(cluster == cluster_id) + cluster_dufo = dufo[cluster == cluster_id] + cluster_nnd = nnd[cluster == cluster_id] + + r_dufo = np.sum(cluster_dufo>0)/all_pts + r_nnd = np.sum(cluster_nnd>0)/all_pts + + if min(r_dufo, r_nnd) > tau1 and max(r_dufo, r_nnd) > tau2: + dynamic[cluster == cluster_id] = cluster_id + return dynamic diff --git a/src/dataset.py b/src/dataset.py index acb514a..778d397 100644 --- a/src/dataset.py +++ b/src/dataset.py @@ -14,14 +14,18 @@ import torch from torch.utils.data import Dataset, DataLoader -import h5py, os, pickle, argparse, sys +import h5py, pickle, argparse from tqdm import tqdm + +import os, sys BASE_DIR = os.path.abspath(os.path.join( os.path.dirname( __file__ ), '..' )) sys.path.append(BASE_DIR) +from src.utils import import_func import warnings warnings.filterwarnings("ignore", category=UserWarning) +# FIXME(Qingwen 2025-08-20): update more pretty here afterward! def collate_fn_pad(batch): num_frames = 2 @@ -73,27 +77,38 @@ def collate_fn_pad(batch): pc1_dynamic_after_mask_ground = torch.nn.utils.rnn.pad_sequence(pc1_dynamic_after_mask_ground, batch_first=True, padding_value=0) res_dict['pc0_dynamic'] = pc0_dynamic_after_mask_ground res_dict['pc1_dynamic'] = pc1_dynamic_after_mask_ground + if 'pch1_dynamic' in batch[0]: + pch1_dynamic_after_mask_ground = [] + for i in range(len(batch)): + pch1_dynamic_after_mask_ground.append(batch[i]['pch1_dynamic'][~batch[i]['gmh1']]) + pch1_dynamic_after_mask_ground = torch.nn.utils.rnn.pad_sequence(pch1_dynamic_after_mask_ground, batch_first=True, padding_value=0) + res_dict['pch1_dynamic'] = pch1_dynamic_after_mask_ground return res_dict + class HDF5Dataset(Dataset): - def __init__(self, directory, n_frames=2, dufo=False, eval = False, leaderboard_version=1): + def __init__(self, directory, n_frames=2, ssl_label=None, eval = False, leaderboard_version=1): ''' - directory: the directory of the dataset - n_frames: the number of frames we use, default is 2: current, next if more then it's the history from current. - dufo: if True, we will read the dynamic cluster label - eval: if True, use the eval index + Args: + directory: the directory of the dataset, the folder should contain some .h5 file and index_total.pkl. + + Following are optional: + * n_frames: the number of frames we use, default is 2: from pc0 to pc1. + * ssl_label: if not None, we will use this label for self-supervised learning + * eval: if True, use the eval index + * leaderboard_version: 1st or 2nd, default is 1. If '2', we will use the index_eval_v2.pkl from assets/docs. ''' super(HDF5Dataset, self).__init__() self.directory = directory if (torch.distributed.is_initialized() and torch.distributed.get_rank() == 0) or not torch.distributed.is_initialized(): - print(f"----[Debug] Loading data with num_frames={n_frames}, eval={eval}, leaderboard_version={leaderboard_version}") + print(f"----[Debug] Loading data with num_frames={n_frames}, ssl_label={ssl_label}, eval={eval}, leaderboard_version={leaderboard_version}") with open(os.path.join(self.directory, 'index_total.pkl'), 'rb') as f: self.data_index = pickle.load(f) self.eval_index = False - self.dufo = dufo - self.n_frames = n_frames + self.ssl_label = import_func(f"src.autolabel.{ssl_label}") if ssl_label is not None else None + self.history_frames = n_frames - 2 if eval: eval_index_file = os.path.join(self.directory, 'index_eval.pkl') @@ -165,11 +180,10 @@ def __getitem__(self, index_): 'pose1': pose1, } - if self.n_frames > 2: + if self.history_frames > 0: past_frames = [] - num_past_frames = self.n_frames - 2 - for i in range(1, num_past_frames + 1): + for i in range(1, self.history_frames + 1): frame_index = index_ - i if frame_index < self.scene_id_bounds[scene_id]["min_index"]: frame_index = self.scene_id_bounds[scene_id]["min_index"] @@ -179,12 +193,14 @@ def __getitem__(self, index_): past_gm = torch.tensor(f[past_timestamp]['ground_mask'][:]) past_pose = torch.tensor(f[past_timestamp]['pose'][:]) - past_frames.append((past_pc, past_gm, past_pose)) + past_frames.append((past_pc, past_gm, past_pose, past_timestamp)) - for i, (past_pc, past_gm, past_pose) in enumerate(past_frames): + for i, (past_pc, past_gm, past_pose, past_timestamp) in enumerate(past_frames): res_dict[f'pch{i+1}'] = past_pc res_dict[f'gmh{i+1}'] = past_gm res_dict[f'poseh{i+1}'] = past_pose + if self.ssl_label is not None: + res_dict[f'pch{i+1}_dynamic'] = torch.tensor(self.ssl_label(f[past_timestamp]).astype('int16')) if 'flow' in f[key]: flow = torch.tensor(f[key]['flow'][:]) @@ -198,9 +214,9 @@ def __getitem__(self, index_): ego_motion = torch.tensor(f[key]['ego_motion'][:]) res_dict['ego_motion'] = ego_motion - if self.dufo: - res_dict['pc0_dynamic'] = torch.tensor(f[key]['label'][:].astype('int16')) - res_dict['pc1_dynamic'] = torch.tensor(f[next_timestamp]['label'][:].astype('int16')) + if self.ssl_label is not None: + res_dict['pc0_dynamic'] = torch.tensor(self.ssl_label(f[key]).astype('int16')) + res_dict['pc1_dynamic'] = torch.tensor(self.ssl_label(f[next_timestamp]).astype('int16')) if self.eval_index: # looks like v2 not follow the same rule as v1 with eval_mask provided diff --git a/src/lossfuncs.py b/src/lossfuncs.py deleted file mode 100644 index cd161d9..0000000 --- a/src/lossfuncs.py +++ /dev/null @@ -1,157 +0,0 @@ -""" -# Created: 2023-07-17 00:00 -# Copyright (C) 2023-now, RPL, KTH Royal Institute of Technology -# Author: Qingwen Zhang (https://kin-zhang.github.io/) -# -# This file is part of DeFlow (https://github.com/KTH-RPL/DeFlow) and SeFlow (https://github.com/KTH-RPL/SeFlow). -# If you find this repo helpful, please cite the respective publication as -# listed on the above website. -# -# Description: Define the loss function for training. -""" -import torch -from assets.cuda.chamfer3D import nnChamferDis -MyCUDAChamferDis = nnChamferDis() -from src.utils.av2_eval import CATEGORY_TO_INDEX, BUCKETED_METACATAGORIES - -# NOTE(Qingwen 24/07/06): squared, so it's sqrt(4) = 2m, in 10Hz the vel = 20m/s ~ 72km/h -# If your scenario is different, may need adjust this TRUNCATED to 80-120km/h vel. -TRUNCATED_DIST = 4 - - -def seflowLoss(res_dict, timer=None): - pc0_label = res_dict['pc0_labels'] - pc1_label = res_dict['pc1_labels'] - - pc0 = res_dict['pc0'] - pc1 = res_dict['pc1'] - - est_flow = res_dict['est_flow'] - - pseudo_pc1from0 = pc0 + est_flow - - unique_labels = torch.unique(pc0_label) - pc0_dynamic = pc0[pc0_label > 0] - pc1_dynamic = pc1[pc1_label > 0] - # fpc1_dynamic = pseudo_pc1from0[pc0_label > 0] - # NOTE(Qingwen): since we set THREADS_PER_BLOCK is 256 - have_dynamic_cluster = (pc0_dynamic.shape[0] > 256) & (pc1_dynamic.shape[0] > 256) - - # first item loss: chamfer distance - # timer[5][1].start("MyCUDAChamferDis") - # raw: pc0 to pc1, est: pseudo_pc1from0 to pc1, idx means the nearest index - est_dist0, est_dist1, _, _ = MyCUDAChamferDis.disid_res(pseudo_pc1from0, pc1) - raw_dist0, raw_dist1, raw_idx0, _ = MyCUDAChamferDis.disid_res(pc0, pc1) - chamfer_dis = torch.mean(est_dist0[est_dist0 <= TRUNCATED_DIST]) + torch.mean(est_dist1[est_dist1 <= TRUNCATED_DIST]) - # timer[5][1].stop() - - # second item loss: dynamic chamfer distance - # timer[5][2].start("DynamicChamferDistance") - dynamic_chamfer_dis = torch.tensor(0.0, device=est_flow.device) - if have_dynamic_cluster: - dynamic_chamfer_dis += MyCUDAChamferDis(pseudo_pc1from0[pc0_label>0], pc1_dynamic, truncate_dist=TRUNCATED_DIST) - # timer[5][2].stop() - - # third item loss: exclude static points' flow - # NOTE(Qingwen): add in the later part on label==0 - static_cluster_loss = torch.tensor(0.0, device=est_flow.device) - - # fourth item loss: same label points' flow should be the same - # timer[5][3].start("SameClusterLoss") - moved_cluster_loss = torch.tensor(0.0, device=est_flow.device) - moved_cluster_norms = torch.tensor([], device=est_flow.device) - for label in unique_labels: - mask = pc0_label == label - if label == 0: - # Eq. 6 in the paper - static_cluster_loss += torch.linalg.vector_norm(est_flow[mask, :], dim=-1).mean() - elif label > 0 and have_dynamic_cluster: - cluster_id_flow = est_flow[mask, :] - cluster_nnd = raw_dist0[mask] - if cluster_nnd.shape[0] <= 0: - continue - - # Eq. 8 in the paper - sorted_idxs = torch.argsort(cluster_nnd, descending=True) - nearby_label = pc1_label[raw_idx0[mask][sorted_idxs]] # nonzero means dynamic in label - non_zero_valid_indices = torch.nonzero(nearby_label > 0) - if non_zero_valid_indices.shape[0] <= 0: - continue - max_idx = sorted_idxs[non_zero_valid_indices.squeeze(1)[0]] - - # Eq. 9 in the paper - max_flow = pc1[raw_idx0[mask][max_idx]] - pc0[mask][max_idx] - - # Eq. 10 in the paper - moved_cluster_norms = torch.cat((moved_cluster_norms, torch.linalg.vector_norm((cluster_id_flow - max_flow), dim=-1))) - - if moved_cluster_norms.shape[0] > 0: - moved_cluster_loss = moved_cluster_norms.mean() # Eq. 11 in the paper - elif have_dynamic_cluster: - moved_cluster_loss = torch.mean(raw_dist0[raw_dist0 <= TRUNCATED_DIST]) + torch.mean(raw_dist1[raw_dist1 <= TRUNCATED_DIST]) - # timer[5][3].stop() - - res_loss = { - 'chamfer_dis': chamfer_dis, - 'dynamic_chamfer_dis': dynamic_chamfer_dis, - 'static_flow_loss': static_cluster_loss, - 'cluster_based_pc0pc1': moved_cluster_loss, - } - return res_loss - -def deflowLoss(res_dict): - pred = res_dict['est_flow'] - gt = res_dict['gt_flow'] - - mask_no_nan = (~gt.isnan() & ~pred.isnan() & ~gt.isinf() & ~pred.isinf()) - - pred = pred[mask_no_nan].reshape(-1, 3) - gt = gt[mask_no_nan].reshape(-1, 3) - - speed = gt.norm(dim=1, p=2) / 0.1 - # pts_loss = torch.norm(pred - gt, dim=1, p=2) - pts_loss = torch.linalg.vector_norm(pred - gt, dim=-1) - - weight_loss = 0.0 - speed_0_4 = pts_loss[speed < 0.4].mean() - speed_mid = pts_loss[(speed >= 0.4) & (speed <= 1.0)].mean() - speed_1_0 = pts_loss[speed > 1.0].mean() - if ~speed_1_0.isnan(): - weight_loss += speed_1_0 - if ~speed_0_4.isnan(): - weight_loss += speed_0_4 - if ~speed_mid.isnan(): - weight_loss += speed_mid - return {'loss': weight_loss} - -# ref from zeroflow loss class FastFlow3DDistillationLoss() -def zeroflowLoss(res_dict): - pred = res_dict['est_flow'] - gt = res_dict['gt_flow'] - mask_no_nan = (~gt.isnan() & ~pred.isnan() & ~gt.isinf() & ~pred.isinf()) - - pred = pred[mask_no_nan].reshape(-1, 3) - gt = gt[mask_no_nan].reshape(-1, 3) - - error = torch.linalg.vector_norm(pred - gt, dim=-1) - # gt_speed = torch.norm(gt, dim=1, p=2) * 10.0 - gt_speed = torch.linalg.vector_norm(gt, dim=-1) * 10.0 - - mins = torch.ones_like(gt_speed) * 0.1 - maxs = torch.ones_like(gt_speed) - importance_scale = torch.max(mins, torch.min(1.8 * gt_speed - 0.8, maxs)) - # error = torch.norm(pred - gt, dim=1, p=2) * importance_scale - error = error * importance_scale - return {'loss': error.mean()} - -# ref from zeroflow loss class FastFlow3DSupervisedLoss() -def ff3dLoss(res_dict): - pred = res_dict['est_flow'] - gt = res_dict['gt_flow'] - classes = res_dict['gt_classes'] - # error = torch.norm(pred - gt, dim=1, p=2) - error = torch.linalg.vector_norm(pred - gt, dim=-1) - is_foreground_class = (classes > 0) # 0 is background, ref: FOREGROUND_BACKGROUND_BREAKDOWN - background_scalar = is_foreground_class.float() * 0.9 + 0.1 - error = error * background_scalar - return {'loss': error.mean()} diff --git a/src/lossfuncs/__init__.py b/src/lossfuncs/__init__.py new file mode 100644 index 0000000..7bf446b --- /dev/null +++ b/src/lossfuncs/__init__.py @@ -0,0 +1,19 @@ +""" +# Created: 2025-08-07 20:12 +# Copyright (C) 2025-now, RPL, KTH Royal Institute of Technology +# Author: Qingwen Zhang (https://kin-zhang.github.io/) +# +# This file is part of +# * OpenSceneFlow (https://github.com/KTH-RPL/OpenSceneFlow) +# * HiMo (https://kin-zhang.github.io/HiMo) +# +# If you find this repo helpful, please cite the respective publication as +# listed on the above website. +# +""" + + + +from .selfsupervise import * +from .supervise import * + diff --git a/src/lossfuncs/selfsupervise.py b/src/lossfuncs/selfsupervise.py new file mode 100644 index 0000000..dccc923 --- /dev/null +++ b/src/lossfuncs/selfsupervise.py @@ -0,0 +1,190 @@ +""" +# Created: 2023-07-17 00:00 +# Updated: 2025-08-07 00:01 +# Copyright (C) 2023-now, RPL, KTH Royal Institute of Technology +# Author: Qingwen Zhang (https://kin-zhang.github.io/) +# +# This file is part of +# * SeFlow (https://github.com/KTH-RPL/SeFlow) +# * HiMo (https://kin-zhang.github.io/HiMo) +# +# If you find this repo helpful, please cite the respective publication as +# listed on the above website. +# +# Description: Define the self-supervised (without GT) loss function for training. +# +""" +import torch +from assets.cuda.chamfer3D import nnChamferDis +MyCUDAChamferDis = nnChamferDis() + +# NOTE(Qingwen 24/07/06): squared, so it's sqrt(4) = 2m, in 10Hz the vel = 20m/s ~ 72km/h +# If your scenario is different, may need adjust this TRUNCATED to 80-120km/h vel. +TRUNCATED_DIST = 4 + +def seflowppLoss(res_dict, timer=None): + pch1_label = res_dict['pch1_labels'] + pc0_label = res_dict['pc0_labels'] + pc1_label = res_dict['pc1_labels'] + + pch1 = res_dict['pch1'] + pc0 = res_dict['pc0'] + pc1 = res_dict['pc1'] + + est_flow = res_dict['est_flow'] + + pseudo_pc1from0 = pc0 + est_flow + pseduo_pch1from0 = pc0 - est_flow + + unique_labels = torch.unique(pc0_label) + pc0_dynamic = pc0[pc0_label > 0] + pc1_dynamic = pc1[pc1_label > 0] + + # fpc1_dynamic = pseudo_pc1from0[pc0_label > 0] + # NOTE(Qingwen): since we set THREADS_PER_BLOCK is 256 + have_dynamic_cluster = (pc0_dynamic.shape[0] > 256) & (pc1_dynamic.shape[0] > 256) + + # first item loss: chamfer distance + # timer[5][1].start("MyCUDAChamferDis") + chamfer_dis = MyCUDAChamferDis(pseudo_pc1from0, pc1, truncate_dist=TRUNCATED_DIST) + MyCUDAChamferDis(pseduo_pch1from0, pch1, truncate_dist=TRUNCATED_DIST) + # timer[5][1].stop() + + # second item loss: dynamic chamfer distance + # timer[5][2].start("DynamicChamferDistance") + dynamic_chamfer_dis = torch.tensor(0.0, device=est_flow.device) + if have_dynamic_cluster: + dynamic_chamfer_dis += MyCUDAChamferDis(pseudo_pc1from0[pc0_label > 0], pc1_dynamic, truncate_dist=TRUNCATED_DIST) + if pch1[pch1_label > 0].shape[0] > 256: + dynamic_chamfer_dis += MyCUDAChamferDis(pseduo_pch1from0[pc0_label > 0], pch1[pch1_label > 0], truncate_dist=TRUNCATED_DIST) + # timer[5][2].stop() + + # third item loss: exclude static points' flow + # NOTE(Qingwen): add in the later part on label==0 + static_cluster_loss = torch.tensor(0.0, device=est_flow.device) + + # fourth item loss: same label points' flow should be the same + # timer[5][3].start("SameClusterLoss") + # raw: pc0 to pc1, est: pseudo_pc1from0 to pc1, idx means the nearest index + raw_dist0, raw_dist1, raw_idx0, _ = MyCUDAChamferDis.disid_res(pc0, pc1) + moved_cluster_loss = torch.tensor(0.0, device=est_flow.device) + moved_cluster_norms = torch.tensor([], device=est_flow.device) + for label in unique_labels: + mask = pc0_label == label + if label == 0: + # Eq. 6 in the SeFlow paper + static_cluster_loss += torch.linalg.vector_norm(est_flow[mask, :], dim=-1).mean() + # NOTE(Qingwen) 2025-04-23: label=1 is dynamic but no cluster id satisfied + elif label > 1 and have_dynamic_cluster: + cluster_id_flow = est_flow[mask, :] + cluster_nnd = raw_dist0[mask] + if cluster_nnd.shape[0] <= 0: + continue + + # Eq. 8 in the SeFlow paper + sorted_idxs = torch.argsort(cluster_nnd, descending=True) + nearby_label = pc1_label[raw_idx0[mask][sorted_idxs]] # nonzero means dynamic in label + non_zero_valid_indices = torch.nonzero(nearby_label > 0) + if non_zero_valid_indices.shape[0] <= 0: + continue + max_idx = sorted_idxs[non_zero_valid_indices.squeeze(1)[0]] + + # Eq. 9 in the SeFlow paper + max_flow = pc1[raw_idx0[mask][max_idx]] - pc0[mask][max_idx] + + # Eq. 10 in the SeFlow paper + moved_cluster_norms = torch.cat((moved_cluster_norms, torch.linalg.vector_norm((cluster_id_flow - max_flow), dim=-1))) + + if moved_cluster_norms.shape[0] > 0: + moved_cluster_loss = moved_cluster_norms.mean() # Eq. 11 in the SeFlow paper + elif have_dynamic_cluster: + moved_cluster_loss = torch.mean(raw_dist0[raw_dist0 <= TRUNCATED_DIST]) + torch.mean(raw_dist1[raw_dist1 <= TRUNCATED_DIST]) + # timer[5][3].stop() + + res_loss = { + 'chamfer_dis': chamfer_dis / 2.0, + 'dynamic_chamfer_dis': dynamic_chamfer_dis / 2.0, + 'static_flow_loss': static_cluster_loss, + 'cluster_based_pc0pc1': moved_cluster_loss, + } + return res_loss + +def seflowLoss(res_dict, timer=None): + pc0_label = res_dict['pc0_labels'] + pc1_label = res_dict['pc1_labels'] + + pc0 = res_dict['pc0'] + pc1 = res_dict['pc1'] + + est_flow = res_dict['est_flow'] + + pseudo_pc1from0 = pc0 + est_flow + + unique_labels = torch.unique(pc0_label) + pc0_dynamic = pc0[pc0_label > 0] + pc1_dynamic = pc1[pc1_label > 0] + # fpc1_dynamic = pseudo_pc1from0[pc0_label > 0] + # NOTE(Qingwen): since we set THREADS_PER_BLOCK is 256 + have_dynamic_cluster = (pc0_dynamic.shape[0] > 256) & (pc1_dynamic.shape[0] > 256) + + # first item loss: chamfer distance + # timer[5][1].start("MyCUDAChamferDis") + # raw: pc0 to pc1, est: pseudo_pc1from0 to pc1, idx means the nearest index + est_dist0, est_dist1, _, _ = MyCUDAChamferDis.disid_res(pseudo_pc1from0, pc1) + raw_dist0, raw_dist1, raw_idx0, _ = MyCUDAChamferDis.disid_res(pc0, pc1) + chamfer_dis = torch.mean(est_dist0[est_dist0 <= TRUNCATED_DIST]) + torch.mean(est_dist1[est_dist1 <= TRUNCATED_DIST]) + # timer[5][1].stop() + + # second item loss: dynamic chamfer distance + # timer[5][2].start("DynamicChamferDistance") + dynamic_chamfer_dis = torch.tensor(0.0, device=est_flow.device) + if have_dynamic_cluster: + dynamic_chamfer_dis += MyCUDAChamferDis(pseudo_pc1from0[pc0_label>0], pc1_dynamic, truncate_dist=TRUNCATED_DIST) + # timer[5][2].stop() + + # third item loss: exclude static points' flow + # NOTE(Qingwen): add in the later part on label==0 + static_cluster_loss = torch.tensor(0.0, device=est_flow.device) + + # fourth item loss: same label points' flow should be the same + # timer[5][3].start("SameClusterLoss") + moved_cluster_loss = torch.tensor(0.0, device=est_flow.device) + moved_cluster_norms = torch.tensor([], device=est_flow.device) + for label in unique_labels: + mask = pc0_label == label + if label == 0: + # Eq. 6 in the paper + static_cluster_loss += torch.linalg.vector_norm(est_flow[mask, :], dim=-1).mean() + # NOTE(Qingwen) 2025-04-23: label=1 is dynamic but no cluster id satisfied + elif label > 1 and have_dynamic_cluster: + cluster_id_flow = est_flow[mask, :] + cluster_nnd = raw_dist0[mask] + if cluster_nnd.shape[0] <= 0: + continue + + # Eq. 8 in the paper + sorted_idxs = torch.argsort(cluster_nnd, descending=True) + nearby_label = pc1_label[raw_idx0[mask][sorted_idxs]] # nonzero means dynamic in label + non_zero_valid_indices = torch.nonzero(nearby_label > 0) + if non_zero_valid_indices.shape[0] <= 0: + continue + max_idx = sorted_idxs[non_zero_valid_indices.squeeze(1)[0]] + + # Eq. 9 in the paper + max_flow = pc1[raw_idx0[mask][max_idx]] - pc0[mask][max_idx] + + # Eq. 10 in the paper + moved_cluster_norms = torch.cat((moved_cluster_norms, torch.linalg.vector_norm((cluster_id_flow - max_flow), dim=-1))) + + if moved_cluster_norms.shape[0] > 0: + moved_cluster_loss = moved_cluster_norms.mean() # Eq. 11 in the paper + elif have_dynamic_cluster: + moved_cluster_loss = torch.mean(raw_dist0[raw_dist0 <= TRUNCATED_DIST]) + torch.mean(raw_dist1[raw_dist1 <= TRUNCATED_DIST]) + # timer[5][3].stop() + + res_loss = { + 'chamfer_dis': chamfer_dis, + 'dynamic_chamfer_dis': dynamic_chamfer_dis, + 'static_flow_loss': static_cluster_loss, + 'cluster_based_pc0pc1': moved_cluster_loss, + } + return res_loss diff --git a/src/lossfuncs/supervise.py b/src/lossfuncs/supervise.py new file mode 100644 index 0000000..b72e7a1 --- /dev/null +++ b/src/lossfuncs/supervise.py @@ -0,0 +1,107 @@ +""" +# Created: 2023-07-17 00:00 +# Copyright (C) 2023-now, RPL, KTH Royal Institute of Technology +# Author: Qingwen Zhang (https://kin-zhang.github.io/) +# +# This file is part of +# * OpenSceneFlow (https://github.com/KTH-RPL/OpenSceneFlow) +# If you find this repo helpful, please cite the respective publication as +# listed on the above website. +# +# Description: Define the supervised (needed GT) loss function for training. +# +""" +import torch +import numpy as np +import os, sys +BASE_DIR = os.path.abspath(os.path.join( os.path.dirname( __file__ ), '../..' )) +sys.path.append(BASE_DIR) + +def deflowLoss(res_dict): + pred = res_dict['est_flow'] + gt = res_dict['gt_flow'] + + mask_no_nan = (~gt.isnan() & ~pred.isnan() & ~gt.isinf() & ~pred.isinf()) + + pred = pred[mask_no_nan].reshape(-1, 3) + gt = gt[mask_no_nan].reshape(-1, 3) + + speed = gt.norm(dim=1, p=2) / 0.1 + pts_loss = torch.linalg.vector_norm(pred - gt, dim=-1) + + weight_loss = 0.0 + for loss_ in [pts_loss[speed < 0.4].mean(), + pts_loss[(speed >= 0.4) & (speed <= 1.0)].mean(), + pts_loss[speed > 1.0].mean()]: + weight_loss += torch.nan_to_num(loss_, nan=0.0) + + return {'loss': weight_loss} + +# designed from MambaFlow: https://github.com/SCNU-RISLAB/MambaFlow +def mambaflowLoss(res_dict): + pred = res_dict['est_flow'] + gt = res_dict['gt_flow'] + mask_no_nan = (~gt.isnan() & ~pred.isnan() & ~gt.isinf() & ~pred.isinf()) + pred = pred[mask_no_nan].reshape(-1, 3) + gt = gt[mask_no_nan].reshape(-1, 3) + + speed = gt.norm(dim=1, p=2) / 0.1 + pts_loss = torch.linalg.vector_norm(pred - gt, dim=-1) + + velocities = speed.cpu().numpy() + + # 计算直方图，返回每个区间的计数和区间边界 + counts, bin_edges = np.histogram(velocities, bins=100, density=False) + + # 计算每个区间的点数占总点数的比例 + total_points = len(velocities) + proportions = counts / total_points + + # 计算每个区间的中心位置，用于绘图 + bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2 + + # 设置占比阈值 + proportion_threshold = 0.01 # 可以根据需要调整这个值 + + # 找出第一个占比小于阈值的柱子 + first_below_threshold = next((i for i, prop in enumerate(proportions) if prop < proportion_threshold), None) + turning_speed = bin_centers[first_below_threshold] + + weight_loss = 0.0 + for loss_ in [pts_loss[speed < turning_speed].mean(), + pts_loss[(speed >= turning_speed) & (speed <= 2)].mean(), + pts_loss[speed > 2].mean()]: + weight_loss += torch.nan_to_num(loss_, nan=0.0) + return {'loss': weight_loss} + +# ref from zeroflow loss class FastFlow3DDistillationLoss() +def zeroflowLoss(res_dict): + pred = res_dict['est_flow'] + gt = res_dict['gt_flow'] + mask_no_nan = (~gt.isnan() & ~pred.isnan() & ~gt.isinf() & ~pred.isinf()) + + pred = pred[mask_no_nan].reshape(-1, 3) + gt = gt[mask_no_nan].reshape(-1, 3) + + error = torch.linalg.vector_norm(pred - gt, dim=-1) + # gt_speed = torch.norm(gt, dim=1, p=2) * 10.0 + gt_speed = torch.linalg.vector_norm(gt, dim=-1) * 10.0 + + mins = torch.ones_like(gt_speed) * 0.1 + maxs = torch.ones_like(gt_speed) + importance_scale = torch.max(mins, torch.min(1.8 * gt_speed - 0.8, maxs)) + # error = torch.norm(pred - gt, dim=1, p=2) * importance_scale + error = error * importance_scale + return {'loss': error.mean()} + +# ref from zeroflow loss class FastFlow3DSupervisedLoss() +def ff3dLoss(res_dict): + pred = res_dict['est_flow'] + gt = res_dict['gt_flow'] + classes = res_dict['gt_classes'] + # error = torch.norm(pred - gt, dim=1, p=2) + error = torch.linalg.vector_norm(pred - gt, dim=-1) + is_foreground_class = (classes > 0) # 0 is background, ref: FOREGROUND_BACKGROUND_BREAKDOWN + background_scalar = is_foreground_class.float() * 0.9 + 0.1 + error = error * background_scalar + return {'loss': error.mean()} diff --git a/src/models/__init__.py b/src/models/__init__.py index 410b773..5f9d9cf 100644 --- a/src/models/__init__.py +++ b/src/models/__init__.py @@ -10,7 +10,7 @@ import warnings warnings.simplefilter(action="ignore", category=FutureWarning) -from .deflow import DeFlow +from .deflow import DeFlow, DeFlowPP from .fastflow3d import FastFlow3D from .nsfp import NSFP diff --git a/src/models/basic/unet.py b/src/models/basic/unet.py index f7f651d..3545324 100644 --- a/src/models/basic/unet.py +++ b/src/models/basic/unet.py @@ -1,7 +1,5 @@ import torch import torch.nn as nn - -from typing import Tuple from . import ConvWithNorms @@ -98,4 +96,147 @@ def forward(self, pc0_B: torch.Tensor, U = self.decoder_step3(T, Bstar) V = self.decoder_step4(U) - return V \ No newline at end of file + return V + +class UNetThreeFrame(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(768, 384, 384) + self.decoder_step2 = UpsampleSkip(384, 192, 192) + self.decoder_step3 = UpsampleSkip(192, 96, 96) + self.decoder_step4 = nn.Conv2d(96, 96, 3, 1, 1) + + def forward(self, pcb0_B: torch.Tensor, 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]}" + + pcb0_F = self.encoder_step_1(pcb0_B) + pcb0_L = self.encoder_step_2(pcb0_F) + pcb0_R = self.encoder_step_3(pcb0_L) + + 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([pcb0_R, pc0_R, pc1_R], + dim=1) # torch.Size([1, 512, 64, 64]), torch.Size([1, 768, 64, 64]) + Lstar = torch.cat([pcb0_L, pc0_L, pc1_L], + dim=1) # torch.Size([1, 256, 128, 128]), torch.Size([1, 384, 128, 128]) + Fstar = torch.cat([pcb0_F, pc0_F, pc1_F], + dim=1) # torch.Size([1, 128, 256, 256]), torch.Size([1, 192, 256, 256]) + Bstar = torch.cat([pcb0_B, pc0_B, pc1_B], + dim=1) # torch.Size([1, 64, 512, 512]), torch.Size([1, 96, 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 + +class ZeroFlowUNetXL(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(64, 128, 3, 2, 1), + ConvWithNorms(128, 128, 3, 1, 1), + ConvWithNorms(128, 128, 3, 1, 1), + ConvWithNorms(128, 128, 3, 1, 1)) + self.encoder_step_2 = 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.encoder_step_3 = nn.Sequential(ConvWithNorms(256, 512, 3, 2, 1), + ConvWithNorms(512, 512, 3, 1, 1), + ConvWithNorms(512, 512, 3, 1, 1), + ConvWithNorms(512, 512, 3, 1, 1), + ConvWithNorms(512, 512, 3, 1, 1), + ConvWithNorms(512, 512, 3, 1, 1)) + self.encoder_step_4 = nn.Sequential(ConvWithNorms(512, 1024, 3, 2, 1), + ConvWithNorms(1024, 1024, 3, 1, 1), + ConvWithNorms(1024, 1024, 3, 1, 1), + ConvWithNorms(1024, 1024, 3, 1, 1), + ConvWithNorms(1024, 1024, 3, 1, 1), + ConvWithNorms(1024, 1024, 3, 1, 1)) + + self.decoder_step1 = UpsampleSkip(2048, 1024, 1024) + self.decoder_step2 = UpsampleSkip(1024, 512, 512) + self.decoder_step3 = UpsampleSkip(512, 256, 256) + self.decoder_step4 = UpsampleSkip(256, 128, 128) + self.decoder_step5 = nn.Conv2d(128, 128, 3, 1, 1) + + def forward(self, pc0_B: torch.Tensor, + pc1_B: torch.Tensor) -> torch.Tensor: + + expected_channels = 64 + 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) + pc0_T = self.encoder_step_4(pc0_R) + + pc1_F = self.encoder_step_1(pc1_B) + pc1_L = self.encoder_step_2(pc1_F) + pc1_R = self.encoder_step_3(pc1_L) + pc1_T = self.encoder_step_4(pc1_R) + + Tstar = torch.cat([pc0_T, pc1_T], + dim=1) # torch.Size([1, 2048, 32, 32]) + Rstar = torch.cat([pc0_R, pc1_R], + dim=1) # torch.Size([1, 1024, 64, 64]) + Lstar = torch.cat([pc0_L, pc1_L], + dim=1) # torch.Size([1, 512, 128, 128]) + Fstar = torch.cat([pc0_F, pc1_F], + dim=1) # torch.Size([1, 256, 256, 256]) + Bstar = torch.cat([pc0_B, pc1_B], + dim=1) # torch.Size([1, 128, 512, 512]) + + S = self.decoder_step1(Tstar, Rstar) + T = self.decoder_step2(S, Lstar) + U = self.decoder_step3(T, Fstar) + V = self.decoder_step4(U, Bstar) + W = self.decoder_step5(V) + + return W \ No newline at end of file diff --git a/src/models/deflow.py b/src/models/deflow.py index d624690..e217801 100644 --- a/src/models/deflow.py +++ b/src/models/deflow.py @@ -12,10 +12,10 @@ import torch.nn as nn import dztimer, torch -from .basic.unet import FastFlow3DUNet +from .basic.unet import FastFlow3DUNet, UNetThreeFrame from .basic.encoder import DynamicEmbedder from .basic.decoder import LinearDecoder, ConvGRUDecoder -from .basic import cal_pose0to1, BaseModel +from .basic import wrap_batch_pcs, BaseModel class DeFlow(BaseModel): def __init__(self, voxel_size = [0.2, 0.2, 6], @@ -45,29 +45,9 @@ def forward(self, batch): output: the predicted flow, pose_flow, and the valid point index of pc0 """ self.timer[0].start("Data Preprocess") - batch_sizes = len(batch["pose0"]) - - pose_flows = [] - transform_pc0s = [] - for batch_id in range(batch_sizes): - selected_pc0 = batch["pc0"][batch_id] - self.timer[0][0].start("pose") - with torch.no_grad(): - if 'ego_motion' in batch: - pose_0to1 = batch['ego_motion'][batch_id] - else: - pose_0to1 = cal_pose0to1(batch["pose0"][batch_id], batch["pose1"][batch_id]) - self.timer[0][0].stop() - - self.timer[0][1].start("transform") - # transform selected_pc0 to pc1 - transform_pc0 = selected_pc0 @ pose_0to1[:3, :3].T + pose_0to1[:3, 3] - self.timer[0][1].stop() - pose_flows.append(transform_pc0 - selected_pc0) - transform_pc0s.append(transform_pc0) - - pc0s = torch.stack(transform_pc0s, dim=0) - pc1s = batch["pc1"] + pcs_dict = wrap_batch_pcs(batch, num_frames=2) + pc0s = pcs_dict['pc0s'] + pc1s = pcs_dict['pc1s'] self.timer[0].stop() self.timer[1].start("Voxelization") @@ -94,7 +74,7 @@ def forward(self, batch): model_res = { "flow": flows, - 'pose_flow': pose_flows, + 'pose_flow': pcs_dict['pose_flows'], "pc0_valid_point_idxes": pc0_valid_point_idxes, "pc0_points_lst": pc0_points_lst, @@ -104,3 +84,75 @@ def forward(self, batch): "num_occupied_voxels": [grid_flow_pseudoimage.size(-1)*grid_flow_pseudoimage.size(-2)] } return model_res + + + +class DeFlowPP(BaseModel): + def __init__(self, voxel_size = [0.2, 0.2, 6], + point_cloud_range = [-51.2, -51.2, -3, 51.2, 51.2, 3], + grid_feature_size = [512, 512], + decoder_option = "gru", + num_iters = 2, + num_frames = 3): + super().__init__() + self.embedder = DynamicEmbedder(voxel_size=voxel_size, + pseudo_image_dims=grid_feature_size, + point_cloud_range=point_cloud_range, + feat_channels=32) + + self.backbone = UNetThreeFrame() + if decoder_option == "gru": + self.head = ConvGRUDecoder(pseudoimage_channels=96, num_iters = num_iters) + else: + self.head = LinearDecoder() + + self.num_frames = num_frames + assert self.num_frames == 3, "DeFlowPP only supports num_frames = 3" + + self.timer = dztimer.Timing() + self.timer.start("Total") + + def forward(self, batch): + """ + input: using the batch from dataloader, which is a dict + Detail: [pc0, pc1, pose0, pose1] + output: the predicted flow, pose_flow, and the valid point index of pc0 + """ + self.timer[0].start("Data Preprocess") + pcs_dict = wrap_batch_pcs(batch, num_frames=self.num_frames) + pc0s = pcs_dict['pc0s'] + pc1s = pcs_dict['pc1s'] + pch1s = pcs_dict['pch1s'] + self.timer[0].stop() + + self.timer[1].start("Voxelization") + pch1_before_pseudoimages, pch1_voxel_infos_lst = self.embedder(pch1s) + pc0_before_pseudoimages, pc0_voxel_infos_lst = self.embedder(pc0s) + pc1_before_pseudoimages, pc1_voxel_infos_lst = self.embedder(pc1s) + self.timer[1].stop() + + self.timer[2].start("Encoder") + grid_flow_pseudoimage = self.backbone(pch1_before_pseudoimages, pc0_before_pseudoimages, + pc1_before_pseudoimages) + self.timer[2].stop() + + self.timer[3].start("Decoder") + flows = self.head( + torch.cat((pch1_before_pseudoimages, pc0_before_pseudoimages, pc1_before_pseudoimages), + dim=1), grid_flow_pseudoimage, pc0_voxel_infos_lst) + self.timer[3].stop() + + model_res = { + "flow": flows, + 'pose_flow': pcs_dict['pose_flows'], + + "pc0_valid_point_idxes": [e["point_idxes"] for e in pc0_voxel_infos_lst], + "pc0_points_lst": [e["points"] for e in pc0_voxel_infos_lst], + + "pc1_valid_point_idxes": [e["point_idxes"] for e in pc1_voxel_infos_lst], + "pc1_points_lst": [e["points"] for e in pc1_voxel_infos_lst], + + 'pch1_valid_point_idxes': [e["point_idxes"] for e in pch1_voxel_infos_lst], + 'pch1_points_lst': [e["points"] for e in pch1_voxel_infos_lst], + } + return model_res \ No newline at end of file diff --git a/src/models/fastflow3d.py b/src/models/fastflow3d.py index e4d5774..f6a3d14 100644 --- a/src/models/fastflow3d.py +++ b/src/models/fastflow3d.py @@ -12,7 +12,7 @@ from .basic.unet import FastFlow3DUNet from .basic.encoder import DynamicEmbedder from .basic.decoder import LinearDecoder -from .basic import cal_pose0to1, BaseModel +from .basic import wrap_batch_pcs, BaseModel class FastFlow3D(BaseModel): @@ -68,31 +68,15 @@ def forward(self, batch, compute_symmetry_y=False): self.timer[0].start("Data Preprocess") - batch_sizes = len(batch["pose0"]) - - pose_flows = [] - transform_pc0s = [] - for batch_id in range(batch_sizes): - selected_pc0 = batch["pc0"][batch_id] - self.timer[0][0].start("pose") - pose_0to1 = cal_pose0to1(batch["pose0"][batch_id], batch["pose1"][batch_id]) - self.timer[0][0].stop() - - self.timer[0][1].start("transform") - # transform selected_pc0 to pc1 - transform_pc0 = selected_pc0 @ pose_0to1[:3, :3].T + pose_0to1[:3, 3] - self.timer[0][1].stop() - pose_flows.append(transform_pc0 - selected_pc0) - transform_pc0s.append(transform_pc0) - - pc0s = torch.stack(transform_pc0s, dim=0) - pc1s = batch["pc1"] + pcs_dict = wrap_batch_pcs(batch, num_frames=2) + pc0s = pcs_dict['pc0s'] + pc1s = pcs_dict['pc1s'] self.timer[0].stop() model_res = self._model_forward(pc0s, pc1s) ret_dict = model_res - ret_dict["pose_flow"] = pose_flows + ret_dict["pose_flow"] = pcs_dict['pose_flows'] 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"] diff --git a/src/models/ssf.py b/src/models/ssf.py index 8fb9799..1c0eeb6 100644 --- a/src/models/ssf.py +++ b/src/models/ssf.py @@ -16,7 +16,7 @@ from .basic.encoder import DynamicVoxelizer from .basic.ssf_module import DynamicScatterVFE from .basic.decoder import SimpleLinearDecoder -from .basic import cal_pose0to1, BaseModel +from .basic import wrap_batch_pcs, BaseModel class SSF(BaseModel): def __init__(self, voxel_size = [0.2, 0.2, 6], @@ -64,31 +64,11 @@ def forward(self, batch): output: the predicted flow, pose_flow, and the valid point index of pc0 """ self.timer[0].start("Data Preprocess") - batch_sizes = len(batch["pose0"]) - - pose_flows = [] - transform_pc0s = [] - for batch_id in range(batch_sizes): - selected_pc0 = batch["pc0"][batch_id] - self.timer[0][0].start("pose") - with torch.no_grad(): - if 'ego_motion' in batch: - pose_0to1 = batch['ego_motion'][batch_id] - else: - pose_0to1 = cal_pose0to1(batch["pose0"][batch_id], batch["pose1"][batch_id]) - self.timer[0][0].stop() - - self.timer[0][1].start("transform") - # transform selected_pc0 to pc1 - transform_pc0 = selected_pc0 @ pose_0to1[:3, :3].T + pose_0to1[:3, 3] - self.timer[0][1].stop() - pose_flows.append(transform_pc0 - selected_pc0) - transform_pc0s.append(transform_pc0) - - pc0s = torch.stack(transform_pc0s, dim=0) + pcs_dict = wrap_batch_pcs(batch, num_frames=2) + pc0s = pcs_dict['pc0s'] pc0s = torch.cat((pc0s, torch.ones((pc0s.size(0), pc0s.size(1),1)).to(pc0s.device) * 0.), dim=2) # indicator_pc0s = torch.ones((pc0s.size(0), pc0s.size(1),1), dtype=torch.int).to(pc0s.device) * 0 - pc1s = batch["pc1"] + pc1s = pcs_dict['pc1s'] pc1s = torch.cat((pc1s, torch.ones((pc1s.size(0), pc1s.size(1),1)).to(pc1s.device) * 1.), dim=2) # indicator_pc1s = torch.ones((pc1s.size(0), pc1s.size(1),1), dtype=torch.int).to(pc1s.device) * 1 pcs_concat = torch.cat((pc0s, pc1s), dim=1) @@ -144,11 +124,12 @@ def forward(self, batch): pc1_points_lst = [e["points"] for e in pc1_voxel_infos_lst] pc0_valid_point_idxes = [e["point_idxes"] for e in pc0_voxel_infos_lst] - pc1_valid_point_idxes = [e["point_idxes"] for e in pc1_voxel_infos_lst] + # since we concat in voxel_infos_dict + pc1_valid_point_idxes = [e["point_idxes"] - pc0s.shape[1] for e in pc1_voxel_infos_lst] model_res = { "flow": flows, - 'pose_flow': pose_flows, + 'pose_flow': pcs_dict['pose_flows'], "pc0_valid_point_idxes": pc0_valid_point_idxes, "pc0_points_lst": pc0_points_lst, diff --git a/src/runner.py b/src/runner.py index 98d5644..e15fd7a 100644 --- a/src/runner.py +++ b/src/runner.py @@ -106,7 +106,7 @@ def __init__(self, cfg, rank, world_size, mode): def _setup_dataloader(self): if self.mode in ['val', 'test', 'eval']: - dataset_path = self.cfg.dataset_path + f"/{self.cfg.av2_mode}" + dataset_path = self.cfg.dataset_path + f"/{self.cfg.data_mode}" is_eval_mode = True else: # 'save' dataset_path = self.cfg.dataset_path diff --git a/src/trainer.py b/src/trainer.py index 53954fa..e2eca7a 100644 --- a/src/trainer.py +++ b/src/trainer.py @@ -27,7 +27,7 @@ from src.utils import import_func from src.utils.mics import weights_init, zip_res from src.utils.av2_eval import write_output_file -from src.models.basic import cal_pose0to1 +from src.models.basic import cal_pose0to1, WarmupCosLR from src.utils.eval_metric import OfficialMetrics, evaluate_leaderboard, evaluate_leaderboard_v2, evaluate_ssf # debugging tools @@ -39,7 +39,25 @@ class ModelWrapper(LightningModule): def __init__(self, cfg, eval=False): super().__init__() - # set grid size + default_self_values = { + "batch_size": 1, + "lr": 2e-4, + "epochs": 3, + "loss_fn": 'deflowLoss', + "add_seloss": None, + "checkpoint": None, + "leaderboard_version": 2, + "supervised_flag": True, + "save_res": False, + "res_name": "default", + "num_frames": 2, + "optimizer": None, + "dataset_path": None, + "data_mode": None, + } + for key, default in default_self_values.items(): + setattr(self, key, cfg.get(key, default)) + if ('voxel_size' in cfg.model.target) and ('point_cloud_range' in cfg.model.target) and not eval and 'point_cloud_range' in cfg: OmegaConf.set_struct(cfg.model.target, True) with open_dict(cfg.model.target): @@ -54,45 +72,33 @@ def __init__(self, cfg, eval=False): abs(int((cfg.model.target.point_cloud_range[1] - cfg.model.target.point_cloud_range[4]) / cfg.model.target.voxel_size[1])), abs(int((cfg.model.target.point_cloud_range[2] - cfg.model.target.point_cloud_range[5]) / cfg.model.target.voxel_size[2]))] + # ---> model + self.point_cloud_range = cfg.model.target.point_cloud_range self.model = instantiate(cfg.model.target) self.model.apply(weights_init) - + if 'pretrained_weights' in cfg and cfg.pretrained_weights is not None: + missing_keys, unexpected_keys = self.model.load_from_checkpoint(cfg.pretrained_weights) + # print(f"Model: {self.model.__class__.__name__}, Number of Frames: {self.num_frames}") + + # ---> loss fn self.loss_fn = import_func("src.lossfuncs."+cfg.loss_fn) if 'loss_fn' in cfg else None - self.add_seloss = cfg.add_seloss if 'add_seloss' in cfg else None - self.cfg_loss_name = cfg.loss_fn if 'loss_fn' in cfg else None - - self.batch_size = int(cfg.batch_size) if 'batch_size' in cfg else 1 - self.lr = cfg.lr if 'lr' in cfg else None - self.lr_scheduler = cfg.lr_scheduler if 'lr_scheduler' in cfg else None - self.epochs = cfg.epochs if 'epochs' in cfg else None + self.cfg_loss_name = cfg.get("loss_fn", None) + # ---> evaluation metric self.metrics = OfficialMetrics() - self.load_checkpoint_path = cfg.checkpoint if 'checkpoint' in cfg else None - + # ---> inference mode + if self.save_res and self.data_mode in ['val', 'valid', 'test']: + self.save_res_path = Path(cfg.dataset_path).parent / "results" / cfg.output + os.makedirs(self.save_res_path, exist_ok=True) + print(f"We are in {cfg.data_mode}, results will be saved in: {self.save_res_path} with version: {self.leaderboard_version} format for online leaderboard.") - self.leaderboard_version = cfg.leaderboard_version if 'leaderboard_version' in cfg else 1 - # NOTE(Qingwen): since we have seflow version which is unsupervised, we need to set the flag to false. - self.supervised_flag = cfg.supervised_flag if 'supervised_flag' in cfg else True - self.save_res = False - if 'av2_mode' in cfg: - self.av2_mode = cfg.av2_mode - self.save_res = cfg.save_res if 'save_res' in cfg else False - - if self.save_res or self.av2_mode == 'test': - self.save_res_path = Path(cfg.dataset_path).parent / "results" / cfg.output - os.makedirs(self.save_res_path, exist_ok=True) - print(f"We are in {cfg.av2_mode}, results will be saved in: {self.save_res_path} with version: {self.leaderboard_version} format for online leaderboard.") - else: - self.av2_mode = None - if 'pretrained_weights' in cfg: - if cfg.pretrained_weights is not None: - self.model.load_from_checkpoint(cfg.pretrained_weights) - - self.dataset_path = cfg.dataset_path if 'dataset_path' in cfg else None - self.vis_name = cfg.res_name if 'res_name' in cfg else 'default' + # self.test_total_num = 0 + if self.data_mode in ['val', 'valid', 'test']: + print(cfg) self.save_hyperparameters() + # FIXME(Qingwen 2025-08-20): update the loss_calculation fn alone to make all things pretty here.... def training_step(self, batch, batch_idx): self.model.timer[4].start("One Scan in model") res_dict = self.model(batch) @@ -102,7 +108,7 @@ def training_step(self, batch, batch_idx): # compute loss total_loss = 0.0 - if self.cfg_loss_name in ['seflowLoss']: + if self.cfg_loss_name in ['seflowLoss', 'seflowppLoss']: loss_items, weights = zip(*[(key, weight) for key, weight in self.add_seloss.items()]) loss_logger = {'chamfer_dis': 0.0, 'dynamic_chamfer_dis': 0.0, 'static_flow_loss': 0.0, 'cluster_based_pc0pc1': 0.0} else: @@ -131,11 +137,15 @@ def training_step(self, batch, batch_idx): if 'pc0_dynamic' in batch: dict2loss['pc0_labels'] = batch['pc0_dynamic'][batch_id][pc0_valid_from_pc2res] dict2loss['pc1_labels'] = batch['pc1_dynamic'][batch_id][pc1_valid_from_pc2res] + if 'pch1_dynamic' in batch and 'pch1_valid_point_idxes' in res_dict: + dict2loss['pch1_labels'] = batch['pch1_dynamic'][batch_id][res_dict['pch1_valid_point_idxes'][batch_id]] # different methods may don't have this in the res_dict if 'pc0_points_lst' in res_dict and 'pc1_points_lst' in res_dict: dict2loss['pc0'] = pc0_points_lst[batch_id] dict2loss['pc1'] = pc1_points_lst[batch_id] + if 'pch1_points_lst' in res_dict: + dict2loss['pch1'] = res_dict['pch1_points_lst'][batch_id] res_loss = self.loss_fn(dict2loss) for i, loss_name in enumerate(loss_items): @@ -144,7 +154,7 @@ def training_step(self, batch, batch_idx): loss_logger[key] += res_loss[key] self.log("trainer/loss", total_loss/batch_sizes, sync_dist=True, batch_size=self.batch_size, prog_bar=True) - if self.add_seloss is not None and self.cfg_loss_name in ['seflowLoss']: + if self.add_seloss is not None and self.cfg_loss_name in ['seflowLoss', 'seflowppLoss']: for key in loss_logger: self.log(f"trainer/{key}", loss_logger[key]/batch_sizes, sync_dist=True, batch_size=self.batch_size) self.model.timer[5].stop() @@ -173,14 +183,22 @@ def train_validation_step_(self, batch, res_dict): pass def configure_optimizers(self): - if self.lr is None: - optimizer = optim.Adam(self.model.parameters(), lr=self.lr) - return optimizer - elif self.lr_scheduler == 'step': - optimizer = optim.Adam(self.model.parameters(), lr=self.lr) - lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=self.trainer.max_epochs//2, gamma=0.1) - return {'optimizer': optimizer, - 'lr_scheduler': lr_scheduler} + optimizers_ = {} + # default Adam + if self.optimizer.name == "AdamW": + optimizers_['optimizer'] = optim.AdamW(self.model.parameters(), lr=self.optimizer.lr, weight_decay=self.optimizer.get("weight_decay", 1e-4)) + else: # if self.optimizer.name == "Adam": + optimizers_['optimizer'] = optim.Adam(self.model.parameters(), lr=self.optimizer.lr) + + if "scheduler" in self.optimizer: + if self.optimizer.scheduler.name == "WarmupCosLR": + optimizers_['lr_scheduler'] = WarmupCosLR(optimizers_['optimizer'], self.optimizer.scheduler.get("min_lr", self.optimizer.lr*0.1), \ + self.optimizer.lr, self.optimizer.scheduler.get("warmup_epochs", 1), self.epochs) + elif self.optimizer.scheduler.name == "StepLR": + optimizers_['lr_scheduler'] = optim.lr_scheduler.StepLR(optimizers_['optimizer'], step_size=self.optimizer.scheduler.get("step_size", self.trainer.max_epochs//3), \ + gamma=self.optimizer.scheduler.get("gamma", 0.1)) + + return optimizers_ def on_train_epoch_start(self): self.time_start_train_epoch = time.time() @@ -191,8 +209,8 @@ def on_train_epoch_end(self): def on_validation_epoch_end(self): self.model.timer.print(random_colors=False, bold=False) - if self.av2_mode == 'test': - print(f"\nModel: {self.model.__class__.__name__}, Checkpoint from: {self.load_checkpoint_path}") + if self.data_mode == 'test': + print(f"\nModel: {self.model.__class__.__name__}, Checkpoint from: {self.checkpoint}") print(f"Test results saved in: {self.save_res_path}, Please run submit command and upload to online leaderboard for results.") if self.leaderboard_version == 1: print(f"\nevalai challenge 2010 phase 4018 submit --file {self.save_res_path}.zip --large --private\n") @@ -204,9 +222,9 @@ def on_validation_epoch_end(self): # wandb.log_artifact(output_file) return - if self.av2_mode == 'val': - print(f"\nModel: {self.model.__class__.__name__}, Checkpoint from: {self.load_checkpoint_path}") - print(f"More details parameters and training status are in checkpoints") + if self.data_mode == 'val': + print(f"\nModel: {self.model.__class__.__name__}, Checkpoint from: {self.checkpoint}") + print(f"More details parameters and training status are in the checkpoint file.") self.metrics.normalize() @@ -223,8 +241,8 @@ def on_validation_epoch_end(self): # Save the dictionaries to a pickle file with open(str(self.save_res_path)+'.pkl', 'wb') as f: pickle.dump((self.metrics.epe_3way, self.metrics.bucketed, self.metrics.epe_ssf), f) - print(f"We already write the {self.vis_name} into the dataset, please run following commend to visualize the flow. Copy and paste it to your terminal:") - print(f"python tools/visualization.py --res_name '{self.vis_name}' --data_dir {self.dataset_path}") + print(f"We already write the {self.res_name} into the dataset, please run following commend to visualize the flow. Copy and paste it to your terminal:") + print(f"python tools/visualization.py --res_name '{self.res_name}' --data_dir {self.dataset_path}") print(f"Enjoy! ^v^ ------ \n") self.metrics = OfficialMetrics() @@ -247,7 +265,7 @@ def eval_only_step_(self, batch, res_dict): else: final_flow[~batch['gm0']] = res_dict['flow'] + pose_flow[~batch['gm0']] - if self.av2_mode == 'val': # since only val we have ground truth flow to eval + if self.data_mode == 'val': # since only val we have ground truth flow to eval gt_flow = batch["flow"] v1_dict = evaluate_leaderboard(final_flow[eval_mask], pose_flow[eval_mask], pc0[eval_mask], \ gt_flow[eval_mask], batch['flow_is_valid'][eval_mask], \ @@ -259,7 +277,7 @@ def eval_only_step_(self, batch, res_dict): self.metrics.step(v1_dict, v2_dict, ssf_dict) # NOTE (Qingwen): Since val and test, we will force set batch_size = 1 - if self.save_res or self.av2_mode == 'test': # test must save data to submit in the online leaderboard. + if self.save_res or self.data_mode == 'test': # test must save data to submit in the online leaderboard. save_pred_flow = final_flow[eval_mask, :3].cpu().detach().numpy() rigid_flow = pose_flow[eval_mask, :3].cpu().detach().numpy() is_dynamic = np.linalg.norm(save_pred_flow - rigid_flow, axis=1, ord=2) >= 0.05 @@ -273,8 +291,10 @@ def run_model_wo_ground_data(self, batch): batch['origin_pc0'] = batch['pc0'].clone() batch['pc0'] = batch['pc0'][~batch['gm0']].unsqueeze(0) batch['pc1'] = batch['pc1'][~batch['gm1']].unsqueeze(0) - if 'pcb0' in batch: - batch['pcb0'] = batch['pcb0'][~batch['gmb0']].unsqueeze(0) + + for i in range(1, self.num_frames-1): + batch[f'pch{i}'] = batch[f'pch{i}'][~batch[f'gmh{i}']].unsqueeze(0) + self.model.timer[12].start("One Scan") res_dict = self.model(batch) self.model.timer[12].stop() @@ -285,7 +305,7 @@ def run_model_wo_ground_data(self, batch): return batch, res_dict def validation_step(self, batch, batch_idx): - if self.av2_mode == 'val' or self.av2_mode == 'test': + if self.data_mode in ['val', 'test']: batch, res_dict = self.run_model_wo_ground_data(batch) self.model.timer[13].start("Eval") self.eval_only_step_(batch, res_dict) @@ -317,13 +337,13 @@ def test_step(self, batch, batch_idx): key = str(batch['timestamp']) scene_id = batch['scene_id'] with h5py.File(os.path.join(self.dataset_path, f'{scene_id}.h5'), 'r+') as f: - if self.vis_name in f[key]: - del f[key][self.vis_name] - f[key].create_dataset(self.vis_name, data=final_flow.cpu().detach().numpy().astype(np.float32)) + if self.res_name in f[key]: + del f[key][self.res_name] + f[key].create_dataset(self.res_name, data=final_flow.cpu().detach().numpy().astype(np.float32)) def on_test_epoch_end(self): self.model.timer.print(random_colors=False, bold=False) - print(f"\n\nModel: {self.model.__class__.__name__}, Checkpoint from: {self.load_checkpoint_path}") + print(f"\n\nModel: {self.model.__class__.__name__}, Checkpoint from: {self.checkpoint}") print(f"We already write the flow_est into the dataset, please run following commend to visualize the flow. Copy and paste it to your terminal:") - print(f"python tools/visualization.py --res_name '{self.vis_name}' --data_dir {self.dataset_path}") + print(f"python tools/visualization.py --res_name '{self.res_name}' --data_dir {self.dataset_path}") print(f"Enjoy! ^v^ ------ \n") diff --git a/src/utils/mics.py b/src/utils/mics.py index c3385b7..795e585 100644 --- a/src/utils/mics.py +++ b/src/utils/mics.py @@ -285,7 +285,7 @@ def __getitem__(self, index): data_dict['pc0'] = f[key]['lidar'][:] data_dict['gm0'] = f[key]['ground_mask'][:] data_dict['pose0'] = f[key]['pose'][:] - for flow_key in self.vis_name + ['dufo_label', 'label']: + for flow_key in self.vis_name + ['dufo', 'label']: if flow_key in f[key]: data_dict[flow_key] = f[key][flow_key][:] diff --git a/tools/visualization.py b/tools/visualization.py index 9cd6651..e1d5e41 100644 --- a/tools/visualization.py +++ b/tools/visualization.py @@ -68,7 +68,7 @@ def check_flow( o3d_vis.update([pcd, pcd1, pcd2, o3d.geometry.TriangleMesh.create_coordinate_frame(size=2)]) def vis( - data_dir: str ="/home/kin/data/av2/preprocess/sensor/mini", + data_dir: str ="/home/kin/data/av2/h5py/demo/val", res_name: str = "flow", # any res_name we write before in HDF5Data start_id: int = 0, point_size: float = 2.0, @@ -124,7 +124,7 @@ def vis( def vis_multiple( - data_dir: str ="/home/kin/data/av2/preprocess/sensor/mini", + data_dir: str ="/home/kin/data/av2/h5py/demo/val", res_name: list = ["flow"], start_id: int = 0, point_size: float = 3.0, diff --git a/tools/write4conf.py b/tools/write4conf.py index 60cc560..bb42431 100644 --- a/tools/write4conf.py +++ b/tools/write4conf.py @@ -17,7 +17,14 @@ def main( model['hyper_parameters']['cfg']['model']['target']['_target_'] = new_path torch.save(model, output_path) +def readmodel( + model_path: str = "/home/kin/model_zoo/seflowpp.ckpt" +): + model = torch.load(model_path) + print(model['hyper_parameters']['cfg'], model['epoch']) + if __name__ == '__main__': start_time = time.time() - fire.Fire(main) + # fire.Fire(main) + fire.Fire(readmodel) print(f"Time used: {time.time() - start_time:.2f} s") \ No newline at end of file diff --git a/train.py b/train.py index 51be0b0..a9a563c 100644 --- a/train.py +++ b/train.py @@ -29,8 +29,8 @@ from src.trainer import ModelWrapper def precheck_cfg_valid(cfg): - if cfg.loss_fn == 'seflowLoss' and cfg.add_seloss is None: - raise ValueError("Please specify the self-supervised loss items for seflowLoss.") + if cfg.loss_fn in ['seflowLoss', 'seflowppLoss'] and (cfg.add_seloss is None or cfg.ssl_label is None): + raise ValueError("Please specify the self-supervised loss items and auto-label source for seflow-series loss.") grid_size = [(cfg.point_cloud_range[3] - cfg.point_cloud_range[0]) * (1/cfg.voxel_size[0]), (cfg.point_cloud_range[4] - cfg.point_cloud_range[1]) * (1/cfg.voxel_size[1]), @@ -57,7 +57,7 @@ def main(cfg): precheck_cfg_valid(cfg) pl.seed_everything(cfg.seed, workers=True) - train_dataset = HDF5Dataset(cfg.train_data, n_frames=cfg.num_frames, dufo=(cfg.loss_fn == 'seflowLoss')) + train_dataset = HDF5Dataset(cfg.train_data, n_frames=cfg.num_frames, ssl_label=cfg.get('ssl_label', None)) train_loader = DataLoader(train_dataset, batch_size=cfg.batch_size, shuffle=True, @@ -76,10 +76,11 @@ def main(cfg): output_dir = HydraConfig.get().runtime.output_dir # overwrite logging folder name for SSL. - if cfg.loss_fn == 'seflowLoss': - cfg.output = cfg.output.replace(cfg.model.name, "seflow") - output_dir = output_dir.replace(cfg.model.name, "seflow") - method_name = "seflow" + if cfg.loss_fn in ['seflowLoss', 'seflowppLoss']: + tmp_ = cfg.loss_fn.split('Loss')[0] + '-' + cfg.model.name + cfg.output = cfg.output.replace(cfg.model.name, tmp_) + output_dir = output_dir.replace(cfg.model.name, tmp_) + method_name = tmp_ else: method_name = cfg.model.name @@ -133,14 +134,15 @@ def main(cfg): print("Initiating wandb and trainer successfully. ^V^ ") print(f"We will use {cfg.gpus} GPUs to train the model. Check the checkpoints in {output_dir} checkpoints folder.") print("Total Train Dataset Size: ", len(train_dataset)) - if cfg.add_seloss is not None and cfg.loss_fn == 'seflowLoss': + if cfg.get('add_seloss', None) is not None and cfg.loss_fn in ['seflowLoss', 'seflowppLoss']: print(f"Note: We are in **self-supervised** training now. No ground truth label is used.") print(f"We will use these loss items in {cfg.loss_fn}: {cfg.add_seloss}") print("-"*40+"\n") # NOTE(Qingwen): search & check: def training_step(self, batch, batch_idx) trainer.fit(model, train_dataloaders = train_loader, val_dataloaders = val_loader, ckpt_path = cfg.checkpoint) - wandb.finish() + if cfg.wandb_mode != "disabled": + wandb.finish() if __name__ == "__main__": main() \ No newline at end of file