finish

app.py

@@ -59,29 +59,13 @@ def label_to_color_image(label):
 
 
 def draw_plot_with_labels(pred_img, seg):
-    fig = plt.figure(figsize=(20, 15))
-
-    grid_spec = gridspec.GridSpec(1, 2, width_ratios=[6, 1])
-
-    plt.subplot(grid_spec[0])
-    plt.imshow(pred_img)
-    plt.axis('off')
+    fig, ax = plt.subplots(1, 1, figsize=(20, 15))
+    ax.imshow(pred_img)
+    ax.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)
-
-    unique_labels = np.unique(seg.numpy().astype("uint8"))
-    ax = plt.subplot(grid_spec[1])
-    plt.imshow(FULL_COLOR_MAP[unique_labels].astype(np.uint8), interpolation="nearest")
-    ax.yaxis.tick_right()
-    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):
@@ -89,7 +73,7 @@ def draw_plot_with_labels(pred_img, seg):
             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)
+            ax.text(x, y, label_name, fontsize=20, color=tuple(color))
 
     return fig