ADE category statistics

handcrafted_solution.py

@@ -121,6 +121,94 @@ def get_uv_depth(vertices, depth):
     vertex_depth = depth[(uv_int[:, 1] , uv_int[:, 0])]
     return uv, vertex_depth
 
+def get_vertices_and_edges_from_two_segmentations(ade_seg_np, gest_seg_np, edge_th = 50.0):
+    '''Get the vertices and edges from the gestalt segmentation mask of the house'''
+    vertices = []
+    connections = []
+    # combined map from ade
+    print(gest_seg_np.shape, ade_seg_np.shape)
+    ade_color0 = np.array([0,0,0]) 
+    ade_mask0 = cv2.inRange(ade_seg_np,  ade_color0-0.5, ade_color0+0.5)
+    ade_color1 = np.array([120,120,120]) 
+    ade_mask1 = cv2.inRange(ade_seg_np,  ade_color1-0.5, ade_color1+0.5)
+    ade_color2 = np.array([180,120,120]) 
+    ade_mask2 = cv2.inRange(ade_seg_np,  ade_color2-0.5, ade_color2+0.5)
+    ade_color3 = np.array([255,9,224]) 
+    ade_mask3 = cv2.inRange(ade_seg_np,  ade_color3-0.5, ade_color3+0.5)
+    ade_mask = cv2.bitwise_or(ade_mask3, ade_mask2)
+    ade_mask = cv2.bitwise_or(ade_mask1, ade_mask)
+    print(ade_mask.any())
+    # Apex
+    apex_color = np.array(gestalt_color_mapping['apex'])
+    apex_mask = cv2.inRange(gest_seg_np,  apex_color-0.5, apex_color+0.5)
+    apex_mask = cv2.bitwise_and(apex_mask, ade_mask)
+    if apex_mask.sum() > 0:
+        output = cv2.connectedComponentsWithStats(apex_mask, 8, cv2.CV_32S)
+        (numLabels, labels, stats, centroids) = output
+        stats, centroids = stats[1:], centroids[1:]
+        
+        for i in range(numLabels-1):
+            vert = {"xy": centroids[i], "type": "apex"}
+            vertices.append(vert)
+    
+    eave_end_color = np.array(gestalt_color_mapping['eave_end_point'])
+    eave_end_mask = cv2.inRange(gest_seg_np,  eave_end_color-0.5, eave_end_color+0.5)
+    eave_end_mask = cv2.bitwise_and(eave_end_mask, ade_mask)
+    if eave_end_mask.sum() > 0:
+        output = cv2.connectedComponentsWithStats(eave_end_mask, 8, cv2.CV_32S)
+        (numLabels, labels, stats, centroids) = output
+        stats, centroids = stats[1:], centroids[1:]
+        
+        for i in range(numLabels-1):
+            vert = {"xy": centroids[i], "type": "eave_end_point"}
+            vertices.append(vert) 
+    
+    print(f'{len(vertices)} vertices detected')
+    # Connectivity
+    apex_pts = []
+    apex_pts_idxs = []
+    for j, v in enumerate(vertices):
+        apex_pts.append(v['xy'])
+        apex_pts_idxs.append(j)
+    apex_pts = np.array(apex_pts)
+            
+    # Ridge connects two apex points
+    for edge_class in ['eave', 'ridge', 'rake', 'valley']:
+        edge_color = np.array(gestalt_color_mapping[edge_class])
+        mask = cv2.morphologyEx(cv2.inRange(gest_seg_np,
+                                            edge_color-0.5,
+                                            edge_color+0.5),
+                                cv2.MORPH_DILATE, np.ones((11, 11)))
+        line_img = np.copy(gest_seg_np) * 0
+        if mask.sum() > 0:
+            output = cv2.connectedComponentsWithStats(mask, 8, cv2.CV_32S)
+            (numLabels, labels, stats, centroids) = output
+            stats, centroids = stats[1:], centroids[1:]
+            edges = []
+            for i in range(1, numLabels):
+                y,x = np.where(labels == i)
+                xleft_idx = np.argmin(x)
+                x_left = x[xleft_idx]
+                y_left = y[xleft_idx]
+                xright_idx = np.argmax(x)
+                x_right = x[xright_idx]
+                y_right = y[xright_idx]
+                edges.append((x_left, y_left, x_right, y_right))
+                cv2.line(line_img, (x_left, y_left), (x_right, y_right), (255, 255, 255), 2)
+            edges = np.array(edges)
+            if (len(apex_pts) < 2) or len(edges) <1:
+                continue
+            pts_to_edges_dist = np.minimum(cdist(apex_pts, edges[:,:2]), cdist(apex_pts, edges[:,2:]))
+            connectivity_mask = pts_to_edges_dist <= edge_th
+            edge_connects = connectivity_mask.sum(axis=0)
+            for edge_idx, edgesum in enumerate(edge_connects):
+                if edgesum>=2:
+                    connected_verts = np.where(connectivity_mask[:,edge_idx])[0]
+                    for a_i, a in enumerate(connected_verts):
+                        for b in connected_verts[a_i+1:]:
+                            connections.append((a, b))
+    return vertices, connections
+
 def get_uv_dept_category(vertices, depth, ade_seg):
     '''Get the depth of the vertices from the depth image'''
     uv = []
@@ -133,11 +221,19 @@ def get_uv_dept_category(vertices, depth, ade_seg):
     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]
+    target_color = set([(120,120,120), (180, 120, 120), (255,9,224)])
+    #filter_ind = [i for i, ele in enumerate(vertex_category) if tuple(ele) in target_color]
+    filter_ind = []
+    for i, ele in enumerate(vertex_category):
+      if tuple(ele) in target_color:
+        filter_ind.append(i)
+
     print(f'retain {len(filter_ind)} idx')
     print(vertex_category[filter_ind])
-    return uv[filter_ind], vertex_depth[filter_ind], vertex_category[filter_ind]
+    #print(vertices)
+    #print(filter_ind)
+    vertices = [vertices[i] for i in filter_ind]
+    return uv[filter_ind], vertex_depth[filter_ind], vertex_category[filter_ind], vertices
 
 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'''
@@ -232,13 +328,15 @@ def predict(entry, visualize=False) -> Tuple[np.ndarray, List[int]]:
         gest_seg_np = np.array(gest_seg).astype(np.uint8)
         # Metric3D
         depth_np = np.array(depth) / 2.5 # 2.5 is the scale estimation coefficient
-        vertices, connections = get_vertices_and_edges_from_segmentation(gest_seg_np, edge_th = 20.)
+        #vertices, connections = get_vertices_and_edges_from_segmentation(gest_seg_np, edge_th = 20.)
+        vertices, connections = get_vertices_and_edges_from_two_segmentations(ade_seg_np, gest_seg_np, edge_th = 20.)
+        
         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)
-        uv, depth_vert, ade_category_vert = get_uv_dept_category(vertices, depth_np, ade_seg_np)
+        uv, depth_vert = get_uv_depth(vertices, depth_np)
+        #uv, depth_vert, ade_category_vert, vertices = 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]