satpalsr commited on
Commit
5831693
1 Parent(s): f5ac2ef

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +63 -0
app.py ADDED
@@ -0,0 +1,63 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import tensorflow as tf
3
+ import gradio as gr
4
+ from huggingface_hub import from_pretrained_keras
5
+ import cv2
6
+
7
+ model = from_pretrained_keras("satpalsr/deeplabv3p-resnet50")
8
+
9
+ colormap = np.array([[0,0,0], [31,119,180], [44,160,44], [44, 127, 125], [52, 225, 143],
10
+ [217, 222, 163], [254, 128, 37], [130, 162, 128], [121, 7, 166], [136, 183, 248],
11
+ [85, 1, 76], [22, 23, 62], [159, 50, 15], [101, 93, 152], [252, 229, 92],
12
+ [167, 173, 17], [218, 252, 252], [238, 126, 197], [116, 157, 140], [214, 220, 252]], dtype=np.uint8)
13
+
14
+ img_size = 512
15
+
16
+ def read_image(image):
17
+ image = tf.convert_to_tensor(image)
18
+ image.set_shape([None, None, 3])
19
+ image = tf.image.resize(images=image, size=[img_size, img_size])
20
+ image = image / 127.5 - 1
21
+ return image
22
+
23
+ def infer(model, image_tensor):
24
+ predictions = model.predict(np.expand_dims((image_tensor), axis=0))
25
+ predictions = np.squeeze(predictions)
26
+ predictions = np.argmax(predictions, axis=2)
27
+ return predictions
28
+
29
+ def decode_segmentation_masks(mask, colormap, n_classes):
30
+ r = np.zeros_like(mask).astype(np.uint8)
31
+ g = np.zeros_like(mask).astype(np.uint8)
32
+ b = np.zeros_like(mask).astype(np.uint8)
33
+ for l in range(0, n_classes):
34
+ idx = mask == l
35
+ r[idx] = colormap[l, 0]
36
+ g[idx] = colormap[l, 1]
37
+ b[idx] = colormap[l, 2]
38
+ rgb = np.stack([r, g, b], axis=2)
39
+ return rgb
40
+
41
+ def get_overlay(image, colored_mask):
42
+ image = tf.keras.preprocessing.image.array_to_img(image)
43
+ image = np.array(image).astype(np.uint8)
44
+ overlay = cv2.addWeighted(image, 0.35, colored_mask, 0.65, 0)
45
+ return overlay
46
+
47
+ def segmentation(input_image):
48
+ image_tensor = read_image(input_image)
49
+ prediction_mask = infer(image_tensor=image_tensor, model=model)
50
+ prediction_colormap = decode_segmentation_masks(prediction_mask, colormap, 20)
51
+ overlay = get_overlay(image_tensor, prediction_colormap)
52
+ return (overlay, prediction_colormap)
53
+
54
+ i = gr.inputs.Image()
55
+ o = [gr.outputs.Image(), gr.outputs.Image()]
56
+
57
+ examples = [["example_image_1.jpg"], ["example_image_2.jpg"], ["example_image_3.jpg"]]
58
+ title = "Human Part Segmentation"
59
+ description = "Upload an image to segment out different human parts."
60
+
61
+ article = "<div style='text-align: center;'><a href='https://twitter.com/SatpalPatawat' target='_blank'>Space by Satpal Singh Rathore</a><br><a href='https://keras.io/examples/vision/deeplabv3_plus/' target='_blank'>Keras example by Soumik Rakshit</a></div>"
62
+ gr.Interface(segmentation, i, o, examples=examples, allow_flagging=False, analytics_enabled=False,
63
+ title=title, description=description, article=article).launch(enable_queue=True)