Results on OASIS data

[rohitrango]

Hi,

Thank you for releasing your code -- great work! 🤩

I'm trying to do deformable registration by training on the OASIS dataset. I've compiled my CSV file by taking all possible pairs from the training set. The first few rows look as follows:

fixed_img_path,moving_img_path,fixed_seg_path,moving_seg_path,fixed_mask_path,moving_mask_path,train
/data/OASIS_OAS1_0001_MR1/aligned_norm.nii.gz,/data/OASIS_OAS1_0002_MR1/aligned_norm.nii.gz,/data/OASIS_OAS1_0001_MR1/aligned_seg35.nii.gz,/data/OASIS_OAS1_0002_MR1/aligned_seg35.nii.gz,None,None,True
/data/OASIS_OAS1_0001_MR1/aligned_norm.nii.gz,/data/OASIS_OAS1_0003_MR1/aligned_norm.nii.gz,/data/OASIS_OAS1_0001_MR1/aligned_seg35.nii.gz,/data/OASIS_OAS1_0003_MR1/aligned_seg35.nii.gz,None,None,True
/data/OASIS_OAS1_0001_MR1/aligned_norm.nii.gz,/data/OASIS_OAS1_0004_MR1/aligned_norm.nii.gz,/data/OASIS_OAS1_0001_MR1/aligned_seg35.nii.gz,/data/OASIS_OAS1_0004_MR1/aligned_seg35.nii.gz,None,None,True
/data/OASIS_OAS1_0001_MR1/aligned_norm.nii.gz,/data/OASIS_OAS1_0005_MR1/aligned_norm.nii.gz,/data/OASIS_OAS1_0001_MR1/aligned_seg35.nii.gz,/data/OASIS_OAS1_0005_MR1/aligned_seg35.nii.gz,None,None,True
/data/OASIS_OAS1_0001_MR1/aligned_norm.nii.gz,/data/OASIS_OAS1_0006_MR1/aligned_norm.nii.gz,/data/OASIS_OAS1_0001_MR1/aligned_seg35.nii.gz,/data/OASIS_OAS1_0006_MR1/aligned_seg35.nii.gz,None,None,True
/data/OASIS_OAS1_0001_MR1/aligned_norm.nii.gz,/data/OASIS_OAS1_0007_MR1/aligned_norm.nii.gz,/data/OASIS_OAS1_0001_MR1/aligned_seg35.nii.gz,/data/OASIS_OAS1_0007_MR1/aligned_seg35.nii.gz,None,None,True
/data/OASIS_OAS1_0001_MR1/aligned_norm.nii.gz,/data/OASIS_OAS1_0009_MR1/aligned_norm.nii.gz,/data/OASIS_OAS1_0001_MR1/aligned_seg35.nii.gz,/data/OASIS_OAS1_0009_MR1/aligned_seg35.nii.gz,None,None,True
/data/OASIS_OAS1_0001_MR1/aligned_norm.nii.gz,/data/OASIS_OAS1_0010_MR1/aligned_norm.nii.gz,/data/OASIS_OAS1_0001_MR1/aligned_seg35.nii.gz,/data/OASIS_OAS1_0010_MR1/aligned_seg35.nii.gz,None,None,True
/data/OASIS_OAS1_0001_MR1/aligned_norm.nii.gz,/data/OASIS_OAS1_0011_MR1/aligned_norm.nii.gz,/data/OASIS_OAS1_0001_MR1/aligned_seg35.nii.gz,/data/OASIS_OAS1_0011_MR1/aligned_seg35.nii.gz,None,None,True

I'm using the following script to train:

python scripts/run.py --job_name oasis_seg --save_dir ./oasis-run-seg --num_keypoints 512 --loss_fn mse --transform_type tps_0 --data_path ./train_oasis_seg.csv --train_dataset csv --run_mode train --backbone truncatedunet --use_amp

But I get a validation Dice score of around 0.65 on the validation set:

which is not so good. I've verified that there are 36 labeled classes in the segmentation (1st channel is background and is ignored).

Training with the dice loss (--loss_fn dice) does not help either.

Have you tried training with the OASIS dataset and have seen different results? Sharing the training scripts / pretrained models would be immensely useful.

Let me know if I'm missing something. Thanks again!

[alanqrwang]

Could you provide some visualizations of the moving, fixed, and aligned images with overlaid keypoints? You can produce these with the --visualize argument.

[rohitrango]

here are some visualizations

[alanqrwang]

Can you also post your training curves? I suspect you could train for longer, because the keypoints aren't spread out very well and some keypoints are off the brain, which is weird. Usually, I train for at least 5000 epochs. Each epoch by default in the code is 32 steps.

[rohitrango]

Are there any tensorboard logs / files that store the training curves? I couldnt find anything in the save dirs.

[alanqrwang]

You can --use_wandb log to wandb.

[rohitrango]

Okay will do!