Update app.py

app.py

@@ -39,28 +39,33 @@ def pil_to_binary_mask(pil_image, threshold=0):
     output_mask = Image.fromarray(mask)
     return output_mask
 
-# utils_mask.py 파일에 있는 get_mask_location 함수를 수정합니다.
+import numpy as np
+from PIL import Image
+
 def get_mask_location(mode, category, parsing, keypoints):
     parsing = np.array(parsing)
     mask = np.zeros_like(parsing)
     
+    print(f"Selected category: {category}")  # 디버깅을 위해 선택된 카테고리 출력
+    
     if category == "상의":
-        # 상의에 해당하는 부분 마스킹 (예: 상의, 드레스 상단부)
+        # 상의에 해당하는 부분만 마스킹 (상체, 팔)
         mask[np.isin(parsing, [5, 6, 7])] = 255
     elif category == "하의":
-        # 하의에 해당하는 부분 마스킹 (예: 바지, 치마)
+        # 하의에 해당하는 부분만 마스킹 (하체)
         mask[np.isin(parsing, [9, 12, 13, 14, 15, 16, 17, 18, 19])] = 255
     elif category == "드레스":
-        # 드레스에 해당하는 부분 마스킹 (상의와 하의 모두 포함)
+        # 드레스에 해당하는 부분 마스킹 (상체와 하체)
         mask[np.isin(parsing, [5, 6, 7, 9, 12, 13, 14, 15, 16, 17, 18, 19])] = 255
     else:
         raise ValueError(f"Unknown category: {category}")
     
-    # 키포인트를 사용하여 마스크 영역을 더 정확하게 조정할 수 있습니다.
-    # 예를 들어, 상의의 경우 어깨와 허리 사이의 영역으로 제한할 수 있습니다.
+    print(f"Mask shape: {mask.shape}, Unique values: {np.unique(mask)}")  # 디버깅을 위해 마스크 정보 출력
     
     mask_gray = Image.fromarray(mask)
     return mask_gray, mask_gray
+    
+
 
 base_path = 'yisol/IDM-VTON'
 example_path = os.path.join(os.path.dirname(__file__), 'example')
@@ -167,15 +172,19 @@ def start_tryon(dict, garm_img, garment_des, is_checked, is_checked_crop, denois
     else:
         human_img = human_img_orig.resize((768,1024))
 
+
     status_message = ""
     if is_checked:
         try:
+            print(f"Processing category: {category}")  # 디버깅을 위해 처리 중인 카테고리 출력
             keypoints = openpose_model(human_img.resize((384,512)))
             model_parse, _ = parsing_model(human_img.resize((384,512)))
             mask, mask_gray = get_mask_location('hd', category, model_parse, keypoints)
             mask = mask.resize((768,1024))
+            print(f"Mask created for category {category}")  # 디버깅을 위해 마스크 생성 확인 메시지 출력
         except Exception as e:
             status_message = f"자동 마스크 생성 중 오류가 발생했습니다: {str(e)}. 기본 마스크를 사용합니다."
+            print(f"Error in mask creation: {str(e)}")  # 디버깅을 위해 오류 메시지 출력
             mask = Image.new('L', (768, 1024), 255)
     else:
         if dict['layers'] and dict['layers'][0]: