Update app.py

app.py

@@ -71,17 +71,24 @@ def label_to_color_image(label):
     if np.max(label) >= len(colormap):
         raise ValueError("label value too large.")
     return colormap[label]
-
-# Create a Gradio interface
+    
+# Gradio 인터페이스 생성
 def image_segmentation(input_img):
-    # Perform image segmentation
+    # 이미지를 넘파이 배열로 변환
+    input_img = np.array(input_img)
+
+    # 이미지 세그멘테이션 수행
     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)
+
+    # 예측된 레이블을 이미지로 변환
+    pred_img = Image.fromarray(pred_label.astype(np.uint8), 'P')
+    pred_img.putpalette(ade_palette())  # 팔레트 설정
 
     return pred_img
 
+# Gradio 인터페이스 생성
 gr.Interface(
     fn=image_segmentation,
     inputs="image",