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
+from PIL import Image, ImageDraw
 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(pred_img, seg):
+def draw_plot_with_labels(pred_img, seg):
     fig = plt.figure(figsize=(20, 15))
 
     grid_spec = gridspec.GridSpec(1, 2, width_ratios=[6, 1])
@@ -66,6 +66,8 @@ def draw_plot(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)
@@ -77,6 +79,18 @@ def draw_plot(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
 
 
@@ -91,7 +105,6 @@ 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
@@ -102,7 +115,7 @@ 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(pred_img, seg)
+    fig = draw_plot_with_labels(pred_img, seg)
     return fig