From b3e93a04d096190f57b486e8bd8fa6631d23596a Mon Sep 17 00:00:00 2001
From: RuslanAgishev <agishrus@fel.cvut.cz>
Date: Thu, 14 Mar 2024 19:14:33 +0100
Subject: [PATCH] trajectory of the footprint in evaluation (fixed for rellis3d
 evaluation)

---
 scripts/eval                       | 14 +++++++++-----
 src/monoforce/datasets/rellis3d.py |  1 +
 src/monoforce/datasets/robingas.py |  1 +
 src/monoforce/datasets/utils.py    |  1 +
 4 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/scripts/eval b/scripts/eval
index 699bde8..6971901 100755
--- a/scripts/eval
+++ b/scripts/eval
@@ -333,11 +333,14 @@ class EvaluatorGeom:
 
         return loss_trans_sum, loss_rot_sum
 
-    def evaluate(self, vis=False, verbose=False):
+    def evaluate(self, vis=False, verbose=False, random_order=False):
         for path in self.data_paths:
             print(f'Evaluation on {os.path.basename(path)}...')
             ds = self.DataClass(path, is_train=False, data_aug_conf=self.data_aug_conf, dphys_cfg=self.dphys_cfg)
-            for i in tqdm(range(len(ds))):
+            sample_range = np.arange(0, len(ds))
+            if random_order:
+                np.random.shuffle(sample_range)
+            for i in tqdm(sample_range):
                 states_true, height = self.get_data(i, ds)
                 trans_diff, rot_diff = self.eval_diff_physics(height, states_true, vis=vis)
                 if rot_diff is not None:
@@ -374,7 +377,8 @@ class EvaluatorGeom:
 class EvaluatorLSS(EvaluatorGeom):
     def __init__(self, dphys_config_path, terrain_encoder_config_path, dataset,
                  # model_path='../config/tb_runs/lss_2024_03_04_09_42_47/train_lss.pt'):
-                 model_path='../config/tb_runs/lss_rellis3d_robingas_dphysics_terrain/train_lss.pt'):
+                 # model_path='../config/tb_runs/lss_rellis3d_robingas_dphysics_terrain/train_lss.pt'):
+                 model_path='../config/tb_runs/lss_rellis3d_2024_03_06_16_07_52/train_lss.pt'):
         super().__init__(dphys_config_path, terrain_encoder_config_path, dataset, model_path)
 
     def load_model(self):
@@ -576,8 +580,8 @@ def evaluate(dataset):
 
 
 def main():
-    dataset = 'robingas'
-    # dataset = 'rellis3d'
+    # dataset = 'robingas'
+    dataset = 'rellis3d'
 
     # vis_data(dataset)
     evaluate(dataset)
diff --git a/src/monoforce/datasets/rellis3d.py b/src/monoforce/datasets/rellis3d.py
index 5bbee66..aee82c5 100644
--- a/src/monoforce/datasets/rellis3d.py
+++ b/src/monoforce/datasets/rellis3d.py
@@ -302,6 +302,7 @@ def get_states_traj(self, id, start_from_zero=False):
             # transform poses to the same coordinate frame as the height map
             Tr = np.linalg.inv(poses[0])
             poses = np.asarray([np.matmul(Tr, p) for p in poses])
+            poses[:, 2, 3] -= self.calib['clearance']
             # count time from 0
             tstamps = traj['stamps']
             tstamps = tstamps - tstamps[0]
diff --git a/src/monoforce/datasets/robingas.py b/src/monoforce/datasets/robingas.py
index 78843d8..f72ef67 100644
--- a/src/monoforce/datasets/robingas.py
+++ b/src/monoforce/datasets/robingas.py
@@ -203,6 +203,7 @@ def get_states_traj(self, i, start_from_zero=False):
             # transform poses to the same coordinate frame as the height map
             Tr = np.linalg.inv(poses[0])
             poses = np.asarray([np.matmul(Tr, p) for p in poses])
+            poses[:, 2, 3] -= self.calib['clearance']
             # count time from 0
             tstamps = traj['stamps']
             tstamps = tstamps - tstamps[0]
diff --git a/src/monoforce/datasets/utils.py b/src/monoforce/datasets/utils.py
index db45fc6..5e5dce3 100644
--- a/src/monoforce/datasets/utils.py
+++ b/src/monoforce/datasets/utils.py
@@ -152,6 +152,7 @@ def load_cam_calib(calib_path):
         transforms = yaml.load(f, Loader=yaml.FullLoader)
     f.close()
     calib['transformations'] = transforms
+    calib['clearance'] = np.abs(np.asarray(calib['transformations']['T_base_link__base_footprint']['data'], dtype=np.float32).reshape((4, 4))[2, 3])
 
     return calib