finish

app.py

@@ -3,7 +3,7 @@ import gradio as gr
 from matplotlib import gridspec
 import matplotlib.pyplot as plt
 import numpy as np
-from PIL import Image, ImageDraw
+from PIL import Image
 import tensorflow as tf
 from transformers import SegformerFeatureExtractor, TFSegformerForSemanticSegmentation
 
@@ -58,7 +58,7 @@ def label_to_color_image(label):
     return colormap[label]
 
 
-def draw_plot_with_labels(pred_img, seg):
+def draw_plot(pred_img, seg):
     fig = plt.figure(figsize=(20, 15))
 
     grid_spec = gridspec.GridSpec(1, 2, width_ratios=[6, 1])
@@ -66,8 +66,6 @@ def draw_plot_with_labels(pred_img, seg):
     plt.subplot(grid_spec[0])
     plt.imshow(pred_img)
     plt.axis('off')
-
-    # 라벨 이름을 추가하기 위한 코드
     LABEL_NAMES = np.asarray(labels_list)
     FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1)
     FULL_COLOR_MAP = label_to_color_image(FULL_LABEL_MAP)
@@ -79,18 +77,6 @@ def draw_plot_with_labels(pred_img, seg):
     plt.yticks(range(len(unique_labels)), LABEL_NAMES[unique_labels])
     plt.xticks([], [])
     ax.tick_params(width=0.0, labelsize=25)
-
-    # 라벨 이름 텍스트 추가
-    draw = ImageDraw.Draw(pred_img)
-    for label, color in enumerate(colormap):
-        mask = seg.numpy() == label
-        if np.any(mask):
-            y, x = np.where(mask)
-            y = np.mean(y).astype(int)
-            x = np.mean(x).astype(int)
-            label_name = LABEL_NAMES[label]
-            draw.text((x, y), label_name, fill=tuple(color), fontsize=20)
-
     return fig
 
 
@@ -105,6 +91,7 @@ def sepia(input_img):
         logits, input_img.size[::-1]
     )  # We reverse the shape of `image` because `image.size` returns width and height.
     seg = tf.math.argmax(logits, axis=-1)[0]
+
     color_seg = np.zeros(
         (seg.shape[0], seg.shape[1], 3), dtype=np.uint8
     )  # height, width, 3
@@ -115,13 +102,13 @@ def sepia(input_img):
     pred_img = np.array(input_img) * 0.5 + color_seg * 0.5
     pred_img = pred_img.astype(np.uint8)
 
-    fig = draw_plot_with_labels(pred_img, seg)
+    fig = draw_plot(pred_img, seg)
     return fig
 
 
 demo = gr.Interface(fn=sepia,
                     inputs=gr.Image(shape=(1024, 1024)),
-                    outputs=['plot'],
+                    outputs=['plot', "text"],
                     examples=["city-1.jpg", "city-2.jpg", "city-3.jpg", "city-4.jpg", "city-5.jpg"],
                     allow_flagging='never')