search

yingkaisha / keras-unet-collection

The Tensorflow, Keras implementation of U-net, V-net, U-net++, UNET 3+, Attention U-net, R2U-net, ResUnet-a, U^2-Net, TransUNET, and Swin-UNET with optional ImageNet-trained backbones.
MIT License
664 stars 184 forks source link

Bug in _model_unet_3plus_2d.py and small suggestion #5

Closed thupalo closed 3 years ago

thupalo commented 3 years ago

File _model_unet_3plus_2d.py line 128 current X = UNET_left(X, filter_num[i_real], stack_num=stack_num_down, activation=activation, pool=pool, batch_norm=batch_norm, name='{}_down{}'.format(name, i_real+1)) should rather be X = UNET_left(X, filter_num_down[i_real], stack_num=stack_num_down, activation=activation, pool=pool, batch_norm=batch_norm, name='{}_down{}'.format(name, i_real+1))

And small enhancement suggestion: I'm writing classification model based on unet_3plus_2d. It would be great to have X_decoder as output from model to be able to construct GCM (controlled by input parameter: return_decoder)

model = unet_3plus_2d(input_shape, .... return_decoder=True, name='unet3plus') then [OUT_stack, X_decoder] = model.output Classification-guided Module (CGM) dropout --> 1-by-1 conv2d --> global-maxpooling --> sigmoid X_CGM = X_decoder[-1] X_CGM = Dropout(rate=0.1)(X_CGM) X_CGM = Conv2D(filter_num_skip[-1], 1, padding='same')(X_CGM) ...

Your opinion?

yingkaisha commented 3 years ago

Hi,

Try pip install keras-unet-collection==0.0.18. The bug should be fixed.

One could build a model without prediction heads through X_decoder = unet_3plus_2d_base(input_tensor, ...), check it out and see if it matches your need.

An example of applying CGM is available: https://github.com/yingkaisha/keras-unet-collection/blob/main/examples/segmentation_unet-three-plus_oxford-iiit.ipynb

Thank you

thupalo commented 3 years ago

Thank for the quick fix, and for maintenance this great set of models. Stay safe.