Update app.py

app.py

@@ -48,6 +48,7 @@ def get_mask_location(mode, category, parsing, keypoints):
     mask = np.zeros_like(parsing)
     
     print(f"Selected category: {category}")
+    print(f"Parsing shape: {parsing.shape}")
     print(f"Unique values in parsing: {np.unique(parsing)}")
     
     if category == "상의":
@@ -71,7 +72,15 @@ def get_mask_location(mode, category, parsing, keypoints):
     print(f"Mask shape: {mask.shape}, Unique values in mask: {np.unique(mask)}")
     print(f"Number of masked pixels: {np.sum(mask == 255)}")
     
-    mask_gray = Image.fromarray(mask)
+    # 마스크 시각화를 위한 코드 추가
+    import matplotlib.pyplot as plt
+    plt.figure(figsize=(10, 10))
+    plt.imshow(mask, cmap='gray')
+    plt.title(f"Mask for {category}")
+    plt.savefig(f"mask_{category}.png")
+    plt.close()
+    
+    mask_gray = Image.fromarray(mask.astype(np.uint8))
     return mask_gray, mask_gray
     
 
@@ -194,14 +203,33 @@ def start_tryon(dict, garm_img, garment_des, is_checked, is_checked_crop, denois
             print(f"Unique values in parsing model output: {np.unique(model_parse)}")
             
             mask, mask_gray = get_mask_location('hd', category, model_parse, keypoints)
+            
+            # 마스크 확인 및 시각화
+            mask_array = np.array(mask)
+            print(f"Mask shape after get_mask_location: {mask_array.shape}")
+            print(f"Unique values in mask after get_mask_location: {np.unique(mask_array)}")
+            print(f"Number of masked pixels after get_mask_location: {np.sum(mask_array == 255)}")
+            
+            plt.figure(figsize=(10, 10))
+            plt.imshow(mask_array, cmap='gray')
+            plt.title(f"Mask after get_mask_location for {category}")
+            plt.savefig(f"mask_after_get_mask_location_{category}.png")
+            plt.close()
+            
             mask = mask.resize((768,1024))
             print(f"Mask created for category {category}")
             
-            # 마스크 확인
-            print(f"Final mask shape: {mask.size}")
-            mask_array = np.array(mask)
-            print(f"Unique values in final mask: {np.unique(mask_array)}")
-            print(f"Number of masked pixels in final mask: {np.sum(mask_array == 255)}")
+            # 최종 마스크 확인
+            mask_array_final = np.array(mask)
+            print(f"Final mask shape: {mask_array_final.shape}")
+            print(f"Unique values in final mask: {np.unique(mask_array_final)}")
+            print(f"Number of masked pixels in final mask: {np.sum(mask_array_final == 255)}")
+            
+            plt.figure(figsize=(10, 10))
+            plt.imshow(mask_array_final, cmap='gray')
+            plt.title(f"Final Mask for {category}")
+            plt.savefig(f"final_mask_{category}.png")
+            plt.close()
             
         except Exception as e:
             status_message = f"자동 마스크 생성 중 오류가 발생했습니다: {str(e)}. 기본 마스크를 사용합니다."