change details

app.py

@@ -10,8 +10,8 @@ MODEL_ID = "nvidia/segformer-b4-finetuned-cityscapes-1024-1024"
 processor = AutoImageProcessor.from_pretrained(MODEL_ID)
 model = AutoModelForSemanticSegmentation.from_pretrained(MODEL_ID)
 
+
 def ade_palette():
-    """ADE20K palette that maps each class to RGB values."""
     return [
         [128, 64, 128], [244, 35, 232], [70, 70, 70], [102, 102, 156], [190, 153, 153],
         [153, 153, 153], [250, 170, 30], [220, 220, 0], [107, 142, 35], [152, 251, 152],
@@ -19,6 +19,7 @@ def ade_palette():
         [0, 60, 100], [0, 80, 100], [0, 0, 230], [119, 11, 32],
     ]
 
+
 labels_list = []
 with open("labels.txt", "r", encoding="utf-8") as fp:
     for line in fp:
@@ -26,6 +27,7 @@ with open("labels.txt", "r", encoding="utf-8") as fp:
 
 colormap = np.asarray(ade_palette(), dtype=np.uint8)
 
+
 def label_to_color_image(label):
     if label.ndim != 2:
         raise ValueError("Expect 2-D input label")
@@ -33,6 +35,7 @@ def label_to_color_image(label):
         raise ValueError("label value too large.")
     return colormap[label]
 
+
 def draw_plot(pred_img, seg_np):
     fig = plt.figure(figsize=(20, 15))
     grid_spec = gridspec.GridSpec(1, 2, width_ratios=[6, 1])
@@ -40,6 +43,7 @@ def draw_plot(pred_img, seg_np):
     plt.subplot(grid_spec[0])
     plt.imshow(pred_img)
     plt.axis('off')
+    plt.title('Segmentation Result', fontsize=20, pad=20)
 
     LABEL_NAMES = np.asarray(labels_list)
     FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1)
@@ -52,8 +56,10 @@ def draw_plot(pred_img, seg_np):
     plt.yticks(range(len(unique_labels)), LABEL_NAMES[unique_labels])
     plt.xticks([], [])
     ax.tick_params(width=0.0, labelsize=25)
+    plt.title('Detected Classes', fontsize=20, pad=20)
     return fig
 
+
 def run_inference(input_img):
     # input: numpy array from gradio -> PIL
     img = Image.fromarray(input_img.astype(np.uint8)) if isinstance(input_img, np.ndarray) else input_img
@@ -63,32 +69,74 @@ def run_inference(input_img):
     inputs = processor(images=img, return_tensors="pt")
     with torch.no_grad():
         outputs = model(**inputs)
-        logits = outputs.logits  # (1, C, h/4, w/4)
+        logits = outputs.logits
 
     # resize to original
     upsampled = torch.nn.functional.interpolate(
         logits, size=img.size[::-1], mode="bilinear", align_corners=False
     )
-    seg = upsampled.argmax(dim=1)[0].cpu().numpy().astype(np.uint8)  # (H,W)
+    seg = upsampled.argmax(dim=1)[0].cpu().numpy().astype(np.uint8)
 
     # colorize & overlay
-    color_seg = colormap[seg]                                # (H,W,3)
+    color_seg = colormap[seg]
     pred_img = (np.array(img) * 0.5 + color_seg * 0.5).astype(np.uint8)
 
     fig = draw_plot(pred_img, seg)
     return fig
 
-demo = gr.Interface(
-    fn=run_inference,
-    inputs=gr.Image(type="numpy", label="Input Image"),
-    outputs=gr.Plot(label="Overlay + Legend"),
-    examples=[
-        "road-2.jpg",
-        "road-3.jpeg",
-    ],
-    flagging_mode="never",
-    cache_examples=False,
-)
+
+# ===== 개선된 인터페이스 (Blocks 사용) =====
+with gr.Blocks(theme=gr.themes.Soft(), title="도시 장면 분할") as demo:
+    gr.Markdown(
+        """
+        # 🏙️ 도시 장면 영상 분할 (City Scene Segmentation)
+        **Cityscapes 데이터셋으로 학습된 SegFormer 모델**을 활용한 도로 및 도시 장면 분할 데모입니다.
+
+        도로, 건물, 차량, 보행자 등 19개 클래스를 자동으로 인식하고 분할합니다.
+        """
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            input_img = gr.Image(
+                type="numpy",
+                label="📷 입력 이미지",
+                height=400
+            )
+            submit_btn = gr.Button(
+                "🎯 분할 실행",
+                variant="primary",
+                size="lg"
+            )
+
+            gr.Markdown("### 📌 예시 이미지")
+            gr.Examples(
+                examples=[
+                    "road-2.jpg",
+                    "road-3.jpeg",
+                ],
+                inputs=input_img,
+                label="도시/도로 장면 샘플"
+            )
+
+        with gr.Column(scale=1):
+            output_plot = gr.Plot(label="✨ 분할 결과 및 범례")
+
+    gr.Markdown(
+        """
+        ---
+        ### 📊 감지 가능한 클래스 (19개)
+        `도로`, `보도`, `건물`, `벽`, `울타리`, `기둥`, `신호등`, `표지판`, `식물`, 
+        `지형`, `하늘`, `사람`, `자전거 탑승자`, `자동차`, `트럭`, `버스`, `기차`, `오토바이`, `자전거`
+        """
+    )
+
+    # 이벤트 핸들러
+    submit_btn.click(
+        fn=run_inference,
+        inputs=input_img,
+        outputs=output_plot
+    )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()