intel / unet

U-Net Biomedical Image Segmentation
Apache License 2.0
303 stars 124 forks source link

Predict multiple output masks in 2D #23

Closed FrancoisMasson1990 closed 3 years ago

FrancoisMasson1990 commented 4 years ago

Hello, Can you give a code example when you mentioned : To predict multiple output masks some modification of the output layer to the model (e.g. more output layers for the sigmoid mask) is required ? Let's say for 3 outputs (0 : background, 1 : edema, 2: non-enhancing tumor)

Also, does the loss function need to be adjusted ?

Thank you in advance

tonyreina commented 3 years ago

Sure. So what you would do is in the prediction layer (last layer) you would just change the number of channels to 3 to predict 3 different masks. For example, https://github.com/IntelAI/unet/blob/master/3D/model.py#L235 there's a variable filters=self.n_cl_out. Just change that to pass 3 filters (or however many output masks you want to predict). You'll need to make sure that the dataloader passes the label in whatever order you've defined (e.g. 0: background, 1: edema, 2:non-enhancing tumor).

The loss function here could be the combination of all 3 Dices. I think in the current code that would essentially be done automatically since the Dice would be run across all 4 dimensions (height, width, depth, output masks). You could also consider calculating the per output channel Dice and doing a weighted average for the loss if you want to target a specific layer prediction.