overlay mask on original image

README.md

@@ -1,8 +1,8 @@
 ---
 title: MobileViT Deeplab Demo
 emoji: 🍕
-colorFrom: red
-colorTo: pink
+colorFrom: black
+colorTo: black
 sdk: gradio
 sdk_version: 3.0.24
 app_file: app.py

app.py

@@ -6,17 +6,21 @@ import torch
 from transformers import MobileViTFeatureExtractor, MobileViTForSemanticSegmentation
 
 model_checkpoint = "apple/deeplabv3-mobilevit-small"
-feature_extractor = MobileViTFeatureExtractor.from_pretrained(model_checkpoint, do_center_crop=False, size=(512, 512))
+feature_extractor = MobileViTFeatureExtractor.from_pretrained(model_checkpoint) #, do_center_crop=False, size=(512, 512))
 model = MobileViTForSemanticSegmentation.from_pretrained(model_checkpoint).eval()
 
 
-# From https://gist.github.com/kaixin96/457cc3d3be699f1f5b2fd4cdb638d4b4
-palette = np.array([
-    [  0,   0,   0], [128,   0,   0], [  0, 128,   0], [128, 128,   0], [  0,   0, 128],
-    [128,   0, 128], [  0, 128, 128], [128, 128, 128], [ 64,   0,   0], [192,   0,   0],
-    [ 64, 128,   0], [192, 128,   0], [ 64,   0, 128], [192,   0, 128], [ 64, 128, 128],
-    [192, 128, 128], [  0,  64,   0], [128,  64,   0], [  0, 192,   0], [128, 192,   0],
-    [ 0, 64, 128]], dtype=np.uint8)
+palette = np.array(
+[
+    [  0,   0,   0], [192,   0,   0], [  0, 192,   0], [192, 192,   0],
+    [  0,   0, 192], [192,   0, 192], [  0, 192, 192], [192, 192, 192],
+    [128,   0,   0], [255,   0,   0], [128, 192,   0], [255, 192,   0],
+    [128,   0, 192], [255,   0, 192], [128, 192, 192], [255, 192, 192],
+    [  0, 128,   0], [192, 128,   0], [  0, 255,   0], [192, 255,   0],
+    [  0, 128, 192]
+],
+dtype=np.uint8)
+
 
 
 def predict(image):
@@ -24,6 +28,11 @@ def predict(image):
         inputs = feature_extractor(image, return_tensors="pt")
         outputs = model(**inputs)
 
+    # Get preprocessed image. The pixel values don't need to be unnormalized
+    # for this particular model.
+    resized = (inputs["pixel_values"].numpy().squeeze().transpose(1, 2, 0)[..., ::-1] * 255).astype(np.uint8)
+
+    # Class predictions for each pixel.
     classes = outputs.logits.argmax(1).squeeze().numpy().astype(np.uint8)
 
     # Super slow method but it works
@@ -32,20 +41,28 @@ def predict(image):
         for x in range(classes.shape[1]):
             colored[y, x] = palette[classes[y, x]]
 
-    # TODO: overlay mask on image?
+    # Resize predictions to input size (not original size).
+    colored = Image.fromarray(colored)
+    colored = colored.resize((resized.shape[1], resized.shape[0]), resample=Image.NEAREST)
+
+    # Keep everything that is not background.
+    mask = (classes != 0) * 255
+    mask = Image.fromarray(mask.astype(np.uint8)).convert("RGB")
+    mask = mask.resize((resized.shape[1], resized.shape[0]), resample=Image.NEAREST)
+
+    # Blend with the input image.
+    resized = Image.fromarray(resized)
+    highlighted = Image.blend(resized, mask, 0.4)
 
-    out_image = Image.fromarray(colored)
-    out_image = out_image.resize((image.shape[1], image.shape[0]), resample=Image.NEAREST)
-    return out_image
+    return colored, highlighted
 
 
 gr.Interface(
     fn=predict,
     inputs=gr.inputs.Image(label="Upload image"),
-    outputs=gr.outputs.Image(),
+    outputs=[gr.outputs.Image(label="Classes"), gr.outputs.Image(label="Highlighted")],
     title="Semantic Segmentation with MobileViT and DeepLabV3",
 ).launch()
 
 
 # TODO: combo box with some example images
-# TODO: combo box with classes to show on the output, if none then do argmax