docs(readme): update opensf link and a source branch.

iyuner · iyuner · commit f0b5ba23b4be · 2025-02-18T14:25:51.000+01:00
* save still core file for reader to quick check.
diff --git a/OpenSceneFlow b/OpenSceneFlow
@@ -1 +1 @@
-Subproject commit dd377faf72337bc0928c9053ed1e0aee2149878b
+Subproject commit bce2b69c03e1cc1576eda2a23a1a75a4eb0d0475
diff --git a/README.md b/README.md
@@ -9,7 +9,7 @@ DeFlow: Decoder of Scene Flow Network in Autonomous Driving
 
 Task: Scene Flow Estimation in Autonomous Driving. 
 
-📜 2025/02/18: Merging all scene flow code to a codebase to update one general repo only. This repo still saved DeFlow README and [cluster slurm files](assets/slurm).
+📜 2025/02/18: Merging all scene flow code to [OpenSceneFLow codebase](https://github.com/KTH-RPL/OpenSceneFlow) for afterward code maintenance. This repo saved README, [cluster slurm files](assets/slurm), and [quick core file](decoder.py) in DeFlow. The old source code branch is also [available here](https://github.com/KTH-RPL/DeFlow/tree/source).
 
 🤗 2024/11/18 16:17: Update model and demo data download link through HuggingFace, personally I found that `wget` from the HuggingFace link is much faster than Zenodo.
 
diff --git a/decoder.py b/decoder.py
@@ -0,0 +1,220 @@
+import torch
+import torch.nn as nn
+from typing import List, Tuple, Dict
+from . import ConvWithNorms
+
+SPLIT_BATCH_SIZE = 512
+
+class MMHeadDecoder(nn.Module):
+
+    def __init__(self, pseudoimage_channels: int = 64):
+        super().__init__()
+
+        self.offset_encoder = nn.Linear(3, 128)
+
+        # FIXME: figure out how to set nheads and num_layers properly
+        # ref: https://pytorch.org/docs/stable/generated/torch.nn.TransformerDecoder.html
+        # https://pytorch.org/docs/stable/generated/torch.nn.TransformerDecoderLayer.html
+        transform_decoder_layers = nn.TransformerDecoderLayer(d_model=128, nhead=4)
+        self.pts_off_transformer = nn.TransformerDecoder(transform_decoder_layers, num_layers=4)
+        
+        self.decoder = nn.Sequential(
+            nn.Linear(pseudoimage_channels*2, 32), nn.GELU(),
+            nn.Linear(32, 3))
+
+    def forward_single(self, before_pseudoimage: torch.Tensor,
+                       after_pseudoimage: torch.Tensor,
+                       point_offsets: torch.Tensor,
+                       voxel_coords: torch.Tensor) -> torch.Tensor:
+        voxel_coords = voxel_coords.long()
+        # assert (voxel_coords[:, 0] == 0).all(), "Z index must be 0"
+
+        # Voxel coords are Z, Y, X, and the pseudoimage is Channel, Y, X
+        # I have confirmed via visualization that these coordinates are correct.
+        after_voxel_vectors = after_pseudoimage[:, voxel_coords[:, 1],
+                                                voxel_coords[:, 2]].T
+        before_voxel_vectors = before_pseudoimage[:, voxel_coords[:, 1],
+                                                  voxel_coords[:, 2]].T
+        
+        # [N, 64] [N, 64] -> [N, 128]
+        concatenated_vectors = torch.cat([before_voxel_vectors, after_voxel_vectors], dim=1)
+        
+        # [N, 128] [N, 128] -> [N, 1, 128]
+        voxel_feature = concatenated_vectors.unsqueeze(1)
+        point_offsets_feature = self.offset_encoder(point_offsets).unsqueeze(1)
+        concatenated_feature = torch.zeros_like(voxel_feature)
+
+        for spilt_range in range(0, concatenated_feature.shape[0], SPLIT_BATCH_SIZE):
+            concatenated_feature[spilt_range:spilt_range+SPLIT_BATCH_SIZE] = self.pts_off_transformer(
+                voxel_feature[spilt_range:spilt_range+SPLIT_BATCH_SIZE],
+                point_offsets_feature[spilt_range:spilt_range+SPLIT_BATCH_SIZE]
+            )
+        
+        flow = self.decoder(concatenated_feature.squeeze(1))
+        return flow
+
+    def forward(
+            self, before_pseudoimages: torch.Tensor,
+            after_pseudoimages: torch.Tensor,
+            voxelizer_infos: List[Dict[str,
+                                       torch.Tensor]]) -> List[torch.Tensor]:
+
+        flow_results = []
+        for before_pseudoimage, after_pseudoimage, voxelizer_info in zip(
+                before_pseudoimages, after_pseudoimages, voxelizer_infos):
+            point_offsets = voxelizer_info["point_offsets"]
+            voxel_coords = voxelizer_info["voxel_coords"]
+            flow = self.forward_single(before_pseudoimage, after_pseudoimage,
+                                       point_offsets, voxel_coords)
+            flow_results.append(flow)
+        return flow_results
+    
+class LinearDecoder(nn.Module):
+
+    def __init__(self, pseudoimage_channels: int = 64):
+        super().__init__()
+
+        self.offset_encoder = nn.Linear(3, 128)
+
+        self.decoder = nn.Sequential(
+            nn.Linear(pseudoimage_channels*4, 32), nn.GELU(),
+            nn.Linear(32, 3))
+
+    def forward_single(self, before_pseudoimage: torch.Tensor,
+                       after_pseudoimage: torch.Tensor,
+                       point_offsets: torch.Tensor,
+                       voxel_coords: torch.Tensor) -> torch.Tensor:
+        voxel_coords = voxel_coords.long()
+        # assert (voxel_coords[:, 0] == 0).all(), "Z index must be 0"
+
+        # Voxel coords are Z, Y, X, and the pseudoimage is Channel, Y, X
+        # I have confirmed via visualization that these coordinates are correct.
+        after_voxel_vectors = after_pseudoimage[:, voxel_coords[:, 1],
+                                                voxel_coords[:, 2]].T
+        before_voxel_vectors = before_pseudoimage[:, voxel_coords[:, 1],
+                                                  voxel_coords[:, 2]].T
+        
+        # [N, 64] [N, 64] -> [N, 128]
+        concatenated_vectors = torch.cat([before_voxel_vectors, after_voxel_vectors], dim=1)
+        
+        # [N, 3] -> [N, 128]
+        point_offsets_feature = self.offset_encoder(point_offsets)
+
+        flow = self.decoder(torch.cat([concatenated_vectors, point_offsets_feature], dim=1))
+        return flow
+
+    def forward(
+            self, before_pseudoimages: torch.Tensor,
+            after_pseudoimages: torch.Tensor,
+            voxelizer_infos: List[Dict[str,
+                                       torch.Tensor]]) -> List[torch.Tensor]:
+
+        flow_results = []
+        for before_pseudoimage, after_pseudoimage, voxelizer_info in zip(
+                before_pseudoimages, after_pseudoimages, voxelizer_infos):
+            point_offsets = voxelizer_info["point_offsets"]
+            voxel_coords = voxelizer_info["voxel_coords"]
+            flow = self.forward_single(before_pseudoimage, after_pseudoimage,
+                                       point_offsets, voxel_coords)
+            flow_results.append(flow)
+        return flow_results
+
+# from https://github.com/weiyithu/PV-RAFT/blob/main/model/update.py
+class ConvGRU(nn.Module):
+    def __init__(self, input_dim=64, hidden_dim=128):
+        super(ConvGRU, self).__init__()
+        self.convz = nn.Conv1d(input_dim+hidden_dim, hidden_dim, 1)
+        self.convr = nn.Conv1d(input_dim+hidden_dim, hidden_dim, 1)
+        self.convq = nn.Conv1d(input_dim+hidden_dim, hidden_dim, 1)
+
+    def forward(self, h, x):
+        hx = torch.cat([h, x], dim=1)
+
+        z = torch.sigmoid(self.convz(hx))
+        r = torch.sigmoid(self.convr(hx))
+        rh_x = torch.cat([r*h, x], dim=1)
+        q = torch.tanh(self.convq(rh_x))
+
+        h = (1 - z) * h + z * q
+        return h
+    
+class ConvGRUDecoder(nn.Module):
+
+    def __init__(self, pseudoimage_channels: int = 64, num_iters: int = 4):
+        super().__init__()
+
+        self.offset_encoder = nn.Linear(3, pseudoimage_channels)
+
+        # NOTE: voxel feature is hidden input, point offset is input, check paper's Fig. 3
+        self.gru = ConvGRU(input_dim=pseudoimage_channels, hidden_dim=pseudoimage_channels*2)
+
+        self.decoder = nn.Sequential(
+            nn.Linear(pseudoimage_channels*3, pseudoimage_channels//2), nn.GELU(),
+            nn.Linear(pseudoimage_channels//2, 3))
+        self.num_iters = num_iters
+
+    def forward_single(self, before_pseudoimage: torch.Tensor,
+                       after_pseudoimage: torch.Tensor,
+                       point_offsets: torch.Tensor,
+                       voxel_coords: torch.Tensor) -> torch.Tensor:
+        voxel_coords = voxel_coords.long()
+        # assert (voxel_coords[:, 0] == 0).all(), "Z index must be 0"
+
+        # Voxel coords are Z, Y, X, and the pseudoimage is Channel, Y, X
+        # I have confirmed via visualization that these coordinates are correct.
+        after_voxel_vectors = after_pseudoimage[:, voxel_coords[:, 1],
+                                                voxel_coords[:, 2]].T
+        before_voxel_vectors = before_pseudoimage[:, voxel_coords[:, 1],
+                                                  voxel_coords[:, 2]].T
+        
+        # [N, 64] [N, 64] -> [N, 128]
+        concatenated_vectors = torch.cat([before_voxel_vectors, after_voxel_vectors], dim=1)
+        
+        # [N, 3] -> [N, 64]
+        point_offsets_feature = self.offset_encoder(point_offsets)
+        
+        # [N, 128] -> [N, 128, 1]
+        concatenated_vectors = concatenated_vectors.unsqueeze(2)
+
+        for itr in range(self.num_iters):
+            concatenated_vectors = self.gru(concatenated_vectors, point_offsets_feature.unsqueeze(2))
+
+        flow = self.decoder(torch.cat([concatenated_vectors.squeeze(2), point_offsets_feature], dim=1))
+        return flow
+
+    def forward(
+            self, before_pseudoimages: torch.Tensor,
+            after_pseudoimages: torch.Tensor,
+            voxelizer_infos: List[Dict[str,
+                                       torch.Tensor]]) -> List[torch.Tensor]:
+
+        flow_results = []
+        for before_pseudoimage, after_pseudoimage, voxelizer_info in zip(
+                before_pseudoimages, after_pseudoimages, voxelizer_infos):
+            point_offsets = voxelizer_info["point_offsets"]
+            voxel_coords = voxelizer_info["voxel_coords"]
+            flow = self.forward_single(before_pseudoimage, after_pseudoimage,
+                                       point_offsets, voxel_coords)
+            flow_results.append(flow)
+        return flow_results
+
+
+class ConvWithNorms(nn.Module):
+
+    def __init__(self, in_num_channels: int, out_num_channels: int,
+                 kernel_size: int, stride: int, padding: int):
+        super().__init__()
+        self.conv = nn.Conv2d(in_num_channels, out_num_channels, kernel_size,
+                              stride, padding)
+        self.batchnorm = nn.BatchNorm2d(out_num_channels)
+        self.nonlinearity = nn.GELU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        conv_res = self.conv(x)
+        if conv_res.shape[2] == 1 and conv_res.shape[3] == 1:
+            # This is a hack to get around the fact that batchnorm doesn't support
+            # 1x1 convolutions
+            batchnorm_res = conv_res
+        else:
+            batchnorm_res = self.batchnorm(conv_res)
+        return self.nonlinearity(batchnorm_res)
diff --git a/deflow.py b/deflow.py
@@ -0,0 +1,113 @@
+
+"""
+# Created: 2023-07-18 15:08
+# 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.
+"""
+
+import torch.nn as nn
+import dztimer, torch
+
+from .basic.unet import FastFlow3DUNet
+from .basic.encoder import DynamicEmbedder
+from .basic.decoder import LinearDecoder, ConvGRUDecoder
+from .basic import cal_pose0to1
+
+class DeFlow(nn.Module):
+    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 = 4):
+        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 = FastFlow3DUNet()
+        if decoder_option == "gru":
+            self.head = ConvGRUDecoder(num_iters = num_iters)
+        elif decoder_option == "linear":
+            self.head = LinearDecoder()
+
+        self.timer = dztimer.Timing()
+        self.timer.start("Total")
+
+    def load_from_checkpoint(self, ckpt_path):
+        ckpt = torch.load(ckpt_path, map_location="cpu")["state_dict"]
+        state_dict = {
+            k[len("model.") :]: v for k, v in ckpt.items() if k.startswith("model.")
+        }
+        print("\nLoading... model weight from: ", ckpt_path, "\n")
+        return self.load_state_dict(state_dict=state_dict, strict=False)
+
+    def forward(self, batch):
+        """
+        input: using the batch from dataloader, which is a dict
+               Detail: [pc0, pc1, pose0, pose1]
+        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"]
+        self.timer[0].stop()
+
+        self.timer[1].start("Voxelization")
+        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(pc0_before_pseudoimages,
+                                            pc1_before_pseudoimages)
+        self.timer[2].stop()
+
+        self.timer[3].start("Decoder")
+        flows = self.head(
+            torch.cat((pc0_before_pseudoimages, pc1_before_pseudoimages),
+                    dim=1), grid_flow_pseudoimage, pc0_voxel_infos_lst)
+        self.timer[3].stop()
+
+        pc0_points_lst = [e["points"] for e in pc0_voxel_infos_lst]
+        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]
+
+        model_res = {
+            "flow": flows,
+            'pose_flow': pose_flows,
+
+            "pc0_valid_point_idxes": pc0_valid_point_idxes,
+            "pc0_points_lst": pc0_points_lst,
+            
+            "pc1_valid_point_idxes": pc1_valid_point_idxes,
+            "pc1_points_lst": pc1_points_lst,
+        }
+        return model_res
