ADE category statistics

handcrafted_solution.py

@@ -121,6 +121,23 @@ def get_uv_depth(vertices, depth):
     vertex_depth = depth[(uv_int[:, 1] , uv_int[:, 0])]
     return uv, vertex_depth
 
+def get_uv_dept_category(vertices, depth, ade_seg):
+    '''Get the depth of the vertices from the depth image'''
+    uv = []
+    for v in vertices:
+        uv.append(v['xy'])
+    uv = np.array(uv)
+    uv_int = uv.astype(np.int32)
+    H, W = depth.shape[:2]
+    uv_int[:, 0] = np.clip( uv_int[:, 0], 0, W-1)
+    uv_int[:, 1] = np.clip( uv_int[:, 1], 0, H-1)
+    vertex_depth = depth[(uv_int[:, 1] , uv_int[:, 0])]
+    vertex_category = ade_seg[(uv_int[:, 1] , uv_int[:, 0])]
+    taget_color = [[120,120,120], [180, 120, 120], [255,9,224]]
+    filter_ind = [i for i, ele in enumerate(vertex_category) if ele in taget_color]
+    print(f'retain {len(filter_ind)} idx')
+    print(vertex_category[filter_ind])
+    return uv[filter_ind], vertex_depth[filter_ind], vertex_category[filter_ind]
 
 def merge_vertices_3d(vert_edge_per_image, th=0.1):
     '''Merge vertices that are close to each other in 3D space and are of same types'''
@@ -202,12 +219,15 @@ def prune_not_connected(all_3d_vertices, connections_3d):
 def predict(entry, visualize=False) -> Tuple[np.ndarray, List[int]]:
     good_entry = convert_entry_to_human_readable(entry)
     vert_edge_per_image = {}
-    for i, (gest, depth, K, R, t) in enumerate(zip(good_entry['gestalt'],
+    for i, (ade, gest, depth, K, R, t) in enumerate(zip(good_entry['ade20k'],
+                                                good_entry['gestalt'],
                                                 good_entry['depthcm'], 
                                                 good_entry['K'],
                                                 good_entry['R'],
                                                 good_entry['t'] 
                                                 )):
+        ade_seg = ade.resize(depth.size)
+        ade_seg_np = np.array(ade_seg).astype(np.uint8)
         gest_seg = gest.resize(depth.size)
         gest_seg_np = np.array(gest_seg).astype(np.uint8)
         # Metric3D
@@ -217,7 +237,8 @@ def predict(entry, visualize=False) -> Tuple[np.ndarray, List[int]]:
             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)
+        #uv, depth_vert = get_uv_depth(vertices, depth_np)
+        uv, depth_vert, ade_category_vert = get_uv_dept_category(vertices, depth_np, ade_seg_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]
@@ -237,6 +258,7 @@ def predict(entry, visualize=False) -> Tuple[np.ndarray, List[int]]:
         print (f'Not enough vertices or connections in the 3D vertices')
         return (good_entry['__key__'], *empty_solution())
     if visualize:
+        print(f"num of est: {len(all_3d_vertices_clean)}, num of gt:{len(good_entry['wf_vertices'])}")
         from hoho.viz3d import plot_estimate_and_gt
         plot_estimate_and_gt(   all_3d_vertices_clean, 
                                 connections_3d_clean,