Update app.py

app.py

@@ -39,6 +39,29 @@ def pil_to_binary_mask(pil_image, threshold=0):
     output_mask = Image.fromarray(mask)
     return output_mask
 
+# utils_mask.py 파일에 있는 get_mask_location 함수를 수정합니다.
+def get_mask_location(mode, category, parsing, keypoints):
+    parsing = np.array(parsing)
+    mask = np.zeros_like(parsing)
+    
+    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}")
+    
+    # 키포인트를 사용하여 마스크 영역을 더 정확하게 조정할 수 있습니다.
+    # 예를 들어, 상의의 경우 어깨와 허리 사이의 영역으로 제한할 수 있습니다.
+    
+    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')
 
@@ -151,8 +174,8 @@ def start_tryon(dict, garm_img, garment_des, is_checked, is_checked_crop, denois
             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))
-        except NotImplementedError:
-            status_message = f"카테고리 {category}에 대한 자동 마스크 생성이 구현되지 않았습니다. 기본 마스크를 사용합니다."
+        except Exception as e:
+            status_message = f"자동 마스크 생성 중 오류가 발생했습니다: {str(e)}. 기본 마스크를 사용합니다."
             mask = Image.new('L', (768, 1024), 255)
     else:
         if dict['layers'] and dict['layers'][0]:
@@ -252,7 +275,6 @@ def start_tryon(dict, garm_img, garment_des, is_checked, is_checked_crop, denois
     else:
         return images[0], mask_gray, status_message
 
-
 garm_list = os.listdir(os.path.join(example_path,"cloth"))
 garm_list_path = [os.path.join(example_path,"cloth",garm) for garm in garm_list]
 
@@ -267,8 +289,6 @@ for ex_human in human_list_path:
     ex_dict['composite'] = None
     human_ex_list.append(ex_dict)
 
-##default human
-
 image_blocks = gr.Blocks(theme="Nymbo/Nymbo_Theme").queue(max_size=12)
 with image_blocks as demo:
     with gr.Column():
@@ -315,11 +335,9 @@ with image_blocks as demo:
     with gr.Column():
         status_message = gr.Textbox(label="상태", interactive=False)
 
-
-
     try_button.click(fn=start_tryon, 
                      inputs=[imgs, garm_img, prompt, is_checked, is_checked_crop, denoise_steps, seed, category], 
                      outputs=[image_out, masked_img, status_message], 
                      api_name='tryon')
 
-image_blocks.launch(auth=("gini","pick"))    
+image_blocks.launch(auth=("gini","pick"))