Improves training process and model initialization

Initializes model weights using Xavier/Glorot initialization for better training.

Adds a check to skip saving patches if they already exist, preventing unnecessary writes.

Introduces a training script for the pointnet model.

Adds progress bar during training loop and shuffles the dataset for better performance.

fast_pointnet.py

@@ -226,6 +226,10 @@ def save_patches_dataset(patches: List[Dict], dataset_dir: str, entry_id: str):
         filename = f"{entry_id}_patch_{i}.pkl"
         filepath = os.path.join(dataset_dir, filename)
         
+        # Skip if file already exists
+        if os.path.exists(filepath):
+            continue
+        
         # Save patch data
         with open(filepath, 'wb') as f:
             pickle.dump(patch, f)
@@ -266,6 +270,22 @@ def train_pointnet(dataset_dir: str, model_save_path: str, epochs: int = 100, ba
     
     # Initialize model with score prediction
     model = FastPointNet(input_dim=7, output_dim=3, max_points=1024, predict_score=True)
+    
+    # Initialize weights using Xavier/Glorot initialization
+    def init_weights(m):
+        if isinstance(m, nn.Conv1d):
+            nn.init.xavier_uniform_(m.weight)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.Linear):
+            nn.init.xavier_uniform_(m.weight)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.BatchNorm1d):
+            nn.init.ones_(m.weight)
+            nn.init.zeros_(m.bias)
+    
+    model.apply(init_weights)
     model.to(device)
     
     # Loss functions

train.py

@@ -12,8 +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 tqdm import tqdm
 
-ds = load_dataset("usm3d/hoho25k", streaming=True, trust_remote_code=True)
+ds = load_dataset("usm3d/hoho25k", cache_dir='/home/skvrnjan/personal/hoho25k', trust_remote_code=True)
+ds = ds.shuffle()
  
 scores_hss = []
 scores_f1 = []
@@ -22,7 +24,7 @@ scores_iou = []
 show_visu = False
 
 idx = 0
-for a in ds['train']:
+for a in tqdm(ds['train'], desc="Processing dataset"):
     #plot_all_modalities(a)
     #pred_vertices, pred_edges = predict_wireframe(a)
     try:

train_pnet.py

@@ -0,0 +1,12 @@
+from fast_pointnet import FastPointNet
+
+if __name__ == "__main__":
+    # Initialize the FastPointNet model
+    model = FastPointNet()
+
+    # Load the dataset
+    dataset = model.load_dataset("/home/skvrnjan/personal/hohocustom/")
+    model_save_path = "/home/skvrnjan/personal/hoho_pnet/"
+
+    # Train the model
+    model.train(dataset, model_save_path, epochs=10)