Tensorflow/Keras: Image Segmentation Model for Image Background Removal Tasks ✂️ (10/2024)
Example Code to Load the Model (Google Colab/Kaggle Notebook):
from huggingface_hub import notebook_login
from huggingface_hub import snapshot_download
import tensorflow as tf
import os
#Realize Autenticação
notebook_login()
#Carregue o modelo do repositório
repo_id = "reidn3r/u2net-image-rembg"
model_path = snapshot_download(repo_id=repo_id)
keras_model_path = os.path.join(model_path, 'u2net-512.keras')
#Carregue o modelo localmente
model = tf.keras.models.load_model(keras_model_path)
#Exibir a arquitetura do modelo
model.summary()
The model was trained on DIS dataset (about 5k images)
The model was implemented based on this article
The architecture:
-
🖼️ Encoder:
- The encoder is responsible for extracting features from the input image, capturing what information is present in the given image to the model.
- It uses pooling operations to reduce spatial dimensionality while preserving the extracted information..
-
📊 Decoder:
- The decoder applies UpSampling operations on the tensor of information returned by the encoder to restore the spatial dimensionality.
- It also uses convolutions combined with information provided by the encoder (through skip connections), locating where the information captured by the encoder is.
-
📈 RSU Block:
- A Residual Block that combines residual connections and concepts from the traditional U-Net architecture to create a block that captures information at multiple resolution scales.
-
🌉 Bridge Blocks:
- Blocks that connect the outputs of the encoder to the decoder. They participate in merging information captured during downsampling in the encoder with fine details (recovered during upsampling in the decoder).