EUNSEO56 commited on
Commit
7995bcb
·
1 Parent(s): 66da7ef
Files changed (1) hide show
  1. app.py +0 -133
app.py DELETED
@@ -1,133 +0,0 @@
1
- import gradio as gr
2
-
3
- from matplotlib import gridspec
4
- import matplotlib.pyplot as plt
5
- import numpy as np
6
- from PIL import Image
7
- import tensorflow as tf
8
- from transformers import SegformerFeatureExtractor, TFSegformerForSemanticSegmentation
9
-
10
- feature_extractor = SegformerFeatureExtractor.from_pretrained(
11
- "nickmuchi/segformer-b4-finetuned-segments-sidewalk"
12
- )
13
- model = TFSegformerForSemanticSegmentation.from_pretrained(
14
- "nickmuchi/segformer-b4-finetuned-segments-sidewalk",
15
- from_pt=True
16
-
17
- )
18
-
19
- def ade_palette():
20
- """ADE20K palette that maps each class to RGB values."""
21
- return [
22
- [204, 87, 92],
23
- [112, 185, 212],
24
- [45, 189, 106],
25
- [234, 123, 67],
26
- [78, 56, 123],
27
- [210, 32, 89],
28
- [90, 180, 56],
29
- [155, 102, 200],
30
- [33, 147, 176],
31
- [255, 183, 76],
32
- [67, 123, 89],
33
- [190, 60, 45],
34
- [134, 112, 200],
35
- [56, 45, 189],
36
- [200, 56, 123],
37
- [87, 92, 204],
38
- [120, 56, 123],
39
- [45, 78, 123],
40
- [156, 200, 56],
41
- [32, 90, 210],
42
- [56, 123, 67],
43
- [180, 56, 123],
44
- [123, 67, 45],
45
- [45, 134, 200],
46
- [67, 56, 123],
47
- [78, 123, 67],
48
- [32, 210, 90],
49
- [45, 56, 189],
50
- [123, 56, 123],
51
- [56, 156, 200],
52
- [189, 56, 45],
53
- [112, 200, 56],
54
- [56, 123, 45],
55
- [200, 32, 90],
56
- [255, 255, 0],
57
- ]
58
-
59
- labels_list = []
60
-
61
- with open(r'labels.txt', 'r') as fp:
62
- for line in fp:
63
- labels_list.append(line[:-1])
64
-
65
- colormap = np.asarray(ade_palette())
66
-
67
-
68
- def label_to_color_image(label):
69
- if label.ndim != 2:
70
- raise ValueError("Expect 2-D input label")
71
-
72
- if np.max(label) >= len(colormap):
73
- raise ValueError("label value too large.")
74
- return colormap[label]
75
-
76
- def draw_plot(pred_img, seg):
77
- fig = plt.figure(figsize=(20, 15))
78
-
79
- grid_spec = gridspec.GridSpec(1, 2, width_ratios=[6, 1])
80
-
81
- plt.subplot(grid_spec[0])
82
- plt.imshow(pred_img)
83
- plt.axis('off')
84
- LABEL_NAMES = np.asarray(labels_list)
85
- FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1)
86
- FULL_COLOR_MAP = label_to_color_image(FULL_LABEL_MAP)
87
-
88
- unique_labels = np.unique(seg.numpy().astype("uint8"))
89
- ax = plt.subplot(grid_spec[1])
90
- plt.imshow(FULL_COLOR_MAP[unique_labels].astype(np.uint8), interpolation="nearest")
91
- ax.yaxis.tick_right()
92
- plt.yticks(range(len(unique_labels)), LABEL_NAMES[unique_labels])
93
- plt.xticks([], [])
94
- ax.tick_params(width=0.0, labelsize=25)
95
- return fig
96
-
97
-
98
- def sepia(input_img):
99
- input_img = Image.fromarray(input_img)
100
-
101
- inputs = feature_extractor(images=input_img, return_tensors="tf")
102
- outputs = model(**inputs)
103
- logits = outputs.logits
104
-
105
- logits = tf.transpose(logits, [0, 2, 3, 1])
106
- logits = tf.image.resize(
107
- logits, input_img.size[::-1]
108
- ) # We reverse the shape of `image` because `image.size` returns width and height.
109
- seg = tf.math.argmax(logits, axis=-1)[0]
110
-
111
- color_seg = np.zeros(
112
- (seg.shape[0], seg.shape[1], 3), dtype=np.uint8
113
- ) # height, width, 3
114
- for label, color in enumerate(colormap):
115
- color_seg[seg.numpy() == label, :] = color
116
-
117
- # Show image + mask
118
- pred_img = np.array(input_img) * 0.5 + color_seg * 0.5
119
- pred_img = pred_img.astype(np.uint8)
120
-
121
- fig = draw_plot(pred_img, seg)
122
- return fig
123
-
124
- demo = gr.Interface(fn=sepia,
125
- inputs=gr.Image(shape=(400, 600)),
126
- outputs=['plot'],
127
- examples=["side-1.jpg", "side-2.jpg", "side-3.jpg", "side-4.jpg", "side-5.jpg", "side-6.jpg"],
128
- allow_flagging='never')
129
-
130
-
131
- demo.launch()
132
-
133
-