Training code availability and how to train on my own dataset

[ssst5114-web]

Hello LesionLocator Team,

First of all, thank you very much for releasing LesionLocator and your excellent CVPR 2025 paper.

I noticed that the current repository mainly provides inference scripts (LesionLocator_segment, LesionLocator_track). However, the paper describes a comprehensive training procedure including pretraining, fine-tuning, and longitudinal training.

Will the full training code be released in the future?

If not, could you please provide some guidance on how to train LesionLocator on a custom dataset (e.g., dataset format, preprocessing requirements, and how to adapt the trainer)?

This would be very helpful for reproducing your results and applying LesionLocator to my own research data.

Thank you very much for your time and for this valuable contribution to the community!

[mrokuss]

Hey @ssst5114-web

Thanks so much for your interest in LesionLocator and for reaching out!

You're absolutely right that, at present, the repository offers only inference scripts (LesionLocator_segment, LesionLocator_track). While we cannot share the full training code at this time, I’m happy to provide guidance to help you replicate and adapt the training process for your own datasets:

1. Dataset & Format

While privacy restrictions prevent us from sharing the in-house dataset (as much as I would like to), we provided a Lesion Dataset with Synthetic Follow-Ups (≈ 700 GB), which can be used for pretraining or in conjunction with your own data. Did you succeed downloading that? It follows the nnU-Net dataset format, which is our standard. You can find the exact specifications here.

2. Pre-Training

The segmentation module was first trained on the task of (promtable) lesion segmentation by initializing with a nnUNet ResEncL checkpoint trained on various segmentation dataset in a MultiTalent fashion. We made the checkpoint publicly available to download here (Dataset619_nativemultistem) and I highly recommend it also for any other segmentation projects ;)

3. Training Strategy & Prompt Simulation

Training essentially involves simulating prompts from ground truth masks:

• For point prompts, we sample a pixel within the lesion and convert it to a small spherical region.
• For box prompts, we compute the bounding box that encloses the lesion, with an added margin.

We mainly adapted the normal nnUNet dataloading pipeline and then simulated a point or box with some jiggle placed in an empty 3D volume (patch) in the dataloader. This can then be passed as a segmentation mask through data augmentation and afterwards concatenated to the input of the network. The logic for converting a segmentation mask into a 3D volumetric prompt is implemented in the prompt handler code.

4. Longitudinal / Tracking Mode

For tracking you have to make sure for your own data that you write a dataloader which draws baseline and follow-up images from the same patient. Prompts are simulated on the baseline scan, the prediction is the mask of the follow-up scan (always single instances). We btw also published a framework for handling patients as a longitudinal adaption of nnUNet which you might want to check out: https://github.com/MIC-DKFZ/LongiSeg

Hope this helps and I'm excited to see what you build with LesionLocator or if you can make it faster since right now this is unfortunately not optimized for speed at all :D Let me know if you run into any obstacles.

Best,
Max