Upload 3 files

feature_solution.py

@@ -338,6 +338,10 @@ def prune_not_connected(all_3d_vertices, connections_3d):
     
     return np.array(new_verts), connected_out
 
+checkpoint_path = "loftr_outdoor.ckpt"
+loftr_model = KF.LoFTR(pretrained=None)
+loftr_model.load_state_dict(torch.load(checkpoint_path)['state_dict'])
+
 def loftr_matcher(gestalt_img_0, gestalt_img1, depth_images):
     import torchvision.transforms as transforms
     rgb_to_gray = transforms.Compose([
@@ -346,7 +350,7 @@ def loftr_matcher(gestalt_img_0, gestalt_img1, depth_images):
         transforms.ToTensor()  # Convert back to tensor
     ])
 
-    device = 'cpu'#torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    device = 'cpu' #torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
     w, h = depth_images.size
     gest_seg_0 = gestalt_img_0.resize(depth_images.size)
@@ -361,7 +365,7 @@ def loftr_matcher(gestalt_img_0, gestalt_img1, depth_images):
     gest_seg_1_tensor = K.image_to_tensor(gest_seg_1_np, False).float().to(device)
     img2 = K.geometry.resize(gest_seg_1_tensor, (int(h/4), int(w/4)))  / 255
 
-    matcher = KF.LoFTR(pretrained="outdoor").to(device)
+    # matcher = KF.LoFTR(pretrained="outdoor").to(device)
 
     input_dict = {
         "image0": img1,
@@ -370,7 +374,7 @@ def loftr_matcher(gestalt_img_0, gestalt_img1, depth_images):
     # print("Input dict shape", input_dict["image0"].shape, input_dict["image1"].shape)
 
     with torch.no_grad():
-        correspondences = matcher(input_dict)
+        correspondences = loftr_model(input_dict)
 
     # mkpts0 = correspondences["keypoints0"].cpu().numpy()
     # mkpts1 = correspondences["keypoints1"].cpu().numpy()
@@ -620,7 +624,7 @@ def predict(entry, visualize=False) -> Tuple[np.ndarray, List[int]]:
         # save_image_with_keypoints(f'keypoints_{j}.png', np.array(good_entry['gestalt'][j]), mkpts_filtered_1, (255, 0, 0))
 
         # Line matching
-        line_0, line_1 = line_matcher(good_entry['gestalt'][i], good_entry['gestalt'][j], good_entry['depthcm'][i]) 
+        # line_0, line_1 = line_matcher(good_entry['gestalt'][i], good_entry['gestalt'][j], good_entry['depthcm'][i]) 
         # save_image_with_lines(f'line_{i}.png', np.array(good_entry['gestalt'][i]), line_0, (255, 0, 0))
         # save_image_with_lines(f'line_{j}.png', np.array(good_entry['gestalt'][j]), line_1, (255, 0, 0))
 

loftr_outdoor.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:21f5bec5968178e8bc8b7633441836fe5de4f47d861dd2cd7dc38e271b0479ec
+size 46341978

script.py

@@ -205,6 +205,13 @@ if __name__ == "__main__":
 							'wf_vertices': pred_vertices.tolist(),
 							'wf_edges': pred_edges
 						})
+	# for i, sample in enumerate(tqdm(dataset)):
+	# 	key, pred_vertices, pred_edges = predict(sample, visualize=False)
+	# 	solution.append({
+	# 					'__key__': key,
+	# 					'wf_vertices': pred_vertices.tolist(),
+	# 					'wf_edges': pred_edges
+	# 				})
 			if i % 100 == 0:
 				# incrementally save the results in case we run out of time
 				print(f"Processed {i} samples")