added cleaner vertex mask

handcrafted_solution.py

@@ -105,15 +105,20 @@ def get_vertices_and_edges_from_segmentation(gest_seg_np, edge_th=50.0):
 
     gest_seg_np = clean_image(gest_seg_np)
     apex_centroids, eave_end_point_centroids, apex_mask, eave_end_point_mask = get_vertices(gest_seg_np)
-    apex_mask = cv2.morphologyEx(apex_mask,
-                                 cv2.MORPH_DILATE, np.ones((11, 11)), iterations=4)
-    eave_end_point_mask = cv2.morphologyEx(eave_end_point_mask,
-                                           cv2.MORPH_DILATE, np.ones((5, 5)), iterations=4)
-    vertex_mask = cv2.bitwise_not(cv2.bitwise_or(apex_mask, eave_end_point_mask))
+
+    vertex_mask = np.zeros_like(apex_mask)
+    for i in apex_centroids:
+        cv2.circle(vertex_mask, np.round(i).astype(np.uint32), 20, (255, ), 40, -1)
+    for i in eave_end_point_centroids:
+        cv2.circle(vertex_mask, np.round(i).astype(np.uint32), 15, (255, ), 30, -1)
+    vertex_mask = np.bitwise_not(vertex_mask)
 
     apex_pts = np.concatenate([apex_centroids, eave_end_point_centroids])
 
+
     for edge_class in ['eave', 'ridge', 'rake', 'valley', 'flashing']:
+        if (len(apex_pts) < 2):
+            break
         edge_color = np.array(gestalt_color_mapping[edge_class])
 
         mask = cv2.inRange(gest_seg_np,
@@ -125,10 +130,10 @@ def get_vertices_and_edges_from_segmentation(gest_seg_np, edge_th=50.0):
 
             mask = cv2.morphologyEx(mask,
                                     cv2.MORPH_DILATE, np.ones((11, 11)), iterations=3)
-            line_img = np.zeros_like(gest_seg_np)
+            # line_img = np.zeros_like(gest_seg_np)
             output = cv2.connectedComponentsWithStats(mask, 8, cv2.CV_32S)
             (numLabels, labels, stats, centroids) = output
-            stats, centroids = stats[1:], centroids[1:]
+            # stats, centroids = stats[1:], centroids[1:]
             edges = []
             for i in range(1, numLabels):
                 y, x = np.where(labels == i)
@@ -139,9 +144,9 @@ def get_vertices_and_edges_from_segmentation(gest_seg_np, edge_th=50.0):
                 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)
+                # 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:
+            if 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
@@ -262,7 +267,7 @@ 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=90.)
+        vertices, connections = get_vertices_and_edges_from_segmentation(gest_seg_np, edge_th=70.)
         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))