Enables voxel model integration for prediction

Integrates the voxel model into the prediction pipeline.
This change allows the system to leverage 3D convolutional
neural networks for improved wireframe prediction.
Also, adds visualization of HSS, F1, and IoU scores.

train.py

@@ -12,9 +12,10 @@ from utils import read_colmap_rec, empty_solution
 
 #from hoho2025.example_solutions import predict_wireframe
 from hoho2025.metric_helper import hss
-from predict import predict_wireframe
+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
 import torch
 
 #ds = load_dataset("usm3d/hoho25k", cache_dir="/media/skvrnjan/sd/hoho25k/", trust_remote_code=True)
@@ -25,23 +26,34 @@ scores_hss = []
 scores_f1 = []
 scores_iou = []
 
-show_visu = True
+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
 
+#voxel_model = load_3dcnn_model(model_path="/home/skvrnjan/personal/hoho_voxel/initial_epoch_100.pth", device=device, predict_score=True)
+voxel_model = None
+
 idx = 0
 for a in tqdm(ds['train'], desc="Processing dataset"):
     #plot_all_modalities(a)
-    #pred_vertices, pred_edges = predict_wireframe(a, pnet_model)
+    #pred_vertices, pred_edges = predict_wireframe_old(a)
+    #pred_vertices, pred_edges = predict_wireframe(a, pnet_model, voxel_model)
     try:
-        pred_vertices, pred_edges = predict_wireframe(a, pnet_model)
+        pred_vertices, pred_edges = predict_wireframe(a, pnet_model, voxel_model)
+        #pred_vertices, pred_edges = predict_wireframe_old(a)
     except:
         pred_vertices, pred_edges = empty_solution()
 
-    if show_visu:
+    score = hss(pred_vertices, pred_edges, a['wf_vertices'], a['wf_edges'], vert_thresh=0.5, edge_thresh=0.5)
+    print(f"Score: {score}")
+    scores_hss.append(score.hss)
+    scores_f1.append(score.f1)
+    scores_iou.append(score.iou)
+
+    if show_visu and score.hss < 0.1:
         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'])
@@ -49,13 +61,7 @@ for a in tqdm(ds['train'], desc="Processing dataset"):
         bpo_cams = plot_bpo_cameras_from_entry_local(None, a)
 
         visu_all = [pcd] + geometries + wireframe + bpo_cams + wireframe2
-        o3d.visualization.draw_geometries(visu_all, window_name="3D Reconstruction")
-
-    score = hss(pred_vertices, pred_edges, a['wf_vertices'], a['wf_edges'], vert_thresh=0.5, edge_thresh=0.5)
-    print(f"Score: {score}")
-    scores_hss.append(score.hss)
-    scores_f1.append(score.f1)
-    scores_iou.append(score.iou)
+        o3d.visualization.draw_geometries(visu_all, window_name=f"3D Reconstruction - HSS: {score.hss:.4f}, F1: {score.f1:.4f}, IoU: {score.iou:.4f}")
 
 for i in range(10):
     print("END OF DATASET")