Update app.py

app.py

@@ -72,45 +72,20 @@ def label_to_color_image(label):
         raise ValueError("label value too large.")
     return colormap[label]
 
-def draw_plot(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')
-    LABEL_NAMES = np.asarray(labels_list)
-    FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1)
-    FULL_COLOR
-
-def get_labels():
-    with open(r'labels.txt', 'r') as fp:
-        labels = []
-        for line in fp:
-            labels.append(line[:-1])
-    return labels
-
-# 레이블 목록 가져오기
-labels = get_labels()
-
-# 버튼 생성
-buttons = []
-for label in labels:
-    button = gr.Button(label)
-    def display_seg(label):
-        seg = np.where(seg == labels.index(label), 1, 0)
-        pred_img = np.array(input_img) * 0.5 + seg * 0.5
-        pred_img = pred_img.astype(np.uint8)
-        fig = draw_plot(pred_img, seg)
-        return fig
-
-    button.click(display_seg, label)
-    buttons.append(button)
-
-demo.components.append(buttons)
-demo.launch()
-
-
+# Create a Gradio interface
+def image_segmentation(input_img):
+    # Perform image segmentation
+    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)
+
+    return pred_img
+
+gr.Interface(
+    fn=image_segmentation,
+    inputs="image",
+    outputs="image"
+).launch()