update

geom_solver.py

@@ -18,10 +18,11 @@ class GeomSolver(object):
 	def __init__(self):
 		self.min_vertices = 18
 		self.kmeans_th = 150
+		self.clr_th = 2.5
 		self.device = 'cuda:0'
 
 
-	def cluster_points(self, point_types, clr_th = 2.5):
+	def cluster_points(self, point_types):
 		point_colors = []
 		for point_type in point_types:
 			point_colors.append(np.array(gestalt_color_mapping[point_type]))
@@ -38,7 +39,7 @@ class GeomSolver(object):
 			gest = self.gests[ki]
 			vert_mask = 0
 			for point_color in point_colors:
-				my_mask = cv2.inRange(gest, point_color-clr_th, point_color+clr_th)
+				my_mask = cv2.inRange(gest, point_color-self.clr_th, point_color+self.clr_th)
 				vert_mask = vert_mask + my_mask
 			vert_mask = (vert_mask > 0).astype(np.uint8)
 
@@ -66,10 +67,11 @@ class GeomSolver(object):
 		criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 200, 0.3)
 		flags = cv2.KMEANS_RANDOM_CENTERS
 		centers = np.zeros((0, 3))
+		assigned_points = []
 		if len(self.xyz[selected_points][dense_pnts]) == 0:
-			return centers
+			return centers, assigned_points
 		for tempi in range(1, 20):
-		    retval, bestLabels, temp_centers = cv2.kmeans(self.xyz[selected_points][dense_pnts].astype(np.float32), tempi, None, criteria, 200,flags)
+		    retval, temp_bestLabels, temp_centers = cv2.kmeans(self.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:
@@ -78,11 +80,16 @@ class GeomSolver(object):
 		        closest_nn = self.kmeans_th
 		    if closest_nn < self.kmeans_th:
 		        break
-		    centers = temp_centers
+		    centers, bestLabels = temp_centers, temp_bestLabels
 		if centers.shape[0] == 0:
-			centers = temp_centers
+			centers, bestLabels = temp_centers, temp_bestLabels
 
-		return centers
+		point_inds = np.arange(self.xyz.shape[0])
+		for ci in range(centers.shape[0]):
+			assigned_inds = point_inds[selected_points][dense_pnts][bestLabels[:,0] == ci]
+			assigned_points.append(assigned_inds)
+
+		return centers, assigned_points
 
 
 	def process_vertices(self):
@@ -112,12 +119,12 @@ class GeomSolver(object):
 		self.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)
 		self.verts = torch.from_numpy(self.xyz.astype(np.float32)).to(self.device)
 
-		centers_apex = self.cluster_points(['apex'])
-		centers_eave = self.cluster_points(['eave_end_point'])
+		centers_apex, assigned_apex = self.cluster_points(['apex'])
+		centers_eave, assigned_eave = self.cluster_points(['eave_end_point'])
 		centers = np.concatenate((centers_apex, centers_eave))
-
-		# image_ids = np.array([p.id for p in self.points3D.values()])
-		# pyt_centers = torch.from_numpy(centers).to(device)
+		self.assigned_points = assigned_apex + assigned_eave
+		self.is_apex = np.zeros((centers.shape[0], )).astype(int)
+		self.is_apex[:centers_apex.shape[0]] = 1
 
 		z_th = centers[:,-1].min() - 50
 		self.wf_center = self.xyz[self.xyz[:,-1] > z_th].mean(0)
@@ -135,17 +142,90 @@ class GeomSolver(object):
 			vertices = np.concatenate((self.vertices, added))
 		self.vertices_aug = vertices
 
+
+	def process_edges(self):
+		N = len(self.gests)
+		image_ids = np.array([p.id for p in self.points3D.values()])
+		center_visibility = [set(np.concatenate([self.points3D[image_ids[pind]].image_ids for pind in ass_item])) for ass_item in self.assigned_points]
+
+		pyt_centers = torch.from_numpy(self.vertices.astype(np.float32)).to(self.device)
+
+		edge_dists = []
+		uvs = []
+		edge_types = {0 : ['eave'], 1 : ['rake', 'valley'], 2 : ['ridge']}
+		for ki in range(N):
+		    gest = self.gests[ki]
+		    edge_masks = {}
+		    per_type_dists = {}
+		    for etype in edge_types:
+		        edge_mask = 0
+		        for edge_class in edge_types[etype]:
+		            edge_color = np.array(gestalt_color_mapping[edge_class])
+		            mask = cv2.morphologyEx(cv2.inRange(gest,
+		                                                edge_color-self.clr_th,
+		                                                edge_color+self.clr_th),
+		                                    cv2.MORPH_DILATE, np.ones((3, 3)))
+		            edge_mask += mask
+		        edge_mask = (edge_mask > 0).astype(np.uint8)
+		        edge_masks[etype] = edge_mask
+		        dist = cv2.distanceTransform(1-edge_mask, cv2.DIST_L2, 3)
+		        per_type_dists[etype] = dist
+		    edge_dists.append(per_type_dists)
+		        
+		    uv = torch.round(self.pyt_cameras[ki].transform_points(pyt_centers)[:, :2]).cpu().numpy().astype(int)
+		    uv_inl = (uv[:, 0] >= 0) * (uv[:, 1] >= 0) * (uv[:, 0] < self.width) * (uv[:, 1] < self.height)
+		    uv = uv[uv_inl]
+		    uvs.append(uv)
+
+		edges = []
+		thresholds_min_mean = {0 : [5, 7], 1 : [15, 25], 2: [30, 1000]}
+		for i in range(pyt_centers.shape[0]):
+		    for j in range(i+1, pyt_centers.shape[0]):
+		        etype = (self.is_apex[i] + self.is_apex[j])
+
+		        points_inter = pyt_centers[i][None] + torch.linspace(0, 1, 20)[:, None].to(self.device) * (pyt_centers[j][None] - pyt_centers[i][None])
+		        min_mean_dist = 1000
+		        all_dists = []
+		        best_ki = -1
+		        best_uvi = -1
+		        for ki in range(N):
+		            cki = self.gestalt_to_colmap_cams[ki]
+		            
+		            if not ( (cki in center_visibility[i]) or (cki in center_visibility[j]) ):
+		                continue
+		        
+		            uvi = torch.round(self.pyt_cameras[ki].transform_points(points_inter)[:, :2]).cpu().numpy().astype(int)
+		            if (uvi <= 0).any() or (uvi[:,0] >= self.width).any() or (uvi[:,1] >= self.height).any():
+		                continue
+		            mean_dist = edge_dists[ki][etype][uvi[:,1], uvi[:,0]].mean()
+		            all_dists.append(mean_dist)
+		            if mean_dist < min_mean_dist:
+		                min_mean_dist = mean_dist
+		                best_ki = ki
+		                best_uvi = uvi
+		                
+		        if best_ki == -1:
+		            continue
+		        ths = thresholds_min_mean[etype]
+		        if min_mean_dist < ths[0] and np.mean(all_dists) < ths[1]:
+		            edges.append((i,j))
+		if len(edges) == 0:
+			edges.append((0, 0))
+		return edges
+
+
 	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
+		edges = self.process_edges()
 
 		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
+		return vertices, 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:
-        vertices = GeomSolver().solve(entry)
-        edges = edges0
+        vertices, edges = GeomSolver().solve(entry)
     except:
         print('ERROR')
         vertices, edges = vertices0, edges0