add network comparison notebook, update dataset and training script, update simulation

Author: dianamindroc

.gitignore

@@ -15,12 +15,19 @@ dist/
 data/
 models/
 
+figures/
+
+simulation/new_sim.py
+simulation/sim_smlm_data.py
+simulation/simulate_smlm.py
+
 scripts/test.py
 scripts/*.def
 scripts/*.sif
 scripts/*.yaml
 scripts/quick_read.py
 scripts/scraps.py
+scripts/test_gpu.py
 
 dataset/SMLMSimulatorStats.py
 

configs/config.yaml

@@ -1,5 +1,5 @@
 dataset:
-  root_folder: '/home/dim26fa/data/sim_dna_tetra_31102024/'
+  root_folder: '/home/dim26fa/data/sim_17022025/samples'
   classes: ['all']
   suffix: '.csv'
   highest_shape: 500
@@ -9,7 +9,7 @@ dataset:
   number_corners_remove: [0,1,2]
   remove_outliers: True
   anisotropy: True
-  anisotropy_factor: 1.5
+  anisotropy_factor: 0.6
   anisotropy_axis: 'z'
   initial_alpha: 1.0
   number_classes: 2
@@ -32,21 +32,22 @@ train:
   scheduler_patience: 20
   gamma: 0.0001 # for other than stepLR use smaller values e.g. 0.0001
   step_size: 50
-  grid_size: 4
   ignore_zeros: False
   augmentation: True
   initial_alpha:
   loss1_improvement_threshold: 30
   min_improvement_epochs: 10
   classifier: True
   pretrained: False
-  ckpt: ''
+  #ckpt: '/home/dim26fa/coding/training/logs_pcn_20241009_225619/best_l1_cd.pth'
+  ckpt: None
   autoencoder: False
   lambda_reg: 0.01
   alpha_density: 0.5
   beta_sparsity: 0.02
-  num_dense: 16384
+  num_dense: 2048
   latent_dim: 1024
+  grid_size: 2
 pcn_config:
   coarse_dim: 16384
   fine_dim: 1024

configs/test_config.yaml

@@ -39,5 +39,6 @@ test:
   #ckpt_path: '/home/dim26fa/coding/training/logs_pcn_20240807_163550/best_l1_cd.pth'
   #ckpt_path: '/home/dim26fa/coding/training/logs_pcn_20240812_124453/best_l1_cd.pth'
   #ckpt_path: '/home/dim26fa/coding/training/logs_pcn_20240814_165920/best_l1_cd.pth'
-  ckpt_path: '/home/dim26fa/coding/training/logs_pcn_20241009_225619/best_l1_cd.pth'
+  #ckpt_path: '/home/dim26fa/coding/training/logs_pcn_20241009_225619/best_l1_cd.pth'
+  ckpt_path: '/home/dim26fa/coding/training/logs_pcn_20241205_012746/best_l1_cd.pth'
   # ckpt_path: '/home/dim26fa/coding/training/logs_pcn_20240819_100525/best_l1_cd.pth'

dataset/SMLMDataset.py

@@ -200,9 +200,21 @@ def __getitem__(self, index):
         filepath = self.filepaths[index]
         if 'csv' in self.suffix:
             # Load data from file
-            arr = np.array(pd.read_csv(filepath))[:, :3]
+            data = pd.read_csv(filepath)[['x', 'y', 'z']]
+            arr = np.array(data)[:, :3]
+            if all(col in data.columns for col in ['rot_x', 'rot_y', 'rot_z', 'rot_w']):
+                # Take first row since rotation should be same for all points
+                rotation = np.array([
+                    data['rot_x'].iloc[0],
+                    data['rot_y'].iloc[0],
+                    data['rot_z'].iloc[0],
+                    data['rot_w'].iloc[0]
+                ])
+            else:
+                rotation = None
         elif 'pts' in self.suffix:
             arr = read_pts_file(filepath)[:, :3]
+            rotation = None
         else:
             raise NotImplementedError("This data type is not implemented at the moment.")
         partial_arr = None
@@ -248,7 +260,10 @@ def __getitem__(self, index):
                   'path': filepath,
                   'label_name': cls,
                   'corner_label': np.array(corner),
-                  'corner_label_name': corner_to_label_name[corner]}
+                  'corner_label_name': corner_to_label_name[corner]
+                  }
+        if rotation is not None:
+            sample['rotation'] = rotation
 
         if self.transform is not None:
             sample = self.transform(sample)