From bec8ddd6ce35aa09a933a4c5943c77de5e869676 Mon Sep 17 00:00:00 2001 From: Kin Date: Mon, 11 Aug 2025 21:58:37 +0200 Subject: [PATCH 01/20] feat(deflowpp): update deflowpp model. * hotfix(ssf): bug fixes in seflow+ssf as ssf concat two pc and forget to - index. --- conf/model/deflowpp.yaml | 11 +++ src/models/__init__.py | 2 +- src/models/basic/unet.py | 147 ++++++++++++++++++++++++++++++++++++++- src/models/deflow.py | 104 ++++++++++++++++++++------- src/models/fastflow3d.py | 26 ++----- src/models/ssf.py | 33 ++------- 6 files changed, 246 insertions(+), 77 deletions(-) create mode 100644 conf/model/deflowpp.yaml 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/src/models/__init__.py b/src/models/__init__.py index 4735f1c..a094526 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..84fce23 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, From 9208a7ebb70fc992cc703f8d79bc0ca1c316d8e6 Mon Sep 17 00:00:00 2001 From: Kin Date: Thu, 21 Aug 2025 00:07:26 +0200 Subject: [PATCH 02/20] feat(autolabel): update more autolabel through himo paper. * add data with lidar_id and lidar_dt * add flow_instances for afterward easy deltaflow update. --- README.md | 34 +++-- assets/README.md | 3 + conf/config.yaml | 5 +- dataprocess/README.md | 6 +- dataprocess/extract_av2.py | 48 ++++--- envsftool.yaml | 7 +- process.py | 254 +++++++++++++++++++++++++++---------- 7 files changed, 251 insertions(+), 106 deletions(-) diff --git a/README.md b/README.md index 41f9f6e..061a5d7 100644 --- a/README.md +++ b/README.md @@ -13,8 +13,8 @@ It is also an official implementation of the following papers (sorted by the tim - **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) - **Flow4D: Leveraging 4D Voxel Network for LiDAR Scene Flow Estimation** *Jaeyeul Kim, Jungwan Woo, Ukcheol Shin, Jean Oh, Sunghoon Im* @@ -134,10 +134,8 @@ Train Flow4D with the leaderboard submit config. [Runtime: Around 18 hours in 4x ```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]" -``` -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 ``` @@ -165,18 +163,30 @@ wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/ssf_best.ckpt wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/ssf_long.ckpt ``` +### Feed-Forward Self-Supervised Model Training -### SeFlow +Train SeFlow/SeFlow++ needed to: +1) process auto-label process. +2) specify the loss function, we set the config of our best model in the leaderboard. -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.] +#### SeFlow ```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" +# [Runtime: Around 11 hours in 4x A100 GPUs.] +python train.py model=deflow 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 ``` -Pretrained weight can be downloaded through: +#### SeFlow++ + ```bash -wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/seflow_best.ckpt +# [Runtime: Around ? hours in ? GPUs.] +python train.py model=deflowpp lr=2e-4 epochs=9 batch_size=16 loss_fn=seflowppLoss +ssl_label=seflowpp_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/seflowpp_best.ckpt ``` ### DeFlow @@ -185,10 +195,8 @@ Train DeFlow with the leaderboard submit config. [Runtime: Around 6-8 hours in 4 ```bash python train.py model=deflow lr=2e-4 epochs=15 batch_size=16 loss_fn=deflowLoss -``` -Pretrained weight can be downloaded through: -```bash +# Pretrained weight can be downloaded through: wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/deflow_best.ckpt ``` diff --git a/assets/README.md b/assets/README.md index 837b01a..74115ed 100644 --- a/assets/README.md +++ b/assets/README.md @@ -100,3 +100,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/conf/config.yaml b/conf/config.yaml index 1e34f76..0ddd779 100644 --- a/conf/config.yaml +++ b/conf/config.yaml @@ -27,8 +27,9 @@ gradient_clip_val: 5.0 # optimizer ==> Adam lr: 2e-6 -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} +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/dataprocess/README.md b/dataprocess/README.md index 893af64..518cdc8 100644 --- a/dataprocess/README.md +++ b/dataprocess/README.md @@ -175,9 +175,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 +186,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..d3731be 100644 --- a/dataprocess/extract_av2.py +++ b/dataprocess/extract_av2.py @@ -40,7 +40,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 +92,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 +135,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 +180,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 +211,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[0], '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 +247,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/envsftool.yaml b/envsftool.yaml index 8ac1925..b8bf407 100644 --- a/envsftool.yaml +++ b/envsftool.yaml @@ -6,6 +6,9 @@ dependencies: - python=3.8 - pytorch::pytorch=2.0.0 - pytorch::torchvision + - pytorch::pytorch-cuda=11.8 + - nvidia/label/cuda-11.8.0::cuda + - nvidia/label/cuda-11.8.0::cuda-toolkit - mkl==2024.0.0 - numba - numpy==1.22 @@ -13,9 +16,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/process.py b/process.py index 076747d..e062155 100644 --- a/process.py +++ b/process.py @@ -3,35 +3,49 @@ # 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): + 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 +57,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 +70,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 From 222e0659fe79019957602196cb39d16ede004395 Mon Sep 17 00:00:00 2001 From: Kin Date: Thu, 21 Aug 2025 00:09:45 +0200 Subject: [PATCH 03/20] feat(seflowpp): the full seflowpp process and train scripts. * tested with demo good, need double check the results on av2 and update the training time & weight link also. --- README.md | 2 +- dataprocess/README.md | 8 +- src/autolabel.py | 58 ++++++++++ src/dataset.py | 50 ++++++--- src/lossfuncs.py | 157 -------------------------- src/lossfuncs/__init__.py | 19 ++++ src/lossfuncs/selfsupervise.py | 194 +++++++++++++++++++++++++++++++++ src/lossfuncs/supervise.py | 107 ++++++++++++++++++ src/trainer.py | 9 +- train.py | 6 +- 10 files changed, 426 insertions(+), 184 deletions(-) create mode 100644 src/autolabel.py delete mode 100644 src/lossfuncs.py create mode 100644 src/lossfuncs/__init__.py create mode 100644 src/lossfuncs/selfsupervise.py create mode 100644 src/lossfuncs/supervise.py diff --git a/README.md b/README.md index 061a5d7..b1878c8 100644 --- a/README.md +++ b/README.md @@ -183,7 +183,7 @@ wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/seflow_best.ckp ```bash # [Runtime: Around ? hours in ? GPUs.] -python train.py model=deflowpp lr=2e-4 epochs=9 batch_size=16 loss_fn=seflowppLoss +ssl_label=seflowpp_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" +python train.py model=deflowpp lr=2e-4 epochs=9 batch_size=16 loss_fn=seflowppLoss +ssl_label=seflowpp_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" num_frames=3 # Pretrained weight can be downloaded through: wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/seflowpp_best.ckpt diff --git a/dataprocess/README.md b/dataprocess/README.md index 518cdc8..a9962d5 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). +- [ ] ZOD (w/o gt): check [here](#zod-dataset). The process script was involved from [HiMo](https://kin-zhang.github.io/HiMo). - [ ] 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: @@ -133,10 +133,10 @@ This directory contains the scripts to preprocess the datasets. 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 +146,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 ``` diff --git a/src/autolabel.py b/src/autolabel.py new file mode 100644 index 0000000..f9a0b5a --- /dev/null +++ b/src/autolabel.py @@ -0,0 +1,58 @@ +""" +# 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['cluster'][:].astype(np.int16)) + # raw seflow: + 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): + 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) > 0.05 and max(r_dufo, r_nnd) > 0.1: + 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..5962ec6 --- /dev/null +++ b/src/lossfuncs/selfsupervise.py @@ -0,0 +1,194 @@ +""" +# 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") + # 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.nanmean(est_dist0[est_dist0 <= TRUNCATED_DIST]) + torch.nanmean(est_dist1[est_dist1 <= TRUNCATED_DIST]) + chamfer_dis += 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") + 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 / 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/trainer.py b/src/trainer.py index fe2cb16..d98acdc 100644 --- a/src/trainer.py +++ b/src/trainer.py @@ -92,6 +92,7 @@ def __init__(self, cfg, eval=False): self.vis_name = cfg.res_name if 'res_name' in cfg else 'default' 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) @@ -101,7 +102,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: @@ -130,11 +131,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: + 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): @@ -143,7 +148,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() diff --git a/train.py b/train.py index 51be0b0..9c64d3b 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.ssl_label) train_loader = DataLoader(train_dataset, batch_size=cfg.batch_size, shuffle=True, From aa642ef77946f828f2060954ae15715753329c13 Mon Sep 17 00:00:00 2001 From: qingwen-berzelius Date: Thu, 21 Aug 2025 12:31:06 +0200 Subject: [PATCH 04/20] fix(trainer): add seflowpp into trainer for folder name. --- train.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/train.py b/train.py index 9c64d3b..af94361 100644 --- a/train.py +++ b/train.py @@ -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 @@ -140,7 +141,8 @@ def main(cfg): # 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 From ff1df2870fa509da92941651f97db8d9739d4a90 Mon Sep 17 00:00:00 2001 From: qingwen-berzelius Date: Thu, 21 Aug 2025 13:14:32 +0200 Subject: [PATCH 05/20] data(zod): update zod extraction scripts. * it could be a good reference for users to extract other data into h5 files etc. --- conf/ground/zod.toml | 31 ++++++++ dataprocess/README.md | 19 ++++- dataprocess/extract_zod.py | 140 +++++++++++++++++++++++++++++++++++++ 3 files changed, 188 insertions(+), 2 deletions(-) create mode 100644 conf/ground/zod.toml create mode 100644 dataprocess/extract_zod.py 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/dataprocess/README.md b/dataprocess/README.md index a9962d5..2651afa 100644 --- a/dataprocess/README.md +++ b/dataprocess/README.md @@ -11,7 +11,7 @@ We've updated the process dataset for: - [x] Argoverse 2.0: check [here](#argoverse-20). The process script Involved from [DeFlow](https://github.com/KTH-RPL/DeFlow). - [x] Waymo: check [here](#waymo-dataset). The process script was involved from [SeFlow](https://github.com/KTH-RPL/SeFlow). - [x] nuScenes: check [here](#nuscenes), The process script was involved from [DeltaFlow](https://github.com/Kin-Zhang/DeltaFlow). -- [ ] ZOD (w/o gt): check [here](#zod-dataset). The process script was involved from [HiMo](https://kin-zhang.github.io/HiMo). +- [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. 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,12 +123,27 @@ 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 diff --git a/dataprocess/extract_zod.py b/dataprocess/extract_zod.py new file mode 100644 index 0000000..e7e29fb --- /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 src.utils.mics 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/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 From 39b6a1819c3ccf8f006938c294eb479c8c5216d5 Mon Sep 17 00:00:00 2001 From: qingwen-berzelius Date: Sat, 23 Aug 2025 18:12:19 +0200 Subject: [PATCH 06/20] !feat(lr): update optimizer to new structure. - update from DeltaFlow. - align some format with copilot comments. * update other submodule repo to align the lr changed. --- README.md | 10 +++--- conf/config.yaml | 4 ++- process.py | 11 +++++++ src/autolabel.py | 10 +++--- src/models/ssf.py | 2 +- src/trainer.py | 84 ++++++++++++++++++++++++++++------------------- train.py | 2 +- 7 files changed, 78 insertions(+), 45 deletions(-) diff --git a/README.md b/README.md index b1878c8..a6bdee8 100644 --- a/README.md +++ b/README.md @@ -133,7 +133,7 @@ mamba activate opensf 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: wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/flow4d_best.ckpt @@ -151,7 +151,7 @@ pip install https://data.pyg.org/whl/torch-2.0.0%2Bcu118/torch_scatter-2.1.2%2Bp 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: @@ -173,7 +173,7 @@ Train SeFlow/SeFlow++ needed to: ```bash # [Runtime: Around 11 hours in 4x A100 GPUs.] -python train.py model=deflow 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" +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 @@ -183,7 +183,7 @@ wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/seflow_best.ckp ```bash # [Runtime: Around ? hours in ? GPUs.] -python train.py model=deflowpp lr=2e-4 epochs=9 batch_size=16 loss_fn=seflowppLoss +ssl_label=seflowpp_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" num_frames=3 +python train.py model=deflowpp optimizer.lr=2e-4 epochs=9 batch_size=16 loss_fn=seflowppLoss +ssl_label=seflowpp_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" num_frames=3 # Pretrained weight can be downloaded through: wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/seflowpp_best.ckpt @@ -194,7 +194,7 @@ wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/seflowpp_best.c 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=15 batch_size=16 loss_fn=deflowLoss +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 diff --git a/conf/config.yaml b/conf/config.yaml index 0ddd779..6fc86ca 100644 --- a/conf/config.yaml +++ b/conf/config.yaml @@ -26,7 +26,9 @@ epochs: 3 gradient_clip_val: 5.0 # optimizer ==> Adam -lr: 2e-6 +optimizer: + name: Adam # [Adam, AdamW] + lr: 1e-3 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: diff --git a/process.py b/process.py index e062155..d6eeca4 100644 --- a/process.py +++ b/process.py @@ -29,6 +29,17 @@ 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] diff --git a/src/autolabel.py b/src/autolabel.py index f9a0b5a..0d54768 100644 --- a/src/autolabel.py +++ b/src/autolabel.py @@ -30,14 +30,16 @@ def shiftClusterid(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['cluster'][:].astype(np.int16)) - # raw seflow: 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): +def seflowpp_auto(input_data, tau1=0.05, tau2=0.1): + """ + 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. + """ dufo = input_data['dufo'][:].astype(np.uint8) cluster = shiftClusterid(input_data['cluster'][:].astype(np.int16)) nnd = input_data['nnd'][:].astype(np.uint8) @@ -53,6 +55,6 @@ def seflowpp_auto(input_data): r_dufo = np.sum(cluster_dufo>0)/all_pts r_nnd = np.sum(cluster_nnd>0)/all_pts - if min(r_dufo, r_nnd) > 0.05 and max(r_dufo, r_nnd) > 0.1: + 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/models/ssf.py b/src/models/ssf.py index 84fce23..1c0eeb6 100644 --- a/src/models/ssf.py +++ b/src/models/ssf.py @@ -125,7 +125,7 @@ def forward(self, batch): pc0_valid_point_idxes = [e["point_idxes"] for e in pc0_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] + pc1_valid_point_idxes = [e["point_idxes"] - pc0s.shape[1] for e in pc1_voxel_infos_lst] model_res = { "flow": flows, diff --git a/src/trainer.py b/src/trainer.py index d98acdc..6912c23 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, + "av2_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,42 +72,28 @@ 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.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 - - - 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) + # ---> inference mode + if self.save_res and self.av2_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.av2_mode}, results will be saved in: {self.save_res_path} with version: {self.leaderboard_version} format for online leaderboard.") - 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 self.save_hyperparameters() # FIXME(Qingwen 2025-08-20): update the loss_calculation fn alone to make all things pretty here.... @@ -177,8 +181,22 @@ def train_validation_step_(self, batch, res_dict): pass def configure_optimizers(self): - optimizer = optim.Adam(self.model.parameters(), lr=self.lr) - return optimizer + 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() diff --git a/train.py b/train.py index af94361..8f6532d 100644 --- a/train.py +++ b/train.py @@ -134,7 +134,7 @@ 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.add_seloss 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") From 0ab8559f8fd8c1b06c2597b32d33a865e3664c60 Mon Sep 17 00:00:00 2001 From: qingwen-berzelius Date: Sat, 23 Aug 2025 18:21:26 +0200 Subject: [PATCH 07/20] dcos(slurm): update slurm script, for reader to easy check how is the trainining setup. --- assets/slurm/0_process.sh | 20 +++++++++++--------- assets/slurm/1_train.sh | 17 +++++++---------- assets/slurm/2_eval.sh | 25 ++++++------------------- 3 files changed, 24 insertions(+), 38 deletions(-) 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..1073745 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 av2_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 av2_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 From 3f6f18e9c451ad85e943973a0f328760128ee1e7 Mon Sep 17 00:00:00 2001 From: qingwen-berzelius Date: Sun, 24 Aug 2025 11:00:38 +0200 Subject: [PATCH 08/20] fix(env): fix some potiential env issue later * add some notes --- assets/README.md | 1 + assets/cuda/chamfer3D/setup.py | 20 +++++++++++++++----- assets/cuda/mmcv/setup.py | 18 +++++++++--------- environment.yaml | 2 ++ 4 files changed, 27 insertions(+), 14 deletions(-) diff --git a/assets/README.md b/assets/README.md index 74115ed..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')" ``` 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/environment.yaml b/environment.yaml index d520844..946145d 100644 --- a/environment.yaml +++ b/environment.yaml @@ -34,6 +34,8 @@ dependencies: - assets/cuda/chamfer3D - assets/cuda/mmcv - open3d==0.18.0 + # - opencv-python-headless==4.9.0.80 # for av2 if there is any on cv2.... + - av==14.0.1 # need to specific this one otherwise av2 will have error - av2==0.3.1 - spconv-cu118==2.3.6 - numpy==1.26.4 From cf7b20ac079122a7d74423dc7b4f11de1a94e4f1 Mon Sep 17 00:00:00 2001 From: qingwen-berzelius Date: Sun, 24 Aug 2025 11:03:02 +0200 Subject: [PATCH 09/20] fix(av2): instance label typo. --- dataprocess/extract_av2.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/dataprocess/extract_av2.py b/dataprocess/extract_av2.py index d3731be..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 @@ -180,7 +182,7 @@ 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) + instances[obj_mask] = dclass[id]+1 else: valid[obj_mask] = 0 return flow, classes, valid, ego1_SE3_ego0, instances @@ -215,7 +217,7 @@ def 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], 'instances': instances[0], + '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) : From d0c7c59e5f5888986d790380fa56b3d806940d8d Mon Sep 17 00:00:00 2001 From: Kin Date: Sun, 24 Aug 2025 11:09:10 +0200 Subject: [PATCH 10/20] fix(process): update key name in new version for seflow-variant process. --- src/utils/mics.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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][:] From fbf0fe7f7a1565124606aa846fd85dc5bc3ab4bc Mon Sep 17 00:00:00 2001 From: Kin Date: Sun, 24 Aug 2025 13:07:51 +0200 Subject: [PATCH 11/20] hotfix(eval): updating num_frames into eval. --- eval.py | 3 ++- save.py | 3 ++- src/trainer.py | 9 +++++---- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/eval.py b/eval.py index bfe7bdd..0c96b8b 100644 --- a/eval.py +++ b/eval.py @@ -48,6 +48,7 @@ def main(cfg): 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.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") @@ -63,7 +64,7 @@ def main(cfg): 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, \ + n_frames=cfg.num_frames, \ eval=True, leaderboard_version=cfg.leaderboard_version), \ batch_size=1, shuffle=False)) wandb.finish() diff --git a/save.py b/save.py index b9195e6..93840ba 100644 --- a/save.py +++ b/save.py @@ -45,6 +45,7 @@ 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, @@ -57,7 +58,7 @@ def main(cfg): # 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() diff --git a/src/trainer.py b/src/trainer.py index 6912c23..2997a94 100644 --- a/src/trainer.py +++ b/src/trainer.py @@ -94,6 +94,7 @@ def __init__(self, cfg, eval=False): 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.") # self.test_total_num = 0 + print(cfg) self.save_hyperparameters() # FIXME(Qingwen 2025-08-20): update the loss_calculation fn alone to make all things pretty here.... @@ -208,7 +209,7 @@ 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}") + 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") @@ -221,8 +222,8 @@ def on_validation_epoch_end(self): 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") + 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() @@ -339,7 +340,7 @@ def test_step(self, batch, batch_idx): 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"Enjoy! ^v^ ------ \n") From 3c82d4825d881073a2750cf100ad990ab183ecef Mon Sep 17 00:00:00 2001 From: Kin Date: Sun, 24 Aug 2025 13:24:15 +0200 Subject: [PATCH 12/20] hotfix(eval/test): for history frames, we update keys' name and it need rm ground mask also. --- src/trainer.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/trainer.py b/src/trainer.py index 2997a94..7389769 100644 --- a/src/trainer.py +++ b/src/trainer.py @@ -290,8 +290,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() From 876cfd45fde1e82d0298057d1cdc874bc1827c97 Mon Sep 17 00:00:00 2001 From: Kin Date: Sun, 24 Aug 2025 14:00:10 +0200 Subject: [PATCH 13/20] small fix on ssl_label to None if under supervise training. --- train.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/train.py b/train.py index 8f6532d..a9a563c 100644 --- a/train.py +++ b/train.py @@ -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, ssl_label=cfg.ssl_label) + 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, @@ -134,7 +134,7 @@ 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 in ['seflowLoss', 'seflowppLoss']: + 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") From e7da73b91b2555616be38a217b115d354d6a9416 Mon Sep 17 00:00:00 2001 From: qingwen-berzelius Date: Thu, 28 Aug 2025 00:22:30 +0200 Subject: [PATCH 14/20] docs(README): update readme from main merge. --- README.md | 66 +++++++++++++++++++++++++++--------------------- src/autolabel.py | 4 ++- 2 files changed, 40 insertions(+), 30 deletions(-) diff --git a/README.md b/README.md index 38dc035..179546a 100644 --- a/README.md +++ b/README.md @@ -11,6 +11,11 @@ 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/todo.todo) ] [ [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* IEEE Transactions on Robotics (**T-RO**) 2025 @@ -134,28 +139,9 @@ And free yourself from trainning, you can download the pretrained weight from [H mamba 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): +### Supervised Training -```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+cu118.html -``` - -Train VoteFlow with the leaderboard submit config. [Runtime: Around 32 hours in 4 x V100 GPUs.] -```bash -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: -```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.] @@ -166,7 +152,7 @@ python train.py model=flow4d optimizer.lr=1e-3 epochs=15 batch_size=8 num_frames wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/flow4d_best.ckpt ``` -### SSF +#### SSF Extra pakcges needed for SSF model: ```bash @@ -190,11 +176,23 @@ wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/ssf_best.ckpt wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/ssf_long.ckpt ``` -### Feed-Forward Self-Supervised Model Training +#### DeFlow + +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) -Train SeFlow/SeFlow++ needed to: +```bash +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 +``` + + +### Self-Supervised Model Training + +Train Feed-forward SSL methods (e.g. SeFlow/SeFlow++), we needed to: 1) process auto-label process. -2) specify the loss function, we set the config of our best model in the leaderboard. +2) specify the loss function, below we set the config to align our best model in the leaderboard. #### SeFlow @@ -216,17 +214,27 @@ python train.py model=deflowpp optimizer.lr=2e-4 epochs=9 batch_size=16 loss_fn= wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/seflowpp_best.ckpt ``` -### DeFlow -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) +### 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 -python train.py model=deflow optimizer.lr=2e-4 epochs=15 batch_size=16 loss_fn=deflowLoss +# Install Pytorch3d +conda install pytorch3d -c pytorch3d + +# Install torch-scatter +pip install torch-scatter -f https://data.pyg.org/whl/torch-2.0.0+cu118.html +``` + +Train VoteFlow with the leaderboard submit config. [Runtime: Around 32 hours in 4 x V100 GPUs.] +```bash +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/deflow_best.ckpt +wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/voteflow_best.ckpt ``` + ## 3. Evaluation You can view Wandb dashboard for the training and evaluation results or upload result to online leaderboard. diff --git a/src/autolabel.py b/src/autolabel.py index 0d54768..736bbb6 100644 --- a/src/autolabel.py +++ b/src/autolabel.py @@ -35,10 +35,12 @@ def seflow_auto(input_data): 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.1): +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)) From 219ddfe061787553c586168664f09818d0ec979e Mon Sep 17 00:00:00 2001 From: qingwen-berzelius Date: Sat, 30 Aug 2025 01:10:51 +0200 Subject: [PATCH 15/20] doc(env): update env in Dockerfile and README. with some print cfg during val. --- Dockerfile | 2 ++ README.md | 45 ++++++++++++------------------- assets/cuda/chamfer3D/__init__.py | 6 ++--- process.py | 2 +- requirements.txt | 3 ++- src/lossfuncs/selfsupervise.py | 20 ++++++-------- src/trainer.py | 3 ++- 7 files changed, 35 insertions(+), 46 deletions(-) diff --git a/Dockerfile b/Dockerfile index 0a50ff6..63035d3 100644 --- a/Dockerfile +++ b/Dockerfile @@ -26,6 +26,8 @@ WORKDIR /home/kin/workspace/OpenSceneFlow # 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/conda install pytorch3d -c pytorch3d +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 # environment for dataprocessing includes data-api diff --git a/README.md b/README.md index 179546a..9f8ba46 100644 --- a/README.md +++ b/README.md @@ -62,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). @@ -120,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. @@ -187,12 +178,11 @@ python train.py model=deflow optimizer.lr=2e-4 epochs=15 batch_size=16 loss_fn=d wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/deflow_best.ckpt ``` +### Feed-Forward Self-Supervised Model Training -### Self-Supervised Model Training - -Train Feed-forward SSL methods (e.g. SeFlow/SeFlow++), we needed to: +Train SeFlow/SeFlow++/VoteFlow needed to: 1) process auto-label process. -2) specify the loss function, below we set the config to align our best model in the leaderboard. +2) specify the loss function, we set the config here for our best model in the leaderboard. #### SeFlow @@ -204,18 +194,7 @@ python train.py model=deflow optimizer.lr=2e-4 epochs=9 batch_size=16 loss_fn=se wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/seflow_best.ckpt ``` -#### SeFlow++ - -```bash -# [Runtime: Around ? hours in ? GPUs.] -python train.py model=deflowpp optimizer.lr=2e-4 epochs=9 batch_size=16 loss_fn=seflowppLoss +ssl_label=seflowpp_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" num_frames=3 - -# Pretrained weight can be downloaded through: -wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/seflowpp_best.ckpt -``` - - -### VoteFLow +#### 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 @@ -235,6 +214,16 @@ wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/voteflow_best.c ``` +#### SeFlow++ + +```bash +# [Runtime: Around 10 hours in 4x A100 GPUs.] for Argoverse 2 +python train.py model=deflowpp optimizer.lr=2e-4 epochs=9 batch_size=8 loss_fn=seflowppLoss +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" num_frames=3 + +# Pretrained weight can be downloaded through: +wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/seflowpp_best.ckpt +``` + ## 3. Evaluation You can view Wandb dashboard for the training and evaluation results or upload result to online leaderboard. 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/process.py b/process.py index d6eeca4..90a6851 100644 --- a/process.py +++ b/process.py @@ -276,5 +276,5 @@ def main( fire.Fire(main) fire.Fire(run_dufocluster) print("\nAlready Finished the main labels: dufo, cluster, ground mask etc.\n") - fire.Fire(run_nnd) + # 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 index cbce091..34c1c97 100644 --- a/requirements.txt +++ b/requirements.txt @@ -13,13 +13,14 @@ omegaconf hydra-core tabulate pathtools -numpy==1.26.0 +numpy==1.26.4 scipy numba pandas scikit-learn h5py hdbscan +av==14.0.1 av2==0.3.4 spconv-cu118==2.3.8 flash-attn==2.6.3 diff --git a/src/lossfuncs/selfsupervise.py b/src/lossfuncs/selfsupervise.py index 5962ec6..dccc923 100644 --- a/src/lossfuncs/selfsupervise.py +++ b/src/lossfuncs/selfsupervise.py @@ -44,15 +44,9 @@ def seflowppLoss(res_dict, timer=None): # 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.nanmean(est_dist0[est_dist0 <= TRUNCATED_DIST]) + torch.nanmean(est_dist1[est_dist1 <= TRUNCATED_DIST]) - chamfer_dis += MyCUDAChamferDis(pseduo_pch1from0, pch1, truncate_dist=TRUNCATED_DIST) + 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 @@ -70,12 +64,14 @@ def seflowppLoss(res_dict, timer=None): # 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 paper + # 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: @@ -84,7 +80,7 @@ def seflowppLoss(res_dict, timer=None): if cluster_nnd.shape[0] <= 0: continue - # Eq. 8 in the paper + # 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) @@ -92,14 +88,14 @@ def seflowppLoss(res_dict, timer=None): continue max_idx = sorted_idxs[non_zero_valid_indices.squeeze(1)[0]] - # Eq. 9 in the paper + # Eq. 9 in the SeFlow paper max_flow = pc1[raw_idx0[mask][max_idx]] - pc0[mask][max_idx] - # Eq. 10 in the paper + # 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 paper + 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() diff --git a/src/trainer.py b/src/trainer.py index 7389769..97ee903 100644 --- a/src/trainer.py +++ b/src/trainer.py @@ -94,7 +94,8 @@ def __init__(self, cfg, eval=False): 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.") # self.test_total_num = 0 - print(cfg) + if self.av2_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.... From 4341de26680fff9eb5a44ed19a658b7da4ab8be4 Mon Sep 17 00:00:00 2001 From: Kin Date: Sat, 30 Aug 2025 01:12:39 +0200 Subject: [PATCH 16/20] docs(README): update README. --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 9f8ba46..3bb7611 100644 --- a/README.md +++ b/README.md @@ -180,7 +180,7 @@ wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/deflow_best.ckp ### Feed-Forward Self-Supervised Model Training -Train SeFlow/SeFlow++/VoteFlow needed to: +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. @@ -218,7 +218,7 @@ wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/voteflow_best.c ```bash # [Runtime: Around 10 hours in 4x A100 GPUs.] for Argoverse 2 -python train.py model=deflowpp optimizer.lr=2e-4 epochs=9 batch_size=8 loss_fn=seflowppLoss +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" num_frames=3 +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 From b27ab76b5267afa34d58b47a611a71211f3723d4 Mon Sep 17 00:00:00 2001 From: Kin Date: Sat, 30 Aug 2025 01:32:35 +0200 Subject: [PATCH 17/20] style(av2_mode): change name av2_mode to data_mode as we have more data involved now. * fix vis_name change to res_name. clean up some codes also. --- README.md | 10 +++++----- assets/slurm/2_eval.sh | 4 ++-- conf/config.yaml | 2 +- conf/eval.yaml | 2 +- conf/save.yaml | 2 +- eval.py | 10 +++++----- save.py | 13 ++++--------- src/runner.py | 2 +- src/trainer.py | 30 +++++++++++++++--------------- tools/visualization.py | 4 ++-- 10 files changed, 37 insertions(+), 42 deletions(-) diff --git a/README.md b/README.md index 3bb7611..71da33b 100644 --- a/README.md +++ b/README.md @@ -102,7 +102,7 @@ If you prefer to build the Docker image by yourself, Check [build-docker-image]( ## 1. Data Preparation -Refer to [dataprocess/README.md](dataprocess/README.md) for dataset download instructions. Currently, we support **Argoverse 2**, **Waymo**, **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). @@ -232,14 +232,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!)** @@ -256,7 +256,7 @@ 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 diff --git a/assets/slurm/2_eval.sh b/assets/slurm/2_eval.sh index 1073745..1a57440 100644 --- a/assets/slurm/2_eval.sh +++ b/assets/slurm/2_eval.sh @@ -12,9 +12,9 @@ cd /proj/berzelius-2023-364/users/x_qinzh/workspace/OpenSceneFlow # ====> leaderboard model -# $PYTHON eval.py wandb_mode=online dataset_path=/proj/berzelius-2023-364/users/x_qinzh/data/av2/autolabel av2_mode=test \ +# $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 \ # save_res=True -$PYTHON eval.py wandb_mode=online dataset_path=/proj/berzelius-2023-364/users/x_qinzh/data/av2/autolabel 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 6fc86ca..a9e68f0 100644 --- a/conf/config.yaml +++ b/conf/config.yaml @@ -28,7 +28,7 @@ gradient_clip_val: 5.0 # optimizer ==> Adam optimizer: name: Adam # [Adam, AdamW] - lr: 1e-3 + 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: 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/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/eval.py b/eval.py index 0c96b8b..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,7 +46,7 @@ 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))) @@ -63,7 +63,7 @@ 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}", \ + 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)) diff --git a/save.py b/save.py index 93840ba..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 @@ -48,19 +48,14 @@ def main(cfg): 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=cfg.num_frames), \ batch_size=1, shuffle=False)) - wandb.finish() if __name__ == "__main__": main() \ No newline at end of file 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 97ee903..4ef319e 100644 --- a/src/trainer.py +++ b/src/trainer.py @@ -53,7 +53,7 @@ def __init__(self, cfg, eval=False): "num_frames": 2, "optimizer": None, "dataset_path": None, - "av2_mode": None, + "data_mode": None, } for key, default in default_self_values.items(): setattr(self, key, cfg.get(key, default)) @@ -88,13 +88,13 @@ def __init__(self, cfg, eval=False): self.metrics = OfficialMetrics() # ---> inference mode - if self.save_res and self.av2_mode in ['val', 'valid', 'test']: + 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.av2_mode}, results will be saved in: {self.save_res_path} with version: {self.leaderboard_version} format for online leaderboard.") + 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.test_total_num = 0 - if self.av2_mode in ['val', 'valid', 'test']: + if self.data_mode in ['val', 'valid', 'test']: print(cfg) self.save_hyperparameters() @@ -209,7 +209,7 @@ 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': + 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: @@ -222,7 +222,7 @@ def on_validation_epoch_end(self): # wandb.log_artifact(output_file) return - if self.av2_mode == 'val': + 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.") @@ -241,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() @@ -265,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], \ @@ -277,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 @@ -305,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) @@ -337,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.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/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, From f0a783e5a1d1828dc5a7390067478f0581dd320c Mon Sep 17 00:00:00 2001 From: Kin Date: Sat, 30 Aug 2025 01:56:41 +0200 Subject: [PATCH 18/20] hotfix(zod): change the path of one lib on zod extraction. test successfully. vis good. --- dataprocess/extract_zod.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/dataprocess/extract_zod.py b/dataprocess/extract_zod.py index e7e29fb..d60ecdd 100644 --- a/dataprocess/extract_zod.py +++ b/dataprocess/extract_zod.py @@ -22,7 +22,7 @@ BASE_DIR = os.path.abspath(os.path.join( os.path.dirname( __file__ ), '..')) sys.path.append(BASE_DIR) -from src.utils.mics import create_reading_index +from dataprocess.misc_data import create_reading_index from linefit import ground_seg GROUNDSEG_config = f"{BASE_DIR}/conf/ground/zod.toml" @@ -52,8 +52,8 @@ def create_group_data(group, pc, pc_dt, pose, pc_id, gm=None, flow_0to1=None, fl # 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 + if cnt % 4 != 0: + continue camera_frame, lidar_frame = frame # img = camera_frame.read() @@ -121,7 +121,7 @@ def process_logs(data_dir: Path, output_dir: Path, nproc: int): # res = list(tqdm(p.imap_unordered(proc, args), total=len(logs), ncols=120)) def main( - dataset_root: str = "/home/kin/data/zod/drives", + 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, From d109526ee92f8827c0c214330ed47d91c83239f3 Mon Sep 17 00:00:00 2001 From: Kin Date: Sat, 30 Aug 2025 11:58:25 +0200 Subject: [PATCH 19/20] revert(env): revert all env to py38 and cu117. here are effort I tried to upgrade and found out some issues: https://github.com/KTH-RPL/OpenSceneFlow/discussions/11 check with env.yaml is good, double check the dockerfile for all methods. --- Dockerfile | 41 ++++++++++++++++++++++++----------------- README.md | 7 +++---- environment.yaml | 25 ++++++++++++------------- envsftool.yaml | 5 ++--- requirements.txt | 33 --------------------------------- src/trainer.py | 2 +- tools/write4conf.py | 9 ++++++++- 7 files changed, 50 insertions(+), 72 deletions(-) delete mode 100644 requirements.txt diff --git a/Dockerfile b/Dockerfile index 63035d3..3fc62ca 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 # install zsh and oh-my-zsh RUN apt update && apt install -y wget git zsh tmux vim g++ @@ -14,25 +21,25 @@ 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 +RUN /opt/conda/bin/conda init zsh && /opt/conda/bin/mamba 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/conda install pytorch3d -c pytorch3d -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 - -# environment for dataprocessing includes data-api -RUN /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 cd /home/kin/workspace/OpenSceneFlow && /opt/conda/bin/mamba env create -f environment.yaml +# VoteFlow and SSF +RUN /opt/conda/envs/opensf/bin/pip install torch-scatter -f https://data.pyg.org/whl/torch-2.0.0+cu117.html +# FastNSF +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 + +# 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 + +RUN cd /home/kin/workspace/OpenSceneFlow && /opt/conda/bin/mamba env create -f envsftool.yaml diff --git a/README.md b/README.md index 71da33b..946d486 100644 --- a/README.md +++ b/README.md @@ -147,9 +147,8 @@ wget https://huggingface.co/kin-zhang/OpenSceneFlow/resolve/main/flow4d_best.ckp 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.] @@ -202,7 +201,7 @@ Extra pakcges needed for VoteFlow, [pytorch3d](https://pytorch3d.org/) (prefer 0 conda install pytorch3d -c pytorch3d # Install torch-scatter -pip install torch-scatter -f https://data.pyg.org/whl/torch-2.0.0+cu118.html +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.] diff --git a/environment.yaml b/environment.yaml index 946145d..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,20 +33,20 @@ dependencies: - assets/cuda/chamfer3D - assets/cuda/mmcv - open3d==0.18.0 - # - opencv-python-headless==4.9.0.80 # for av2 if there is any on cv2.... - - av==14.0.1 # need to specific this one otherwise av2 will have error - - 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 b8bf407..887457a 100644 --- a/envsftool.yaml +++ b/envsftool.yaml @@ -6,9 +6,8 @@ dependencies: - python=3.8 - pytorch::pytorch=2.0.0 - pytorch::torchvision - - pytorch::pytorch-cuda=11.8 - - 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 - mkl==2024.0.0 - numba - numpy==1.22 diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index 34c1c97..0000000 --- a/requirements.txt +++ /dev/null @@ -1,33 +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.4 -scipy -numba -pandas -scikit-learn -h5py -hdbscan -av==14.0.1 -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/src/trainer.py b/src/trainer.py index 4ef319e..e2eca7a 100644 --- a/src/trainer.py +++ b/src/trainer.py @@ -137,7 +137,7 @@ 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: + 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 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 From 5a7c53e6df0259e82c9adb0dcb130381fe3ea774 Mon Sep 17 00:00:00 2001 From: Kin Date: Sat, 30 Aug 2025 14:28:19 +0200 Subject: [PATCH 20/20] docs(Dockerfile): update dockerfile * checked locally on all methods' training is good. docs update README --- Dockerfile | 17 ++++++++++------- README.md | 12 ++++++------ 2 files changed, 16 insertions(+), 13 deletions(-) diff --git a/Dockerfile b/Dockerfile index 3fc62ca..55d295e 100644 --- a/Dockerfile +++ b/Dockerfile @@ -9,7 +9,7 @@ RUN curl -o ~/miniforge3.sh -LO https://github.com/conda-forge/miniforge/release 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 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++ @@ -21,7 +21,6 @@ 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 && /opt/conda/bin/mamba init zsh # change to conda env ENV PATH /opt/conda/bin:$PATH @@ -31,15 +30,19 @@ 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 cd /home/kin/workspace/OpenSceneFlow && /opt/conda/bin/mamba env create -f environment.yaml -# VoteFlow and SSF +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 -# FastNSF 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 +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 -RUN cd /home/kin/workspace/OpenSceneFlow && /opt/conda/bin/mamba 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 && /opt/conda/bin/conda clean -ya diff --git a/README.md b/README.md index 946d486..ad357a0 100644 --- a/README.md +++ b/README.md @@ -14,7 +14,7 @@ It is also an official implementation of the following papers (sorted by the tim - **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/todo.todo) ] [ [Project](https://github.com/Kin-Zhang/DeltaFlow) ] +[ 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* @@ -72,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 @@ -95,7 +95,7 @@ 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. @@ -127,7 +127,7 @@ 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 ``` ### Supervised Training @@ -260,8 +260,8 @@ Note: The leaderboard result in DeFlow&SeFlow main paper is [version 1](https:// ```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