finish

app.py

@@ -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,17 @@ 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)
+
+    # 라벨 이름 텍스트 추가
+    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]
+            plt.text(x, y, label_name, fontsize=20, color=tuple(color))
+
     return fig
 
 
@@ -91,7 +104,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
@@ -108,7 +120,7 @@ def sepia(input_img):
 
 demo = gr.Interface(fn=sepia,
                     inputs=gr.Image(shape=(1024, 1024)),
-                    outputs=['plot', "text"],
+                    outputs=['plot'],
                     examples=["city-1.jpg", "city-2.jpg", "city-3.jpg", "city-4.jpg", "city-5.jpg"],
                     allow_flagging='never')