update

geom_solver.py

@@ -16,11 +16,11 @@ def my_empty_solution():
 class GeomSolver(object):
 
 	def __init__(self):
-		self.min_vertices = 12
-		self.mean_vertices = 18 
-		self.max_vertices = 25
-		self.kmeans_th = 150
-		self.point_dist_th = 25
+		self.min_vertices = 10
+		self.mean_vertices = 18
+		self.max_vertices = 30
+		self.kmeans_th = 200
+		self.point_dist_th = 50
 		self.th_min_support = 3
 		self.clr_th = 2.5
 		self.device = 'cuda:0'
@@ -72,11 +72,14 @@ class GeomSolver(object):
 
 		criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 200, 0.3)
 		flags = cv2.KMEANS_RANDOM_CENTERS
+		point_inds = np.arange(self.xyz.shape[0])
 		centers = np.zeros((0, 3))
 		assigned_points = []
 		if len(self.xyz[selected_points][dense_pnts]) == 0:
 			return centers, assigned_points
-		for tempi in range(1, 20):
+		if len(self.xyz[selected_points][dense_pnts]) == 1:
+			return self.xyz[selected_points][dense_pnts], [point_inds[selected_points][dense_pnts]]
+		for tempi in range(1, 30):
 		    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=2, radius=1.2*self.kmeans_th) 
@@ -90,7 +93,6 @@ class GeomSolver(object):
 		if centers.shape[0] == 0:
 			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]
@@ -98,6 +100,12 @@ class GeomSolver(object):
 				continue
 			centers_selected.append(centers[ci])
 			assigned_points.append(assigned_inds)
+		if len(centers_selected) == 0:
+			print("Not centers with enough support!")
+			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
 		centers_selected = np.stack(centers_selected)
 		return centers_selected, assigned_points
 
@@ -143,7 +151,8 @@ class GeomSolver(object):
 		self.vertices = centers
 		nvert = centers.shape[0]
 		# desired_vertices = (self.xyz[:,-1] > z_th).sum() // 300
-		desired_vertices = 2*nvert
+		desired_vertices = 2*nvert 
+		# desired_vertices = self.mean_vertices
 		if desired_vertices < self.min_vertices:
 			desired_vertices = self.mean_vertices
 		if desired_vertices > self.max_vertices:
@@ -206,6 +215,7 @@ class GeomSolver(object):
 
 		edges = []
 		thresholds_min_mean = {0 : [5, 7], 1 : [9, 25], 2: [30, 1000]}
+		# thresholds_min_mean = {0 : [1, 7], 1 : [1, 25], 2: [1, 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])