Fix shape mismatch with different input sized UnetR training script

[louisfb01]

Following #52 .

Training script crashes with our contrasts other than T1w, especially a problem when trying to train multi-contrast. The issue is related to a shape mismatch between what the model predicts and our data. It happens with one random dimension (between x,y,z) and random images where the reformatted image patch is not matching the expected shape with one additional dimension. Happens in the RandWeightedCropd function in the transforms.

Example output: ValueError: image and weight map spatial shape mismatch: (75, 74, 300) vs (74, 74, 300).

To reproduce:

1. Implement the Unetr branch and follow the steps in README file.
2. run python monai_main.py --dataset_json_name dataset.

To test:

• [ ] Try with minimal transforms to see if there's a bug.
• [x] replace RandWeightedCropd with RandCropByPosNegLabeld.

Causes shape mismatch in 1D (flattened matrix). Does not fix.

• [x] #55
• [x] #56
• [ ] TBD if not fixed

[jcohenadad]

thank you for opening this issue!

Implement the Unetr branch and follow the steps in README file.

Can you please use a hyperlink (easier for us to access) and also point to the actual branch commit (because a branch can evolve)

same thing for the README --> hyperlink to the commit of the README you used

[jcohenadad]

Note: better to use the "view file" feature, instead of the "diff" (easier to read): https://github.com/sct-pipeline/contrast-agnostic-softseg-spinalcord/blob/deeee90dd1f44d6e7ab796af03f99563cf58f1fc/monai_main.py