clean up

handcrafted_solution.py

@@ -457,8 +457,22 @@ def prune_not_connected(all_3d_vertices, connections_3d):
     
     return np.array(new_verts), connected_out
 
-
-def predict(entry, visualize=False) -> Tuple[np.ndarray, List[int]]:
+def uv_to_v3d(uv, depth_vert, K, R, t):
+  # Normalize the uv to the camera intrinsics
+  xy_local = np.ones((len(uv), 3))
+  xy_local[:, 0] = (uv[:, 0] - K[0,2]) / K[0,0]
+  xy_local[:, 1] = (uv[:, 1] - K[1,2]) / K[1,1]
+  # Get the 3D vertices
+  vertices_3d_local = depth_vert[...,None] * (xy_local/np.linalg.norm(xy_local, axis=1)[...,None])
+  world_to_cam = np.eye(4)
+  world_to_cam[:3, :3] = R
+  world_to_cam[:3, 3] = t.reshape(-1)
+  cam_to_world =  np.linalg.inv(world_to_cam)
+  vertices_3d = cv2.transform(cv2.convertPointsToHomogeneous(vertices_3d_local), cam_to_world)
+  vertices_3d = cv2.convertPointsFromHomogeneous(vertices_3d).reshape(-1, 3)
+  return vertices_3d
+
+def predict(entry, visualize=False, depth_scale=2.5, ) -> Tuple[np.ndarray, List[int]]:
     good_entry = convert_entry_to_human_readable(entry)
     points3D = good_entry['points3d']
     vert_edge_per_image = {}
@@ -477,10 +491,7 @@ def predict(entry, visualize=False) -> Tuple[np.ndarray, List[int]]:
         elif i==2: # only visualize view 0,1
           continue 
         '''
-
-        depth_scale = 2.5
         
-
         ade_seg = ade.resize(depth.size)
         ade_seg_np = np.array(ade_seg).astype(np.uint8)
         gest_seg = gest.resize(depth.size)
@@ -495,26 +506,14 @@ def predict(entry, visualize=False) -> Tuple[np.ndarray, List[int]]:
             print (f'Not enough vertices ({len(vertices)}) or connections ({len(connections)}) in image {i}')
             vert_edge_per_image[i] = np.empty((0, 2)), [], np.empty((0, 3))
             continue
-        
-
-        
+                
         #uv, depth_vert = get_uv_depth(vertices, depth_np)         
         sfm_depth_np = get_SfM_depth(points3D, depth_np, gest_seg_np, K, R, t, 10)
         uv, depth_vert = get_smooth_uv_depth(vertices, depth_np, gest_seg_np, sfm_depth_np)
         #uv, depth_vert = get_smooth_uv_depth(vertices, depth_np, gest_seg_np, None)        
 
-        # Normalize the uv to the camera intrinsics
-        xy_local = np.ones((len(uv), 3))
-        xy_local[:, 0] = (uv[:, 0] - K[0,2]) / K[0,0]
-        xy_local[:, 1] = (uv[:, 1] - K[1,2]) / K[1,1]
-        # Get the 3D vertices
-        vertices_3d_local = depth_vert[...,None] * (xy_local/np.linalg.norm(xy_local, axis=1)[...,None])
-        world_to_cam = np.eye(4)
-        world_to_cam[:3, :3] = R
-        world_to_cam[:3, 3] = t.reshape(-1)
-        cam_to_world =  np.linalg.inv(world_to_cam)
-        vertices_3d = cv2.transform(cv2.convertPointsToHomogeneous(vertices_3d_local), cam_to_world)
-        vertices_3d = cv2.convertPointsFromHomogeneous(vertices_3d).reshape(-1, 3)
+        vertices_3d = uv_to_v3d(uv, depth_vert, K, R, t)        
+        
         vert_edge_per_image[i] = vertices, connections, vertices_3d
     all_3d_vertices, connections_3d = merge_vertices_3d(vert_edge_per_image, 3.0) # TODO: 3cm looks too small    
     #print(f'after merge, {len(all_3d_vertices)} 3d vertices and {len(connections_3d)} 3d connections')