Fix model card metadata and update documentation

README.md

@@ -1,178 +1,237 @@
-# SlicerStrokeSegmentator
-
-SlicerStrokeSegmentator is a 3D Slicer extension for automated chronic ischemic stroke lesion segmentation from T1-weighted brain MRI using nnU-Net v2.
-
-The extension performs local inference using a five-fold ensemble of pretrained models and displays predicted lesion segmentations directly in 3D Slicer.
+---
+language:
+- en
+tags:
+- medical-imaging
+- stroke
+- lesion-segmentation
+- brain-mri
+- nnunet
+- 3d-slicer
+- pytorch
+pipeline_tag: image-segmentation
+---
+
+# Chronic Stroke Lesion Segmentation — nnU-Net v2
+
+Pretrained nnU-Net v2 models for automated chronic ischemic stroke lesion segmentation from T1-weighted brain MRI.
+
+The models can be used with the existing NNUNet extension for 3D Slicer, allowing researchers to generate lesion segmentations and visualize the results directly in Slicer.
 
 ![SlicerStrokeSegmentator screenshot](SlicerStrokeSegmentator/Resources/Screenshots/StrokeSegmentator.png)
 
+The model repository includes five trained checkpoints, the nnU-Net configuration files, and the custom BCE+Dice trainer required for model loading.
+
 ## Features
 
-- Automated chronic ischemic stroke lesion segmentation.
-- T1-weighted brain MRI input.
-- Five-fold nnU-Net v2 ensemble.
-- Local pretrained model directory selection.
-- Automatic verification of model checkpoints.
-- Model directory saved between Slicer sessions.
-- Segmentation results displayed directly in 3D Slicer.
-- CPU-based inference.
+- Automated chronic stroke lesion segmentation from T1-weighted MRI.
+- nnU-Net v2 3D full-resolution architecture.
+- Five-fold ensemble inference.
+- Custom BCE+Dice trainer.
+- Publicly available pretrained model weights.
+- Integration with the existing 3D Slicer NNUNet extension.
+- Local inference without uploading patient images to a remote server.
+
+## Model configuration
+
+| Parameter | Value |
+|---|---|
+| Framework | nnU-Net v2 |
+| Architecture configuration | 3d_fullres |
+| Trainer | nnUNetTrainerBCEDice |
+| Plans | nnUNetPlans |
+| Target voxel spacing | 1 × 1 × 1 mm³ |
+| Patch size | 128 × 128 × 128 |
+| Training batch size | 2 |
+| Intensity normalization | ZScoreNormalization |
+| Ensemble | Five folds (0–4) |
+| Input modality | T1-weighted MRI |
+
+The target voxel spacing and normalization are specified in the nnU-Net plans. These preprocessing operations do not replace spatial registration.
 
 ## Requirements
 
-- 3D Slicer (tested with version 5.12.4 on Linux).
-- Compatible nnU-Net v2 installation in Slicer's Python environment.
-- T1-weighted MRI registered to MNI152 1 mm isotropic space.
-- Pretrained model weights obtained separately from the authors.
-- Sufficient RAM and disk space for inference.
+- 3D Slicer with the NNUNet extension.
+- Compatible nnU-Net v2 installation.
+- PyTorch.
+- T1-weighted brain MRI prepared consistently with the training data, including registration to MNI152 1 mm space.
+- Sufficient system memory and disk space.
 
-The extension currently uses CPU inference. Processing time depends on the computer hardware and input data.
+Tested environment: 3D Slicer 5.12.4 on Linux, with nnU-Net v2 2.8.1.
 
-## Installation
+Successful inference has been demonstrated in this environment. Windows and macOS compatibility have not yet been independently verified.
 
-The extension can be loaded into 3D Slicer from its source directory.
+## Installation and model setup
 
-1. Download or clone this repository.
-2. Open 3D Slicer.
-3. Load the SlicerStrokeSegmentator scripted module.
-4. Open the Stroke Segmentator module under Segmentation.
-5. Obtain the pretrained model weights from the authors.
-6. Extract the weights into a local directory.
-7. Click "Select Local Model Directory".
-8. Select the nnU-Net results root directory containing Dataset001_StrokeT1w.
+### 1. Install the NNUNet extension
 
-The extension uses ~/stroke_model as its default model directory.
+Open 3D Slicer, navigate to the Extension Manager, and install the existing NNUNet extension.
 
-Installation through the Slicer Extension Manager is subject to acceptance into the Slicer ExtensionsIndex.
+Restart Slicer if requested.
 
-## Model Directory Structure
+### 2. Download the pretrained model
 
-The expected directory structure is:
+Download the model repository:
 
-    stroke_model/
-      Dataset001_StrokeT1w/
-        nnUNetTrainerImprovedLossCheckpoints__nnUNetPlans__3d_fullres/
-          dataset.json
-          plans.json
-          fold_0/
-            checkpoint_final.pth
-          fold_1/
-            checkpoint_final.pth
-          fold_2/
-            checkpoint_final.pth
-          fold_3/
-            checkpoint_final.pth
-          fold_4/
-            checkpoint_final.pth
+https://huggingface.co/parvezamu/stroke-lesion-segmentation
 
-Select the stroke_model directory, not an individual trainer or fold directory.
+The published model checkpoints are stored under:
 
-## Usage
+```text
+Dataset001_StrokeT1w/
+  nnUNetTrainer__nnUNetPlans__3d_fullres/
+    fold_0/
+    fold_1/
+    fold_2/
+    fold_3/
+    fold_4/
+```
 
-1. Open the Stroke Segmentator module.
-2. Load a T1-weighted MRI registered to MNI152 1 mm space.
-3. Select the MRI in the Input T1w selector.
-4. Select the local pretrained model directory if needed.
-5. Verify that all required model checkpoints are present.
-6. Click "Segment Stroke Lesion".
-7. Wait for inference to finish.
-8. The predicted lesion segmentation will appear in the Slicer scene.
+Each fold contains a `checkpoint_final.pth` file.
 
-The extension uses the custom nnU-Net trainer:
+The repository also provides the nnU-Net configuration metadata, including `dataset.json` and `plans.json`.
 
-    nnUNetTrainerImprovedLossCheckpoints
+### 3. Install the custom trainer
 
-All five model folds are combined during prediction.
+Download:
 
-## Public Testing Dataset - ATLAS v3.0 (629 Cases)
+[custom_trainer/nnUNetTrainerBCEDice.py](https://huggingface.co/parvezamu/stroke-lesion-segmentation/blob/main/custom_trainer/nnUNetTrainerBCEDice.py)
 
-A collection of 629 publicly available ATLAS v3.0 T1-weighted brain MRI scans and their corresponding ground-truth stroke lesion masks is provided for testing SlicerStrokeSegmentator.
+Copy the trainer into the `nnunetv2/training/nnUNetTrainer/` directory of the Python environment used by the Slicer NNUNet extension.
 
-The collection includes:
+The custom trainer inherits from the standard nnU-Net trainer and defines the BCE+Dice loss implementation. No additional custom loss file is required.
 
-- 329 newly released ATLAS v3.0 cases.
-- 300 cases from the former ATLAS v2.0 validation dataset.
+For the tested Linux installation, the destination was:
+
+```text
+Slicer-5.12.4-linux-amd64/
+  lib/Python/lib/python3.12/site-packages/
+    nnunetv2/training/nnUNetTrainer/
+      nnUNetTrainerBCEDice.py
+```
+
+The exact location depends on the Slicer installation and operating system.
+
+### 4. Prepare the local model directory
+
+The checkpoint metadata identifies the trainer as `nnUNetTrainerBCEDice`.
+
+After downloading the model, rename the local trainer directory only from:
+
+```text
+nnUNetTrainer__nnUNetPlans__3d_fullres
+```
+
+to:
+
+```text
+nnUNetTrainerBCEDice__nnUNetPlans__3d_fullres
+```
+
+Arrange the local model files as follows:
+
+```text
+stroke_model/
+  Dataset001_StrokeT1w/
+    nnUNetTrainerBCEDice__nnUNetPlans__3d_fullres/
+      dataset.json
+      plans.json
+      fold_0/
+        checkpoint_final.pth
+      fold_1/
+        checkpoint_final.pth
+      fold_2/
+        checkpoint_final.pth
+      fold_3/
+        checkpoint_final.pth
+      fold_4/
+        checkpoint_final.pth
+```
 
-These 629 cases were not included in the original 655-case ATLAS v2.0 training dataset used to develop the segmentation model.
+This is a local installation adjustment. The original Hugging Face checkpoint directory remains unchanged.
 
-The images have been registered to MNI152 space for use with the segmentation pipeline. Users should verify that voxel spacing, image geometry, and preprocessing are compatible with the model.
+### 5. Run inference in 3D Slicer
 
-### Dataset Download
+1. Open 3D Slicer and launch the NNUNet extension.
+2. Load a T1-weighted MRI prepared in the required MNI152 space.
+3. Configure the local model directory and select the appropriate model.
+4. Select the available trained folds for ensemble inference.
+5. Run segmentation.
+6. Review the resulting lesion segmentation in the Slicer scene.
 
-[Download ATLAS v3.0 - 629 Cases (Google Drive)](https://drive.google.com/drive/folders/1kgHygdB5PSDqj0ScFCj-nxuHOd5gPvXk?usp=sharing)
+The predicted segmentation should be visually inspected, particularly for small lesions, atypical anatomy, and registration errors.
 
-The dataset is organised into two folders:
+## Input image preparation
 
-    ATLAS_v3.0/
-      input/
-        T1-weighted MRI scans
-      masks/
-        Ground-truth stroke lesion masks
+The model was developed for T1-weighted brain MRI prepared in MNI152 space.
 
-The input folder contains the MRI scans, while the masks folder contains their corresponding reference lesion segmentations.
+nnU-Net automatically performs configuration-dependent operations such as voxel resampling and Z-score intensity normalization during inference.
 
-### Testing Procedure
+However, resampling is not equivalent to anatomical registration.
 
-1. Download a T1-weighted MRI from the input folder.
-2. Open 3D Slicer and select the Stroke Segmentator module.
-3. Load the MRI into Slicer.
-4. Select the pretrained model directory.
-5. Click "Segment Stroke Lesion".
-6. Wait until segmentation is complete.
-7. Load the corresponding ground-truth lesion mask from the masks folder.
-8. Compare the predicted segmentation with the reference mask.
+Users should ensure that their input images have undergone spatial preparation consistent with the training pipeline. Differences in registration, skull stripping, image orientation, field of view, or intensity characteristics may affect segmentation performance.
 
-Suggested evaluation metrics include:
+## Public testing dataset — ATLAS v3.0 collection
+
+A testing collection containing 629 ATLAS cases is available separately:
+
+[Download ATLAS testing collection (Google Drive)](https://drive.google.com/drive/folders/1kgHygdB5PSDqj0ScFCj-nxuHOd5gPvXk?usp=sharing)
+
+The collection is described as containing:
+
+- 329 newly released ATLAS v3.0 cases.
+- 300 cases from the former ATLAS v2.0 validation dataset.
+
+These cases are separate from the original 655-case ATLAS v2.0 training dataset used to develop the segmentation model.
+
+The testing collection contains T1-weighted MRI images and corresponding reference lesion masks. Users should verify image geometry, preprocessing, case correspondence, and applicable dataset access and redistribution conditions before use.
+
+### Suggested evaluation metrics
 
 - Dice similarity coefficient (DSC).
 - Lesion-wise precision, recall, and F1-score.
-- Hausdorff distance (HD95).
+- 95th-percentile Hausdorff distance (HD95).
 - Lesion volume difference.
 - Segmentation failure rate.
 
-These metrics require a separate evaluation procedure and are not automatically calculated by the extension.
+These metrics require a separate evaluation procedure and are not automatically calculated by the pretrained model.
 
-### Original Dataset Sources
+### Original dataset source
 
-ATLAS dataset download portal:
-https://fcon_1000.projects.nitrc.org/indi/retro/atlas_download.html
+[ATLAS dataset download portal](https://fcon_1000.projects.nitrc.org/indi/retro/atlas_download.html)
 
-ISLES'26 dataset information:
-https://isles-26.grand-challenge.org/dataset/
+Please acknowledge the original ATLAS dataset publications and comply with their applicable terms.
 
-The official ISLES'26 challenge evaluates algorithms using native-space MRI. The MNI152-registered testing collection described here is intended for testing this extension and is not an official ISLES'26 submission.
+## Model availability
 
-Users must comply with the original dataset's applicable access, attribution, and redistribution conditions.
+The pretrained model checkpoints and custom trainer are publicly available:
 
-## Model Access
+https://huggingface.co/parvezamu/stroke-lesion-segmentation
 
-Pretrained model weights are not included in this repository.
+The associated research manuscript is currently in preparation.
 
-The associated research manuscript is currently under preparation. Model weights may be provided to authorised researchers upon request, subject to applicable access conditions.
-
-Contact:
+For research enquiries:
 
 - pahm409@aucklanduni.ac.nz
 - parvezamu@gmail.com
 
-The repository will be updated with information about model availability following publication arrangements.
-
-## Research and Clinical Use
+## Research and clinical use
 
-This software is intended for research purposes only.
+For research use only.
 
-It is not established as a clinically validated diagnostic tool and should not be used independently for clinical diagnosis, treatment decisions, or patient management.
+This model is not established as a clinically validated diagnostic tool and should not be used independently for clinical diagnosis, treatment decisions, or patient management.
 
-Segmentation accuracy may vary depending on MRI acquisition, preprocessing, registration quality, lesion characteristics, and differences between training and evaluation datasets.
+Segmentation performance may vary with MRI acquisition parameters, lesion size, anatomical changes, registration quality, and differences between training and evaluation datasets.
 
 ## Citation
 
 Ahmad et al. Manuscript in preparation.
 
-Auckland Bioengineering Institute,
-University of Auckland,
-New Zealand.
+Auckland Bioengineering Institute  
+University of Auckland  
+New Zealand
 
-Please cite the associated publication when it becomes available and acknowledge the original ATLAS dataset publications when using the testing data.
+Please cite the associated publication when available and acknowledge the original ATLAS dataset publications.
 
 ## Acknowledgements
 
@@ -180,8 +239,13 @@ HRC grant 21/144 IMPRESS.
 
 ## License
 
-The extension source code is released under the Apache License, Version 2.0.
 
-See LICENSE.txt for the complete licence terms.
 
-Pretrained model weights are distributed separately and may have different access and licensing conditions.
+The separate SlicerStrokeSegmentator extension source code is released under the Apache License, Version 2.0, as specified in its source repository. That license should not automatically be assumed to apply to the pretrained model weights.
+
+### Successful inference in 3D Slicer
+
+The following screenshot shows the predicted chronic stroke lesion segmentation
+overlaid on the T1-weighted MRI (`A207_T1w`) in axial, coronal, and sagittal views.
+
+![Five-fold chronic stroke lesion segmentation using nnU-Net v2 in 3D Slicer](SlicerStrokeSegmentator/Resources/Screenshots/StrokeSegmentator.png)