5Grains commited on
Commit
d916009
·
1 Parent(s): 448a801

Delete app.py

Browse files
Files changed (1) hide show
  1. app.py +0 -111
app.py DELETED
@@ -1,111 +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
- # Use a pipeline as a high-level helper
9
- from transformers import pipeline
10
-
11
- pipe = pipeline("image-segmentation", model="nvidia/segformer-b0-finetuned-cityscapes-512-1024")
12
- # Load model directly
13
- from transformers import AutoFeatureExtractor, SegformerForSemanticSegmentation
14
-
15
- extractor = AutoFeatureExtractor.from_pretrained("nvidia/segformer-b0-finetuned-cityscapes-512-1024")
16
- model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-cityscapes-512-1024")
17
- def ade_palette():
18
- """ADE20K palette that maps each class to RGB values."""
19
- return [
20
- [255, 0, 0],
21
- [255, 127, 127],
22
- [255, 127, 0],
23
- [255, 255, 0],
24
- [127, 255, 0],
25
- [0, 255, 0],
26
- [127, 255, 127],
27
- [0, 255, 127],
28
- [0, 255, 255],
29
- [0, 127, 255],
30
- [0, 0, 255],
31
- [127, 127, 255],
32
- [127, 0, 255],
33
- [255, 0, 255],
34
- [255, 0, 127],
35
- [0, 0, 0],
36
- [127, 127, 127],
37
- [255, 255, 255]
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.astype(np.uint8)) # Convert to uint8 before displaying
63
- plt.axis('off')
64
-
65
- LABEL_NAMES = np.asarray(labels_list)
66
- FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1)
67
- FULL_COLOR_MAP = label_to_color_image(FULL_LABEL_MAP)
68
-
69
- unique_labels = np.unique(seg.numpy().astype("uint8"))
70
- ax = plt.subplot(grid_spec[1])
71
- plt.imshow(FULL_COLOR_MAP[unique_labels].astype(np.uint8), interpolation="nearest")
72
- ax.yaxis.tick_right()
73
- plt.yticks(range(len(unique_labels)), LABEL_NAMES[unique_labels])
74
- plt.xticks([], [])
75
- ax.tick_params(width=0.0, labelsize=25)
76
- return fig
77
-
78
- def sepia(input_img):
79
- input_img = Image.fromarray(input_img)
80
-
81
- inputs = extractor(images=input_img, return_tensors="tf")
82
- outputs = model(**inputs)
83
- logits = outputs.logits
84
-
85
- logits = tf.transpose(logits, [0, 2, 3, 1])
86
- logits = tf.image.resize(
87
- logits, input_img.size[::-1]
88
- ) # We reverse the shape of `image` because `image.size` returns width and height.
89
- seg = tf.math.argmax(logits, axis=-1)[0]
90
-
91
- color_seg = np.zeros(
92
- (seg.shape[0], seg.shape[1], 3), dtype=np.uint8
93
- ) # height, width, 3
94
- for label, color in enumerate(colormap):
95
- color_seg[seg.numpy() == label, :] = color
96
-
97
- # Show image + mask
98
- pred_img = np.array(input_img) * 0.5 + color_seg * 0.5
99
- pred_img = pred_img.astype(np.uint8)
100
-
101
- fig = draw_plot(pred_img, seg)
102
- return fig
103
-
104
- demo = gr.Interface(fn=sepia,
105
- inputs=gr.Image(shape=(400, 600)),
106
- outputs=['plot'],
107
- examples=["person-1.jpg", "person-2.jpg", "person-3.jpg", "person-4.jpg", "person-5.jpg", ],
108
- allow_flagging='never')
109
-
110
-
111
- demo.launch()