Re-train ivadomed models with nnUNet

[hermancollin]

We need to re-train our models using nnUNet and compare performance with ivadomed. There are already scripts to convert microscopy BIDS dataset to nnunetv2 format (e.g. https://github.com/axondeepseg/model_seg_rabbit_axon-myelin_bf).

Eventually, we would also like to aggregate all datasets and train on everything at once.

[hermancollin]

@mathieuboudreau I'm thinking of the timeline for this. More and more people are showing interest in the nnunet models, especially the unmyelinated axon one. Instead of directing them to the other repo, it would be nice to implement everything in ADS. Should we wait for the ivadomed models to be re-trained to start the transition, or should we already start to add support for nnunet?

[mathieuboudreau]

should we wait for the ivadomed models to be re-trained to start the transition,

Do you have a ballpark idea how long retraining the models would take? eg 1 day / 1 week / 1 month / etc?

[hermancollin]

For every dataset, it should take 1-2 days. Some models could be trained in parallel, so everything could be done within a week.

[mathieuboudreau]

If that's the case, and you expect it to go smoothly, I'd say lets just wait until we have them retrained on nnunet and then update ADS as a whole accordingly afterwards. That's a short enough timeframe that we can point our users to the temp repo and tell them that integrating it in ADS proper is in the works and coming out soon

[hermancollin]

Ok so here's an outline:

• TEM model I already started uploading scripts on the default-TEM-model repo. There is a preprocessing script to convert the data and a script for training. This is more or less what will be required for every following model. Training is ongoing.
• SEM model here: https://github.com/axondeepseg/default-SEM-model
• BF model here: https://github.com/axondeepseg/default-BF-model
• BF human model here: https://github.com/axondeepseg/model_seg_human_axon-myelin_bf

[hermancollin]

Currently the only model left to train is the wakehealth model. Once this is done, we can combine the 4 datasets mentioned earlier with the VCU dataset (already in nnunet format with a model trained) and get one big model.

Here are some thoughts. The ivadomed models were each <100 Mb, which is in retrospect a big advantage over nnunet. We already somewhat discussed this and we agreed bigger models weren't a problem, but nnunet models will be heavier. A single nnunet model is >200 Mb, and all our results so far indicate that ensembling 3-5 folds is beneficial. This means each ensembled model would effectively be 600 Mb to >1 Gb large (more than all ivadomed models).

At this point, I would wait to see if the model trained on all datasets can be used instead of individual ones. When we have results over all 5 testing sets, we'll see if it's necessary to upload individual nnunet models. This could mean a 5x size reduction in total download for the user.

[hermancollin]

Update: the models were all retrained a few months ago. It is now time to upload everything on github. We don't know yet if these models will be downloaded by default by ADS, but they will still be accessible in the future. Here is a list of all dedicated models and their respective repo where we will put the checkpoints as release assets.

• TEM (TEM1 in the paper): https://github.com/axondeepseg/default-TEM-model
• SEM (SEM1 in the paper): https://github.com/axondeepseg/default-SEM-model
• BF default (BF1 in the paper): https://github.com/axondeepseg/default-BF-model
• BF VCU (BF2 in the paper): https://github.com/axondeepseg/model_seg_rabbit_axon-myelin_bf
• BF wakehealth (BF3 in the paper): originally here but migrated in this new repo for nnunetv2 https://github.com/axondeepseg/default-wakehealth-model
• CARS (CARS1 in the paper): https://github.com/axondeepseg/default-CARS-model

As a sidenote, the generalist models (both BF generalist and full generalist) will be available here: https://github.com/axondeepseg/model_seg_generalist. The full generalist model folds are already uploaded.

[hermancollin]

I'll also add the fold_all individual checkpoints in the same releases. Checklist with the corresponding release so I don't miss anything:

dedicated models

• TEM1 - https://github.com/axondeepseg/default-TEM-model/releases/tag/r20240403
• SEM1 - https://github.com/axondeepseg/default-SEM-model/releases/tag/r20240403
• BF1 - https://github.com/axondeepseg/default-BF-model/releases/tag/r20240405
• BF2 - https://github.com/axondeepseg/model_seg_rabbit_axon-myelin_bf/releases/tag/r20240405
• BF3 - https://github.com/axondeepseg/default-wakehealth-model/releases/tag/r20240408
• CARS1 - https://github.com/axondeepseg/default-CARS-model/releases/tag/r20240403

---

generalist models

• BF_AGG - light version only in r20240416
• BF_FULL - ensemble in r20240224 and light in r20240416

The full ensembled model with 5 folds is suffixed with _ensemble and the fold_all model with _light.