Generates edge patches from predicted vertices

Implements edge patch generation based on predicted vertices instead of ground truth.

It finds the closest ground truth vertex for each predicted vertex within a distance threshold and propagates ground truth connections to predicted vertices.
This facilitates training using predicted wireframes.

Also, the IoU threshold for positive patch overlap is increased to 1.0, effectively disabling it.

Enables generation of edge datasets from predicted vertices.

Adds visualization of ground truth vertices and connections in blue to the edge patch visualization.

predict.py

@@ -16,14 +16,16 @@ 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
+from scipy.spatial.distance import cdist
+from scipy.optimize import linear_sum_assignment
 
 GENERATE_DATASET = False
 DATASET_DIR = '/home/skvrnjan/personal/hohocustom/'
 #DATASET_DIR = '/mnt/personal/skvrnjan/hohocustom'
 
-GENERATE_DATASET_EDGES = False
-EDGES_DATASET_DIR = '/home/skvrnjan/personal/hohocustom_edges/'
-#EDGES_DATASET_DIR = '/mnt/personal/skvrnjan/hohocustom_edges'
+GENERATE_DATASET_EDGES = True
+#EDGES_DATASET_DIR = '/home/skvrnjan/personal/hohocustom_edges/'
+EDGES_DATASET_DIR = '/mnt/personal/skvrnjan/hohocustom_edges'
 
 def convert_entry_to_human_readable(entry):
     out = {}
@@ -723,12 +725,41 @@ def visu_pcloud_and_preds(colmap_rec, extracted_ids, extracted_points, extracted
         ball_radius = 1.0  # meters
         plotter.show(title=f"Extracted Points within {ball_radius}m radius - Spheres at group means")
 
-def generate_edge_patches(frame):
-    vertices = frame['wf_vertices']
-    connections = frame['wf_edges']
-
-    vertices = np.array(vertices) if vertices else np.empty((0, 3))
-
+def generate_edge_patches(frame, pred_vertices):
+    gt_vertices = np.array(frame['wf_vertices']) if frame['wf_vertices'] else np.empty((0, 3))
+    gt_connections = frame['wf_edges']
+    
+    vertices = np.array(pred_vertices) if pred_vertices is not None and len(pred_vertices) > 0 else np.empty((0, 3))
+    
+    # Find closest GT vertex for each predicted vertex
+    connections = []
+    if len(vertices) > 0 and len(gt_vertices) > 0:
+        # For each GT vertex, find the closest predicted vertex
+        gt_to_pred_mapping = {}
+        for gt_idx, gt_vertex in enumerate(gt_vertices):
+            # Calculate distances from this GT vertex to all predicted vertices
+            distances = np.linalg.norm(vertices - gt_vertex, axis=1)
+            
+            # Find the closest predicted vertex
+            closest_pred_idx = np.argmin(distances)
+            closest_distance = distances[closest_pred_idx]
+            
+            # Only map if within distance threshold
+            distance_threshold = 1.5
+            if closest_distance <= distance_threshold:
+                gt_to_pred_mapping[gt_idx] = closest_pred_idx
+        
+        # Propagate GT connections to predicted vertices
+        for gt_connection in gt_connections:
+            gt_start, gt_end = gt_connection
+            if gt_start in gt_to_pred_mapping and gt_end in gt_to_pred_mapping:
+                pred_start = gt_to_pred_mapping[gt_start]
+                pred_end = gt_to_pred_mapping[gt_end]
+                connections.append((pred_start, pred_end))
+        
+        print(f"Matched {len(gt_to_pred_mapping)} GT vertices to predicted vertices")
+        print(f"Propagated {len(connections)} connections from GT to predicted vertices")
+    
     positive_patches = []
     negative_patches = []
 
@@ -897,7 +928,7 @@ def generate_edge_patches(frame):
                     union_volume = current_volume + pos_volume - overlap_volume
                     if union_volume > 0:
                         iou = overlap_volume / union_volume
-                        if iou > 0.5:  # 0.2 IoU threshold
+                        if iou > 1.:  # 0.2 IoU threshold
                             has_overlap = True
                             break
                 
@@ -954,7 +985,7 @@ def generate_edge_patches(frame):
     all_patches = positive_patches + negative_patches
 
     # Visualize edge patches
-    if True:  # Set to True to enable visualization
+    if False:  # Set to True to enable visualization
         # Create plotter
         plotter = pv.Plotter()
         
@@ -965,6 +996,24 @@ def generate_edge_patches(frame):
             whole_cloud["colors"] = gray_colors
             plotter.add_mesh(whole_cloud, scalars="colors", rgb=True, point_size=3, render_points_as_spheres=True)
         
+        # Add GT vertices and connections in blue
+        gt_vertices = np.array(frame['wf_vertices']) if frame['wf_vertices'] else np.empty((0, 3))
+        gt_connections = frame['wf_edges']
+        
+        if len(gt_vertices) > 0:
+            # Add GT vertices as blue spheres
+            for gt_vertex in gt_vertices:
+                gt_sphere = pv.Sphere(radius=0.15, center=gt_vertex)
+                plotter.add_mesh(gt_sphere, color='blue', opacity=0.8)
+            
+            # Add GT connections as blue lines
+            for gt_connection in gt_connections:
+                gt_start_idx, gt_end_idx = gt_connection
+                if gt_start_idx < len(gt_vertices) and gt_end_idx < len(gt_vertices):
+                    gt_line_points = np.array([gt_vertices[gt_start_idx], gt_vertices[gt_end_idx]])
+                    gt_line = pv.Line(gt_line_points[0], gt_line_points[1])
+                    plotter.add_mesh(gt_line, color='blue', line_width=8)
+        
         # Visualize each patch
         for patch_idx, patch in enumerate(all_patches):
             # Use green for positive (edge), red for negative (non-edge)
@@ -1008,7 +1057,7 @@ def generate_edge_patches(frame):
         # Set title based on label distribution
         positive_count = sum(1 for patch in all_patches if patch['label'] == 1)
         negative_count = sum(1 for patch in all_patches if patch['label'] == 0)
-        title = f"Edge Patches - Positive (Green): {positive_count}, Negative (Red): {negative_count}"
+        title = f"Edge Patches - Positive (Green): {positive_count}, Negative (Red): {negative_count}, GT (Blue)"
         
         plotter.show(title=title)
 
@@ -1332,11 +1381,6 @@ def predict_wireframe(entry, pnet_model, voxel_model, pnet_class_model, config)
     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'])
-        return empty_solution()
-
     vert_edge_per_image = {}
     idxs_points = []
     all_connections = []
@@ -1467,13 +1511,20 @@ def predict_wireframe(entry, pnet_model, voxel_model, pnet_class_model, config)
     non_zero_mask = ~np.all(np.isclose(predicted_vertices, [0.0, 0.0, 0.0]), axis=1)
     valid_indices = np.where(non_zero_mask)[0]
 
+    # Filter vertices to only include non-zero ones
+    filtered_vertices = predicted_vertices[valid_indices]
+
+    #patches = generate_edge_patches(good_entry, filtered_vertices)
+    if GENERATE_DATASET_EDGES:
+        patches = generate_edge_patches(good_entry, filtered_vertices)
+        save_patches_dataset_class(patches, EDGES_DATASET_DIR, good_entry['order_id'])
+
+        return empty_solution()
+
     if len(valid_indices) == 0:
         print("No valid predicted vertices found")
         return empty_solution()
 
-    # Filter vertices to only include non-zero ones
-    filtered_vertices = predicted_vertices[valid_indices]
-
     # Create mapping from old indices to new indices
     old_to_new_mapping = {old_idx: new_idx for new_idx, old_idx in enumerate(valid_indices)}
 

train.py

@@ -19,9 +19,9 @@ 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 = []
@@ -31,11 +31,11 @@ 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
+#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
@@ -46,7 +46,7 @@ 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, pnet_class_model, config)
+    #pred_vertices, pred_edges = predict_wireframe(a, pnet_model, voxel_model, pnet_class_model, config)
     try:
         pred_vertices, pred_edges = predict_wireframe(a, pnet_model, voxel_model, pnet_class_model, config)
         #pred_vertices, pred_edges = predict_wireframe_old(a)
@@ -70,8 +70,8 @@ for a in tqdm(ds['train'], desc="Processing dataset"):
         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
+    #if idx >= 100:  # Limit to first 10 samples for testing
+    #    break
 
 for i in range(10):
     print("END OF DATASET")