update

geom_solver.py

@@ -17,8 +17,9 @@ class GeomSolver(object):
 
 	def __init__(self):
 		self.min_vertices = 18
-		self.kmeans_th = 150
-		self.point_dist_th = 25
+		self.kmeans_th = 130
+		self.point_dist_th = 60
+		self.th_min_support = 3
 		self.clr_th = 2.5
 		self.device = 'cuda:0'
 		self.return_edges = False
@@ -63,8 +64,8 @@ class GeomSolver(object):
 		selected_points[visible_counts < 1] = False
 
 		pnts = torch.from_numpy(self.xyz[selected_points].astype(np.float32))[None]
-		bdists, inds, nn = ball_query(pnts, pnts, K=3, radius=30) 
-		dense_pnts = (bdists[0] > 0).sum(1) == 2
+		bdists, inds, nn = ball_query(pnts, pnts, K=3, radius=40) 
+		dense_pnts = (bdists[0] > 0).sum(1) == 3
 
 		criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 200, 0.3)
 		flags = cv2.KMEANS_RANDOM_CENTERS
@@ -75,7 +76,7 @@ class GeomSolver(object):
 		for tempi in range(1, 20):
 		    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) 
+		    bdists, inds, nn = ball_query(cpnts, cpnts, K=2, radius=1.2*self.kmeans_th) 
 		    if bdists.max() > 0:
 		        closest_nn = (bdists[bdists>0].min()**0.5).item()
 		    else:
@@ -87,11 +88,15 @@ class GeomSolver(object):
 			centers, bestLabels = temp_centers, temp_bestLabels
 
 		point_inds = np.arange(self.xyz.shape[0])
+		centers_selected = []
 		for ci in range(centers.shape[0]):
 			assigned_inds = point_inds[selected_points][dense_pnts][bestLabels[:,0] == ci]
+			if len(assigned_inds) < self.th_min_support:
+				continue
+			centers_selected.append(centers[ci])
 			assigned_points.append(assigned_inds)
-
-		return centers, assigned_points
+		centers_selected = np.stack(centers_selected)
+		return centers_selected, assigned_points
 
 
 	def process_vertices(self):