Improves edge prediction with class PointNet

The best settings so far:

Mean HSS: 0.2896
Mean F1: 0.3718
Mean IoU: 0.2466
{'vertex_threshold': 0.4, 'edge_threshold': 0.6, 'only_predicted_connections': False}

predict.py

@@ -15,14 +15,15 @@ from fast_pointnet import save_patches_dataset, predict_vertex_from_patch
 from fast_voxel import predict_vertex_from_patch_voxel
 import time
 from fast_pointnet_class import save_patches_dataset as save_patches_dataset_class
+from fast_pointnet_class import predict_class_from_patch
 
 GENERATE_DATASET = False
 DATASET_DIR = '/home/skvrnjan/personal/hohocustom/'
 #DATASET_DIR = '/mnt/personal/skvrnjan/hohocustom'
 
-GENERATE_DATASET_EDGES = True
-#EDGES_DATASET_DIR = '/home/skvrnjan/personal/hohocustom_edges/'
-EDGES_DATASET_DIR = '/mnt/personal/skvrnjan/hohocustom_edges'
+GENERATE_DATASET_EDGES = False
+EDGES_DATASET_DIR = '/home/skvrnjan/personal/hohocustom_edges/'
+#EDGES_DATASET_DIR = '/mnt/personal/skvrnjan/hohocustom_edges'
 
 def convert_entry_to_human_readable(entry):
     out = {}
@@ -953,7 +954,7 @@ def generate_edge_patches(frame):
     all_patches = positive_patches + negative_patches
 
     # Visualize edge patches
-    if False:  # Set to True to enable visualization
+    if True:  # Set to True to enable visualization
         # Create plotter
         plotter = pv.Plotter()
         
@@ -1013,6 +1014,95 @@ def generate_edge_patches(frame):
 
     return all_patches
 
+def generate_edge_patches_forward(frame, pred_vertices):
+    vertices = pred_vertices
+
+    cylinder_radius = 0.5
+
+    colmap = frame['colmap_binary']
+
+    # Create 6D point cloud from COLMAP data
+    colmap_points_6d = []
+    for pid, p3D in colmap.points3D.items():
+        # Combine xyz coordinates and RGB color
+        point_6d = np.concatenate([p3D.xyz, p3D.color / 255.0])  # Normalize color to [0,1]
+        colmap_points_6d.append(point_6d)
+
+    colmap_points_6d = np.array(colmap_points_6d) if colmap_points_6d else np.empty((0, 6))
+    
+    colmap_points_6d[:, 3:] = colmap_points_6d[:, 3:] * 2 - 1
+
+    # Extract 3D coordinates for faster vectorized operations
+    colmap_points_3d = colmap_points_6d[:, :3]
+
+    forward_patches = []
+
+    # For each vertex pair, create a patch without label
+    for i in range(len(vertices)):
+        for j in range(i + 1, len(vertices)):
+            start_vertex = vertices[i]
+            end_vertex = vertices[j]
+            
+            # Create line vector from start to end
+            line_vector = end_vertex - start_vertex
+            line_length = np.linalg.norm(line_vector)
+            
+            # Normalize line vector
+            line_direction = line_vector / line_length
+            
+            # Extend the line by 25 cm (0.25 meters) on both ends for more context
+            extension_length = 0.25  # 25 cm in meters
+            extended_start = start_vertex - extension_length * line_direction
+            extended_end = end_vertex + extension_length * line_direction
+            extended_line_length = line_length + 2 * extension_length
+            
+            # Vectorized distance calculation
+            # Vector from extended start to all points
+            start_to_points = colmap_points_3d - extended_start[np.newaxis, :]
+            
+            # Project onto line direction to get distance along extended line
+            projection_lengths = np.dot(start_to_points, line_direction)
+            
+            # Filter points within extended line segment bounds
+            within_bounds = (projection_lengths >= 0) & (projection_lengths <= extended_line_length)
+            
+            # Find closest points on extended line segment for all points
+            closest_points_on_line = extended_start[np.newaxis, :] + projection_lengths[:, np.newaxis] * line_direction[np.newaxis, :]
+            
+            # Calculate perpendicular distances from points to line
+            perpendicular_distances = np.linalg.norm(colmap_points_3d - closest_points_on_line, axis=1)
+            
+            # Find points within cylinder
+            within_cylinder = within_bounds & (perpendicular_distances <= cylinder_radius)
+            
+            if np.sum(within_cylinder) <= 10:
+                continue
+
+            points_in_cylinder = colmap_points_6d[within_cylinder]
+            point_indices_in_cylinder = np.where(within_cylinder)[0]
+            
+            # Center the patch at the midpoint of the original line (not extended)
+            line_midpoint = (start_vertex + end_vertex) / 2
+            
+            # Shift points to center around origin
+            points_centered = points_in_cylinder.copy()
+            points_centered[:, :3] -= line_midpoint
+            
+            # Create edge patch without label
+            edge_patch = {
+                'patch_6d': points_centered,
+                'connection': (i, j),
+                'line_start': start_vertex - line_midpoint,
+                'line_end': end_vertex - line_midpoint,
+                'cylinder_radius': cylinder_radius,
+                'point_indices': point_indices_in_cylinder,
+                'center': line_midpoint
+            }
+            
+            forward_patches.append(edge_patch)
+
+    return forward_patches
+
 def calculate_cylinder_overlap_volume(cyl1, cyl2):
     """
     Calculate the intersection volume between two cylinders using numpy vectorization.
@@ -1113,21 +1203,28 @@ def calculate_cylinder_overlap_volume(cyl1, cyl2):
     
     return max(0.0, overlap_volume)
 
-def predict_wireframe(entry, pnet_model, voxel_model) -> Tuple[np.ndarray, List[int]]:
+def predict_wireframe(entry, pnet_model, voxel_model, pnet_class_model, config) -> Tuple[np.ndarray, List[int]]:
     """
     Predict 3D wireframe from a dataset entry.
     """
     good_entry = convert_entry_to_human_readable(entry)
     colmap_rec = good_entry['colmap_binary']
 
+    vertex_threshold = config.get('vertex_threshold', 0.5)
+    edge_threshold = config.get('edge_threshold', 0.5)
+    only_predicted_connections = config.get('only_predicted_connections', False)
+
     if GENERATE_DATASET_EDGES:
         patches = generate_edge_patches(good_entry)
-        save_patches_dataset_class(patches, EDGES_DATASET_DIR, good_entry['order_id'])
+        #save_patches_dataset_class(patches, EDGES_DATASET_DIR, good_entry['order_id'])
         return empty_solution()
 
     vert_edge_per_image = {}
     idxs_points = []
     all_connections = []
+
+    print(f"Processing {len(good_entry['gestalt'])} images")
+
     for i, (gest, depth, K, R, t, img_id, ade_seg) in enumerate(zip(good_entry['gestalt'],
                                                 good_entry['depth'], 
                                                 good_entry['K'],
@@ -1222,7 +1319,7 @@ def predict_wireframe(entry, pnet_model, voxel_model) -> Tuple[np.ndarray, List[
 
         #visu_patch_and_pred(patch, pred_vertex, pred_dist, pred_class)
 
-        if pred_class > 0.5:
+        if pred_class > vertex_threshold:
             predicted_vertices.append(pred_vertex)
         else:
             predicted_vertices.append(np.array([0.0, 0.0, 0.0]))  # Append a zero vertex if not predicted
@@ -1273,9 +1370,26 @@ def predict_wireframe(entry, pnet_model, voxel_model) -> Tuple[np.ndarray, List[
 
     #print(f"Filtered vertices: {len(filtered_vertices)} from {len(predicted_vertices)}")
     #print(f"Filtered connections: {len(filtered_connections)} from {len(connections)}")
+    
+    forward_patches = generate_edge_patches_forward(good_entry, filtered_vertices)
+    new_connections = []
+    if len(forward_patches) > 0:
+        for patch in forward_patches:
+            start_idx, end_idx = patch['connection']
+            
+            pred_class, pred_score = predict_class_from_patch(pnet_class_model, patch, device='cuda')
+            if pred_score > edge_threshold:
+                new_connections.append((start_idx, end_idx))
 
     predicted_vertices = np.array(filtered_vertices)
-    connections = filtered_connections
+    
+    if only_predicted_connections:
+        connections = new_connections
+    else:
+        connections = filtered_connections + new_connections
+
+    # Remove duplicates from connections
+    connections = list(set(connections))
 
     return predicted_vertices, connections 
 

train.py

@@ -16,11 +16,12 @@ from predict import predict_wireframe, predict_wireframe_old
 from tqdm import tqdm
 from fast_pointnet import load_pointnet_model
 from fast_voxel import load_3dcnn_model
+from fast_pointnet_class import load_pointnet_model as load_pointnet_class_model
 import torch
 
-#ds = load_dataset("usm3d/hoho25k", cache_dir="/media/skvrnjan/sd/hoho25k/", trust_remote_code=True)
-ds = load_dataset("usm3d/hoho25k", cache_dir="/mnt/personal/skvrnjan/hoho25k/", trust_remote_code=True)
-ds = ds.shuffle()
+ds = load_dataset("usm3d/hoho25k", cache_dir="/media/skvrnjan/sd/hoho25k/", trust_remote_code=True)
+#ds = load_dataset("usm3d/hoho25k", cache_dir="/mnt/personal/skvrnjan/hoho25k/", trust_remote_code=True)
+#ds = ds.shuffle()
  
 scores_hss = []
 scores_f1 = []
@@ -30,19 +31,24 @@ show_visu = False
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-#pnet_model = load_pointnet_model(model_path="/home/skvrnjan/personal/hoho_pnet/initial_epoch_100.pth", device=device, predict_score=True)
-pnet_model = None
+pnet_model = load_pointnet_model(model_path="/home/skvrnjan/personal/hoho_pnet/initial_epoch_100.pth", device=device, predict_score=True)
+#pnet_model = None
+
+pnet_class_model = load_pointnet_class_model(model_path="/home/skvrnjan/personal/hoho_pnet_edges_v2/initial_epoch_100.pth", device=device)
+#pnet_class_model = None
 
 #voxel_model = load_3dcnn_model(model_path="/home/skvrnjan/personal/hoho_voxel/initial_epoch_100.pth", device=device, predict_score=True)
 voxel_model = None
 
+config = {'vertex_threshold': 0.4, 'edge_threshold': 0.7, 'only_predicted_connections': False}
+
 idx = 0
 for a in tqdm(ds['train'], desc="Processing dataset"):
     #plot_all_modalities(a)
     #pred_vertices, pred_edges = predict_wireframe_old(a)
-    #pred_vertices, pred_edges = predict_wireframe(a, pnet_model, voxel_model)
+    #pred_vertices, pred_edges = predict_wireframe(a, pnet_model, voxel_model, pnet_class_model)
     try:
-        pred_vertices, pred_edges = predict_wireframe(a, pnet_model, voxel_model)
+        pred_vertices, pred_edges = predict_wireframe(a, pnet_model, voxel_model, pnet_class_model, config)
         #pred_vertices, pred_edges = predict_wireframe_old(a)
     except:
         pred_vertices, pred_edges = empty_solution()
@@ -53,7 +59,7 @@ for a in tqdm(ds['train'], desc="Processing dataset"):
     scores_f1.append(score.f1)
     scores_iou.append(score.iou)
 
-    if show_visu and score.hss < 0.1:
+    if show_visu:
         colmap = read_colmap_rec(a['colmap_binary'])
         pcd, geometries = plot_reconstruction_local(None, colmap, points=True, cameras=True, crop_outliers=True)
         wireframe = plot_wireframe_local(None, a['wf_vertices'], a['wf_edges'], a['wf_classifications'])
@@ -63,9 +69,14 @@ for a in tqdm(ds['train'], desc="Processing dataset"):
         visu_all = [pcd] + geometries + wireframe + bpo_cams + wireframe2
         o3d.visualization.draw_geometries(visu_all, window_name=f"3D Reconstruction - HSS: {score.hss:.4f}, F1: {score.f1:.4f}, IoU: {score.iou:.4f}")
 
+    idx += 1
+    if idx >= 100:  # Limit to first 10 samples for testing
+        break
+
 for i in range(10):
     print("END OF DATASET")
 print(f"Mean HSS: {np.mean(scores_hss):.4f}")
 print(f"Mean F1: {np.mean(scores_f1):.4f}")
 print(f"Mean IoU: {np.mean(scores_iou):.4f}")
+print(config)