Matthijs Hollemans commited on
Commit
6a36cd0
β€’
1 Parent(s): c304fb7

overlay mask on original image

Browse files
Files changed (2) hide show
  1. README.md +2 -2
  2. app.py +31 -14
README.md CHANGED
@@ -1,8 +1,8 @@
1
  ---
2
  title: MobileViT Deeplab Demo
3
  emoji: πŸ•
4
- colorFrom: red
5
- colorTo: pink
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, do_center_crop=False, size=(512, 512))
10
  model = MobileViTForSemanticSegmentation.from_pretrained(model_checkpoint).eval()
11
 
12
 
13
- # From https://gist.github.com/kaixin96/457cc3d3be699f1f5b2fd4cdb638d4b4
14
- palette = np.array([
15
- [ 0, 0, 0], [128, 0, 0], [ 0, 128, 0], [128, 128, 0], [ 0, 0, 128],
16
- [128, 0, 128], [ 0, 128, 128], [128, 128, 128], [ 64, 0, 0], [192, 0, 0],
17
- [ 64, 128, 0], [192, 128, 0], [ 64, 0, 128], [192, 0, 128], [ 64, 128, 128],
18
- [192, 128, 128], [ 0, 64, 0], [128, 64, 0], [ 0, 192, 0], [128, 192, 0],
19
- [ 0, 64, 128]], dtype=np.uint8)
 
 
 
 
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
- # TODO: overlay mask on image?
 
 
 
 
 
 
 
 
 
 
 
36
 
37
- out_image = Image.fromarray(colored)
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