V3

src/parcellation.py

@@ -94,10 +94,8 @@ def run_inference(input_file, only_face_cropping, only_skull_stripping):
 
     # Load the pre-trained models from the fixed "model/" folder
     print("Loading models...")
-    cnet, ssnet, pnet_c, pnet_s, pnet_a, hnet_c, hnet_a = load_model("model/", device=device)
+    cnet, ssnet, pnet, hnet = load_model("model/", device=device)
     print("Models loaded successfully.")
-    # cnet, ssnet, pnet_a, hnet_c, hnet_a = load_model("model/", device=device)
-
     # --- Processing Flow (based on the original parcellation.py) ---
     # 1. Load the input image, convert to canonical orientation, and remove extra dimensions
     print("Loading and preprocessing the input image...")
@@ -115,7 +113,7 @@ def run_inference(input_file, only_face_cropping, only_skull_stripping):
 
     # 3. Cropping
     print("Cropping the input image...")
-    cropped, out_filename = cropping(opt.o, basename, odata, data, cnet, device)
+    cropped, shift, out_filename = cropping(opt.o, basename, odata, data, cnet, device)
     print("Cropping completed.")
     if only_face_cropping:
         pass
@@ -123,18 +121,18 @@ def run_inference(input_file, only_face_cropping, only_skull_stripping):
     else:
         # 4. Skull stripping
         print("Performing skull stripping...")
-        stripped, shift, out_filename = stripping(opt.o, basename, cropped, odata, data, ssnet, device)
+        stripped, out_filename = stripping(opt.o, basename, cropped, odata, data, ssnet, shift, device)
         print("Skull stripping completed.")
         if only_skull_stripping:
             pass
         else:
             # 5. Parcellation
             print("Starting parcellation...")
-            parcellated = parcellation(stripped, pnet_c, pnet_s, pnet_a, device)
+            parcellated = parcellation(stripped, pnet, device)
             print("Parcellation completed.")
             # 6. Separate into hemispheres
             print("Separating hemispheres...")
-            separated = hemisphere(stripped, hnet_c, hnet_a, device)
+            separated = hemisphere(stripped, hnet, device)
             print("Hemispheres separated.")
             # 7. Postprocessing
             print("Postprocessing the parcellated data...")

src/utils/cropping.py

@@ -1,31 +1,44 @@
 import numpy as np
 import torch
 from scipy.ndimage import binary_closing
+from scipy import ndimage
 
 from utils.functions import normalize, reimburse_conform
 
 
 def crop(voxel, model, device):
     """
-    Crops the given voxel data using the provided model and device.
+    Apply a neural network-based cropping operation on 3D voxel data.
+
+    This function slides a 3-slice window across the input volume along the first axis
+    and predicts a binary mask for each slice using the given model. The outputs are then
+    aggregated into a full 3D prediction volume.
 
     Args:
-        voxel (numpy.ndarray): The input voxel data to be cropped, expected to be of shape (N, 256, 256).
-        model (torch.nn.Module): The PyTorch model used for cropping.
-        device (torch.device): The device (CPU or GPU) on which the computation will be performed.
+        voxel (numpy.ndarray): Input 3D array of shape (N, 256, 256). The first dimension
+            corresponds to the slice index (typically coronal or sagittal).
+        model (torch.nn.Module): The trained PyTorch model that predicts binary masks
+            for each input slice triplet.
+        device (torch.device): The device (CPU, CUDA, or MPS) on which inference will run.
 
     Returns:
-        torch.Tensor: The cropped output tensor of shape (256, 256, 256).
+        torch.Tensor: The predicted 3D binary mask of shape (256, 256, 256).
     """
+    # Pad the input volume by one slice at each end to allow 3-slice context
+    voxel = np.pad(voxel, [(1, 1), (0, 0), (0, 0)], "constant", constant_values=voxel.min())
     model.eval()
+
     with torch.inference_mode():
-        output = torch.zeros(256, 256, 256).to(device)
-        for i, v in enumerate(voxel):
-            image = v.reshape(1, 1, 256, 256)
-            image = torch.tensor(image).to(device)
-            x_out = torch.sigmoid(model(image)).detach()
-            output[i] = x_out
-        return output.reshape(256, 256, 256)
+        box = torch.zeros(256, 256, 256)
+
+        # Iterate through each target slice and predict using a 3-slice input context
+        for i in range(1, 257):
+            image = np.stack([voxel[i - 1], voxel[i], voxel[i + 1]])
+            image = torch.tensor(image.reshape(1, 3, 256, 256)).to(device)
+            x_out = torch.sigmoid(model(image)).detach().cpu()
+            box[i - 1] = x_out
+
+        return box.reshape(256, 256, 256)
 
 
 def closing(voxel):
@@ -59,7 +72,7 @@ def cropping(output_dir, basename, odata, data, cnet, device):
         numpy.ndarray: The cropped medical imaging data.
     """
     voxel = data.get_fdata().astype("float32")
-    voxel = normalize(voxel)
+    voxel = normalize(voxel, "cropping")
 
     coronal = voxel.transpose(1, 2, 0)
     sagittal = voxel
@@ -72,4 +85,18 @@ def cropping(output_dir, basename, odata, data, cnet, device):
 
     out_filename = reimburse_conform(output_dir, basename, "cropped", odata, data, out_e)
 
-    return cropped, out_filename
+    # Compute center of mass for the masked brain
+    x, y, z = map(int, ndimage.center_of_mass(out_e))
+
+    # Compute shifts required to center the brain
+    xd = 128 - x
+    yd = 120 - y
+    zd = 128 - z
+
+    # Translate (roll) the image to center the brain region
+    cropped = np.roll(cropped, (xd, yd, zd), axis=(0, 1, 2))
+
+    # Crop out boundary padding to reduce size and focus on the centered brain
+    cropped = cropped[16:-16, 16:-16, 16:-16]
+
+    return cropped, (xd, yd, zd), out_filename

src/utils/functions.py

@@ -5,9 +5,13 @@ import numpy as np
 from nibabel import processing
 
 
-def normalize(voxel):
+def normalize(voxel, mode):
     nonzero = voxel[voxel > 0]
-    voxel = np.clip(voxel, 0, np.mean(nonzero) + np.std(nonzero) * 2)
+    if mode in ["cropping", "stripping"]:
+        clip = 2
+    elif mode in ["parcellation", "hemisphere"]:
+        clip = 3
+    voxel = np.clip(voxel, 0, np.mean(nonzero) + np.std(nonzero) * clip)
     voxel = (voxel - np.min(voxel)) / (np.max(voxel) - np.min(voxel))
     voxel = (voxel * 2) - 1
     return voxel.astype("float32")

src/utils/hemisphere.py

@@ -1,92 +1,104 @@
 import torch
 from scipy.ndimage import binary_dilation
-
+import numpy as np
 from utils.functions import normalize
 
 
-def separate(voxel, model, device, mode):
+def separate(voxel, model, device):
     """
-    Separates the voxel data based on the specified mode and processes it using the given model.
+    Perform slice-wise inference using a hemisphere separation model.
+
+    This function runs a 2.5D neural network across slices of a 3D input volume.
+    Each slice is processed in the context of its immediate neighbors (previous
+    and next slices) to improve spatial coherence. The model outputs a
+    three-class probability map distinguishing background, left hemisphere,
+    and right hemisphere regions.
 
     Args:
-        voxel (list or numpy.ndarray): The input voxel data to be processed.
-        model (torch.nn.Module): The neural network model used for processing the voxel data.
-        device (torch.device): The device (CPU or GPU) on which the model and data are loaded.
-        mode (str): The mode of separation, either 'c' for coronal or 'a' for axial.
+        voxel (numpy.ndarray): Input voxel data of shape (N, 224, 224).
+        model (torch.nn.Module): Trained hemisphere segmentation model (U-Net architecture).
+        device (torch.device): Computational device (CPU, CUDA, or MPS).
 
     Returns:
-        torch.Tensor: The processed output tensor with shape (stack[0], 3, stack[1], stack[2]).
+        torch.Tensor: A tensor of shape (224, 3, 224, 224) containing softmax
+        probabilities for each class at every voxel.
     """
-    if mode == "c":
-        # Set the stack dimensions for coronal mode
-        stack = (224, 192, 192)
-    elif mode == "a":
-        # Set the stack dimensions for axial mode
-        stack = (192, 224, 192)
-
-    # Set the model to evaluation mode
     model.eval()
 
-    # Disable gradient calculation for inference
+    # Pad the volume by one slice on both ends to provide full 3-slice context
+    voxel = np.pad(voxel, [(1, 1), (0, 0), (0, 0)], "constant", constant_values=voxel.min())
+
     with torch.inference_mode():
-        # Initialize an output tensor with the specified stack dimensions
-        output = torch.zeros(stack[0], 3, stack[1], stack[2]).to(device)
+        # Output tensor for storing model predictions (class probabilities)
+        box = torch.zeros(224, 3, 224, 224)
 
-        # Iterate over each slice in the voxel data
-        for i, v in enumerate(voxel):
-            # Reshape the slice and convert it to a tensor
-            image = torch.tensor(v.reshape(1, 1, stack[1], stack[2]))
-            # Move the tensor to the specified device
-            image = image.to(device)
-            # Perform a forward pass through the model and apply softmax
-            x_out = torch.softmax(model(image), 1).detach()
-            # Store the output in the corresponding slice of the output tensor
-            output[i] = x_out
+        # Iterate slice-by-slice along the first axis
+        for i in range(1, 225):
+            image = np.stack([voxel[i - 1], voxel[i], voxel[i + 1]])
+            image = torch.tensor(image.reshape(1, 3, 224, 224)).to(device)
 
-        # Return the processed output tensor
-        return output
+            # Model inference with softmax normalization across classes
+            x_out = torch.softmax(model(image), dim=1).detach().cpu()
+            box[i - 1] = x_out
 
+        # Return complete 3D probability map
+        return box.reshape(224, 3, 224, 224)
 
-def hemisphere(voxel, hnet_c, hnet_a, device):
+
+def hemisphere(voxel, hnet, device):
     """
-    Processes a voxel image to separate and dilate hemispheres using neural networks.
+    Perform hemisphere separation on a brain MRI volume using a deep learning model.
+
+    The function predicts left and right hemisphere regions from a normalized
+    3D MRI volume using multi-view inference (coronal and transverse planes).
+    Predictions from both orientations are fused to improve robustness. The final
+    label map is post-processed using binary dilation to smooth and expand hemisphere
+    boundaries, ensuring anatomical continuity.
 
     Args:
-        voxel (torch.Tensor): The input voxel image tensor.
-        hnet_c (torch.nn.Module): The neural network model for coronal separation.
-        hnet_a (torch.nn.Module): The neural network model for transverse separation.
-        device (torch.device): The device to run the neural networks on (e.g., 'cpu' or 'cuda').
+        voxel (numpy.ndarray): Input 3D brain volume to be separated into hemispheres.
+        hnet (torch.nn.Module): Trained hemisphere segmentation model.
+        device (torch.device): Target device for computation (e.g., 'cuda', 'cpu').
 
     Returns:
-        numpy.ndarray: The processed and dilated mask of the hemispheres.
+        numpy.ndarray: A 3D integer array representing the hemisphere mask:
+            - 0: Background
+            - 1: Left hemisphere
+            - 2: Right hemisphere
     """
-    # Normalize the voxel data
-    voxel = normalize(voxel)
+    # Normalize voxel intensities for inference
+    voxel = normalize(voxel, "hemisphere")
 
-    # Transpose the voxel data for coronal and transverse views
+    # Prepare different anatomical orientations for inference
     coronal = voxel.transpose(1, 2, 0)
     transverse = voxel.transpose(2, 1, 0)
 
-    # Separate the coronal and transverse views using the respective models
-    out_c = separate(coronal, hnet_c, device, "c").permute(1, 3, 0, 2)
-    out_a = separate(transverse, hnet_a, device, "a").permute(1, 3, 2, 0)
+    # Perform inference for both coronal and transverse orientations
+    out_c = separate(coronal, hnet, device).permute(1, 3, 0, 2)
+    out_a = separate(transverse, hnet, device).permute(1, 3, 2, 0)
 
-    # Combine the outputs from both views
+    # Fuse both outputs by summing class probabilities
     out_e = out_c + out_a
 
-    # Get the final output by taking the argmax along the first dimension
-    out_e = torch.argmax(out_e, 0).cpu().numpy()
+    # Determine final class labels (0, 1, or 2) by selecting the most probable class
+    out_e = torch.argmax(out_e, dim=0).cpu().numpy()
 
-    # Clear the CUDA cache
+    # Release any residual GPU memory
     torch.cuda.empty_cache()
 
-    # Perform binary dilation on the mask for class 1
-    dilated_mask_1 = binary_dilation(out_e == 1, iterations=5).astype("int16")
+    # --------------------------
+    # Post-processing step: binary dilation
+    # --------------------------
+
+    # First, dilate the left hemisphere (class 1)
+    dilated_mask_1 = binary_dilation(out_e == 1, iterations=1).astype("int16")
+    # Preserve right hemisphere voxels from the original prediction
     dilated_mask_1[out_e == 2] = 2
 
-    # Perform binary dilation on the mask for class 2
-    dilated_mask_2 = binary_dilation(dilated_mask_1 == 2, iterations=5).astype("int16") * 2
+    # Then, dilate the right hemisphere (class 2) symmetrically
+    dilated_mask_2 = binary_dilation(dilated_mask_1 == 2, iterations=1).astype("int16") * 2
+    # Restore left hemisphere voxels to prevent overwriting
     dilated_mask_2[dilated_mask_1 == 1] = 1
 
-    # Return the final dilated mask
+    # Return the final dilated and fused hemisphere mask
     return dilated_mask_2

src/utils/load_model.py

@@ -8,70 +8,65 @@ from utils.network import UNet
 
 def load_model(model_dir, device):
     """
-    This function loads multiple pre-trained models and sets them to evaluation mode.
-    The models loaded are:
-    1. CNet: A U-Net model for some specific task.
-    2. SSNet: Another U-Net model for a different task.
-    3. PNet coronal: A U-Net model for coronal plane predictions.
-    4. PNet sagittal: A U-Net model for sagittal plane predictions.
-    5. PNet axial: A U-Net model for axial plane predictions.
-    6. HNet coronal: A U-Net model for coronal plane predictions with different input/output channels.
-    7. HNet axial: A U-Net model for axial plane predictions with different input/output channels.
+    Load and initialize the pretrained neural network models required for the OpenMAP-T1 pipeline.
 
-    Parameters:
-    opt (object): An options object containing model paths.
-    device (torch.device): The device on which to load the models (CPU or GPU).
+    This function loads four U-Net–based models from the specified pretrained model directory.
+    Each model is moved to the target device (CPU, CUDA, or MPS) and set to evaluation mode.
+
+    Models loaded:
+        1. **CNet (Cropping Network)** — Performs face cropping and brain localization.
+        2. **SSNet (Skull Stripping Network)** — Removes non-brain tissues from MRI scans.
+        3. **PNet (Parcellation Network)** — Predicts fine-grained anatomical labels across 142 regions.
+        4. **HNet (Hemisphere Network)** — Segments the brain into hemispheric masks (left/right/other).
+
+    Args:
+        opt (argparse.Namespace): Parsed command-line arguments containing the pretrained model directory path (`opt.m`).
+        device (torch.device): Target device on which to load models (e.g., `torch.device('cuda')`).
 
     Returns:
-    tuple: A tuple containing all the loaded models.
+        tuple:
+            A tuple containing four initialized and evaluation-ready models:
+            (cnet, ssnet, pnet, hnet).
     """
-    # Unzip the Model.zip file
     model_zip_path = os.path.join(model_dir, "model.zip")
     with zipfile.ZipFile(model_zip_path, "r") as zip_ref:
         zip_ref.extractall(model_dir)
-
-    # Load CNet model
-    cnet = UNet(1, 1)
-    cnet.load_state_dict(torch.load(os.path.join(model_dir, "CNet", "CNet.pth"), weights_only=True))
+    # --------------------------
+    # Load CNet (Cropping Network)
+    # --------------------------
+    # Input: 3-channel (neighboring slices), Output: 1-channel binary mask
+    cnet = UNet(3, 1)
+    print(os.path.join(model_dir, "model", "CNet", "CNet.pth"))
+    cnet.load_state_dict(torch.load(os.path.join(model_dir, "model", "CNet", "CNet.pth"), weights_only=True))
     cnet.to(device)
     cnet.eval()
 
-    # Load SSNet model
-    ssnet = UNet(1, 1)
-    ssnet.load_state_dict(torch.load(os.path.join(model_dir, "SSNet", "SSNet.pth"), weights_only=True))
+    # ------------------------------
+    # Load SSNet (Skull Stripping Network)
+    # ------------------------------
+    # Input: 3-channel (neighboring slices), Output: 1-channel brain mask
+    ssnet = UNet(3, 1)
+    ssnet.load_state_dict(torch.load(os.path.join(model_dir, "model", "SSNet", "SSNet.pth"), weights_only=True))
     ssnet.to(device)
     ssnet.eval()
 
-    # Load PNet coronal model
-    pnet_c = UNet(3, 142)
-    pnet_c.load_state_dict(torch.load(os.path.join(model_dir, "PNet", "coronal.pth"), weights_only=True))
-    pnet_c.to(device)
-    pnet_c.eval()
-
-    # Load PNet sagittal model
-    pnet_s = UNet(3, 142)
-    pnet_s.load_state_dict(torch.load(os.path.join(model_dir, "PNet", "sagittal.pth"), weights_only=True))
-    pnet_s.to(device)
-    pnet_s.eval()
-
-    # Load PNet axial model
-    pnet_a = UNet(3, 142)
-    pnet_a.load_state_dict(torch.load(os.path.join(model_dir, "PNet", "axial.pth"), weights_only=True))
-    pnet_a.to(device)
-    pnet_a.eval()
-
-    # Load HNet coronal model
-    hnet_c = UNet(1, 3)
-    hnet_c.load_state_dict(torch.load(os.path.join(model_dir, "HNet", "coronal.pth"), weights_only=True))
-    hnet_c.to(device)
-    hnet_c.eval()
+    # -----------------------------
+    # Load PNet (Parcellation Network)
+    # -----------------------------
+    # Input: 4 channels (multi-modal or augmented context), Output: 142 anatomical regions
+    pnet = UNet(4, 142)
+    pnet.load_state_dict(torch.load(os.path.join(model_dir, "model", "PNet", "PNet.pth"), weights_only=True))
+    pnet.to(device)
+    pnet.eval()
 
-    # Load HNet axial model
-    hnet_a = UNet(1, 3)
-    hnet_a.load_state_dict(torch.load(os.path.join(model_dir, "HNet", "axial.pth"), weights_only=True))
-    hnet_a.to(device)
-    hnet_a.eval()
+    # -----------------------------
+    # Load HNet (Hemisphere Network)
+    # -----------------------------
+    # Input: 3 channels, Output: 3-class hemisphere mask (left, right, background)
+    hnet = UNet(3, 3)
+    hnet.load_state_dict(torch.load(os.path.join(model_dir, "model", "HNet", "HNet.pth"), weights_only=True))
+    hnet.to(device)
+    hnet.eval()
 
-    # Return all loaded models
-    return cnet, ssnet, pnet_c, pnet_s, pnet_a, hnet_c, hnet_a
-    # return cnet, ssnet, pnet_a, hnet_c, hnet_a
+    # Return all loaded, device-initialized, and evaluation-ready models
+    return cnet, ssnet, pnet, hnet

src/utils/parcellation.py

@@ -4,102 +4,133 @@ import torch
 from utils.functions import normalize
 
 
-def parcellate(voxel, model, device, mode):
+def parcellate(
+    voxel: np.ndarray,
+    model: torch.nn.Module,
+    device: torch.device,
+    mode: str,
+    n_classes: int = 142,
+) -> torch.Tensor:
     """
-    Parcellates a given voxel volume using a specified model and mode.
+    Perform 2.5D neural network inference for brain parcellation along a specific anatomical plane.
+
+    The function processes a 3D volume slice by slice using a 3-slice context window (previous,
+    current, next). An additional constant-valued fourth channel encodes the orientation mode
+    (Axial, Coronal, or Sagittal), allowing the network to distinguish the processing plane.
 
     Args:
-        voxel (numpy.ndarray): The input voxel volume to be parcellated.
-        model (torch.nn.Module): The neural network model used for parcellation.
-        device (torch.device): The device (CPU or GPU) on which the model is run.
-        mode (str): The mode of parcellation. Can be 'c', 's', or 'a', which determines the stack dimensions.
+        voxel (numpy.ndarray): 3D voxel data of shape (N, 224, 224), representing a single anatomical view.
+        model (torch.nn.Module): The trained PyTorch parcellation model.
+        device (torch.device): Device for inference (CPU, CUDA, or MPS).
+        mode (str): The anatomical plane used for inference. Must be one of {'Axial', 'Coronal', 'Sagittal'}.
+        n_classes (int, optional): Number of output anatomical labels. Defaults to 142.
 
     Returns:
-        torch.Tensor: The parcellated voxel volume.
+        torch.Tensor: A tensor of shape (224, n_classes, 224, 224) containing softmax probabilities
+        for each class at each voxel position.
     """
-    if mode == "c":
-        stack = (224, 192, 192)
-    elif mode == "s":
-        stack = (192, 224, 192)
-    elif mode == "a":
-        stack = (192, 224, 192)
-
-    # Set the model to evaluation mode
     model.eval()
-
-    # Pad the voxel volume to handle edge cases
-    voxel = np.pad(voxel, [(1, 1), (0, 0), (0, 0)], "constant", constant_values=voxel.min())
-
-    # Disable gradient calculation for inference
+    voxel = voxel.astype(np.float32)
+
+    # Set the constant value for the 4th channel to encode plane orientation
+    if mode == "Axial":
+        section_value = 1.0
+    elif mode == "Coronal":
+        section_value = -1.0
+    elif mode == "Sagittal":
+        section_value = 0.0
+    else:
+        raise ValueError("mode must be one of {'Axial','Coronal','Sagittal'}")
+
+    # Pad one slice on both ends to safely allow 3-slice context
+    voxel_pad = np.pad(
+        voxel,
+        [(1, 1), (0, 0), (0, 0)],
+        mode="constant",
+        constant_values=float(voxel.min()),
+    )
+
+    # Initialize a container for the network outputs (CPU for accumulation)
+    box = torch.empty((224, n_classes, 224, 224), dtype=torch.float32, device="cpu")
+
+    # Inference loop: iterate over slices and feed triplets to the model
     with torch.inference_mode():
-        # Initialize an empty tensor to store the parcellation results
-        box = torch.zeros(stack[0], 142, stack[1], stack[2])
+        for i in range(1, 225):
+            prev_ = voxel_pad[i - 1]
+            curr_ = voxel_pad[i]
+            next_ = voxel_pad[i + 1]
 
-        # Iterate over each slice in the stack dimension
-        for i in range(1, stack[0] + 1):
-            # Stack three consecutive slices to form the input image
-            image = np.stack([voxel[i - 1], voxel[i], voxel[i + 1]])
-            image = torch.tensor(image.reshape(1, 3, stack[1], stack[2]))
-            image = image.to(device)
+            # Build 4-channel input (3 context slices + orientation encoding)
+            four_ch = np.empty((4, 224, 224), dtype=np.float32)
+            four_ch[0] = prev_
+            four_ch[1] = curr_
+            four_ch[2] = next_
+            four_ch[3].fill(section_value)
 
-            # Perform the forward pass through the model and apply softmax
-            x_out = torch.softmax(model(image), 1).detach().cpu()
+            inp = torch.from_numpy(four_ch).unsqueeze(0).to(device)
 
-            # Store the output in the corresponding slice of the box tensor
-            box[i - 1] = x_out
+            # Model inference with softmax normalization
+            logits = model(inp)
+            probs = torch.softmax(logits, dim=1)
 
-        # Reshape the box tensor to the desired output shape
-        return box.reshape(stack[0], 142, stack[1], stack[2])
+            # Store softmax output for this slice
+            box[i - 1] = probs
 
+    return box
 
-def parcellation(voxel, pnet_c, pnet_s, pnet_a, device):
+
+def parcellation(voxel, pnet, device):
     """
-    Perform parcellation on the given voxel data using provided neural networks for coronal, sagittal, and axial views.
+    Perform full 3D brain parcellation by aggregating predictions across multiple anatomical planes.
+
+    The function normalizes the input MRI volume, generates three differently oriented representations
+    (coronal, sagittal, axial), and performs 2.5D inference on each using a shared parcellation network.
+    The resulting probability maps are fused by summation and converted into a discrete segmentation map
+    via argmax over anatomical classes.
 
     Args:
-        voxel (torch.Tensor): The input 3D voxel data to be parcellated.
-        pnet_c (torch.nn.Module): The neural network model for coronal view parcellation.
-        pnet_s (torch.nn.Module): The neural network model for sagittal view parcellation.
-        pnet_a (torch.nn.Module): The neural network model for axial view parcellation.
-        device (torch.device): The device (CPU or GPU) to perform computations on.
+        voxel (numpy.ndarray): Input 3D brain volume (float array).
+        pnet (torch.nn.Module): Trained parcellation network (U-Net or similar architecture).
+        device (torch.device): Device on which inference will be executed (CPU or GPU).
 
     Returns:
-        numpy.ndarray: The parcellated output as a numpy array.
+        numpy.ndarray: Final 3D parcellation map (integer label image) with voxel-wise anatomical labels.
     """
-    # Normalize the voxel data
-    voxel = normalize(voxel)
+    # Normalize input intensities for network inference
+    voxel = normalize(voxel, "parcellation")
 
-    # Prepare the voxel data for different views
+    # Prepare three anatomical views for 2.5D inference
     coronal = voxel.transpose(1, 2, 0)
     sagittal = voxel
     axial = voxel.transpose(2, 1, 0)
 
-    # Perform parcellation for the coronal view
-    print("Performing parcellation for coronal view...")
-    out_c = parcellate(coronal, pnet_c, device, "c").permute(1, 3, 0, 2)
+    # ------------------------
+    # Coronal view inference
+    # ------------------------
+    out_c = parcellate(coronal, pnet, device, "Coronal").permute(1, 3, 0, 2)
     torch.cuda.empty_cache()
-    print("Parcellation for coronal view completed.")
 
-    # Perform parcellation for the sagittal view
-    print("Performing parcellation for sagittal view...")
-    out_s = parcellate(sagittal, pnet_s, device, "s").permute(1, 0, 2, 3)
+    # ------------------------
+    # Sagittal view inference
+    # ------------------------
+    out_s = parcellate(sagittal, pnet, device, "Sagittal").permute(1, 0, 2, 3)
     torch.cuda.empty_cache()
-    print("Parcellation for sagittal view completed.")
 
-    # Combine the results from coronal and sagittal views
+    # Fuse coronal and sagittal predictions
     out_e = out_c + out_s
     del out_c, out_s
-    print("Combining results from coronal and sagittal views...")
-    # Perform parcellation for the axial view
-    out_a = parcellate(axial, pnet_a, device, "a").permute(1, 3, 2, 0)
+
+    # ------------------------
+    # Axial view inference
+    # ------------------------
+    out_a = parcellate(axial, pnet, device, "Axial").permute(1, 3, 2, 0)
     torch.cuda.empty_cache()
-    print("Parcellation for axial view completed.")
 
-    # Combine the results from all views
-    out_e = out_a + out_e
+    # Combine outputs from all three anatomical orientations
+    out_e = out_e + out_a
     del out_a
 
-    # Get the final parcellated output by taking the argmax
+    # Convert probability maps to final integer labels
     parcellated = torch.argmax(out_e, 0).numpy()
 
     return parcellated

src/utils/stripping.py

@@ -7,96 +7,92 @@ from utils.functions import normalize, reimburse_conform
 
 def strip(voxel, model, device):
     """
-    Applies a given model to a 3D voxel array and returns the processed output.
+    Perform slice-wise inference using the brain stripping model.
+
+    This function processes the input 3D volume slice by slice (along the first axis),
+    using a three-slice context window for each prediction. The output is a 3D mask
+    representing the brain region.
 
     Args:
-        voxel (numpy.ndarray): A 3D numpy array of shape (256, 256, 256) representing the input voxel data.
-        model (torch.nn.Module): A PyTorch model to be used for processing the voxel data.
-        device (torch.device): The device (CPU or GPU) on which the model and data should be loaded.
+        voxel (numpy.ndarray): Input voxel data of shape (N, 224, 224), typically
+            a single anatomical orientation (e.g., coronal or sagittal view).
+        model (torch.nn.Module): The trained PyTorch brain stripping model.
+        device (torch.device): Device used for inference (CPU, CUDA, or MPS).
 
     Returns:
-        torch.Tensor: A 3D tensor of shape (256, 256, 256) containing the processed output.
+        torch.Tensor: A tensor of shape (224, 224, 224) representing the predicted
+        binary brain mask.
     """
-    # Set the model to evaluation mode
     model.eval()
 
-    # Disable gradient calculation for inference
-    with torch.inference_mode():
-        # Initialize an empty tensor to store the output
-        output = torch.zeros(256, 256, 256).to(device)
-
-        # Iterate over each slice in the voxel data
-        for i, v in enumerate(voxel):
-            # Reshape the slice to match the model's input dimensions
-            image = v.reshape(1, 1, 256, 256)
+    # Pad one slice on both ends to ensure valid 3-slice context at the boundaries
+    voxel = np.pad(voxel, [(1, 1), (0, 0), (0, 0)], "constant", constant_values=voxel.min())
 
-            # Convert the numpy array to a PyTorch tensor and move it to the specified device
-            image = torch.tensor(image).to(device)
-
-            # Apply the model to the input image and apply the sigmoid activation function
-            x_out = torch.sigmoid(model(image)).detach()
+    with torch.inference_mode():
+        box = torch.zeros(224, 224, 224)
 
-            # Store the output in the corresponding slice of the output tensor
-            output[i] = x_out
+        # Perform model inference for each slice using a 3-slice context
+        for i in range(1, 225):
+            image = np.stack([voxel[i - 1], voxel[i], voxel[i + 1]])
+            image = torch.tensor(image.reshape(1, 3, 224, 224)).to(device)
+            x_out = torch.sigmoid(model(image)).detach().cpu()
+            box[i - 1] = x_out
 
-        # Reshape the output tensor to the original voxel dimensions and return it
-        return output.reshape(256, 256, 256)
+        # Return as a 3D mask tensor
+        return box.reshape(224, 224, 224)
 
 
-def stripping(output_dir, basename, voxel, odata, data, ssnet, device):
+def stripping(output_dir, basename, voxel, odata, data, ssnet, shift, device):
     """
-    Perform brain stripping on a given voxel using a specified neural network.
+    Perform full 3D brain stripping using a deep learning model.
 
-    This function normalizes the input voxel, applies brain stripping in three anatomical planes
-    (coronal, sagittal, and axial), and combines the results to produce a final stripped brain image.
-    The stripped image is then centered and cropped.
+    This function applies a neural network-based skull-stripping algorithm to
+    isolate the brain region from a 3D MRI volume. It performs inference along
+    three anatomical orientations—coronal, sagittal, and axial—and fuses the
+    predictions to obtain a robust binary mask. The mask is then applied to the
+    input image, recentred, and saved.
 
     Args:
-        voxel (numpy.ndarray): The input 3D voxel data to be stripped.
-        data (nibabel.Nifti1Image): The original neuroimaging data.
-        ssnet (torch.nn.Module): The neural network model used for brain stripping.
-        device (torch.device): The device on which the neural network model is loaded (e.g., CPU or GPU).
+        output_dir (str): Directory where intermediate and final results will be saved.
+        basename (str): Base name of the current case (used for file naming).
+        voxel (numpy.ndarray): Input 3D voxel data (preprocessed MRI image).
+        odata (nibabel.Nifti1Image): Original NIfTI image before preprocessing.
+        data (nibabel.Nifti1Image): Preprocessed NIfTI image used for model input.
+        ssnet (torch.nn.Module): Trained brain stripping network.
+        shift (tuple[int, int, int]): The (x, y, z) offsets applied previously during cropping.
+        device (torch.device): Device used for inference (CPU, CUDA, or MPS).
 
     Returns:
-        tuple: A tuple containing:
-            - stripped (numpy.ndarray): The stripped and processed brain image.
-            - (xd, yd, zd) (tuple of int): The shifts applied to center the brain image in the x, y, and z directions.
+        numpy.ndarray: The skull-stripped 3D brain volume.
     """
-    # Normalize the input voxel data
-    voxel = normalize(voxel)
+    # Preserve original intensity data for later restoration
+    original = voxel.copy()
+
+    # Normalize the voxel intensities for model input
+    voxel = normalize(voxel, "stripping")
 
-    # Prepare the voxel data in three anatomical planes: coronal, sagittal, and axial
+    # Prepare data in three anatomical orientations
     coronal = voxel.transpose(1, 2, 0)
     sagittal = voxel
     axial = voxel.transpose(2, 1, 0)
 
-    # Apply the brain stripping model to each plane
-    out_c = strip(coronal, ssnet, device).permute(2, 0, 1)
-    out_s = strip(sagittal, ssnet, device)
-    out_a = strip(axial, ssnet, device).permute(2, 1, 0)
+    # Apply the model along each anatomical plane
+    out_c = strip(coronal, ssnet, device).permute(2, 0, 1)  # coronal → native orientation
+    out_s = strip(sagittal, ssnet, device)  # sagittal
+    out_a = strip(axial, ssnet, device).permute(2, 1, 0)  # axial → native orientation
 
-    # Combine the results from the three planes and threshold the output
+    # Fuse predictions by averaging across the three planes and apply threshold
     out_e = ((out_c + out_s + out_a) / 3) > 0.5
     out_e = out_e.cpu().numpy()
 
-    # Multiply the original data by the thresholded output to get the stripped brain image
-    stripped = data.get_fdata().astype("float32") * out_e
+    # Apply the binary mask to extract the brain region
+    stripped = original * out_e
 
-    out_filename = reimburse_conform(output_dir, basename, "stripped", odata, data, out_e)
-
-    # Calculate the center of mass of the stripped brain image
-    x, y, z = map(int, ndimage.center_of_mass(out_e))
-
-    # Calculate the shifts needed to center the brain image
-    xd = 128 - x
-    yd = 120 - y
-    zd = 128 - z
+    # Restore the mask to the original conformed geometry
+    # Pad to original full size and reverse the previously applied shift
+    out_e = np.pad(out_e, [(16, 16), (16, 16), (16, 16)], "constant", constant_values=0)
+    out_e = np.roll(out_e, (-shift[0], -shift[1], -shift[2]), axis=(0, 1, 2))
 
-    # Apply the shifts to center the brain image
-    stripped = np.roll(stripped, (xd, yd, zd), axis=(0, 1, 2))
-
-    # Crop the centered brain image
-    stripped = stripped[32:-32, 16:-16, 32:-32]
+    out_filename = reimburse_conform(output_dir, basename, "stripped", odata, data, out_e)
 
-    # Return the stripped brain image and the shifts applied
-    return stripped, (xd, yd, zd), out_filename
+    return stripped, out_filename