Filters segmented point clouds by depth and masks

Improves the extraction of segmented point clouds by adding filtering based on depth and masks for apex and eave end points. This reduces noise and inaccuracies in the point cloud data. Additionally, switches the training loop to use the validation set.

predict.py

@@ -595,16 +595,111 @@ def extract_segmented_pcloud(gest_seg_np, colmap_rec, img_id_substring, ade_seg,
     # Filter points that fall within the apex or eave_end masks
     filtered_points_xyz = []
     filtered_point_idxs = []
+    filtered_points_color = []
 
+    # Apex
+    apex_color = np.array(gestalt_color_mapping['apex'])
+    apex_mask = cv2.inRange(gest_seg_np, apex_color-10., apex_color+10.)
+    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(1, numLabels):
+            cur_mask = labels == i
+            # Dilate the current mask to make it slightly larger
+            kernel = np.ones((5, 5), np.uint8)
+            cur_mask = cv2.dilate(cur_mask.astype(np.uint8), kernel, iterations=2).astype(bool)
+            color = np.random.rand(3)
+            # Create boolean mask for points in current apex mask and house mask
+            valid_points_mask = cur_mask[uv[:, 1], uv[:, 0]] & house_mask[uv[:, 1], uv[:, 0]]
+            
+            for z in range(5):
+                if np.sum(valid_points_mask) < 5:
+                    cur_mask = cv2.dilate(cur_mask.astype(np.uint8), kernel, iterations=1).astype(bool)
+                    valid_points_mask = cur_mask[uv[:, 1], uv[:, 0]] & house_mask[uv[:, 1], uv[:, 0]]
+                else:
+                    break
+                # 
+            if np.any(valid_points_mask):
+                # Get indices of valid points
+                valid_point_indices = valid_indices[valid_points_mask]
+                
+                # Get 3D points in camera coordinates for depth filtering
+                valid_world_points = points_xyz_world[valid_point_indices]
+                valid_cam_points = points_cam[valid_point_indices]
+                
+                # Compute depths (Z coordinates in camera space)
+                depths = valid_cam_points[:, 2]
+                
+                # Find minimum depth and filter points within min_depth + 2 meters
+                if len(depths) > 0:
+                    min_depth = np.min(depths)
+                    depth_filter = depths <= (min_depth + 2.0)
+                    
+                    # Apply depth filter
+                    final_valid_indices = valid_point_indices[depth_filter]
+                    
+                    # Add corresponding points to filtered lists
+                    filtered_points_xyz.extend(points_xyz_world[final_valid_indices])
+                    filtered_point_idxs.extend(points_idxs[final_valid_indices])
+                    filtered_points_color.extend([color] * np.sum(depth_filter))
+
+
+    # Eave end
+    eave_end_color = np.array(gestalt_color_mapping['eave_end_point'])
+    eave_end_mask = cv2.inRange(gest_seg_np, eave_end_color-10, eave_end_color+10)
+    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(1, numLabels):
+            cur_mask = labels == i
+            kernel = np.ones((5,5), np.uint8)
+            cur_mask = cv2.dilate(cur_mask.astype(np.uint8), kernel, iterations=2).astype(bool)
+            color = np.random.rand(3)
+            valid_points_mask = cur_mask[uv[:, 1], uv[:, 0]] & house_mask[uv[:, 1], uv[:, 0]]
+            
+            for z in range(5):
+                if np.sum(valid_points_mask) < 5:
+                    cur_mask = cv2.dilate(cur_mask.astype(np.uint8), kernel, iterations=1).astype(bool)
+                    valid_points_mask = cur_mask[uv[:, 1], uv[:, 0]] & house_mask[uv[:, 1], uv[:, 0]]
+                else:
+                    break
+            if np.any(valid_points_mask):
+                # Get indices of valid points
+                valid_point_indices = valid_indices[valid_points_mask]
+                
+                # Get 3D points in camera coordinates for depth filtering
+                valid_world_points = points_xyz_world[valid_point_indices]
+                valid_cam_points = points_cam[valid_point_indices]
+                
+                # Compute depths (Z coordinates in camera space)
+                depths = valid_cam_points[:, 2]
+                
+                # Find minimum depth and filter points within min_depth + 2 meters
+                if len(depths) > 0:
+                    min_depth = np.min(depths)
+                    depth_filter = depths <= (min_depth + 2.0)
+                    
+                    # Apply depth filter
+                    final_valid_indices = valid_point_indices[depth_filter]
+                    
+                    # Add corresponding points to filtered lists
+                    filtered_points_xyz.extend(points_xyz_world[final_valid_indices])
+                    filtered_point_idxs.extend(points_idxs[final_valid_indices])
+                    filtered_points_color.extend([color] * np.sum(depth_filter))
+
+    '''
     for i, (u, v) in enumerate(uv):
     # Check if this projected point falls within the combined maskvalid_indices
         if combined_mask[v, u] > 0 and house_mask[v, u] > 0:
             original_idx = valid_indices[i]  # Get original index
             filtered_points_xyz.append(points_xyz_world[original_idx])
             filtered_point_idxs.append(points_idxs[original_idx])
-
+    '''
     filtered_points_xyz = np.array(filtered_points_xyz) if filtered_points_xyz else np.empty((0, 3))
     filtered_point_idxs = np.array(filtered_point_idxs) if filtered_point_idxs else np.empty((0,))
+    filtered_points_color = np.array(filtered_points_color) if filtered_points_color else np.empty((0, 3))
 
     '''
     depth_fitted, depth_sparse, _, col_img = get_fitted_dense_depth(depth, colmap_rec, img_id_substring, ade_seg, K, R, t)
@@ -663,7 +758,7 @@ def extract_segmented_pcloud(gest_seg_np, colmap_rec, img_id_substring, ade_seg,
     # Filtered COLMAP points in red
     if len(filtered_points_xyz) > 0:
         pcd_filtered.points = o3d.utility.Vector3dVector(filtered_points_xyz)
-        pcd_filtered.colors = o3d.utility.Vector3dVector(np.full((len(filtered_points_xyz), 3), [1.0, 0.0, 0.0]))
+        pcd_filtered.colors = o3d.utility.Vector3dVector(np.array(filtered_points_color))
 
     # Segmented depth points in blue
     if len(segmented_points_3d) > 0:

train.py

@@ -22,7 +22,7 @@ scores_iou = []
 show_visu = False
 
 idx = 0
-for a in ds['train']:
+for a in ds['validation']:
     #plot_all_modalities(a)
     #pred_vertices, pred_edges = predict_wireframe(a)
     try: