qubvel / segmentation_models

Segmentation models with pretrained backbones. Keras and TensorFlow Keras.
MIT License
4.73k stars 1.03k forks source link

May I know how to get all available pure backbone model ? #519

Open HandsomeBrotherShuaiLi opened 2 years ago

HandsomeBrotherShuaiLi commented 2 years ago

Hi, Thanks for your great work and it helps me a lot for the segmentation tasks. I find that there are so many implemented classic backbone models in your framework. So may I know how to import these great backbones in our own script via segmentation_models lib? For instance: from segmentation_models.backbones import resnet18 etc.

Thanks ~

mehdi2020 commented 2 years ago

Hi, you can write: model = Unet(backbone_name='resnet18', ...) https://segmentation-models.readthedocs.io/en/latest/tutorial.html

HandsomeBrotherShuaiLi commented 2 years ago

Hi, you can write: model = Unet(backbone_name='resnet18', ...) https://segmentation-models.readthedocs.io/en/latest/tutorial.html

Hi, thanks for your reply. But my purpose is that I don't want encoder-decoder framework, I just want borrow all classification backbone models from this library without any decoder, that is to say, I only want the downsampling encoder model without upsampling decoder.

HandsomeBrotherShuaiLi commented 2 years ago

I wrote some ugly codes to implement my idea:

import segmentation_models as sm

sm.set_framework('tf.keras')
import tensorflow.keras as keras
keras.backend.set_image_data_format('channels_last')
backbone = "vgg16"
framework = sm.Unet(backend=backbone,  
                    encoder_weights='imagenet',
                    activation="sigmoid",
                    classes=1,
                    input_shape=(256, 256, 3))
encoder_features = sm.Backbones.get_feature_layers(backbone.lower(), n=4)
final_downsampling_feature = framework.get_layer(name=encoder_features[0]).output
model = keras.models.Model(framework.input, final_downsampling_feature)
model.summary()

The model's structure is like:

Model: "model_1"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_1 (InputLayer)         [(None, 256, 256, 3)]     0         
_________________________________________________________________
block1_conv1 (Conv2D)        (None, 256, 256, 64)      1792      
_________________________________________________________________
block1_conv2 (Conv2D)        (None, 256, 256, 64)      36928     
_________________________________________________________________
block1_pool (MaxPooling2D)   (None, 128, 128, 64)      0         
_________________________________________________________________
block2_conv1 (Conv2D)        (None, 128, 128, 128)     73856     
_________________________________________________________________
block2_conv2 (Conv2D)        (None, 128, 128, 128)     147584    
_________________________________________________________________
block2_pool (MaxPooling2D)   (None, 64, 64, 128)       0         
_________________________________________________________________
block3_conv1 (Conv2D)        (None, 64, 64, 256)       295168    
_________________________________________________________________
block3_conv2 (Conv2D)        (None, 64, 64, 256)       590080    
_________________________________________________________________
block3_conv3 (Conv2D)        (None, 64, 64, 256)       590080    
_________________________________________________________________
block3_pool (MaxPooling2D)   (None, 32, 32, 256)       0         
_________________________________________________________________
block4_conv1 (Conv2D)        (None, 32, 32, 512)       1180160   
_________________________________________________________________
block4_conv2 (Conv2D)        (None, 32, 32, 512)       2359808   
_________________________________________________________________
block4_conv3 (Conv2D)        (None, 32, 32, 512)       2359808   
_________________________________________________________________
block4_pool (MaxPooling2D)   (None, 16, 16, 512)       0         
_________________________________________________________________
block5_conv1 (Conv2D)        (None, 16, 16, 512)       2359808   
_________________________________________________________________
block5_conv2 (Conv2D)        (None, 16, 16, 512)       2359808   
_________________________________________________________________
block5_conv3 (Conv2D)        (None, 16, 16, 512)       2359808   
=================================================================
Total params: 14,714,688
Trainable params: 14,714,688
Non-trainable params: 0
_________________________________________________________________
​