Update app.py

app.py

@@ -72,27 +72,33 @@ def label_to_color_image(label):
         raise ValueError("label value too large.")
     return colormap[label]
     
-# Gradio 인터페이스 생성
+# Gradio 인터페이스에 사용될 이미지 세그멘테이션 함수를 정의합니다.
 def image_segmentation(input_img):
-    # 이미지를 넘파이 배열로 변환
-    input_img = np.array(input_img)
+    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 = Image.fromarray(pred_label.astype(np.uint8), 'P')
-    pred_img.putpalette(ade_palette())  # 팔레트 설정
+    pred_img = label_to_color_image(pred_label)
 
     return pred_img
 
+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 인터페이스 생성
-gr.Interface(
+iface = gr.Interface(
     fn=image_segmentation,
     inputs="image",
     outputs="image"
-).launch()
+)
 
+# 인터페이스 실행
+iface.launch()