update the script to convert old weight to new version.

Kin-Zhang · Kin-Zhang · commit 4d4eecdce9f1 · 2024-09-10T15:24:45.000+02:00
write4conf.py file, todo update model weights later.
diff --git a/README.md b/README.md
@@ -159,11 +159,11 @@ python save.py checkpoint=/home/kin/seflow_best.ckpt dataset_path=/home/kin/data
 # The output of above command will be like:
 Model: DeFlow, Checkpoint from: /home/kin/model_zoo/v2/seflow_best.ckpt
 We already write the flow_est into the dataset, please run following commend to visualize the flow. Copy and paste it to your terminal:
-python tools/scene_flow.py --flow_mode 'seflow_best' --data_dir /home/kin/data/av2/preprocess_v2/sensor/vis
+python tools/visualization.py --res_name 'seflow_best' --data_dir /home/kin/data/av2/preprocess_v2/sensor/vis
 Enjoy! ^v^ ------ 
 
 # Then run the command in the terminal:
-python tools/scene_flow.py --flow_mode 'seflow_best' --data_dir /home/kin/data/av2/preprocess_v2/sensor/vis
+python tools/visualization.py --res_name 'seflow_best' --data_dir /home/kin/data/av2/preprocess_v2/sensor/vis
 ```
 
 https://github.com/user-attachments/assets/f031d1a2-2d2f-4947-a01f-834ed1c146e6
diff --git a/conf/model/deflow.yaml b/conf/model/deflow.yaml
@@ -1,7 +1,7 @@
 name: deflow
 
 target:
-  _target_: src.models.deflow.DeFlow
+  _target_: src.models.DeFlow
   decoder_option: gru # choices: [linear, gru]
   num_iters: 4
   voxel_size: ${voxel_size}
diff --git a/conf/model/fastflow3d.yaml b/conf/model/fastflow3d.yaml
@@ -1,7 +1,7 @@
 name: fastflow3d
 
 target:
-  _target_: src.models.fastflow3d.FastFlow3D
+  _target_: src.models.FastFlow3D
   voxel_size: ${voxel_size}
   point_cloud_range: ${point_cloud_range}
   
diff --git a/conf/save.yaml b/conf/save.yaml
@@ -1,4 +1,4 @@
-dataset_path: /home/kin/data/av2/preprocess_v2/sensor/vis
+dataset_path: /home/kin/data/av2/preprocess_v2/demo/sensor/val
 checkpoint: /home/kin/model_zoo/seflow_best.ckpt
 res_name: # if None will directly be the `model_name.ckpt` in checkpoint path
 
diff --git a/dataprocess/README.md b/dataprocess/README.md
@@ -140,7 +140,7 @@ Or you can run testing file to visualize the data.
 
 ```bash
 # view gt flow
-python3 tools/scene_flow.py --data_dir /home/kin/data/av2/preprocess/sensor/mini --flow_mode flow
+python tools/visualization.py --data_dir /home/kin/data/av2/preprocess/sensor/mini --res_name flow
 
-python3 tools/scene_flow.py --data_dir /home/kin/data/waymo/preprocess/val --flow_mode flow
+python tools/visualization.py --data_dir /home/kin/data/waymo/preprocess/val --res_name flow
 ```
diff --git a/save.py b/save.py
@@ -22,7 +22,7 @@
 from src.trainer import ModelWrapper
 from src.utils import bc
 
-@hydra.main(version_base=None, config_path="conf", config_name="vis")
+@hydra.main(version_base=None, config_path="conf", config_name="save")
 def main(cfg):
     pl.seed_everything(cfg.seed, workers=True)
     output_dir = HydraConfig.get().runtime.output_dir
diff --git a/src/dataset.py b/src/dataset.py
@@ -149,12 +149,6 @@ def __getitem__(self, index_):
             pc1 = torch.tensor(f[next_timestamp]['lidar'][:][:,:3])
             gm1 = torch.tensor(f[next_timestamp]['ground_mask'][:])
             pose1 = torch.tensor(f[next_timestamp]['pose'][:])
-            # if pc0[~gm0].shape[0] == 0:
-            #     print(f"\nWarning: empty point cloud! scene_id: [{scene_id}, {timestamp}], Check the data!\n")
-            #     return self.__getitem__(index_+1)
-            # elif pc1[~gm1].shape[0] == 0:
-            #     print(f"\nWarning: empty point cloud! scene_id: [{scene_id}, {next_timestamp}], Check the data!\n")
-            #     return self.__getitem__(index_-1)
             res_dict = {
                 'scene_id': scene_id,
                 'timestamp': key,
diff --git a/src/models/__init__.py b/src/models/__init__.py
@@ -0,0 +1,2 @@
+from .deflow import DeFlow
+from .fastflow3d import FastFlow3D
diff --git a/src/trainer.py b/src/trainer.py
@@ -214,7 +214,7 @@ def on_validation_epoch_end(self):
 
         if self.save_res:
             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/scene_flow.py --flow_mode '{self.vis_name}' --data_dir {self.dataset_path}")
+            print(f"python tools/visualization.py --res_name '{self.vis_name}' --data_dir {self.dataset_path}")
             print(f"Enjoy! ^v^ ------ \n")
         
     def eval_only_step_(self, batch, res_dict):
@@ -310,5 +310,5 @@ 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"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/scene_flow.py --flow_mode '{self.vis_name}' --data_dir {self.dataset_path}")
+        print(f"python tools/visualization.py --res_name '{self.vis_name}' --data_dir {self.dataset_path}")
         print(f"Enjoy! ^v^ ------ \n")
diff --git a/tools/README.md b/tools/README.md
@@ -12,11 +12,11 @@ run `tools/scene_flow.py` to view the scene flow dataset with ground truth flow.
 
 ```bash
 # view gt flow
-python3 tools/scene_flow.py --data_dir /home/kin/data/av2/preprocess/sensor/mini --flow_mode flow
+python3 tools/scene_flow.py --data_dir /home/kin/data/av2/preprocess/sensor/mini --res_name flow
 
 # view est flow
-python3 tools/scene_flow.py --data_dir /home/kin/data/av2/preprocess/sensor/mini --flow_mode deflow_best
-python3 tools/scene_flow.py --data_dir /home/kin/data/av2/preprocess/sensor/mini --flow_mode seflow_best
+python3 tools/scene_flow.py --data_dir /home/kin/data/av2/preprocess/sensor/mini --res_name deflow_best
+python3 tools/scene_flow.py --data_dir /home/kin/data/av2/preprocess/sensor/mini --res_name seflow_best
 ```
 
 Demo Effect (press `SPACE` to stop and start in the visualization window):
diff --git a/tools/visualization.py b/tools/visualization.py
@@ -26,12 +26,12 @@
 
 def check_flow(
     data_dir: str ="/home/kin/data/av2/preprocess/sensor/mini",
-    flow_mode: str = "flow", # "flow", "flow_est"
+    res_name: str = "flow", # "flow", "flow_est"
     start_id: int = -1,
     point_size: float = 3.0,
 ):
-    dataset = HDF5Data(data_dir, vis_name=flow_mode, flow_view=True)
-    o3d_vis = MyVisualizer(view_file=VIEW_FILE, window_title=f"view {'ground truth flow' if flow_mode == 'flow' else f'{flow_mode} flow'}, `SPACE` start/stop")
+    dataset = HDF5Data(data_dir, vis_name=res_name, flow_view=True)
+    o3d_vis = MyVisualizer(view_file=VIEW_FILE, window_title=f"view {'ground truth flow' if res_name == 'flow' else f'{res_name} flow'}, `SPACE` start/stop")
 
     opt = o3d_vis.vis.get_render_option()
     opt.background_color = np.asarray([80/255, 90/255, 110/255])
@@ -58,18 +58,18 @@ def check_flow(
 
         pcd2 = o3d.geometry.PointCloud()
         # pcd2.points = o3d.utility.Vector3dVector(pc0[:, :3] + pose_flow) # if you want to check pose_flow
-        pcd2.points = o3d.utility.Vector3dVector(pc0[:, :3] + data[flow_mode][~gm0])
+        pcd2.points = o3d.utility.Vector3dVector(pc0[:, :3] + data[res_name][~gm0])
         pcd2.paint_uniform_color([0.0, 0.0, 1.0]) # blue: pc0 + 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",
-    flow_mode: str = "flow", # "flow", "flow_est"
+    res_name: str = "flow", # "flow", "flow_est"
     start_id: int = -1,
     point_size: float = 2.0,
 ):
-    dataset = HDF5Data(data_dir, vis_name=flow_mode, flow_view=True)
-    o3d_vis = MyVisualizer(view_file=VIEW_FILE, window_title=f"view {'ground truth flow' if flow_mode == 'flow' else f'{flow_mode} flow'}, `SPACE` start/stop")
+    dataset = HDF5Data(data_dir, vis_name=res_name, flow_view=True)
+    o3d_vis = MyVisualizer(view_file=VIEW_FILE, window_title=f"view {'ground truth flow' if res_name == 'flow' else f'{res_name} flow'}, `SPACE` start/stop")
 
     opt = o3d_vis.vis.get_render_option()
     # opt.background_color = np.asarray([216, 216, 216]) / 255.0
@@ -94,8 +94,8 @@ def vis(
         pose_flow = pc0[:, :3] @ ego_pose[:3, :3].T + ego_pose[:3, 3] - pc0[:, :3]
         
         pcd = o3d.geometry.PointCloud()
-        if flow_mode in ['dufo_label', 'label']:
-            labels = data[flow_mode]
+        if res_name in ['dufo_label', 'label']:
+            labels = data[res_name]
             pcd_i = o3d.geometry.PointCloud()
             for label_i in np.unique(labels):
                 pcd_i.points = o3d.utility.Vector3dVector(pc0[labels == label_i][:, :3])
@@ -104,9 +104,9 @@ def vis(
                 else:
                     pcd_i.paint_uniform_color(color_map[label_i % len(color_map)])
                 pcd += pcd_i
-        elif flow_mode in data:
+        elif res_name in data:
             pcd.points = o3d.utility.Vector3dVector(pc0[:, :3])
-            flow = data[flow_mode] - pose_flow # ego motion compensation here.
+            flow = data[res_name] - pose_flow # ego motion compensation here.
             flow_color = flow_to_rgb(flow) / 255.0
             is_dynamic = np.linalg.norm(flow, axis=1) > 0.1
             flow_color[~is_dynamic] = [1, 1, 1]
diff --git a/tools/write4conf.py b/tools/write4conf.py
@@ -0,0 +1,20 @@
+import torch
+import fire, time
+
+def main(
+    # download from: https://zenodo.org/records/12632962
+    model_path: str = "/home/kin/model_zoo/v2/seflow_best.ckpt", 
+    # new output weight file
+    output_path: str = "/home/kin/model_zoo/v3/seflow_best.ckpt",
+):
+    model = torch.load(model_path)
+    model_name = model['hyper_parameters']['cfg']['model']['name']
+    old_path = model['hyper_parameters']['cfg']['model']['target']['_target_']
+    new_path = old_path.replace(f"scripts.network.models.{model_name}", "src.models")
+    model['hyper_parameters']['cfg']['model']['target']['_target_'] = new_path
+    torch.save(model, output_path)
+
+if __name__ == '__main__':
+    start_time = time.time()
+    fire.Fire(main)
+    print(f"Time used: {time.time() - start_time:.2f} s")

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-dataset_path: /home/kin/data/av2/preprocess_v2/sensor/vis`
	`1`	`+dataset_path: /home/kin/data/av2/preprocess_v2/demo/sensor/val`
`2`	`2`	`checkpoint: /home/kin/model_zoo/seflow_best.ckpt`
`3`	`3`	res_name: # if None will directly be the `model_name.ckpt` in checkpoint path
`4`	`4`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+from .deflow import DeFlow`
	`2`	`+from .fastflow3d import FastFlow3D`