Merge pull request #226 from KMarshallX/pre-release-2_0-0910

KMarshallX · web-flow · commit c7a112d7f163 · 2025-09-10T15:50:32.000+10:00
Pre release 2 0 0910
diff --git a/config/adapt_config.py b/config/adapt_config.py
@@ -36,7 +36,7 @@
 adapt_parser.add_argument('--patch_size', nargs=3, type=int, default=(64, 64, 64),
                         help='Expected size for training (x y z)')
 # optimizer type, available: [sgd, adam]
-adapt_parser.add_argument('--optimizer', type=str, default="adam", help='available: [sgd, adam]')
+adapt_parser.add_argument('--optimizer', type=str, default="adam", help='available: [sgd, adam, adamw]')
 # loss metric type, available: [bce, dice, tver]
 adapt_parser.add_argument('--loss_metric', type=str, default="tver", help="available: [bce, dice, tver]")
 
diff --git a/config/angiboost_config.py b/config/angiboost_config.py
@@ -31,7 +31,7 @@
 # expected size of the training patch
 angiboost_parser.add_argument('--osz', nargs=3, type=int, default=(64, 64, 64),
                         help='Expected size of the training patch (x y z)')
-angiboost_parser.add_argument('--optimizer', type=str, default="adam", help="available: [sgd, adam]")
+angiboost_parser.add_argument('--optimizer', type=str, default="adam", help="available: [sgd, adam, adamw]")
 angiboost_parser.add_argument('--loss_metric', type=str, default="tver", help="available: [bce, dice, tver]")
 
 # Optimizer tuning
diff --git a/config/boost_config.py b/config/boost_config.py
@@ -30,7 +30,7 @@
 # expected size of the training patch
 boost_parser.add_argument('--patch_size', nargs=3, type=int, default=(64, 64, 64),
                         help='Expected size of the training patch (x y z)')
-boost_parser.add_argument('--optimizer', type=str, default="adam", help="available: [sgd, adam]")
+boost_parser.add_argument('--optimizer', type=str, default="adam", help="available: [sgd, adam, adamw]")
 boost_parser.add_argument('--loss_metric', type=str, default="tver", help="available: [bce, dice, tver]")
 
 # Optimizer tuning
diff --git a/config/train_config.py b/config/train_config.py
@@ -29,7 +29,7 @@
 # expected patch size for training
 train_parser.add_argument('--output_size', nargs=3, type=int, default=(64, 64, 64), help='Expected patch size for training (x y z)')
 # optimizer type, available: [sgd, adam]
-train_parser.add_argument('--optimizer', type=str, default="adam", help='available: [sgd, adam]')
+train_parser.add_argument('--optimizer', type=str, default="adam", help='available: [sgd, adam, adamw]')
 # loss metric type, available: [bce, dice, tver]
 train_parser.add_argument('--loss_metric', type=str, default="tver", help="available: [bce, dice, tver]")
 
diff --git a/environment.yml b/environment.yml
@@ -10,6 +10,12 @@ dependencies:
   - pytorch-cuda=12.4
   - python=3.10
   - pip
+  - jupyter          
+  - notebook         
+  - jupyterlab       
+  - ipython          
+  - ipykernel        
+  - ipywidgets  
   - pip:
     - numpy==1.26.4
     - scipy==1.15.2
@@ -23,3 +29,4 @@ dependencies:
     - antspyx==0.4.2
     - connected-components-3d==3.13.0
     - osfclient==0.0.5
+    - jupyterlmod==5.3.0
diff --git a/library/aug_utils.py b/library/aug_utils.py
@@ -375,24 +375,29 @@ def __init__(self, mode: str = 'spatial'):
 
         self.mode = mode
 
-    def _blur(self, std: float = 0.85) -> tio.RandomBlur:
-        return tio.RandomBlur(std=std)
+    def _blur(self, p: float = 1, std: float = 0.85) -> tio.RandomBlur:
+        return tio.RandomBlur(std=std, p=p)
 
-    def _bias(self, coefficients: float = 0.15, order: int = 3) -> tio.RandomBiasField:
-        return tio.RandomBiasField(coefficients=coefficients, order=order)
+    def _bias(self, p: float = 1, coefficients: float = 0.15, order: int = 3) -> tio.RandomBiasField:
+        return tio.RandomBiasField(coefficients=coefficients, order=order, p=p)
 
-    def _noise(self, mean: float = 0, std: float = 0.008) -> tio.RandomNoise:
-        return tio.RandomNoise(mean=mean, std=std)
+    def _noise(self, p: float = 1, mean: float = 0, std: float = 0.008) -> tio.RandomNoise:
+        return tio.RandomNoise(mean=mean, std=std, p=p)
 
-    def _flip(self, axes: Union[Tuple[int, ...], int], probability: float = 1.0) -> tio.RandomFlip:
-        return tio.RandomFlip(axes=axes, flip_probability=probability)
+    def _flip(self, p: float = 1, axes: Union[Tuple[int, ...], int] = (0, 1, 2), probability: float = 1.0) -> tio.RandomFlip:
+        # Randomly choose a single axis from the available axes
+        if isinstance(axes, tuple):
+            random_axis = int(np.random.choice(axes))
+        else:
+            random_axis = axes
+        return tio.RandomFlip(axes=random_axis, flip_probability=probability, p=p)
 
-    def _elastic_deform(self, num_control_points: int = 9, 
+    def _elastic_deform(self, p: float = 1, num_control_points: int = 9, 
                         max_displacement: int = 7, 
                         locked_borders: int = 2) -> tio.RandomElasticDeformation:
         return tio.RandomElasticDeformation(num_control_points=num_control_points, 
                                             max_displacement=max_displacement, 
-                                            locked_borders=locked_borders)
+                                            locked_borders=locked_borders, p=p)
 
     def __call__(self, image_batch: torch.Tensor, seg_batch: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
         """
@@ -416,36 +421,42 @@ def __call__(self, image_batch: torch.Tensor, seg_batch: torch.Tensor) -> Tuple[
                 self._blur(),
                 self._bias(),
                 self._noise(),
-                self._flip(axes=(0, 1)),
-                self._elastic_deform(),
-            ])
-        elif self.mode == 'random':
-            transforms = tio.OneOf([
-                self._blur(),
-                self._bias(),
-                self._noise(),
-                self._flip(axes=(0, 1)),
+                self._flip(axes=(0, 1, 2)),
                 self._elastic_deform()
             ])
+        elif self.mode == 'random':
+            transforms = tio.OneOf({
+                self._blur() : 0.1,
+                self._bias() : 0.1,
+                self._noise() : 0.1,
+                self._flip(axes=(0, 1, 2)) : 0.35,
+                self._elastic_deform() : 0.35
+            })
         elif self.mode == 'spatial':
             transforms = tio.Compose([
-                self._flip(axes=(0, 1)),
+                self._flip(axes=(0, 1, 2)),
                 self._elastic_deform()
             ])
         elif self.mode == 'intensity':
-            transforms = tio.Compose([
+            transforms = tio.OneOf([
                 self._blur(),
                 self._bias(),
                 self._noise()
             ])
         elif self.mode == 'off':
             # No augmentation, return original subject
-            return subject_batch['image'].data, subject_batch['label'].data # type: ignore
+            return subject_batch['image'].data.unsqueeze(1), subject_batch['label'].data.unsqueeze(1) # type: ignore
         else:
             raise ValueError(f"Unsupported mode '{self.mode}' for TorchIO augmentations")
         
         # Apply the transform to the subject batch
         transformed_subject = transforms(subject_batch)
+        
+        # # Track which transform was applied (for OneOf modes)
+        # # TESTING BLOCK, UNCOMMENT THIS FOR DEBUGGING
+        # if self.mode in ['random', 'intensity']:
+        #     applied_transforms = [str(transform) for transform in transformed_subject.history]
+        #     print(f"Applied transforms in {self.mode} mode: {applied_transforms[-1] if applied_transforms else 'None'}")
 
         # Extract image and label tensors
         image_tensor = transformed_subject['image'].data.unsqueeze(1) # type: ignore
diff --git a/library/data_loaders.py b/library/data_loaders.py
@@ -114,9 +114,9 @@ def _zooming(self, img_crop: np.ndarray, seg_crop: np.ndarray) -> Tuple[np.ndarr
         """
         if img_crop.shape != self.patch_size:
             zoom_factors = tuple(float(out_dim) / float(crop_dim) for out_dim, crop_dim in zip(self.patch_size, img_crop.shape))
-            img_crop = scind.zoom(img_crop, zoom_factors, order=3, mode='nearest')
+            img_crop = scind.zoom(img_crop, zoom_factors, order=0, mode='nearest')
             seg_crop = scind.zoom(seg_crop, zoom_factors, order=0, mode='nearest')
-        return img_crop, seg_crop
+        return img_crop.astype(np.float32), seg_crop.astype(np.int8)
 
     def __repr__(self) -> str:
         return (f"SingleChannelLoader(\n"
@@ -164,7 +164,7 @@ def __iter__(self):
                 img_batch[self.batch_multiplier] = large_img_crop
                 seg_batch[self.batch_multiplier] = large_seg_crop  
                 # Yield the cropped and resized patches as a batch of pytorch tensors
-                yield torch.from_numpy(img_batch).float(), torch.from_numpy(seg_batch).long()
+                yield torch.from_numpy(img_batch).float(), torch.from_numpy(seg_batch).ceil().int()
 
             except Exception as e:
                 logger.error(f"Error generating patch {i}: {e}")
diff --git a/library/loss_func.py b/library/loss_func.py
@@ -150,7 +150,8 @@ def choose_optimizer(optim_name: str, model_params: Any, lr: float) -> Any:
     """
     optimizer_registry = {
         'sgd': lambda: torch.optim.SGD(model_params, lr),
-        'adam': lambda: torch.optim.Adam(model_params, lr)
+        'adam': lambda: torch.optim.Adam(model_params, lr),
+        'adamw': lambda: torch.optim.AdamW(model_params, lr)
     }
     
     if optim_name not in optimizer_registry:
diff --git a/library/module_utils.py b/library/module_utils.py
@@ -399,7 +399,7 @@ def process_single_image(
             resized_image = self._resize_image(image_array, target_size)
             
             # Standardize image
-            standardized_image = standardiser(resized_image)
+            standardized_image = normaliser(resized_image)
             
             # Create patches
             patches = patchify(standardized_image, (64, 64, 64), 64)
@@ -431,7 +431,7 @@ def predict_all_images(
         threshold: float,
         connect_threshold: int,
         save_mip: bool = False,
-        save_probability: bool = False
+        save_probability: bool = True
     ) -> None:
         """
         Process all images in the input directory.
diff --git a/library/train_utils.py b/library/train_utils.py
@@ -232,7 +232,7 @@ def _train_epoch(
             total_lr += optimizer.param_groups[0]['lr']
             num_batches += 1
 
-        logger.info(f"Data loading time: {data_loading_time:.2f}s, Model training time: {model_training_time:.2f}s")  # TEST
+        logger.info(f"\nData loading time: {data_loading_time:.2f}s, Model training time: {model_training_time:.2f}s")  # TEST
         return total_loss / num_batches, total_lr / num_batches
 
     def train_model(
@@ -269,7 +269,7 @@ def train_model(
             # Log progress
             tqdm.write(
                 f'Epoch [{epoch+1}/{self.num_epochs}], '
-                f'Loss: {avg_loss:.4f}, LR: {avg_lr:.8f}'
+                f'Loss: {avg_loss:.8f}, LR: {avg_lr:.8f}'
             )
         
         # Save model
diff --git a/train_debug.py b/train_debug.py
@@ -0,0 +1,103 @@
+#!/usr/bin/env python3
+
+"""
+Training the chosen model
+
+Editor: Marshall Xu
+Last Edited: 01/08/2025
+"""
+
+import logging
+import config.train_config as train_config
+from library import preprocess_procedure, Trainer
+
+# Set up logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+
+def run_training():
+    """Execute the complete training pipeline."""
+    config = train_config.args
+
+    # Determine input paths based on preprocessing mode
+    input_path = config.image_path if config.prep_mode == 4 else config.preprocessed_path
+
+    logger.info("Training session will start shortly..")
+    logger.info("Parameters Info:")
+    logger.info("*" * 61)
+    logger.info(f"Input image path: {config.image_path}, Segmentation path: {config.label_path}, Prep_mode: {config.prep_mode}")
+    logger.info(f"Epoch number: {config.epochs}, Learning rate: {config.learning_rate}")
+
+    # Preprocessing procedure
+    preprocess_procedure(config.image_path, config.preprocessed_path, config.prep_mode)
+
+    # Initialize and run training
+    trainer = Trainer(
+        loss_name=config.loss_metric, model_name=config.model,
+        input_channels=config.input_channel, output_channels=config.output_channel, filter_count=config.filters,
+        optimizer_name=config.optimizer, learning_rate=config.learning_rate, optimizer_gamma=config.optimizer_gamma, num_epochs=config.epochs,
+        batch_multiplier=config.batch_multiplier, patch_size=tuple(config.output_size), augmentation_mode=config.augmentation_mode,
+        crop_low_thresh=config.crop_low_thresh
+    )
+
+    # Training loop
+    trainer.train(input_path, config.label_path, config.output_model)
+
+def run_debugging(input_path, label_path, output_model, epochs=1000, learning_rate=1e-3, augmentation_mode='off', preprocessed_path=None):
+    """Execute the complete training pipeline."""
+    config = train_config.args
+
+    # Determine input paths based on preprocessing mode
+    input_path = input_path if config.prep_mode == 4 else preprocessed_path
+
+    logger.info("Training session will start shortly..")
+    logger.info("Parameters Info:")
+    logger.info("*" * 61)
+    logger.info(f"Input image path: {input_path}, Segmentation path: {label_path}, Prep_mode: {config.prep_mode}")
+    logger.info(f"Augmentation mode: {augmentation_mode}")
+    logger.info(f"Epoch number: {epochs}, Learning rate: {learning_rate}")
+
+    # Preprocessing procedure
+    preprocess_procedure(input_path, preprocessed_path, config.prep_mode) #type: ignore
+
+    # Initialize and run training
+    trainer = Trainer(
+        loss_name=config.loss_metric, model_name=config.model,
+        input_channels=config.input_channel, output_channels=config.output_channel, filter_count=config.filters,
+        optimizer_name=config.optimizer, learning_rate=learning_rate, optimizer_gamma=config.optimizer_gamma, num_epochs=epochs,
+        batch_multiplier=config.batch_multiplier, patch_size=tuple(config.output_size), augmentation_mode=augmentation_mode,
+        crop_low_thresh=config.crop_low_thresh
+    )
+
+    # Training loop
+    trainer.train(input_path, label_path, output_model) #type: ignore
+
+
+
+if __name__ == "__main__":
+    input_path = "/scratch/project/simvascmri/data/single_data/prep_160um_skull/"
+    label_path = "/scratch/project/simvascmri/data/single_data/tta_seg_160um/"
+    output_model = "./result_models/DEBUG_NewVB_train_aug_off_ep1k_0826"
+
+    run_debugging(input_path, label_path, output_model, epochs=1000, learning_rate=1e-3, augmentation_mode='off')
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+

Original file line number	Diff line number	Diff line change
`@@ -150,7 +150,8 @@ def choose_optimizer(optim_name: str, model_params: Any, lr: float) -> Any:`
`150`	`150`	`"""`
`151`	`151`	`optimizer_registry = {`
`152`	`152`	`'sgd': lambda: torch.optim.SGD(model_params, lr),`
`153`		`- 'adam': lambda: torch.optim.Adam(model_params, lr)`
	`153`	`+ 'adam': lambda: torch.optim.Adam(model_params, lr),`
	`154`	`+ 'adamw': lambda: torch.optim.AdamW(model_params, lr)`
`154`	`155`	`}`
`155`	`156`
`156`	`157`	`if optim_name not in optimizer_registry:`