update

geom_solver.py

@@ -1,16 +1,15 @@
 import numpy as np
 from pytorch3d.ops import ball_query
 from helpers import *
-from handcrafted_solution import *
+from handcrafted_solution import convert_entry_to_human_readable
+import cv2
 import hoho
 import itertools
 import torch
 from pytorch3d.renderer import PerspectiveCameras
-
+from hoho.color_mappings import gestalt_color_mapping
 
 def my_empty_solution():
-	'''Return a minimal valid solution, i.e. 2 vertices and 1 edge.'''
-	# return np.zeros((1,3)), [(0, 0)]
 	return np.zeros((18,3)), [(0, 0)]
 
 
@@ -19,54 +18,28 @@ class GeomSolver(object):
 	def __init__(self):
 		self.min_vertices = 18
 		self.kmeans_th = 150
-		self.clr_th = 2.5
 		self.device = 'cuda:0'
 
-	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]
-
-		cameras, images, points3D = human_entry['cameras'], human_entry['images'], human_entry['points3d']
-		colmap_cameras_tf = list(human_entry['images'].keys())
-		xyz = np.stack([p.xyz for p in points3D.values()])
-		color = np.stack([p.rgb for p in points3D.values()])
-
-		gestalt_camcet = np.stack([eye for eye, target, up, fov in itertools.starmap(hoho.Rt_to_eye_target, zip(*[human_entry[k] for k in 'ade20k K R t'.split()]))])
-		col_camcet = np.stack([eye for eye, target, up, fov in col_cams])
-		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 cluster_points(self, point_types, clr_th = 2.5):
+		point_colors = []
+		for point_type in point_types:
+			point_colors.append(np.array(gestalt_color_mapping[point_type]))
 
-		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)
-		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=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(self.device)
-
-		apex_color = np.array(gestalt_color_mapping['apex'])
-		eave_end_color = np.array(gestalt_color_mapping['eave_end_point'])
-
-		dist_points = np.zeros((xyz.shape[0], ))
-		visible_counts = np.zeros((xyz.shape[0], ), dtype=int)
+		dist_points = np.zeros((self.verts.shape[0], ))
+		visible_counts = np.zeros((self.verts.shape[0], ), dtype=int)
 		proj_uv = []
-		for ki in range(N):
-			if broken_cams[ki]:
+		for ki in range(len(self.gestalt_to_colmap_cams)):
+			if self.broken_cams[ki]:
 				proj_uv.append(([], []))
 				continue
-			cki = gestalt_to_colmap_cams[ki]
+			cki = self.gestalt_to_colmap_cams[ki]
 
-			gest = np.array(human_entry['gestalt'][ki])
-			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
+			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)
+				vert_mask = vert_mask + my_mask
 			vert_mask = (vert_mask > 0).astype(np.uint8)
 
 			dist = cv2.distanceTransform(1-vert_mask, cv2.DIST_L2, 3)
@@ -74,11 +47,9 @@ class GeomSolver(object):
 			ndist = np.zeros_like(dist)
 			ndist = cv2.normalize(dist, ndist, 0, 1.0, cv2.NORM_MINMAX)
 
-			in_this_image = np.array([cki in p.image_ids for p in points3D.values()])
-			# tempind = 2103
-			# print(in_this_image[tempind-1], cki, points3D[tempind].image_ids)
-			uv = torch.round(pyt_cameras[ki].transform_points(verts)[:, :2]).cpu().numpy().astype(int)
-			uv_inl = (uv[:, 0] >= 0) * (uv[:, 1] >= 0) * (uv[:, 0] < width) * (uv[:, 1] < height) * in_this_image
+			in_this_image = np.array([cki in p.image_ids for p in self.points3D.values()])
+			uv = torch.round(self.pyt_cameras[ki].transform_points(self.verts)[:, :2]).cpu().numpy().astype(int)
+			uv_inl = (uv[:, 0] >= 0) * (uv[:, 1] >= 0) * (uv[:, 0] < self.width) * (uv[:, 1] < self.height) * in_this_image
 			proj_uv.append((uv, uv_inl))
 			uv = uv[uv_inl]
 
@@ -88,17 +59,17 @@ class GeomSolver(object):
 		selected_points = (dist_points / (visible_counts + 1e-6)) <= 10
 		selected_points[visible_counts < 1] = False
 
-
-		pnts = torch.from_numpy(xyz[selected_points].astype(np.float32))[None]
+		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
 
-
 		criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 200, 0.3)
 		flags = cv2.KMEANS_RANDOM_CENTERS
-		centers = None
+		centers = np.zeros((0, 3))
+		if len(self.xyz[selected_points][dense_pnts]) == 0:
+			return centers
 		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)
+		    retval, 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:
@@ -108,23 +79,56 @@ class GeomSolver(object):
 		    if closest_nn < self.kmeans_th:
 		        break
 		    centers = temp_centers
-		if centers is None:
+		if centers.shape[0] == 0:
 			centers = temp_centers
-		# image_ids = np.array([p.id for p in points3D.values()])
+
+		return centers
+
+
+	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]
+
+		cameras, images, self.points3D = human_entry['cameras'], human_entry['images'], human_entry['points3d']
+		colmap_cameras_tf = list(human_entry['images'].keys())
+		self.xyz = np.stack([p.xyz for p in self.points3D.values()])
+		color = np.stack([p.rgb for p in self.points3D.values()])
+		self.gests = [np.array(gest0) for gest0 in human_entry['gestalt']]
+
+		gestalt_camcet = np.stack([eye for eye, target, up, fov in itertools.starmap(hoho.Rt_to_eye_target, zip(*[human_entry[k] for k in 'ade20k K R t'.split()]))])
+		col_camcet = np.stack([eye for eye, target, up, fov in col_cams])
+		self.gestalt_to_colmap_cams = [colmap_cameras_tf[np.argmin(((gcam - col_camcet)**2).sum(1)**0.5)] for gcam in gestalt_camcet]
+		self.broken_cams = np.array([np.min(((gcam - col_camcet)**2).sum(1)**0.5) for gcam in gestalt_camcet]) > 300
+
+		self.height, self.width = cameras[1].height, cameras[1].width
+		N = len(self.gestalt_to_colmap_cams)
+		K = to_K(*human_entry['cameras'][1].params)[None].repeat(N, 0)
+		R = np.stack([quaternion_to_rotation_matrix(human_entry['images'][self.gestalt_to_colmap_cams[ind]].qvec) for ind in range(N)])
+		T = np.stack([human_entry['images'][self.gestalt_to_colmap_cams[ind]].tvec for ind in range(N)])
+
+		R = np.linalg.inv(R)
+		image_size = torch.Tensor([self.height, self.width]).repeat(N, 1)
+		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 = 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)
 
-		z_th = centers[:,-1].min() - 10
-		self.wf_center = xyz[xyz[:,-1] > z_th].mean(0)
+		z_th = centers[:,-1].min() - 50
+		self.wf_center = self.xyz[self.xyz[:,-1] > z_th].mean(0)
 		self.wf_center[-1] = centers[:, -1].mean()
-		self.with_broken_cams = broken_cams.any()
 
 		self.vertices = centers
-
-		if self.with_broken_cams:
-			vertices = self.vertices
+		if self.broken_cams.any():
+			vertices = centers
 			print("There are broken cams.")
 		else:
-			nvert = self.vertices.shape[0]
+			nvert = centers.shape[0]
 			# added_one = (self.min_vertices * self.wf_center - self.vertices.sum(0)) / (self.min_vertices - nvert)
 			added_one = self.wf_center
 			added = added_one[None].repeat(self.min_vertices - nvert,0)

helpers.py

@@ -1,6 +1,8 @@
 import numpy as np
 from PIL import Image as PImage
 import io
+from scipy.spatial.distance import cdist
+from scipy.optimize import linear_sum_assignment
 
 
 def to_K(f, cx, cy):
@@ -21,9 +23,6 @@ def quaternion_to_rotation_matrix(qvec):
     return R
    
 
-from scipy.spatial.distance import cdist
-from scipy.optimize import linear_sum_assignment
-
 def preregister_mean_std(verts_to_transform, target_verts, single_scale=True):
     mu_target = target_verts.mean(axis=0)
     mu_in = verts_to_transform.mean(axis=0)

my_solution.py

@@ -39,6 +39,7 @@ def predict(entry, visualize=False) -> Tuple[np.ndarray, List[int]]:
         vertices = GeomSolver().solve(entry)
         edges = edges0
     except:
+        print('ERROR')
         vertices, edges = vertices0, edges0
 
     if vertices.shape[0] < vertices0.shape[0]:
@@ -46,11 +47,6 @@ def predict(entry, visualize=False) -> Tuple[np.ndarray, List[int]]:
         verts_new[:vertices.shape[0]] = vertices
         vertices = verts_new
 
-    # if len(vertices) == 1:
-    #     # print("Added one more vertex")
-    #     vertices = np.concatenate((vertices, np.zeros((1,3))))
-    #     edges = [(0,0)]
-
     if (len(edges) < 1) and (len(vertices) >= 2):
         # print("Added only edges")
        edges = edges0