han-byeol commited on
Commit
6e69a3d
1 Parent(s): 763e28a

Delete app.py

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