add final config/run/plot script for pcd_to_pose

martius-lab · Mar 25, 2024 · 6a8cc7d · 6a8cc7d
1 parent 8bb3239
commit 6a8cc7d
Show file tree

Hide file tree

Showing 4 changed files with 149 additions and 21 deletions.
diff --git a/hitchhiking_rotations/cfgs/cfg_pcd_to_pose.py b/hitchhiking_rotations/cfgs/cfg_pcd_to_pose.py
@@ -17,7 +17,7 @@ def get_cfg_pcd_to_pose(device):
     return {
         "verbose": False,
         "batch_size": 32,
-        "epochs": 100,
+        "epochs": 300,
         "training_data": {
             "_target_": "hitchhiking_rotations.datasets.PointCloudDataset",
             "mode": "train",
@@ -47,7 +47,7 @@ def get_cfg_pcd_to_pose(device):
             "metrics": ["l1", "l2", "geodesic_distance", "chordal_distance"],
         },
         "trainers": {
-            "r9_l1": {
+            "r9_svd_l1": {
                 **shared_trainer_cfg,
                 **{
                     "preprocess_target": "${u:passthrough}",
@@ -57,7 +57,7 @@ def get_cfg_pcd_to_pose(device):
                     "model": "${model9}",
                 },
             },
-            "r9_l2": {
+            "r9_svd_l2": {
                 **shared_trainer_cfg,
                 **{
                     "preprocess_target": "${u:passthrough}",
@@ -67,7 +67,7 @@ def get_cfg_pcd_to_pose(device):
                     "model": "${model9}",
                 },
             },
-            "r9_geodesic_distance": {
+            "r9_svd_geodesic_distance": {
                 **shared_trainer_cfg,
                 **{
                     "preprocess_target": "${u:passthrough}",
@@ -77,7 +77,7 @@ def get_cfg_pcd_to_pose(device):
                     "model": "${model9}",
                 },
             },
-            "r9_chordal_distance": {
+            "r9_svd_chordal_distance": {
                 **shared_trainer_cfg,
                 **{
                     "preprocess_target": "${u:passthrough}",
@@ -87,7 +87,67 @@ def get_cfg_pcd_to_pose(device):
                     "model": "${model9}",
                 },
             },
+            "r9_l1": {
+                **shared_trainer_cfg,
+                **{
+                    "preprocess_target": "${u:flatten}",
+                    "postprocess_pred_loss": "${u:flatten}",
+                    "postprocess_pred_logging": "${u:procrustes_to_rotmat}",
+                    "loss": "${u:l1}",
+                    "model": "${model9}",
+                },
+            },
+            "r9_l2": {
+                **shared_trainer_cfg,
+                **{
+                    "preprocess_target": "${u:flatten}",
+                    "postprocess_pred_loss": "${u:flatten}",
+                    "postprocess_pred_logging": "${u:procrustes_to_rotmat}",
+                    "loss": "${u:l2}",
+                    "model": "${model9}",
+                },
+            },
+            "r9_geodesic_distance": {
+                **shared_trainer_cfg,
+                **{
+                    "preprocess_target": "${u:passthrough}",
+                    "postprocess_pred_loss": "${u:n_3x3}",
+                    "postprocess_pred_logging": "${u:procrustes_to_rotmat}",
+                    "loss": "${u:geodesic_distance}",
+                    "model": "${model9}",
+                },
+            },
+            "r9_chordal_distance": {
+                **shared_trainer_cfg,
+                **{
+                    "preprocess_target": "${u:passthrough}",
+                    "postprocess_pred_loss": "${u:n_3x3}",
+                    "postprocess_pred_logging": "${u:procrustes_to_rotmat}",
+                    "loss": "${u:chordal_distance}",
+                    "model": "${model9}",
+                },
+            },
             "r6_l1": {
+                **shared_trainer_cfg,
+                **{
+                    "preprocess_target": "${u:rotmat_to_gramschmidt_f}",
+                    "postprocess_pred_loss": "${u:flatten}",
+                    "postprocess_pred_logging": "${u:gramschmidt_to_rotmat}",
+                    "loss": "${u:l1}",
+                    "model": "${model6}",
+                },
+            },
+            "r6_l2": {
+                **shared_trainer_cfg,
+                **{
+                    "preprocess_target": "${u:rotmat_to_gramschmidt_f}",
+                    "postprocess_pred_loss": "${u:flatten}",
+                    "postprocess_pred_logging": "${u:gramschmidt_to_rotmat}",
+                    "loss": "${u:l2}",
+                    "model": "${model6}",
+                },
+            },
+            "r6_gso_l1": {
                 **shared_trainer_cfg,
                 **{
                     "preprocess_target": "${u:passthrough}",
@@ -97,7 +157,7 @@ def get_cfg_pcd_to_pose(device):
                     "model": "${model6}",
                 },
             },
-            "r6_l2": {
+            "r6_gso_l2": {
                 **shared_trainer_cfg,
                 **{
                     "preprocess_target": "${u:passthrough}",
@@ -107,7 +167,7 @@ def get_cfg_pcd_to_pose(device):
                     "model": "${model6}",
                 },
             },
-            "r6_geodesic_distance": {
+            "r6_gso_geodesic_distance": {
                 **shared_trainer_cfg,
                 **{
                     "preprocess_target": "${u:passthrough}",
@@ -117,7 +177,7 @@ def get_cfg_pcd_to_pose(device):
                     "model": "${model6}",
                 },
             },
-            "r6_chordal_distance": {
+            "r6_gso_chordal_distance": {
                 **shared_trainer_cfg,
                 **{
                     "preprocess_target": "${u:passthrough}",
@@ -127,6 +187,16 @@ def get_cfg_pcd_to_pose(device):
                     "model": "${model6}",
                 },
             },
+            "quat_c_geodesic_distance": {
+                **shared_trainer_cfg,
+                **{
+                    "preprocess_target": "${u:passthrough}",
+                    "postprocess_pred_loss": "${u:quaternion_to_rotmat}",
+                    "postprocess_pred_logging": "${u:quaternion_to_rotmat}",
+                    "loss": "${u:geodesic_distance}",
+                    "model": "${model4}",
+                },
+            },
             "quat_c_chordal_distance": {
                 **shared_trainer_cfg,
                 **{
@@ -177,6 +247,46 @@ def get_cfg_pcd_to_pose(device):
                     "model": "${model4}",
                 },
             },
+            "quat_rf_cosine_distance": {
+                **shared_trainer_cfg,
+                **{
+                    "preprocess_target": "${u:rotmat_to_quaternion_rand_flip}",
+                    "postprocess_pred_loss": "${u:passthrough}",
+                    "postprocess_pred_logging": "${u:quaternion_to_rotmat}",
+                    "loss": "${u:cosine_distance}",
+                    "model": "${model4}",
+                },
+            },
+            "quat_rf_l2": {
+                **shared_trainer_cfg,
+                **{
+                    "preprocess_target": "${u:rotmat_to_quaternion_rand_flip}",
+                    "postprocess_pred_loss": "${u:passthrough}",
+                    "postprocess_pred_logging": "${u:quaternion_to_rotmat}",
+                    "loss": "${u:l2}",
+                    "model": "${model4}",
+                },
+            },
+            "quat_rf_l1": {
+                **shared_trainer_cfg,
+                **{
+                    "preprocess_target": "${u:rotmat_to_quaternion_rand_flip}",
+                    "postprocess_pred_loss": "${u:passthrough}",
+                    "postprocess_pred_logging": "${u:quaternion_to_rotmat}",
+                    "loss": "${u:l1}",
+                    "model": "${model4}",
+                },
+            },
+            "quat_rf_l2_dp": {
+                **shared_trainer_cfg,
+                **{
+                    "preprocess_target": "${u:rotmat_to_quaternion_rand_flip}",
+                    "postprocess_pred_loss": "${u:passthrough}",
+                    "postprocess_pred_logging": "${u:quaternion_to_rotmat}",
+                    "loss": "${u:l2_dp}",
+                    "model": "${model4}",
+                },
+            },
             "rotvec_l1": {
                 **shared_trainer_cfg,
                 **{

diff --git a/scripts/run_all.py b/scripts/run_all.py
@@ -1,15 +1,18 @@
 from hitchhiking_rotations import HITCHHIKING_ROOT_DIR
 import os
+import sys
 
-p = os.path.join(HITCHHIKING_ROOT_DIR, "scripts", "train.py")
 
-for seed in range(10):
-    os.system(f"python3 {p} --experiment cube_image_to_pose --seed {seed}")
+exps = ["pcd_to_pose", "cube_image_to_pose", "pose_to_cube_image"]
+assert len(sys.argv) == 2, "Please provide the experiment name as an argument."
+assert sys.argv[1] in exps, f"Experiment name should be one of {exps}"
 
+exp = sys.argv[1]
+p = os.path.join(HITCHHIKING_ROOT_DIR, "scripts", "train.py")
 for seed in range(10):
-    os.system(f"python3 {p} --experiment pose_to_cube_image --seed {seed}")
-
-
-os.system("python3 " + str(os.path.join(HITCHHIKING_ROOT_DIR, "visu", "figure_19.py")))
-os.system("python3 " + str(os.path.join(HITCHHIKING_ROOT_DIR, "visu", "figure_12a.py")))
-os.system("python3 " + str(os.path.join(HITCHHIKING_ROOT_DIR, "visu", "figure_12b.py")))
+    os.system(f"python3 {p} --experiment {exp} --seed {seed}")
+if exp in exps[:2]:
+    os.system("python3 " + str(os.path.join(HITCHHIKING_ROOT_DIR, "visu", "figure_19.py", exp)))
+    os.system("python3 " + str(os.path.join(HITCHHIKING_ROOT_DIR, "visu", "figure_12a.py", exp)))
+else:
+    os.system("python3 " + str(os.path.join(HITCHHIKING_ROOT_DIR, "visu", "figure_12b.py")))
diff --git a/visu/figure_12a.py b/visu/figure_12a.py
@@ -1,4 +1,5 @@
 import os
+import sys
 import numpy as np
 import seaborn as sns
 import matplotlib.pyplot as plt
@@ -10,9 +11,14 @@
 
 selected_metric = "geodesic_distance"
 training_metric = "chordal_distance"
+exps = ["pcd_to_pose", "cube_image_to_pose"]
 
+assert len(sys.argv) == 2, "Please provide the experiment name as an argument."
+assert sys.argv[1] in exps, f"Experiment name should be one of {exps}"
 
-files = [str(s) for s in Path(os.path.join(HITCHHIKING_ROOT_DIR, "results", "cube_image_to_pose")).rglob("*result.npy")]
+exp = sys.argv[1]
+
+files = [str(s) for s in Path(os.path.join(HITCHHIKING_ROOT_DIR, "results", exp)).rglob("*result.npy")]
 results = [np.load(file, allow_pickle=True) for file in files]
 
 df_res = {}
@@ -71,6 +77,6 @@
 plt.ylabel("")
 plt.tight_layout()
 
-out_p = os.path.join(HITCHHIKING_ROOT_DIR, "results", "cube_image_to_pose", "figure_12a.pdf")
+out_p = os.path.join(HITCHHIKING_ROOT_DIR, "results", exp, "figure_12a.pdf")
 plt.savefig(out_p)
 plt.show()
diff --git a/visu/figure_19.py b/visu/figure_19.py
@@ -1,4 +1,5 @@
 import os
+import sys
 import numpy as np
 import seaborn as sns
 import matplotlib.pyplot as plt
@@ -7,14 +8,22 @@
 import pandas as pd
 from hitchhiking_rotations.utils import RotRep
 
+
+exps = ["pcd_to_pose", "cube_image_to_pose"]
+
+assert len(sys.argv) == 2, "Please provide the experiment name as an argument."
+assert sys.argv[1] in exps, f"Experiment name should be one of {exps}"
+
+exp = sys.argv[1]
+
 plt.figure(figsize=(14, 14))
 plt.style.use(os.path.join(HITCHHIKING_ROOT_DIR, "assets", "prettyplots.mplstyle"))
 sns.set_style("whitegrid")
 plt.rcParams.update({"font.size": 11})
 
 for j, selected_metric in enumerate(["geodesic_distance", "chordal_distance"]):
     files = [
-        str(s) for s in Path(os.path.join(HITCHHIKING_ROOT_DIR, "results", "cube_image_to_pose")).rglob("*result.npy")
+        str(s) for s in Path(os.path.join(HITCHHIKING_ROOT_DIR, "results", exp)).rglob("*result.npy")
     ]
     results = [np.load(file, allow_pickle=True) for file in files]
 
@@ -113,7 +122,7 @@
     # plt.xscale("log")
     plt.tight_layout()
 
-out_p = os.path.join(HITCHHIKING_ROOT_DIR, "results", "cube_image_to_pose", f"figure_19_combined.pdf")
+out_p = os.path.join(HITCHHIKING_ROOT_DIR, "results", exp, f"figure_19_combined.pdf")
 
 plt.savefig(out_p)