uniGradICON icon indicating copy to clipboard operation
uniGradICON copied to clipboard

How to train on my own dataset?

Open supgy opened this issue 9 months ago • 9 comments

HI author, I have some CT and MRI images of the same body part, how can I train the model use my own data set. I notice that you used
self.imgs = torch.load(datapath) to load dataset, do I need to convert the NII image into a tensor first? Hope your reply~

supgy avatar Mar 21 '25 08:03 supgy

https://icon.readthedocs.io/en/move-in-constricon/medical_training.html

Here is our best effort tutorial on fine-tuning. It's almost ready to merge into master, so if you run into any issues let me know so that I can fix them!

HastingsGreer avatar Mar 21 '25 14:03 HastingsGreer

For MR to CT registration, a good start would be to use the following for model.py:

model.py

import unigradicon

input_shape = [1, 1, 175, 175, 175]

def make_network():

return unigradicon.get_multigradicon()

HastingsGreer avatar Mar 21 '25 14:03 HastingsGreer

Wow,I'm so happy to receive your patient reply. I will try it on my dataset. If there are good experimental results, I will show them to you

supgy avatar Mar 21 '25 14:03 supgy

Ah- a quick correction to my comment. Specifically for ct mr, It should be import unigradicon import icon_registration as icon ...

return unigradicon.get_multigradicon(loss_fn=icon.losses.SquaredLNCC(sigma=5))

HastingsGreer avatar Mar 22 '25 00:03 HastingsGreer

Hello author, I have successfully completed the experiment on my dataset based on the tutorial you provided. Now I want to use the generated transform.hdf5 to calculate the Jacobian determinant, but I don't know how to do it correctly. My code has always had errors, can you give me some advice?Here is my code:

import itk transform_path = "transforms/patient100transform.hdf5" transform = itk.transformread(transform_path) fixed_image = itk.imread("ct_test/patient100.nii.gz", itk.F)

displacement_field = itk.transform_to_displacement_field_filter(fixed,transform = transform[0]) jacobian_image = itk.displacement_field_jacobian_determinant_filter(displacement_field) min_jacobian = itk.minimum_maximum_image_filter(jacobian_image).GetMinimum() jacobian_array = itk.array_from_image(jacobian_image) folding_rate = (jacobian_array < 0).sum() / jacobian_array.size print(f"Min Jacobian: {min_jacobian}, Folding Rate: {folding_rate}")

supgy avatar Mar 26 '25 09:03 supgy

You can check out this tutorial

In addition, we have included this function as a CLI in the latest PyPI package. If you pip install the latest unigradicon Python package, you can compute the Jacobian via unigradicon-jacobian --fixed=[Your_fixed_image] --transform=trans.hdf5 --jacob=jacobian.nii.gz

lintian-a avatar May 25 '25 18:05 lintian-a

Hi, thanks for creating this great fine-tuning tutorial! I'm fine-tuning on a local dataset which I'm loading from nifti. I'm just querying what format is expected for the input data to unigradicon model?

From the code it seems that you load nifti using 'itk.imread'. From my experience (see code below), ITK sets negative x/y values for direction/offset when loading from nifti, but doesn't flip the image data array. This is because it expects the data to be in RAS+ and wants it in LPS+. In my example, the nifti data is stored in RAS+ anyway, so I just reverse these negative values when loading.

# Load using nibabel.
img = nib.load(filepath)
print('nifti')
print(img.affine)

# Load using itk.
img = itk.imread(filepath)
print('itk')
print(img.GetDirection(), img.GetSpacing(), img.GetOrigin())

Output:

$ nifti
[[   1.171875     0.           0.        -300.       ]
 [   0.           1.171875     0.        -142.8999939]
 [   0.           0.           2.        -332.       ]
 [   0.           0.           0.           1.       ]]
itk
itkMatrixD33 ([[-1.0, 0.0, 0.0], [0.0, -1.0, 0.0], [0.0, 0.0, 1.0]]) itkVectorD3 ([1.17188, 1.17188, 2]) itkPointD3 ([300, 142.9, -332])

I think unigradicon doesn't make use of direction/offset anyway when training, just the image data? So ITK's behaviour here won't change the loaded data. However, using np.array(img) transposes the image axes - see below. So I believe the array presented as input to the network is actually in +SPL coords? This is assuming that your training data was stored in LPS+ coords? I think this is important as I'll need to call np.transpose(data) on my nibabel-loaded data before passing for fine-tuning to get best results.

# Load using nibabel.
img = nib.load(filepath)
data = img.get_fdata()
print('nifti')
print(data.shape)

# Load using itk.
img = itk.imread(filepath)
data = np.array(img)
print('itk')
print(data.shape)

Output:

nifti
(512, 512, 250)
itk
(250, 512, 512)

Thanks! Brett

clarkbab avatar Aug 18 '25 06:08 clarkbab

Hi! Do you have a pipeline using unigradicon already, and just want to improve its performance? if so, one option that’s very safe is to just use that pipeline to generate the preprocessed training data so that there’s an exact match.

The second best option is to just heavily augment with all possible axis flips and transposes during training, so that the final combination used in evaluation is in distribution. This is the approach we took when training unigradicon, where it made a lot if sense because we had no way to predict what crazy combination of image libraries the end user would be using. its also the approach we found to be most performant for lung registration in the GradICON paper- the network is hugely overparametrized, so fighting overfitting in this way is all upside.

Finally, it is technically possible to carefully reason through all the transposes and data conversions as you are suggesting, but its something of a last resort- it always takes me a lot of coffee and two days of focus to get into the mindset. This is needed for e.g. learn2reg submission, or writing/debugging code like icon’s itk_wrapper --Hastings Greer 919 525 6808 -- github.com/HastingsGreer

On Mon, Aug 18, 2025 at 2:30 AM Brett Clark @.***> wrote:

clarkbab left a comment (uncbiag/uniGradICON#44) https://github.com/uncbiag/uniGradICON/issues/44#issuecomment-3195302519

Hi, thanks for creating this great fine-tuning tutorial! I'm fine-tuning on a local dataset which I'm loading from nifti. I'm just querying what format is expected for the input data to unigradicon model?

From the code it seems that you load nifti using 'itk.imread' https://github.com/uncbiag/uniGradICON/blob/6c2167707442e4cb4b11449542aa9c32fbb460ef/src/unigradicon/__init__.py#L368. From my experience (see code below), ITK sets negative x/y values for direction/offset when loading from nifti, but doesn't flip the image data array. This is because it expects the data to be in RAS+ and wants it in LPS+. In my example, the nifti data is stored in RAS+ anyway, so I just reverse these negative values when loading.

Load using nibabel.

img = nib.load(filepath) print('nifti') print(img.affine)

Load using itk.

img = itk.imread(filepath) print('itk') print(img.GetDirection(), img.GetSpacing(), img.GetOrigin())

Output:

$ nifti [[ 1.171875 0. 0. -300. ] [ 0. 1.171875 0. -142.8999939] [ 0. 0. 2. -332. ] [ 0. 0. 0. 1. ]] itk itkMatrixD33 ([[-1.0, 0.0, 0.0], [0.0, -1.0, 0.0], [0.0, 0.0, 1.0]]) itkVectorD3 ([1.17188, 1.17188, 2]) itkPointD3 ([300, 142.9, -332])

I think unigradicon doesn't make use of direction/offset anyway when training, just the image data? So ITK's behaviour here won't change the loaded data. However, using np.array(img) transposes the image axes https://github.com/uncbiag/ICON/blob/37c2c13f4f353f705cc4c31931e4c89ad46212e4/src/icon_registration/itk_wrapper.py#L52

  • see below. So I believe the array presented as input to the network is actually in +SPL coords? This is assuming that your training data was stored in LPS+ coords? I think this is important as I'll need to call np.transpose(data) on my nibabel-loaded data before passing for fine-tuning to get best results.

Load using nibabel.

img = nib.load(filepath) data = img.get_fdata() print('nifti') print(data.shape)

Load using itk.

img = itk.imread(filepath) data = np.array(img) print('itk') print(data.shape)

Output:

nifti (512, 512, 250) itk (250, 512, 512)

Thanks! Brett

— Reply to this email directly, view it on GitHub https://github.com/uncbiag/uniGradICON/issues/44#issuecomment-3195302519, or unsubscribe https://github.com/notifications/unsubscribe-auth/ABUCOAPQUEVMU3VD23DOAH33OFXGXAVCNFSM6AAAAABZPMFREWVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZTCOJVGMYDENJRHE . You are receiving this because you commented.Message ID: @.***>

HastingsGreer avatar Aug 18 '25 11:08 HastingsGreer

Thanks @HastingsGreer for the quick reply. Yes, I already have a data processing pipeline which I'm using to evaluate UGI on my local dataset, and have fine-tuned a model using some of this local data. Perhaps orientation doesn't matter so much, as you say, the model has been trained on a variety of datasets and image orientations.

I have another quick question about data augmentation. Is it expected that we apply this in "make_batch"? If so, I think it would be worth mentioning this in the tutorial, as I think the metrics we see on Tensorboard are just training loss (no validation loss) - so could potentially be overfitting to fine-tuning dataset without aug?

clarkbab avatar Aug 18 '25 23:08 clarkbab