update UI

app.py

@@ -715,6 +715,7 @@
 # if __name__ == "__main__":
 #     demo.launch(server_name="0.0.0.0", server_port=7860)
 
+
 import os
 import sys
 
@@ -743,7 +744,6 @@ from transformers import pipeline
 
 from huggingface_hub import hf_hub_download
 
-# Show where diffusers3 is imported from (helps diagnose import collisions on Spaces)
 import diffusers3
 print("[BOOT] diffusers3 loaded from:", getattr(diffusers3, "__file__", "<?>"), flush=True)
 
@@ -823,10 +823,6 @@ def apply_parsing_white_mask_to_person_cv2(
     person_pil: Image.Image,
     parsing_img: Image.Image
 ) -> np.ndarray:
-    """
-    person_pil(RGB) 크기에 parsing_img(L) 마스크를 맞춰서
-    흰색(255) 영역만 person을 남기고 나머지는 흰색 배경으로 만드는 함수.
-    """
     person_rgb = np.array(person_pil.convert("RGB"), dtype=np.uint8)
 
     mask = np.array(parsing_img.convert("L"), dtype=np.uint8)
@@ -843,11 +839,6 @@ def apply_parsing_white_mask_to_person_cv2(
     return result_bgr
 
 
-from typing import Optional, Tuple
-import numpy as np
-from PIL import Image
-
-
 def clean_and_smooth_parsing_mask(
     parsing_img: Image.Image,
     *,
@@ -858,10 +849,6 @@ def clean_and_smooth_parsing_mask(
     morph_iters: int = 1,
     blur_ksize: int = 0,
 ) -> Image.Image:
-    """
-    Clean small white blobs and smooth boundaries on a grayscale (0/255) PIL mask.
-    White=foreground, Black=background.
-    """
     if not isinstance(parsing_img, Image.Image):
         raise TypeError("parsing_img must be a PIL.Image.Image")
 
@@ -870,11 +857,6 @@ def clean_and_smooth_parsing_mask(
 
     mask = np.where(arr >= white_threshold, 255, 0).astype(np.uint8)
 
-    try:
-        import cv2
-    except ImportError as e:
-        raise ImportError("This function requires opencv-python (cv2).") from e
-
     num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(mask, connectivity=8)
 
     keep = np.zeros_like(mask)
@@ -1031,10 +1013,10 @@ def make_depth_from_parsing_edges(parsing_img: Image.Image) -> Image.Image:
 
     depth_img = _edges_from_parsing(parsing_img)
 
-    inverted_depth = cv2.bitwise_not(depth_img)
-    contours, _ = cv2.findContours(inverted_depth, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+#     inverted_depth = cv2.bitwise_not(depth_img)
+    contours, _ = cv2.findContours(depth_img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
 
-    filled_depth = inverted_depth.copy()
+    filled_depth = depth_img.copy()
     cv2.drawContours(filled_depth, contours, -1, (255), thickness=cv2.FILLED)
 
     filled_depth = cv2.resize(filled_depth, (W, H), interpolation=cv2.INTER_AREA)
@@ -1138,10 +1120,17 @@ def get_pipe_and_device() -> Tuple[StableDiffusionXLControlNetImg2ImgPipeline, s
     return pipe, device, dtype
 
 
+# ✅ UI 표기 → 내부 extractor category 문자열 매핑
+_UI_TO_EXTRACTOR_CATEGORY = {
+    "Upper-body": "Upper-cloth",
+    "Lower-body": "Bottom",
+    "Dress": "Dress",
+}
+
+
 def run_one(paths: Paths, prompt: str, steps: int = DEFAULT_STEPS, category: str = "Dress"):
     """
-    Returns:
-      images(list[PIL]), mask_pil(PIL), depth_map(PIL), person_pil(PIL), garment_pil(PIL), garment_mask_pil(PIL)
+    category: UI에서 넘어오는 값(Upper-body/Lower-body/Dress)
     """
     global H, W
     pipe, device, _dtype = get_pipe_and_device()
@@ -1149,9 +1138,11 @@ def run_one(paths: Paths, prompt: str, steps: int = DEFAULT_STEPS, category: str
 
     H, W = compute_hw_from_person(paths.person_path)
 
-    # ✅ UI에서 받은 category 그대로 extractor에 전달
+    # ✅ UI category를 extractor가 기대하는 문자열로 변환
+    extractor_category = _UI_TO_EXTRACTOR_CATEGORY.get(category, "Dress")
+
     res = run_simple_extractor(
-        category=category,
+        category=extractor_category,
         input_path=os.path.abspath(paths.person_path),
         model_restore=schp_ckpt,
     )
@@ -1246,7 +1237,8 @@ def run_one(paths: Paths, prompt: str, steps: int = DEFAULT_STEPS, category: str
         "depth:", depth_map.size,
         "garment:", garment_pil.size,
         "gmask:", garment_mask_pil.size,
-        "category:", category,
+        "ui_category:", category,
+        "extractor_category:", extractor_category,
         flush=True
     )
 
@@ -1309,7 +1301,6 @@ def infer_web(person_fp, sketch_fp, style_fp, prompt, steps, seed, category):
     if person_fp is None or style_fp is None:
         raise gr.Error("person / style 이미지는 필수입니다. (sketch는 선택)")
 
-    # ✅ category는 UI 라디오에서 들어오며 기본값은 "Dress"
     if category not in ("Upper-body", "Lower-body", "Dress"):
         raise gr.Error(f"Invalid category: {category}")
 
@@ -1336,23 +1327,22 @@ def infer_web(person_fp, sketch_fp, style_fp, prompt, steps, seed, category):
 with gr.Blocks(title="VISTA Demo (HF Spaces)") as demo:
     gr.Markdown("## VISTA Demo\nperson / style 필수, sketch(guide)는 선택입니다.")
 
-    # ✅ Markdown 아래 행에 토글(라디오) 추가: 기본값 Dress
+    # ✅ UI 표기는 Upper-body/Lower-body/Dress 유지 (기본 Dress)
     category_toggle = gr.Radio(
-        choices=["Upper-body", "Lower-body", "Dress"],
+        choices=["Dress", "Upper-body", "Lower-body"],
         value="Dress",
         label="Category",
         interactive=True,
     )
 
+    # ✅ 한 행에 Person / Style / Output 배치
     with gr.Row():
         person_in = gr.Image(label="Person Image (required)", type="filepath")
         style_in = gr.Image(label="Style Image (required)", type="filepath")
+        out_img = gr.Image(label="Output", type="pil")
 
     with gr.Accordion("Sketch / Guide (optional)", open=False):
-        sketch_in = gr.Image(
-            label="Sketch / Guide",
-            type="filepath"
-        )
+        sketch_in = gr.Image(label="Sketch / Guide", type="filepath")
 
     with gr.Row():
         prompt_in = gr.Textbox(label="Prompt", value="upper garment", lines=2)
@@ -1361,7 +1351,7 @@ with gr.Blocks(title="VISTA Demo (HF Spaces)") as demo:
 
     run_btn = gr.Button("Run")
 
-    out_img = gr.Image(label="Output", type="pil")
+    # 파일 다운로드는 Output 아래(다음 행)에 두는 게 일반적으로 보기 좋음
     out_file = gr.File(label="Download result.png")
 
     gr.Markdown("### Debug Visualizations (mask/depth/etc)")
@@ -1384,3 +1374,5 @@ demo.queue()
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)
 
+
+