Spaces:
Runtime error
Runtime error
Matthijs Hollemans
commited on
Commit
β’
6a36cd0
1
Parent(s):
c304fb7
overlay mask on original image
Browse files
README.md
CHANGED
@@ -1,8 +1,8 @@
|
|
1 |
---
|
2 |
title: MobileViT Deeplab Demo
|
3 |
emoji: π
|
4 |
-
colorFrom:
|
5 |
-
colorTo:
|
6 |
sdk: gradio
|
7 |
sdk_version: 3.0.24
|
8 |
app_file: app.py
|
|
|
1 |
---
|
2 |
title: MobileViT Deeplab Demo
|
3 |
emoji: π
|
4 |
+
colorFrom: black
|
5 |
+
colorTo: black
|
6 |
sdk: gradio
|
7 |
sdk_version: 3.0.24
|
8 |
app_file: app.py
|
app.py
CHANGED
@@ -6,17 +6,21 @@ import torch
|
|
6 |
from transformers import MobileViTFeatureExtractor, MobileViTForSemanticSegmentation
|
7 |
|
8 |
model_checkpoint = "apple/deeplabv3-mobilevit-small"
|
9 |
-
feature_extractor = MobileViTFeatureExtractor.from_pretrained(model_checkpoint
|
10 |
model = MobileViTForSemanticSegmentation.from_pretrained(model_checkpoint).eval()
|
11 |
|
12 |
|
13 |
-
|
14 |
-
|
15 |
-
[ 0, 0, 0], [
|
16 |
-
[
|
17 |
-
[
|
18 |
-
[
|
19 |
-
[ 0,
|
|
|
|
|
|
|
|
|
20 |
|
21 |
|
22 |
def predict(image):
|
@@ -24,6 +28,11 @@ def predict(image):
|
|
24 |
inputs = feature_extractor(image, return_tensors="pt")
|
25 |
outputs = model(**inputs)
|
26 |
|
|
|
|
|
|
|
|
|
|
|
27 |
classes = outputs.logits.argmax(1).squeeze().numpy().astype(np.uint8)
|
28 |
|
29 |
# Super slow method but it works
|
@@ -32,20 +41,28 @@ def predict(image):
|
|
32 |
for x in range(classes.shape[1]):
|
33 |
colored[y, x] = palette[classes[y, x]]
|
34 |
|
35 |
-
#
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
36 |
|
37 |
-
|
38 |
-
out_image = out_image.resize((image.shape[1], image.shape[0]), resample=Image.NEAREST)
|
39 |
-
return out_image
|
40 |
|
41 |
|
42 |
gr.Interface(
|
43 |
fn=predict,
|
44 |
inputs=gr.inputs.Image(label="Upload image"),
|
45 |
-
outputs=gr.outputs.Image(),
|
46 |
title="Semantic Segmentation with MobileViT and DeepLabV3",
|
47 |
).launch()
|
48 |
|
49 |
|
50 |
# TODO: combo box with some example images
|
51 |
-
# TODO: combo box with classes to show on the output, if none then do argmax
|
|
|
6 |
from transformers import MobileViTFeatureExtractor, MobileViTForSemanticSegmentation
|
7 |
|
8 |
model_checkpoint = "apple/deeplabv3-mobilevit-small"
|
9 |
+
feature_extractor = MobileViTFeatureExtractor.from_pretrained(model_checkpoint) #, do_center_crop=False, size=(512, 512))
|
10 |
model = MobileViTForSemanticSegmentation.from_pretrained(model_checkpoint).eval()
|
11 |
|
12 |
|
13 |
+
palette = np.array(
|
14 |
+
[
|
15 |
+
[ 0, 0, 0], [192, 0, 0], [ 0, 192, 0], [192, 192, 0],
|
16 |
+
[ 0, 0, 192], [192, 0, 192], [ 0, 192, 192], [192, 192, 192],
|
17 |
+
[128, 0, 0], [255, 0, 0], [128, 192, 0], [255, 192, 0],
|
18 |
+
[128, 0, 192], [255, 0, 192], [128, 192, 192], [255, 192, 192],
|
19 |
+
[ 0, 128, 0], [192, 128, 0], [ 0, 255, 0], [192, 255, 0],
|
20 |
+
[ 0, 128, 192]
|
21 |
+
],
|
22 |
+
dtype=np.uint8)
|
23 |
+
|
24 |
|
25 |
|
26 |
def predict(image):
|
|
|
28 |
inputs = feature_extractor(image, return_tensors="pt")
|
29 |
outputs = model(**inputs)
|
30 |
|
31 |
+
# Get preprocessed image. The pixel values don't need to be unnormalized
|
32 |
+
# for this particular model.
|
33 |
+
resized = (inputs["pixel_values"].numpy().squeeze().transpose(1, 2, 0)[..., ::-1] * 255).astype(np.uint8)
|
34 |
+
|
35 |
+
# Class predictions for each pixel.
|
36 |
classes = outputs.logits.argmax(1).squeeze().numpy().astype(np.uint8)
|
37 |
|
38 |
# Super slow method but it works
|
|
|
41 |
for x in range(classes.shape[1]):
|
42 |
colored[y, x] = palette[classes[y, x]]
|
43 |
|
44 |
+
# Resize predictions to input size (not original size).
|
45 |
+
colored = Image.fromarray(colored)
|
46 |
+
colored = colored.resize((resized.shape[1], resized.shape[0]), resample=Image.NEAREST)
|
47 |
+
|
48 |
+
# Keep everything that is not background.
|
49 |
+
mask = (classes != 0) * 255
|
50 |
+
mask = Image.fromarray(mask.astype(np.uint8)).convert("RGB")
|
51 |
+
mask = mask.resize((resized.shape[1], resized.shape[0]), resample=Image.NEAREST)
|
52 |
+
|
53 |
+
# Blend with the input image.
|
54 |
+
resized = Image.fromarray(resized)
|
55 |
+
highlighted = Image.blend(resized, mask, 0.4)
|
56 |
|
57 |
+
return colored, highlighted
|
|
|
|
|
58 |
|
59 |
|
60 |
gr.Interface(
|
61 |
fn=predict,
|
62 |
inputs=gr.inputs.Image(label="Upload image"),
|
63 |
+
outputs=[gr.outputs.Image(label="Classes"), gr.outputs.Image(label="Highlighted")],
|
64 |
title="Semantic Segmentation with MobileViT and DeepLabV3",
|
65 |
).launch()
|
66 |
|
67 |
|
68 |
# TODO: combo box with some example images
|
|