[Semantic Segmentation] - Add SegFormer Architecture, Weight Conversion Script and Presets

[DavidLandup0]

This PR adds:

• SegFormerBackbone and presets
• Preprocessor flow
• SegFormerImageSegmenter and presets for Cityscapes and ADE20k (B0...B5 each)
• Conversion script
• Tests

Basic Usage

preprocessor = keras_hub.models.ImageSegmenterPreprocessor.from_preset("segformer_b0_ade20k_512")
segmenter = keras_hub.models.SegFormerImageSegmenter.from_preset("segformer_b0_ade20k_512")
segmenter(np.random.rand(1, 512, 512, 3))

End-to-end example with preprocessor:

import urllib.request 
from PIL import Image 
import numpy as np
import keras_hub

preprocessor = keras_hub.models.ImageSegmenterPreprocessor.from_preset("segformer_b0_ade20k_512")
segmenter = keras_hub.models.SegFormerImageSegmenter.from_preset("segformer_b0_ade20k_512")

img_url = "https://www.vanorohotel.com/wp-content/uploads/2021/07/drz-vanoro_6737.jpg"  
urllib.request.urlretrieve(img_url, "image.png") 
  
img = np.array(Image.open("image.png").resize((512, 512)))
img = np.expand_dims(img, 0)
inputs = preprocessor(img)
outs = segmenter(inputs)

With Image Converter

converter = keras_hub.layers.ImageConverter(image_size=(512, 512))
preprocessor = keras_hub.models.ImageSegmenterPreprocessor.from_preset("segformer_b0_ade20k_512", image_converter=converter)
segmenter = keras_hub.models.SegFormerImageSegmenter.from_preset("segformer_b0_ade20k_512")

Training Pipeline Example

A few examples in the notebook below:

• Instantiation of Backbone and Segmenter with MiT Encoder
• Running on input images
• Training pipeline with TFDS on Oxford IIIT Pets as an example

https://colab.research.google.com/drive/1EBNg6nPKx_KzyRuQQtHZ_PG_Nsf2pAg2#scrollTo=V9Ub4NHKCx9e

After a few minutes of training from scratch (both encoder and segmenter):

[DavidLandup0]

Found the issue - a transpose call shuffling the order of a latent in the encoder incorrectly. I'll get the presets up on Kaggle now

[divyashreepathihalli]

Amazing!! the PR looks great. Just a few comments.

[divyashreepathihalli]

how about transformations to y - which would be masks in the training set. If the images are resized/transformed, so should the masks

[DavidLandup0]

Good point. In this case, since we don't want to apply either normalization or rescaling to the masks, the image converter would just resize the images. Do we have a standard practice of defining a default converter or let the user decide this?

[divyashreepathihalli]

for segformer, I think we should add the default resizing of mask here if we are performing resizing on images.

[divyashreepathihalli]

deeplabv3 code did something similar

[DavidLandup0]

Added - if a converter is present, it resizes both the images and the masks. It also rescales and normalizes the images to match the HuggingFace preprocessor

[divyashreepathihalli]

wouldn't this also rescale the masks if rescale is part of image_converter?

[DavidLandup0]

Yes, it would. The ImageConverter should only resize - rescaling and normalization is enabled in the preprocessor by default, since it's required for the weights to function as intended. This should at the very least be documented, or ideally set as the default to avoid confusion.

If I add a default ImageConverter with just resizing and export the SegFormer model as a task again - the ImageConverter's from_preset() would restore it to that state as well? In that case, it may make the most sense to export the SegFormer models with the ImageConverter configs already baked in. Thoughts?