Add functionality to evaluate top K assignments based on cost

Author: aneeshc12

ablations_over_360.sh

@@ -1,11 +1,38 @@
+#!/bin/bash --login
+#SBATCH -n 16
+#SBATCH -A research
+#SBATCH --mem-per-cpu=2048
+#SBATCH --time=72:00:00
+#SBATCH --mincpus=10
+#SBATCH --gres=gpu:1
+#SBATCH [email protected]
+#SBATCH --mail-type=ALL
+#SBATCH -w gnode041
+
+# module add cuda/8.0
+# module add cudnn/7-cuda-8.0
+
 # Define arrays for possible values
 DOWNSAMPLE_VOXEL_SIZES=(0 0.01 0.001)
 GLOB_DIST_VALUES=(1 1.5 2.0)
 LOCAL_DIST_VALUES=(0.1 0.4 0.8 1.2)
-FPFH_VOXEL_VALUES=(0.02 0.05 0.1 0.2)
+FPFH_VOXEL_VALUES=(0.2)
 USE_MESH=0
 
+# conda init bash
+conda activate reid
+# /home2/$USER/Change_detection/download_weights.sh
+mkdir -p /scratch/$USER/360_runs/
+
+cd /home2/$USER/Change_detection
+
 # Iterate over values
+for i in 1 2 3 4; do
+echo $i
+done;
+
+python runthrough_360.py
+
 for DOWNSAMPLE_VOXEL_SIZE in "${DOWNSAMPLE_VOXEL_SIZES[@]}"; do
     for GLOB_DIST in "${GLOB_DIST_VALUES[@]}"; do
         for LOCAL_DIST in "${LOCAL_DIST_VALUES[@]}"; do
@@ -25,17 +52,20 @@ for DOWNSAMPLE_VOXEL_SIZE in "${DOWNSAMPLE_VOXEL_SIZES[@]}"; do
                 python runthrough_360.py \
                     --sam-checkpoint-path=/scratch/$USER/sam_vit_h_4b8939.pth \
                     --ram-pretrained-path=/scratch/$USER/ram_swin_large_14m.pth \
-                    --test-folder-path=/home2/$USER/Documents/aneesh_thread/Change_detection/360_zip \
-                    --memory-save-path=/scratch/$USER/vin-experiments/360_runs/memory_${ABLATION_NAME}.pcd \
+                    --test-folder-path=/home2/$USER/Change_detection/360_zip \
+                    --memory-save-path=/scratch/$USER/360_runs/memory_${ABLATION_NAME}.pcd \
                     --down-sample-voxel-size=${DOWNSAMPLE_VOXEL_SIZE} \
                     --no-create-ext-mesh \
-                    --save-results-path=/scratch/$USER/vin-experiments/360_runs/results_${ABLATION_NAME}.json \
+                    --save-results-path=/scratch/$USER/360_runs/results_${ABLATION_NAME}.json \
                     --fpfh-global-dist-factor=${GLOB_DIST} \
                     --fpfh-local-dist-factor=${LOCAL_DIST} \
-                    --fpfh-voxel-size=${FPFH_VOXEL} \
-                    > /scratch/$USER/vin-experiments/360_runs/run_${ABLATION_NAME}.txt
+                    --fpfh-voxel-size=${FPFH_VOXEL} 
+                    # > /scratch/$USER/360_runs/run_${ABLATION_NAME}.txt
+
+                echo "Test done!"
+                exit 0
 
             done
         done
     done
-done
+done

fpfh/fpfh_register.py

@@ -52,4 +52,4 @@ def register_point_clouds(source, target, voxel_size, global_dist_factor = 1.5,
         o3d.pipelines.registration.TransformationEstimationPointToPoint()
     )
 
-    return result_icp.transformation
+    return result_icp.transformation, result_icp.fitness

jcbb.py

@@ -3,6 +3,8 @@
 import torch
 from copy import copy
 
+import time
+
 class JCBB():
     def __init__(self, cosine_similarities, R_matrices):
         self.cosine_similarities = cosine_similarities
@@ -19,8 +21,12 @@ def __init__(self, cosine_similarities, R_matrices):
     # get all assignment paths
     def get_assignments(self):
         self.leaf_nodes = []
-        JCBBNode(self, self.num_detected, 0, [], [i for i in range(self.num_memory)])
+        JCBBNode(self, self.num_detected, 0, 0, [], [i for i in range(self.num_memory)])
+        self.sort_assignments()
         return self.leaf_nodes
+    
+    def sort_assignments(self):
+        self.leaf_nodes = sorted(self.leaf_nodes, key=lambda x: x[1], reverse=True)
 
 """
 Node in the JCBB tree
@@ -30,34 +36,48 @@ def get_assignments(self):
 assignment : j
 path : array of assignments before this (index is i, value is j)
 choices_left : list of indices left to assign
+
+calculate the product of similarities as we go down (add log costs)
 """
-def JCBBNode(tree, max_depth, depth, path, choices_left):
+def JCBBNode(tree, max_depth, similarity_sum, depth, path, choices_left):
     if depth == max_depth:
-        tree.leaf_nodes.append(path)
+        tree.leaf_nodes.append([path, similarity_sum/max_depth])
         return
 
     # recurse through 
     for c in choices_left:
+        # add the cosine similarity of this assignment to the cost
+        new_sum = similarity_sum + np.log(tree.cosine_similarities[depth][c])
+
         child_path = copy(path)
-        
+
         # let this child be assignment (depth+1, c), remove c from choices
         child_path.append(c)
         child_choices_left = [d for d in choices_left if d != c]
 
         JCBBNode(
             tree,
             max_depth,
+            new_sum,
             depth+1,
             child_path,
             child_choices_left
         )
 
 if __name__ ==  '__main__':
-    j = JCBB(np.zeros((3,4)), None)
+    sim = np.random.uniform(size=(4,20))
+
+    start = time.time()
+
+    j = JCBB(sim, None)
     paths = j.get_assignments()
     # for p in paths:
     #     for k, v in p.items():
     #         print(k,v, sep="|", end=" ")
     #     print()
 
-    print(paths)
+    print("Time taken: ", time.time()-start)
+
+    # print(sim)
+
+    # print(paths)

object_memory.py

@@ -1039,7 +1039,7 @@ def consolidate_memory(self, bounding_box_threshold=0.3,  occlusion_overlap_thre
     def localise(self, image_path, depth_image_path, testname="", save_point_clouds=False,
                  outlier_removal_config=None, 
                  fpfh_global_dist_factor = 1.5, fpfh_local_dist_factor = 0.4, 
-                 fpfh_voxel_size = 0.05):
+                 fpfh_voxel_size = 0.05, topK=5):
         """
         Given an image and a corresponding depth image in an unknown frame, consult the stored memory
         and output a pose in the world frame of the point clouds stored in memory.
@@ -1097,52 +1097,48 @@ def localise(self, image_path, depth_image_path, testname="", save_point_clouds=
         j = JCBB(cosine_similarities, R_matrices)
         assns = j.get_assignments()
 
-        # calculate paths for all assingments, pick the best one
-        best_assignment = assns[0]
-        min_cost = 1e11
+        # only consider the top K assingments
+        assns_to_consider = [assn[0] for assn in assns[:topK]]
 
-        for assn in assns:
-            cost = 0
+        print("Phrases: ", detected_phrases)
+        print(cosine_similarities)
+        print("Assignments being considered: ", assns_to_consider)
+
+        assn_data = [ [assn, None, None] for assn in assns_to_consider ]
 
-            ### COST FUNCTION ###
+        # go through all top K assingments, record ICP costs
+        for assn_num, assn in enumerate(assns_to_consider):
+            # use ALL object pointclouds together
+            all_detected_points = []
+            all_memory_points = []
             for i,j in enumerate(assn):
-                cost += (1 - cosine_similarities[i,j])      
-            cost = np.log(cost) * 1./len(assn) # normalized product of cosine DIFFERENCES
-            # NOTE: diving by length of assignments doesn't matter as all are same rn
+                all_detected_points.append(detected_pointclouds[i])
+                all_memory_points.append(self.memory[j].pcd)
+            all_detected_points = np.concatenate(all_detected_points, axis=-1)
+            all_memory_points = np.concatenate(all_memory_points, axis=-1)
+
+            # centering all the pointclouds
+            detected_mean = np.mean(all_detected_points, axis=-1)
+            memory_mean = np.mean(all_memory_points, axis=-1)
+            all_detected_pcd = o3d.geometry.PointCloud()
+            all_detected_pcd.points = o3d.utility.Vector3dVector(all_detected_points.T - detected_mean)
+            all_memory_pcd = o3d.geometry.PointCloud()
+            all_memory_pcd.points = o3d.utility.Vector3dVector(all_memory_points.T - memory_mean)
+            
+            # remove outliers from detected pcds
+            all_detected_pcd_filtered, _ = all_detected_pcd.remove_radius_outlier(nb_points=outlier_removal_config["radius_nb_points"],
+                                                            radius=outlier_removal_config["radius"])
 
-            # get min cost
-            if cost < min_cost:
-                min_cost = cost
-                best_assignment = assn
+            transform, fitness = register_point_clouds(all_detected_pcd_filtered, all_memory_pcd, 
+                                            voxel_size = fpfh_voxel_size, global_dist_factor = fpfh_global_dist_factor, 
+                                            local_dist_factor = fpfh_local_dist_factor)
 
-        print("Phrases: ", detected_phrases)
-        print(cosine_similarities)
-        print("Assignment: ", best_assignment)
-
-        # use ALL object pointclouds together
-        all_detected_points = []
-        all_memory_points = []
-        for i,j in enumerate(best_assignment):
-            all_detected_points.append(detected_pointclouds[i])
-            all_memory_points.append(self.memory[j].pcd)
-        all_detected_points = np.concatenate(all_detected_points, axis=-1)
-        all_memory_points = np.concatenate(all_memory_points, axis=-1)
-
-        # centering all the pointclouds
-        detected_mean = np.mean(all_detected_points, axis=-1)
-        memory_mean = np.mean(all_memory_points, axis=-1)
-        all_detected_pcd = o3d.geometry.PointCloud()
-        all_detected_pcd.points = o3d.utility.Vector3dVector(all_detected_points.T - detected_mean)
-        all_memory_pcd = o3d.geometry.PointCloud()
-        all_memory_pcd.points = o3d.utility.Vector3dVector(all_memory_points.T - memory_mean)
-        
-        # remove outliers from detected pcds
-        all_detected_pcd_filtered, _ = all_detected_pcd.remove_radius_outlier(nb_points=outlier_removal_config["radius_nb_points"],
-                                                        radius=outlier_removal_config["radius"])
+            assn_data[assn_num] = [assn, transform, fitness]
 
-        transform = register_point_clouds(all_detected_pcd_filtered, all_memory_pcd, 
-                                          voxel_size = fpfh_voxel_size, global_dist_factor = fpfh_global_dist_factor, 
-                                          local_dist_factor = fpfh_local_dist_factor)
+        best_assn = max(assn_data, key=lambda x: x[-1])
+
+        assn = best_assn[0]
+        transform = best_assn[1]
 
         R = copy.copy(transform[:3,:3])
         t = copy.copy(transform[:3, 3])
@@ -1152,7 +1148,9 @@ def localise(self, image_path, depth_image_path, testname="", save_point_clouds=
 
         localised_pose = np.concatenate((tAvg, qAvg))
 
-        return localised_pose
+        # moved objects will have indices that are not present in the first row of assn
+
+        return localised_pose, assn
 
 @dataclass
 class LocalArgs:
@@ -1188,8 +1186,8 @@ class LocalArgs:
     with open(poses_json_path, 'r') as f:
         poses = json.load(f)
 
-    for target in range(1,9): # all of them
-    # for target in [6,7,8]: # sanity check
+    # for target in range(1,9): # all of them
+    for target in [6,7,8]: # sanity check
         target_num = target
         target_pose = None
         for i, view in enumerate(poses["views"]):