MIC-DKFZ / nnUNet

Apache License 2.0
5.63k stars 1.71k forks source link

Improve segmentation in a specific region #2230

Closed Laurabc closed 1 month ago

Laurabc commented 3 months ago

I am interested in improving the segmentation quality in a specific region of my dataset.

  1. In first place, I was thinking in giving more weight to the loss function in the pixels in this region, how can I do this?

  2. Other thing I am considering is region-based training, but I get the following error: UnboundLocalError: local variable 'region_labels' referenced before assignment. I may be using it wrong, I'm following this https://github.com/MIC-DKFZ/nnUNet/blob/master/documentation/region_based_training.md. I believe my dataset.json is correct, but I am not completely sure about how my dataset should be, is it ok if my labels have 0, 1, 2 in each pixel (background, foreground, specific region I want to improve, respectively)?

  3. Is there any other easy way that I'm overlooking?

sten2lu commented 3 months ago

HI @Laurabc,

First, regarding the general approach:

I encourage you to oversample specific regions during nnU-Net training to increase their weighting during training. For this to work, you would need to specify precisely which regions to oversample, e.g., for each case in some automated fashion. However, you would need to have some specific way to find these regions.

Therefore, a good way of doing this would be region-based training, as nnU-Net will then sample often from the specific_region label by default due to its oversampling foreground strategy.

To do this, you can simply create a version of the dataset.json like this with according label maps in labelsTr.

...
"labels": {
    "background": 0,
    "foreground": [1, 2],
    "specific_region": [2],
},
"regions_class_order": [1, 2],
...

Best regards,

Carsten

Laurabc commented 3 months ago

Hi, @sten2lu I'm not sure how I should oversample. I'm working with CT scans, with annotations of an organ, but there is a part of the organ where the segmentation works worse. How do I oversample?

Regarding the region-based, yes, that's what I tried before opening the issue, and I got "UnboundLocalError: local variable 'region_labels' referenced before assignment". My guess was that the labels I provided weren't in the correct format, I'm using .nrrd files with 0,1, and 2 in each pixel, is this ok? Preprocessing works, the problem is in training.

anenbergb commented 3 months ago

@Laurabc the UnboundLocalError: local variable 'region_labels' referenced before assignment error can be resolved by installing batchgeneratorv2 locally. See issue https://github.com/MIC-DKFZ/batchgeneratorsv2/issues/1

sten2lu commented 1 month ago

Hi @Laurabc,

has your issue been resolved?

Best regards,

Carsten

Laurabc commented 1 month ago

@sten2lu yes