update

geom_solver.py

@@ -16,12 +16,14 @@ def my_empty_solution():
 
 class GeomSolver(object):
 
-	def __init__(self, entry):
+	def __init__(self):
 		self.min_vertices = 18
+		self.kmeans_th = 150
+		self.clr_th = 2.5
+		self.device = 'cuda:0'
 
-		human_entry = convert_entry_to_human_readable(entry)
-		self.human_entry = human_entry
-
+	def process_vertices(self):
+		human_entry = self.human_entry
 		col_cams = [hoho.Rt_to_eye_target(human_entry['ade20k'][0], to_K(*human_entry['cameras'][1].params), quaternion_to_rotation_matrix(colmap_img.qvec), colmap_img.tvec) for colmap_img in human_entry['images'].values()]
 		eye, target, up, fov = col_cams[0]
 
@@ -35,23 +37,19 @@ class GeomSolver(object):
 		gestalt_to_colmap_cams = [colmap_cameras_tf[np.argmin(((gcam - col_camcet)**2).sum(1)**0.5)] for gcam in gestalt_camcet]
 		broken_cams = np.array([np.min(((gcam - col_camcet)**2).sum(1)**0.5) for gcam in gestalt_camcet]) > 300
 
-		# def get_vertices(self):
-		clr_th = 2.5
-		device = 'cuda:0'
+
 		height = cameras[1].height
 		width = cameras[1].width
 		N = len(gestalt_to_colmap_cams)
 		K = to_K(*human_entry['cameras'][1].params)[None].repeat(N, 0)
-		# print(gestalt_to_colmap_cams, N, human_entry['images'])
-		# print(sorted(human_entry['images'].keys()))
 		R = np.stack([quaternion_to_rotation_matrix(human_entry['images'][gestalt_to_colmap_cams[ind]].qvec) for ind in range(N)])
 		T = np.stack([human_entry['images'][gestalt_to_colmap_cams[ind]].tvec for ind in range(N)])
 
 		R = np.linalg.inv(R)
-		image_size=torch.Tensor([height, width]).repeat(N, 1)
-		pyt_cameras = PerspectiveCameras(device=device, R=R, T=T, in_ndc=False, focal_length=K[:, 0, :1], principal_point=K[:, :2, 2], image_size=image_size)
+		image_size  = torch.Tensor([height, width]).repeat(N, 1)
+		pyt_cameras = PerspectiveCameras(device=self.device, R=R, T=T, in_ndc=False, focal_length=K[:, 0, :1], principal_point=K[:, :2, 2], image_size=image_size)
 
-		verts = torch.from_numpy(xyz.astype(np.float32)).to(device)
+		verts = torch.from_numpy(xyz.astype(np.float32)).to(self.device)
 
 		apex_color = np.array(gestalt_color_mapping['apex'])
 		eave_end_color = np.array(gestalt_color_mapping['eave_end_point'])
@@ -66,8 +64,8 @@ class GeomSolver(object):
 			cki = gestalt_to_colmap_cams[ki]
 
 			gest = np.array(human_entry['gestalt'][ki])
-			apex_mask = cv2.inRange(gest,  apex_color-clr_th, apex_color+clr_th)
-			eave_end_mask = cv2.inRange(gest,  eave_end_color-clr_th, eave_end_color+clr_th)
+			apex_mask = cv2.inRange(gest,  apex_color-self.clr_th, apex_color+self.clr_th)
+			eave_end_mask = cv2.inRange(gest,  eave_end_color-self.clr_th, eave_end_color+self.clr_th)
 			vert_mask = apex_mask + eave_end_mask
 			vert_mask = (vert_mask > 0).astype(np.uint8)
 
@@ -99,16 +97,15 @@ class GeomSolver(object):
 		criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 200, 0.3)
 		flags = cv2.KMEANS_RANDOM_CENTERS
 		centers = None
-		kmeans_th = 150
-		for tempi in range(1,11):
+		for tempi in range(1, 20):
 		    retval, bestLabels, temp_centers = cv2.kmeans(xyz[selected_points][dense_pnts].astype(np.float32), tempi, None, criteria, 200,flags)
 		    cpnts = torch.from_numpy(temp_centers.astype(np.float32))[None]
 		    bdists, inds, nn = ball_query(cpnts, cpnts, K=1, radius=100) 
 		    if bdists.max() > 0:
 		        closest_nn = (bdists[bdists>0].min()**0.5).item()
 		    else:
-		        closest_nn = kmeans_th
-		    if closest_nn < kmeans_th:
+		        closest_nn = self.kmeans_th
+		    if closest_nn < self.kmeans_th:
 		        break
 		    centers = temp_centers
 		if centers is None:
@@ -123,8 +120,6 @@ class GeomSolver(object):
 
 		self.vertices = centers
 
-
-	def get_vertices(self, visualize=False):
 		if self.with_broken_cams:
 			vertices = self.vertices
 			print("There are broken cams.")
@@ -134,9 +129,20 @@ class GeomSolver(object):
 			added_one = self.wf_center
 			added = added_one[None].repeat(self.min_vertices - nvert,0)
 			vertices = np.concatenate((self.vertices, added))
+		self.vertices_aug = vertices
+
+	def solve(self, entry, visualize=False):
+		human_entry = convert_entry_to_human_readable(entry)
+		self.human_entry = human_entry
+		self.process_vertices()
+		vertices = self.vertices_aug
+
 		if visualize:
 			from hoho.viz3d import plot_estimate_and_gt
 			plot_estimate_and_gt(vertices, [(0,0)], self.human_entry['wf_vertices'], self.human_entry['wf_edges'])
+		
 		return vertices
 
+		
+
 

helpers.py

@@ -119,7 +119,7 @@ def my_compute_WED(pd_vertices, pd_edges, gt_vertices, gt_edges, cv_ins=-1/2, cv
     # Step 5: Calculation of WED
     WED = translation_costs + deletion_costs + insertion_costs + deletion_edge_costs + insertion_edge_costs
     print(translation_costs, deletion_costs, insertion_costs, deletion_edge_costs, insertion_edge_costs)
-    
+
     if normalized:
         total_length_of_gt_edges = np.linalg.norm((gt_vertices[gt_edges[:, 0]] - gt_vertices[gt_edges[:, 1]]), axis=1).sum()
         WED = WED / total_length_of_gt_edges

my_solution.py

@@ -36,8 +36,7 @@ def predict(entry, visualize=False) -> Tuple[np.ndarray, List[int]]:
     # return (entry['__key__'], *my_empty_solution())
     vertices0, edges0 = my_empty_solution()
     try:
-        solver = GeomSolver(entry)
-        vertices = solver.get_vertices()
+        vertices = GeomSolver().solve(entry)
         edges = edges0
     except:
         vertices, edges = vertices0, edges0