1

app.py

@@ -8,6 +8,28 @@ import numpy as np
 feature_extractor = SegformerFeatureExtractor.from_pretrained("nvidia/segformer-b0-finetuned-cityscapes-512-1024")
 model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-cityscapes-512-1024")
 
+colors = np.array([
+    [255, 0, 0],      # 빨강
+    [255, 228, 0],    # 노랑
+    [171, 242, 0],    # 연두
+    [0, 216, 255],    # 하늘
+    [0, 0, 255],      # 파랑
+    [255, 0, 221],    # 핑크
+    [116, 116, 116],  # 회색
+    [95, 0, 255],     # 보라
+    [255, 94, 0],     # 주황
+    [71, 200, 62],    # 초록
+    [153, 0, 76],     # 마젠타
+    [67, 116, 217],   # 애매한하늘+파랑
+    [153, 112, 0],    # 겨자
+    [87, 129, 0],     # 녹색
+    [255, 169, 169],  # 분홍분홍
+    [35, 30, 183],    # 어두운 파랑
+    [225, 186, 133],  # 살색
+    [206, 251, 201],  # 연한초록
+    [165, 102, 255]   # 애매한보라
+], dtype=np.uint8)
+
 def segment_image(image):
     # 이미지를 처리하고 모델에 전달하기
     inputs = feature_extractor(images=image, return_tensors="pt")
@@ -18,9 +40,11 @@ def segment_image(image):
     result = logits.argmax(dim=1)[0]
     result = result.cpu().detach().numpy()
 
-    result_image = Image.fromarray(result.astype(np.uint8), mode="P").convert("RGB")
+    # 색상 팔레트 적용
+    result_color = colors[result]
+
+    result_image = Image.fromarray(result_color.astype(np.uint8))
 
-    # 결과 배열 반환
     return result_image
 
 # Gradio 인터페이스 정의