x

train.py

@@ -6,6 +6,8 @@ import tempfile,zipfile
 import io
 import open3d as o3d
 import os
+import argparse # Added for command-line arguments
+import numpy as np # Make sure numpy is imported if not already implicitly
 
 from visu import plot_reconstruction_local, plot_wireframe_local, plot_bpo_cameras_from_entry_local, _plotly_rgb_to_normalized_o3d_color
 from utils import read_colmap_rec, empty_solution
@@ -21,8 +23,31 @@ from fast_pointnet_class_10d import load_pointnet_model as load_pointnet_class_m
 import torch
 import time
 
-#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)
+# --- Argument Parsing ---
+parser = argparse.ArgumentParser(description="Train and evaluate HoHo model with custom config.")
+parser.add_argument('--vertex_threshold', type=float, default=0.4, help='Vertex threshold for prediction.')
+parser.add_argument('--edge_threshold', type=float, default=0.625, help='Edge threshold for prediction.')
+parser.add_argument('--only_predicted_connections', type=lambda x: (str(x).lower() == 'true'), default=True, help='Use only predicted connections (True/False).')
+parser.add_argument('--max_samples', type=int, default=100, help='Maximum number of samples to process.')
+parser.add_argument('--results_dir', type=str, default="results", help='Directory to save result files.')
+
+
+args = parser.parse_args()
+
+# --- Configuration from Arguments ---
+config = {
+    'vertex_threshold': args.vertex_threshold,
+    'edge_threshold': args.edge_threshold,
+    'only_predicted_connections': args.only_predicted_connections
+}
+print(f"Running with configuration: {config}")
+
+# Create results directory if it doesn't exist
+os.makedirs(args.results_dir, exist_ok=True)
+
+
+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 = []
@@ -45,7 +70,6 @@ pnet_class_model = load_pointnet_class_model(model_path="pnet_class.pth", device
 #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.6, 'only_predicted_connections': False}
 
 idx = 0
 prediction_times = []
@@ -60,9 +84,13 @@ for a in tqdm(ds['train'], desc="Processing dataset"):
         end_time = time.time()
         prediction_time = end_time - start_time
         prediction_times.append(prediction_time)
-        mean_time = np.mean(prediction_times)
-        print(f"Prediction time: {prediction_time:.4f} seconds, Mean time: {mean_time:.4f} seconds")  
-    except:
+        if prediction_times: # ensure not empty before calculating mean
+            mean_time = np.mean(prediction_times)
+            print(f"Prediction time: {prediction_time:.4f} seconds, Mean time: {mean_time:.4f} seconds")  
+        else:
+            print(f"Prediction time: {prediction_time:.4f} seconds")
+    except Exception as e: # Catch specific exceptions if possible, or log the error
+        print(f"Error during prediction: {e}")
         pred_vertices, pred_edges = empty_solution()
 
     score = hss(pred_vertices, pred_edges, a['wf_vertices'], a['wf_edges'], vert_thresh=0.5, edge_thresh=0.5)
@@ -82,13 +110,51 @@ 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
+    if idx >= args.max_samples:
+        print(f"Reached max_samples limit: {args.max_samples}")
         break
 
-for i in range(10):
+for i in range(10): # This loop seems to be for console output spacing
     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)
 
+mean_hss_val = np.mean(scores_hss) if scores_hss else 0.0
+mean_f1_val = np.mean(scores_f1) if scores_f1 else 0.0
+mean_iou_val = np.mean(scores_iou) if scores_iou else 0.0
+
+
+print(f"Mean HSS: {mean_hss_val:.4f}")
+print(f"Mean F1: {mean_f1_val:.4f}")
+print(f"Mean IoU: {mean_iou_val:.4f}")
+print(f"Final Config: {config}")
+if prediction_times:
+    print(f"Overall Mean Prediction Time: {np.mean(prediction_times):.4f} seconds")
+
+
+# --- Writing results to a file ---
+vt_str = str(config['vertex_threshold']).replace('.', 'p')
+et_str = str(config['edge_threshold']).replace('.', 'p')
+opc_str = str(config['only_predicted_connections'])
+
+results_filename = f"results_vt{vt_str}_et{et_str}_opc{opc_str}_samples{args.max_samples}.txt"
+results_filepath = os.path.join(args.results_dir, results_filename)
+
+with open(results_filepath, 'w') as f:
+    f.write(f"Configuration: {config}\n")
+    f.write(f"Max Samples Processed: {args.max_samples}\n")
+    f.write(f"Mean HSS: {mean_hss_val:.4f}\n")
+    f.write(f"Mean F1: {mean_f1_val:.4f}\n")
+    f.write(f"Mean IoU: {mean_iou_val:.4f}\n")
+    if prediction_times:
+        f.write(f"Overall Mean Prediction Time: {np.mean(prediction_times):.4f} seconds\n")
+    f.write("\nIndividual HSS Scores:\n")
+    for s_hss in scores_hss:
+        f.write(f"{s_hss:.4f}\n")
+    f.write("\nIndividual F1 Scores:\n")
+    for s_f1 in scores_f1:
+        f.write(f"{s_f1:.4f}\n")
+    f.write("\nIndividual IoU Scores:\n")
+    for s_iou in scores_iou:
+        f.write(f"{s_iou:.4f}\n")
+
+
+print(f"Results saved to {results_filepath}")