some depth maps are nov retrieved from the pointcloud

handcrafted_solution.py

@@ -5,7 +5,9 @@ from collections import defaultdict
 from typing import Tuple, List
 
 import cv2
+import hoho
 import numpy as np
+import scipy.interpolate as si
 from PIL import Image as PImage
 from hoho.color_mappings import gestalt_color_mapping
 from hoho.read_write_colmap import read_cameras_binary, read_images_binary, read_points3D_binary
@@ -243,14 +245,19 @@ def get_vertices_and_edges_from_segmentation(gest_seg_np, *, point_radius=30, ma
 
     missed_vertices = []
     if len(ridge_edges) > 0 and len(rake_edges) > 0:
-
         inferred_vertices = infer_missing_vertices(ridge_edges, rake_edges)
         missed_vertices = get_missed_vertices(vertices, inferred_vertices, **kwargs)
         vertices = np.concatenate([vertices, missed_vertices])
 
     vertices = KDTree(vertices)
 
-    for edge_class in ['eave', 'valley', 'flashing', 'step_flashing', 'hip']:
+    for edge_class in ['eave',
+                       'step_flashing',
+                       'flashing',
+                       'post',
+                       'valley',
+                       'hip',
+                       'transition_line']:
         class_edges, class_directions = get_lines_and_directions(gest_seg_np, edge_class,
                                                                  rho=rho,
                                                                  theta=theta,
@@ -334,6 +341,7 @@ def get_vertices_and_edges_from_segmentation(gest_seg_np, *, point_radius=30, ma
 def get_uv_depth(vertices, depth):
     '''Get the depth of the vertices from the depth image'''
 
+    # depth[depth > 5000] = np.inf
     uv = np.array([v['xy'] for v in vertices])
     uv_int = uv.astype(np.int32)
     H, W = depth.shape[:2]
@@ -380,12 +388,10 @@ def merge_vertices_3d(vert_edge_per_image, merge_th=0.1, **kwargs):
         for vv in v:
             already_there.add(vv)
     old_idx_to_new = {}
-    count = 0
-    for idxs in merged:
+    for count, idxs in enumerate(merged):
         new_vertices.append(all_3d_vertices[idxs].mean(axis=0))
         for idx in idxs:
             old_idx_to_new[idx] = count
-        count += 1
     # print (connections_3d)
     new_vertices = np.array(new_vertices)
     # print (connections_3d)
@@ -422,49 +428,109 @@ def prune_not_connected(all_3d_vertices, connections_3d):
 
 
 def predict(entry, visualize=False, scale_estimation_coefficient=2.5, **kwargs) -> Tuple[np.ndarray, List[int]]:
-    good_entry = convert_entry_to_human_readable(entry)
-    if 'gestalt' not in good_entry or 'depthcm' not in good_entry or 'K' not in good_entry or 'R' not in good_entry or 't' not in good_entry:
+    if 'gestalt' not in entry or 'depthcm' not in entry or 'K' not in entry or 'R' not in entry or 't' not in entry:
         print('Missing required fields in the entry')
-        return (good_entry['__key__'], *empty_solution())
+        return (entry['__key__'], *empty_solution())
+    entry = hoho.decode(entry)
+
     vert_edge_per_image = {}
-    for i, (gest, depth, K, R, t) in enumerate(zip(good_entry['gestalt'],
-                                                   good_entry['depthcm'],
-                                                   good_entry['K'],
-                                                   good_entry['R'],
-                                                   good_entry['t']
-                                                   )):
-        gest_seg = gest.resize(depth.size)
-        gest_seg_np = np.array(gest_seg).astype(np.uint8)
-        # Metric3D
-        depth_np = np.array(depth) / scale_estimation_coefficient
-        vertices, connections = get_vertices_and_edges_from_segmentation(gest_seg_np, **kwargs)
-        if (len(vertices) < 2) or (len(connections) < 1):
-            print(f'Not enough vertices or 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)
-        # 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)
+    image_dict = {}
+    for k, v in entry["images"].items():
+        image_dict[v.name] = v
+    for i, (gest, depthcm, K, R, t, imagekey) in enumerate(zip(entry['gestalt'],
+                                                               entry['depthcm'],
+                                                               entry['K'],
+                                                               entry['R'],
+                                                               entry['t'],
+                                                               entry['__imagekey__']
+                                                               )):
+
+        try:
+            gest_seg = gest.resize(depthcm.size)
+            gest_seg_np = np.array(gest_seg).astype(np.uint8)
+            vertices, connections = get_vertices_and_edges_from_segmentation(gest_seg_np, **kwargs)
+            if (len(vertices) < 2) or (len(connections) < 1):
+                print(f'Not enough vertices or connections in image {i}')
+                vert_edge_per_image[i] = np.empty((0, 2)), [], np.empty((0, 3))
+                continue
+            belonging_points = []
+            for i in image_dict[imagekey].point3D_ids[np.where(image_dict[imagekey].point3D_ids != -1)]:
+                belonging_points.append(entry["points3d"][i])
+
+            if len(belonging_points) < 1:
+                print(f'No 3D points in image {i}')
+                vert_edge_per_image[i] = np.empty((0, 2)), [], np.empty((0, 3))
+                raise KeyError
+            projected2d, _ = cv2.projectPoints(np.array([i.xyz for i in belonging_points]), R, t, K, 0)
+            important = np.where(np.all(projected2d >= 0, axis=2))
+            # Normalize the uv to the camera intrinsics
+            world_to_cam = np.eye(4)
+            world_to_cam[:3, :3] = R
+            world_to_cam[:3, 3] = t
+
+            homo_belonging_points = cv2.convertPointsToHomogeneous(np.array([i.xyz for i in belonging_points]))
+            depth = cv2.convertPointsFromHomogeneous(cv2.transform(homo_belonging_points, world_to_cam))
+            depth = np.array([i[0][2] for i in depth])
+            depth = depth[important[0]]
+            projected2d = projected2d[important]
+            if len(depth) < 1:
+                print(f'No 3D points in image {i}')
+                vert_edge_per_image[i] = np.empty((0, 2)), [], np.empty((0, 3))
+                raise KeyError
+            # print(projected2d.shape, depth.shape)
+
+            interpolator = si.NearestNDInterpolator(projected2d, depth)
+
+            vertex_coordinates = np.array([v['xy'] for v in vertices])
+            xi, yi = vertex_coordinates[:, 0], vertex_coordinates[:, 1]
+            depth_vert = interpolator(xi, yi)
+            xy_local = np.ones((len(vertex_coordinates), 3))
+            xy_local[:, 0] = (vertex_coordinates[:, 0] - K[0, 2]) / K[0, 0]
+            xy_local[:, 1] = (vertex_coordinates[:, 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)
+
+        except KeyError:
+            gest_seg = gest.resize(depthcm.size)
+            gest_seg_np = np.array(gest_seg).astype(np.uint8)
+            # Metric3D
+            depth_np = np.array(depthcm) / scale_estimation_coefficient
+            vertices, connections = get_vertices_and_edges_from_segmentation(gest_seg_np, **kwargs)
+            if (len(vertices) < 2) or (len(connections) < 1):
+                print(f'Not enough vertices or 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)
+            # 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)
+
         vert_edge_per_image[i] = vertices, connections, vertices_3d
     all_3d_vertices, connections_3d = merge_vertices_3d(vert_edge_per_image, **kwargs)
-    all_3d_vertices_clean, connections_3d_clean = prune_not_connected(all_3d_vertices, connections_3d)
+    all_3d_vertices_clean, connections_3d_clean = all_3d_vertices, connections_3d
+    # all_3d_vertices_clean, connections_3d_clean = prune_not_connected(all_3d_vertices, connections_3d)
     if (len(all_3d_vertices_clean) < 2) or len(connections_3d_clean) < 1:
         print(f'Not enough vertices or connections in the 3D vertices')
-        return (good_entry['__key__'], *empty_solution())
+        return (entry['__key__'], *empty_solution())
     if visualize:
         from hoho.viz3d import plot_estimate_and_gt
         plot_estimate_and_gt(all_3d_vertices_clean,
                              connections_3d_clean,
-                             good_entry['wf_vertices'],
-                             good_entry['wf_edges'])
-    return good_entry['__key__'], all_3d_vertices_clean, connections_3d_clean
+                             entry['wf_vertices'],
+                             entry['wf_edges'])
+    return entry['__key__'], all_3d_vertices_clean, connections_3d_clean

script.py

@@ -132,9 +132,9 @@ if __name__ == "__main__":
                                        point_radius=25,
                                        max_angle=15,
                                        extend=30,
-                                       merge_th=3.0,
-                                       min_missing_distance=300.0,
-                                       scale_estimation_coefficient=4.5,
+                                       merge_th=100.0,
+                                       min_missing_distance=30000000.0,
+                                       scale_estimation_coefficient=2.54,
                                        ))
         
         for i, result in enumerate(tqdm(results)):