Update app.py

app.py

@@ -72,33 +72,16 @@ def label_to_color_image(label):
         raise ValueError("label value too large.")
     return colormap[label]
     
-# Gradio 인터페이스에 사용될 이미지 세그멘테이션 함수를 정의합니다.
-def image_segmentation(input_img):
-    input_img = np.array(input_img)  # 이미지를 넘파이 배열로 변환
+def draw_plot(pred_img, seg, show_seg=False):
+    fig = plt.figure(figsize=(20, 15))
 
-    # 이미지 세그멘테이션 수행
-    inputs = feature_extractor(images=input_img, return_tensors="pt")
-    outputs = model(**inputs)
-    pred_label = np.argmax(outputs.logits[0].numpy(), axis=0)
-    pred_img = label_to_color_image(pred_label)
+    grid_spec = gridspec.GridSpec(1, 2, width_ratios=[6, 1])
 
-    return pred_img
+    plt.subplot(grid_spec[0])
+    plt.imshow(pred_img)
+    plt.axis('off')
 
-def label_to_color_image(label):
-    if label.ndim != 2:
-        raise ValueError("Expect 2-D input label")
-
-    if np.max(label) >= len(colormap):
-        raise ValueError("label value too large.")
-    return colormap[label]
-
-# Gradio 인터페이스 생성
-iface = gr.Interface(
-    fn=image_segmentation,
-    inputs="image",
-    outputs="image"
-)
-
-# 인터페이스 실행
-iface.launch()
+    if show_seg:
+        unique_labels = np.unique(seg.numpy().astype("uint8"))
+        ax = plt.subplot(grid_spec[1])