Karin0616 commited on
Commit
b89aec9
โ€ข
1 Parent(s): 37bdf9b

i cant do any more...

Browse files
Files changed (1) hide show
  1. app.py +70 -33
app.py CHANGED
@@ -15,7 +15,7 @@ model = TFSegformerForSemanticSegmentation.from_pretrained(
15
  "nvidia/segformer-b5-finetuned-cityscapes-1024-1024"
16
  )
17
 
18
- def palette():
19
 
20
  return [
21
  [204, 87, 92], # road (Reddish)
@@ -45,7 +45,8 @@ labels_list = []
45
  with open(r'labels.txt', 'r') as fp:
46
  for line in fp:
47
  labels_list.append(line[:-1])
48
- colormap = np.asarray(palette())
 
49
 
50
  def label_to_color_image(label):
51
  if label.ndim != 2:
@@ -55,34 +56,70 @@ def label_to_color_image(label):
55
  raise ValueError("label value too large.")
56
  return colormap[label]
57
 
58
- # Gradio ์ธํ„ฐํŽ˜์ด์Šค ์ •์˜
59
- iface = gr.Interface(
60
- fn=predict_segmentation,
61
- inputs=gr.Image(shape=(564,846)),
62
- outputs="image",
63
- examples=["city1.jpg","city2.jpg","city3.jpg"],
64
- )
65
- iface.launch()
66
-
67
- # ์ด๋ฏธ์ง€ ์„ธ๊ทธ๋ฉ˜ํ…Œ์ด์…˜ ํ•จ์ˆ˜ ์ •์˜
68
- def predict_segmentation(image, model):
69
- # ์ด๋ฏธ์ง€ ๋ณ€ํ™˜
70
- image = Image.fromarray(image.astype('uint8'), 'RGB')
71
- image = image.resize((1024, 1024)) # ๋ชจ๋ธ์˜ ์ž…๋ ฅ ํฌ๊ธฐ์— ๋งž๊ฒŒ ์กฐ์ ˆ
72
- image_array = tf.keras.preprocessing.image.img_to_array(image)
73
- image_array = tf.expand_dims(image_array, 0)
74
-
75
- # ๋ชจ๋ธ ์ถ”๋ก 
76
- predictions = model(image_array)["output_0"]
77
-
78
- # ๋ ˆ์ด๋ธ”๋ณ„ ์ƒ‰์ƒ ๋งคํ•‘
79
- segmented_image = tf.zeros_like(predictions)
80
- for label, color in label_colors.items():
81
- mask = tf.reduce_all(tf.equal(predictions, color), axis=-1, keepdims=True)
82
- for i in range(3):
83
- segmented_image += tf.cast(mask, tf.float32) * tf.constant(color[i], dtype=tf.float32)
84
-
85
- # ์ด๋ฏธ์ง€ ๋ฆฌํ„ด
86
- segmented_image = tf.cast(segmented_image, tf.uint8)
87
- segmented_image = tf.image.resize(segmented_image, [image.height, image.width])
88
- return segmented_image.numpy()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
15
  "nvidia/segformer-b5-finetuned-cityscapes-1024-1024"
16
  )
17
 
18
+ def ade_palette():
19
 
20
  return [
21
  [204, 87, 92], # road (Reddish)
 
45
  with open(r'labels.txt', 'r') as fp:
46
  for line in fp:
47
  labels_list.append(line[:-1])
48
+
49
+ colormap = np.asarray(ade_palette())
50
 
51
  def label_to_color_image(label):
52
  if label.ndim != 2:
 
56
  raise ValueError("label value too large.")
57
  return colormap[label]
58
 
59
+ def draw_plot(pred_img, seg):
60
+ fig = plt.figure(figsize=(20, 15))
61
+
62
+ grid_spec = gridspec.GridSpec(1, 2, width_ratios=[6, 1])
63
+
64
+ plt.subplot(grid_spec[0])
65
+ plt.imshow(pred_img)
66
+ plt.axis('off')
67
+ LABEL_NAMES = np.asarray(labels_list)
68
+ FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1)
69
+ FULL_COLOR_MAP = label_to_color_image(FULL_LABEL_MAP)
70
+
71
+ unique_labels = np.unique(seg.numpy().astype("uint8"))
72
+
73
+ ax = plt.subplot(grid_spec[1])
74
+ plt.imshow(FULL_COLOR_MAP[unique_labels].astype(np.uint8), interpolation="nearest")
75
+
76
+ ax.yaxis.tick_left()
77
+ plt.yticks(range(len(unique_labels)), LABEL_NAMES[unique_labels])
78
+ plt.xticks([], [])
79
+ ax.tick_params(width=0.0, labelsize=27)
80
+ return fig
81
+
82
+ def sepia(input_img):
83
+ input_img = Image.fromarray(input_img)
84
+
85
+ inputs = feature_extractor(images=input_img, return_tensors="tf")
86
+ outputs = model(**inputs)
87
+ logits = outputs.logits
88
+
89
+ logits = tf.transpose(logits, [0, 2, 3, 1])
90
+ logits = tf.image.resize(
91
+ logits, input_img.size[::-1]
92
+ ) # We reverse the shape of `image` because `image.size` returns width and height.
93
+ seg = tf.math.argmax(logits, axis=-1)[0]
94
+
95
+ color_seg = np.zeros(
96
+ (seg.shape[0], seg.shape[1], 3), dtype=np.uint8
97
+ ) # height, width, 3
98
+ for label, color in enumerate(colormap):
99
+ color_seg[seg.numpy() == label, :] = color
100
+
101
+ # Show image + mask
102
+ pred_img = np.array(input_img) * 0.5 + color_seg * 0.5
103
+ pred_img = pred_img.astype(np.uint8)
104
+
105
+ fig = draw_plot(pred_img, seg)
106
+ return fig
107
+
108
+
109
+ demo = gr.Interface(fn=sepia,
110
+ inputs=gr.Image(shape=(564,846)),
111
+ outputs=['plot'],
112
+ live=True,
113
+ examples=["city1.jpg","city2.jpg","city3.jpg"],
114
+ allow_flagging='never',
115
+ title="City Image Segmentation Model",
116
+ theme="huggingfacedark",
117
+ description=["This model is a high-performance city image segmentation model based on the Segformer architecture provided by NVIDIA. Specifically, the "segformer-b5" model, trained on the Cityscapes dataset, excels at performing intricate segmentation on high-resolution images of 1024x1024 pixels. It accurately identifies various urban elements such as roads, buildings, pedestrians, providing visually rich segmentation results.",
118
+ "This is a machine learning activity project at Kyunggi University."],
119
+
120
+
121
+ )
122
+
123
+
124
+ demo.launch()
125
+