Feat: Add finetuning code & dice loss support for IO

.github/workflows/build_and_deploy.yml

@@ -0,0 +1,169 @@
+name: Build and Deploy
+
+on:
+  workflow_dispatch:
+    inputs:
+      upload_to_testpypi:
+        description: 'Upload to TestPyPI'
+        required: true
+        default: 'false'
+        type: boolean
+      reason:
+        description: 'Reason for manual build'
+        required: false
+        default: 'Manual testing'
+
+jobs:
+  # Add a check job that verifies organization membership
+  check-membership:
+    runs-on: [self-hosted, linux]
+    if: github.event_name == 'workflow_dispatch'
+    steps:
+      - name: Check organization membership
+        uses: actions/github-script@v6
+        with:
+          script: |
+            try {
+              const isOrgMember = await github.rest.orgs.checkMembershipForUser({
+                org: 'uncbiag',
+                username: context.actor
+              });
+              if (isOrgMember.status !== 204) {
+                core.setFailed('Only members of uncbiag organization can manually trigger this workflow');
+                process.exit(1);
+              }
+            } catch (error) {
+              core.setFailed('Only members of uncbiag organization can manually trigger this workflow');
+            }
+
+  build:
+    runs-on: [self-hosted, linux]
+    needs: [check-membership]
+    if: always() && (github.event_name != 'workflow_dispatch' || success())
+
+    permissions:
+      contents: read
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v2
+        with:
+          ref: main
+
+      - name: Set up Python 3.8
+        uses: actions/setup-python@v2
+        with:
+          python-version: 3.8
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install build twine pytest
+
+      - name: Build package
+        run: python -m build
+
+      - name: Get package name and version
+        id: pkg_info
+        run: |
+          # Extract package name and version from setup.cfg
+          PKG_NAME=$(grep -m 1 "name = " setup.cfg | cut -d "=" -f 2 | tr -d '[:space:]')
+          PKG_VERSION=$(grep -m 1 "version = " setup.cfg | cut -d "=" -f 2 | tr -d '[:space:]')
+          echo "pkg_name=$PKG_NAME" >> $GITHUB_OUTPUT
+          echo "pkg_version=$PKG_VERSION" >> $GITHUB_OUTPUT
+
+      - name: Install package from local build
+        run: |
+          # Find and install the wheel file
+          WHEEL_FILE=$(find dist -name "*.whl" | head -n 1)
+          python -m pip install $WHEEL_FILE
+
+      - name: Download test data
+        run: |
+          wget https://www.hgreer.com/assets/slicer_mirror/RegLib_C01_1.nrrd
+          wget https://www.hgreer.com/assets/slicer_mirror/RegLib_C01_2.nrrd
+          wget https://www.hgreer.com/assets/RegLib_C01_1_foreground_mask.nii.gz
+
+      - name: Run unit tests with GPU
+        run:
+          python -m unittest discover
+
+      - name: Test CLI with CPU
+        env:
+          CUDA_VISIBLE_DEVICES: ""  # Force CPU execution
+        run: |
+          echo "Running all tests on CPU"
+
+          # All commands running on CPU
+          unigradicon-register --fixed=RegLib_C01_2.nrrd --fixed_modality=mri --moving=RegLib_C01_1.nrrd --moving_modality=mri \
+              --transform_out=trans.hdf5 --warped_moving_out=warped_C01_1.nrrd --io_iterations=None
+
+          unigradicon-register --fixed=RegLib_C01_2.nrrd --fixed_modality=mri --moving=RegLib_C01_1.nrrd --moving_modality=mri \
+              --transform_out=trans.hdf5 --warped_moving_out=warped_C01_1.nrrd --io_iterations=3
+
+          unigradicon-warp --fixed=RegLib_C01_2.nrrd --moving=RegLib_C01_1.nrrd \
+              --transform=trans.hdf5 --warped_moving_out=warped_2_C01_1.nrrd --nearest_neighbor
+
+          unigradicon-warp --fixed=RegLib_C01_2.nrrd --moving=RegLib_C01_1_foreground_mask.nii.gz \
+              --transform=trans.hdf5 --warped_moving_out=warped_2_C01_1.nrrd --nearest_neighbor
+
+          unigradicon-jacobian --fixed=RegLib_C01_2.nrrd --transform=trans.hdf5 --jacob=jacobian.nii.gz
+
+      - name: Upload distributions
+        uses: actions/upload-artifact@v4
+        with:
+          name: release-dists
+          path: dist/
+
+  testpypi-publish:
+    runs-on: ubuntu-latest
+    needs: [build]
+    if: github.event.inputs.upload_to_testpypi == 'true'
+
+    permissions:
+      id-token: write
+
+    environment:
+      name: pypi
+
+    outputs:
+      pkg_name: ${{ steps.set_pkg_info.outputs.pkg_name }}
+      pkg_version: ${{ steps.set_pkg_info.outputs.pkg_version }}
+
+    steps:
+      - name: Retrieve release distributions
+        uses: actions/download-artifact@v4
+        with:
+          name: release-dists
+          path: dist/
+
+      # Get package info from the build artifact
+      - name: Set package info
+        id: set_pkg_info
+        run: |
+          # Extract package name and version from a wheel file
+          WHEEL_FILE=$(find dist -name "*.whl" | head -n 1)
+          WHEEL_BASENAME=$(basename "$WHEEL_FILE")
+          PKG_NAME=$(echo "$WHEEL_BASENAME" | cut -d'-' -f1)
+          PKG_VERSION=$(echo "$WHEEL_BASENAME" | cut -d'-' -f2)
+          echo "pkg_name=$PKG_NAME" >> $GITHUB_OUTPUT
+          echo "pkg_version=$PKG_VERSION" >> $GITHUB_OUTPUT
+
+      - name: Publish release distributions to test PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          repository-url: https://test.pypi.org/legacy/
+
+      - name: Verify installation from TestPyPI
+        run: |
+          python -m pip install --index-url https://test.pypi.org/simple/ --no-deps ${{ steps.set_pkg_info.outputs.pkg_name }}==${{ steps.set_pkg_info.outputs.pkg_version }}
+          # Install dependencies from PyPI (not TestPyPI)
+          python -m pip install ${{ steps.set_pkg_info.outputs.pkg_name }}==${{ steps.set_pkg_info.outputs.pkg_version }} --no-index --find-links dist/
+
+          wget https://www.hgreer.com/assets/slicer_mirror/RegLib_C01_1.nrrd
+          wget https://www.hgreer.com/assets/slicer_mirror/RegLib_C01_2.nrrd
+          unigradicon-register --fixed=RegLib_C01_2.nrrd --fixed_modality=mri --moving=RegLib_C01_1.nrrd --moving_modality=mri \
+              --transform_out=trans_testpypi.hdf5 --warped_moving_out=warped_testpypi.nrrd --io_iterations=None
+
+
+

.github/workflows/gpu-test-action.yml

@@ -20,12 +20,15 @@ jobs:
     - name: Install dependencies
       run: |
         pip install -r requirements.txt
-
         pip install -e .
+        pip install pytest
 
     - name: fast test with unittest
       run: |
         python -m unittest -k CPU
     - name: GPU test with unittest
       run: |
         python -m unittest discover
+    - name: pytest finetuning suite
+      run: |
+        pytest tests/finetuning/ -v

.github/workflows/test_readme_works.yml

@@ -36,4 +36,5 @@ jobs:
         unigradicon-warp --fixed=RegLib_C01_2.nrrd --moving=RegLib_C01_1_foreground_mask.nii.gz \
             --transform=trans.hdf5 --warped_moving_out=warped_2_C01_1.nrrd --nearest_neighbor
         unigradicon-jacobian --fixed=RegLib_C01_2.nrrd --transform=trans.hdf5 --jacob=jacobian.nii.gz
+        unigradicon-finetune --help
 

.github/workflows/unit_tests.yml

@@ -0,0 +1,27 @@
+name: unit-tests
+
+on:
+  pull_request:
+  push:
+    branches: main
+
+jobs:
+  pytest:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ['3.9', '3.11', '3.12']
+
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install package and test deps
+      run: |
+        pip install -e .
+        pip install pytest
+    - name: Run pytest
+      run: pytest tests/finetuning/ -v

.gitignore

@@ -1,2 +1,7 @@
 *.pyc
 *.egg-info
+datasets/
+network_weights/
+results/
+configs/learn2reg/l2r_*.json
+__pycache__/

README.md

@@ -41,7 +41,7 @@ source unigradicon_virtualenv/bin/activate
 pip install unigradicon
 ```
 
-To register one pair of image
+To register one pair of images
 ```
 wget https://www.hgreer.com/assets/slicer_mirror/RegLib_C01_1.nrrd
 wget https://www.hgreer.com/assets/slicer_mirror/RegLib_C01_2.nrrd
@@ -51,7 +51,7 @@ unigradicon-register --fixed=RegLib_C01_2.nrrd --fixed_modality=mri --moving=Reg
 
 To register without instance optimization (IO)
 ```
-unigradicon-register --fixed=RegLib_C01_2.nrrd --fixed_modality=mri --moving=RegLib_C01_1.nrrd --moving_modality=mri --transform_out=trans.hdf5 --warped_moving_out=warped_C01_1.nrrd --io_iterations None
+unigradicon-register --fixed=RegLib_C01_2.nrrd --fixed_modality=mri --moving=RegLib_C01_1.nrrd --moving_modality=mri --transform_out=trans.hdf5 --warped_moving_out=warped_C01_1.nrrd --io_iterations 'None'
 ```
 
 To use a different similarity measure in the IO. We currently support three similarity measures
@@ -62,43 +62,92 @@ To use a different similarity measure in the IO. We currently support three simi
 unigradicon-register --fixed=RegLib_C01_2.nrrd --fixed_modality=mri --moving=RegLib_C01_1.nrrd --moving_modality=mri --transform_out=trans.hdf5 --warped_moving_out=warped_C01_1.nrrd --io_iterations 50 --io_sim lncc2
 ```
 
-To load specific model weight in the inference. We currently support uniGradICON and multiGradICON.
+To change the instance optimization learning rate, use `--io_lr` (default is `0.00002`; must be positive). You can try increasing `--io_lr` while lowering `--io_iterations`; this may reduce runtime while giving similar registration quality.
+```
+unigradicon-register --fixed=RegLib_C01_2.nrrd --fixed_modality=mri --moving=RegLib_C01_1.nrrd --moving_modality=mri --transform_out=trans.hdf5 --warped_moving_out=warped_C01_1.nrrd --io_iterations 50 --io_lr 0.00002
+```
+
+To load specific model weights during inference. We currently support uniGradICON and multiGradICON.
 ```
 unigradicon-register --fixed=RegLib_C01_2.nrrd --fixed_modality=mri --moving=RegLib_C01_1.nrrd --moving_modality=mri --transform_out=trans.hdf5 --warped_moving_out=warped_C01_1.nrrd --model multigradicon
 ```
 
-To mask out the background using the provided segmentation before registration (segmentations for both moving and fixed images are necessary for accurate registration):
+### Masks, Segmentations, and Backward Compatibility
+
+Starting in 1.0.6, binary masks and label-map segmentations have separate CLI arguments:
+
+| Argument | Purpose |
+|----------|---------|
+| `--fixed_mask`, `--moving_mask` | Binary ROI masks for input masking and/or loss-function masking |
+| `--fixed_segmentation`, `--moving_segmentation` | Label maps for Dice loss |
+| `--input_masking` | Apply masks to image intensities before registration |
+| `--loss_function_masking` | Use masks inside the similarity loss |
+| `--dice_loss_weight` | Enable Dice loss using segmentations |
+
+For backward compatibility with 1.0.5 commands, if masks are not provided and segmentations are provided without Dice loss, the segmentations are still accepted as masks. This compatibility path is deprecated; prefer `--fixed_mask` and `--moving_mask` for masking.
+
+To mask out the background before registration, provide binary masks for both moving and fixed images and enable input masking:
 ```
-unigradicon-register --fixed=RegLib_C01_2.nrrd --fixed_modality=mri --fixed_segmentation=[fixed_image_segmentation_file_name] --moving=RegLib_C01_1.nrrd --moving_modality=mri --moving_segmentation=[moving_image_segmentation_file_name] --transform_out=trans.hdf5 --warped_moving_out=warped_C01_1.nrrd --io_iterations None
+unigradicon-register --fixed=RegLib_C01_2.nrrd --fixed_modality=mri --fixed_mask=[fixed_image_mask_file_name] --moving=RegLib_C01_1.nrrd --moving_modality=mri --moving_mask=[moving_image_mask_file_name] --transform_out=trans.hdf5 --warped_moving_out=warped_C01_1.nrrd --io_iterations 'None' --input_masking
 ```
 
-To apply loss function masking using the provided segmentations in the IO:
+To apply loss function masking in the IO, provide binary masks and set `--loss_function_masking`:
 ```
-unigradicon-register --fixed=RegLib_C01_2.nrrd --fixed_modality=mri --fixed_segmentation=[fixed_image_segmentation_file_name] --moving=RegLib_C01_1.nrrd --moving_modality=mri --moving_segmentation=[moving_image_segmentation_file_name] --transform_out=trans.hdf5 --warped_moving_out=warped_C01_1.nrrd --io_iterations 50 --io_sim lncc2 --loss_function_masking
+unigradicon-register --fixed=RegLib_C01_2.nrrd --fixed_modality=mri --fixed_mask=[fixed_image_mask_file_name] --moving=RegLib_C01_1.nrrd --moving_modality=mri --moving_mask=[moving_image_mask_file_name] --transform_out=trans.hdf5 --warped_moving_out=warped_C01_1.nrrd --io_iterations 50 --io_sim lncc2 --loss_function_masking
 ```
 
-This loss function ensures proper intensity adjustments for registration tasks requiring mass conservation, utilizing the change of variables rule from integration. To be effective, images must be in an intensity space where conservation holds. This loss function is specifically valid for CT modality, where -1000 HU represents air. To apply determinant-based intensity correction during registration in the IO:
+Input masking and loss masking are independent. If you want masks used for both input preprocessing and loss masking, pass both flags:
 ```
-unigradicon-register --fixed=RegLib_C01_2.nrrd --fixed_modality=mri --fixed_segmentation=[fixed_image_segmentation_file_name] --moving=RegLib_C01_1.nrrd --moving_modality=mri --moving_segmentation=[moving_image_segmentation_file_name] --transform_out=trans.hdf5 --warped_moving_out=warped_C01_1.nrrd --io_iterations 50 --io_sim lncc2 --intensity_conservation_loss
+unigradicon-register ... --input_masking --loss_function_masking
+```
+
+Intensity conservation loss ensures proper intensity adjustments for registration tasks requiring mass conservation, utilizing the change of variables rule from integration. This loss function is specifically valid for CT modality, where -1000 HU represents air. To apply this loss function, set the modality to CT and enable the intensity conservation loss flag:
+```
+unigradicon-register --fixed=RegLib_C01_2.nrrd --fixed_modality=ct --moving=RegLib_C01_1.nrrd --moving_modality=ct --transform_out=trans.hdf5 --warped_moving_out=warped_C01_1.nrrd --io_iterations 50 --io_sim lncc2 --intensity_conservation_loss
+```
+
+To optimize using Dice loss for improved anatomical structure alignment, provide segmentations and set a Dice loss weight. When enabled, the system converts the segmentations to one-hot encoding, warps them along with the images, and adds a weighted Dice loss term to the optimization objective. This encourages better alignment of corresponding anatomical structures between images.
+
+The total loss becomes:  
+`L_total = lambda * L_inverse_consistency + L_similarity + dice_loss_weight * L_dice`.
+
+This feature is particularly useful for organ registration, brain structure alignment, and other tasks where anatomical correspondence is critical. Note that the model expects segmentations to be single-channel images with the same shape as the input images. The segmentations are automatically converted to one-hot encoding. At least one **non-zero** label id must appear in **both** segmentations (excluding background `0`); otherwise registration fails with a clear error.
+
+```
+unigradicon-register --fixed=RegLib_C01_2.nrrd --fixed_modality=mri --fixed_segmentation=[fixed_image_segmentation_file_name] --moving=RegLib_C01_1.nrrd --moving_modality=mri --moving_segmentation=[moving_image_segmentation_file_name] --transform_out=trans.hdf5 --warped_moving_out=warped_C01_1.nrrd --io_iterations 50 --io_sim lncc2 --dice_loss_weight 0.1
+```
+
+To use custom network weights (e.g., after [finetuning](src/unigradicon/finetuning/README.md)):
+```
+unigradicon-register --fixed=fixed.nii.gz --fixed_modality=mri --moving=moving.nii.gz --moving_modality=mri --transform_out=trans.hdf5 --warped_moving_out=warped.nii.gz --network_weights /path/to/network_weights_final.trch
+```
+
+If you customized preprocessing during finetuning, pass the same values at inference with `--ct_window` (for CT) or `--quantile_range` (for MRI). The resulting intensity range after clamping must have max greater than min (for example, avoid equal `ct_window` bounds or a degenerate quantile range).
+```
+unigradicon-register --fixed=fixed.nii.gz --fixed_modality=mri --moving=moving.nii.gz --moving_modality=mri --transform_out=trans.hdf5 --quantile_range 0.0 0.99 --network_weights /path/to/network_weights_final.trch
+```
+
+To finetune on your own data, see the [finetuning guide](./src/unigradicon/finetuning/README.md):
+```
+unigradicon-finetune --config /path/to/config.yaml
 ```
 
 To warp an image
 ```
-unigradicon-warp --fixed [fixed_image_file_name] --moving [moving_image_file_name]  --transform trans.hdf5 --warped_moving_out warped.nii.gz --linear
+unigradicon-warp --fixed [fixed_image_file_name] --moving [moving_image_file_name] --transform trans.hdf5 --warped_moving_out warped.nii.gz --linear
 ```
 To warp a label map
 
 ```
-unigradicon-warp --fixed [fixed_image_file_name] --moving [moving_image_segmentation_file_name]  --transform trans.hdf5 --warped_moving_out warped_seg.nii.gz --nearest_neighbor
+unigradicon-warp --fixed [fixed_image_file_name] --moving [moving_image_segmentation_file_name] --transform trans.hdf5 --warped_moving_out warped_seg.nii.gz --nearest_neighbor
 ```
 
 ### 👉 Inference via colab notebook
 We provide a [colab notebook](https://colab.research.google.com/drive/1O4F0j_ZaRHxyFJ-pfApl2E1_47-wzxMj#offline=true&sandboxMode=true) where the users can directly access and visualize the output of uniGradICON network. 
 
 ### 👉 Inference via Slicer Extension
 
-A Slicer extensions is available [here](https://github.com/uncbiag/SlicerUniGradICON?tab=readme-ov-file). It is an official Slicer Extension and can be installed via the Slicer Extension Manager. This requires Slicer >=5.7.0. Please make sure to install the Slicer PyTorch extension before as uniGradICON depends on it.
-
+A Slicer extension is available [here](https://github.com/uncbiag/SlicerUniGradICON?tab=readme-ov-file). It is an official Slicer Extension and can be installed via the Slicer Extension Manager. This requires Slicer >=5.7.0. Please make sure to install the Slicer PyTorch extension first, since uniGradICON depends on it.
 
 ## Training and testing data
 
@@ -145,7 +194,8 @@ A Slicer extensions is available [here](https://github.com/uncbiag/SlicerUniGrad
         <td>578,888</td>
         <td>Inter-pat.</td>
         <td>MRI</td>
-    </tr>
+    </tr>r
+```
     <tr>
         <td>4.</td>
         <td>L2R-Abdomen</td>
@@ -262,11 +312,13 @@ A Slicer extensions is available [here](https://github.com/uncbiag/SlicerUniGrad
     </tr>
 </table>
 
-## Get involved
+## Finetuning
 
-Our goal is to continuously improve the `uniGradICON` model, e.g., by training on more datasets with additional diversity. Feel free to point us to datasets that should be included or let us know if you want to help with future developments.
+You can finetune the `uniGradICON` model on your own data using the [finetuning guide](./src/unigradicon/finetuning/README.md).
 
+## Get involved
 
+Our goal is to continuously improve the `uniGradICON` model, e.g., by training on more datasets with additional diversity. Feel free to point us to datasets that should be included or let us know if you want to help with future developments.
 
 ## Visualization
 
@@ -299,4 +351,3 @@ If you find this repository useful, please consider citing:
   organization={Springer}
 }
 ```
-