init fusion lcd orin config

loop_gt.py

@@ -0,0 +1,130 @@
+import argparse
+import torch
+from torch.utils.data import Dataset
+import pykitti
+import os
+from sklearn.neighbors import KDTree
+import pickle
+import numpy as np
+skip_frame=50
+
+class KITTILoader3DPosesOnlyLoopPositives(Dataset):
+
+    def __init__(self, dir, sequence, poses, positive_range=5., negative_range=25., hard_range=None):
+        super(KITTILoader3DPosesOnlyLoopPositives, self).__init__()
+
+        self.positive_range = positive_range
+        self.negative_range = negative_range
+        self.hard_range = hard_range
+        self.dir = dir
+        self.sequence = sequence
+
+        if int(sequence) > 21:
+            self.poses = np.load(poses)
+        else:
+            calib = np.genfromtxt(os.path.join(dir, 'sequences', sequence, 'calib.txt'))[:, 1:]
+            T_cam_velo = np.reshape(calib[4], (3, 4))
+            T_cam_velo = np.vstack([T_cam_velo, [0, 0, 0, 1]])
+            poses2 = []
+            with open(poses, 'r') as f:
+                for x in f:
+                    x = x.strip().split()
+                    x = [float(v) for v in x]
+                    pose = np.zeros((4, 4))
+                    pose[0, 0:4] = np.array(x[0:4])
+                    pose[1, 0:4] = np.array(x[4:8])
+                    pose[2, 0:4] = np.array(x[8:12])
+                    pose[3, 3] = 1.0
+                    pose = np.linalg.inv(T_cam_velo) @ (pose @ T_cam_velo)
+                    poses2.append(pose)
+            self.poses = np.stack(poses2)
+        self.kdtree = KDTree(self.poses[:, :3, 3])
+
+    def __len__(self):
+        return len(self.poses)
+
+    def __getitem__(self, idx):
+
+        x = self.poses[idx, 0, 3]
+        y = self.poses[idx, 1, 3]
+        z = self.poses[idx, 2, 3]
+        r0 = self.poses[idx, :3, :3]
+        rs = self.poses[:, :3, :3]
+        dr = np.linalg.inv(r0) @ rs.swapaxes(0, 2)
+        angle = np.arccos(np.clip((np.trace(dr) - 1) / 2, -1, 1))
+        angle = angle * 180 / np.pi
+
+        idx_angle = np.where(angle < 99999)[0]
+        anchor_pose = torch.tensor([x, y, z])
+
+        indices = self.kdtree.query_radius(anchor_pose.unsqueeze(0).numpy(), self.positive_range, sort_results=True, return_distance=True)
+        indices = [indices[0][0], indices[1][0]]
+        min_range = max(0, idx - skip_frame)
+        max_range = min(idx + skip_frame, len(self.poses))
+        positive_idxs = list(set(indices[0]) & set(idx_angle) - set(range(min_range, max_range)))
+        loop_angle = angle[positive_idxs]
+        reverse = 0
+        if len(loop_angle) > 0:
+            reverse=np.sum(loop_angle>90)
+            if min(loop_angle) > 90:
+                reverse = -1*reverse
+        positive_idxs.sort()
+        num_loop = len(positive_idxs)
+
+        indices = self.kdtree.query_radius(anchor_pose.unsqueeze(0).numpy(), self.negative_range)
+        indices = set(indices[0])
+        negative_idxs = set(range(len(self.poses))) - indices
+        negative_idxs = list(negative_idxs)
+        negative_idxs.sort()
+
+        hard_idxs = None
+        if self.hard_range is not None:
+            inner_indices = self.kdtree.query_radius(anchor_pose.unsqueeze(0).numpy(), self.hard_range[0])
+            outer_indices = self.kdtree.query_radius(anchor_pose.unsqueeze(0).numpy(), self.hard_range[1])
+            hard_idxs = set(outer_indices[0]) - set(inner_indices[0])
+            pass
+
+        return num_loop, positive_idxs, negative_idxs, hard_idxs, reverse
+
+
+if __name__ == '__main__':
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--root_folder', default='/home/adlab36/chenyouyuan/FUSIONLCD',
+                        help='dataset directory')
+    args = parser.parse_args()
+    base_dir = args.root_folder
+    for sequence in ['00', '05', '06', '07', '08', '09', '50', '54', '55', '56', '59']:#'00', '05', '06', '07', '08', '09', '50', '54', '55', '56', '59', '120205','130405'
+        if int(sequence) < 50:
+            poses_file = base_dir + "/sequences/" + sequence + "/poses.txt"
+        elif int(sequence)<100:
+            poses_file = base_dir + "/sequences/" + sequence + "/poses.npy"
+        else:
+            pass
+            
+        
+        dataset = KITTILoader3DPosesOnlyLoopPositives(base_dir, sequence, poses_file, 4, 15, [8, 15])
+        lc_gt = []
+        lc_gt_file = os.path.join(base_dir, 'sequences', sequence, 'loop_GT_4m.pickle')
+        loop_pairs = []
+        loop_files = []
+        for i in range(len(dataset)):
+            sample, pos, neg, hard, reverse = dataset[i]
+            if sample > 0.:
+                loop_files.append([i, reverse])
+                sample_dict = {}
+                sample_dict['idx'] = i
+                sample_dict['positive_idxs'] = pos
+                for p in pos:
+                    if i < p:
+                        loop_pairs.append([i, p])
+                # sample_dict['negative_idxs'] = neg
+                # sample_dict['hard_idxs'] = hard
+                lc_gt.append(sample_dict)
+        loop_files = np.array(loop_files)
+        num_reverse_file = int(np.sum(loop_files[:, 1]<0))
+        num_reverse_pairs = int(np.sum(np.abs(loop_files[:, 1])))/2
+        with open(lc_gt_file, 'wb') as f:
+            pickle.dump(lc_gt, f)
+        print('Sequence %02d done,%05d files, %05d files with loop, %05d[%.4f] files only has reverse loop, %05d loop pairs, %05d[%.4f] reverse loop' %
+              (int(sequence), len(dataset), len(loop_files), num_reverse_file, num_reverse_file / len(loop_files), len(loop_pairs),num_reverse_pairs,num_reverse_pairs/len(loop_pairs)))