Update app.py

app.py

@@ -104,7 +104,7 @@ def preprocess_mask(mask):
     mask = torch.from_numpy(mask)
     return mask
     
- def load_img(nparr, gray: bool = False):
+def load_img(nparr, gray: bool = False):
     # alpha_channel = None
     # nparr = np.frombuffer(img_bytes, np.uint8)
     if gray:
@@ -131,7 +131,8 @@ def model_process(image, mask):
     image_pil = Image.fromarray(image)
     mask_pil = Image.fromarray(mask).convert("L")
     print(f'image_pil_ = {image_pil.shape}')
-    print(f'mask_pil_ = {mask_pil.shape}')  
+    print(f'mask_pil_ = {mask_pil.shape}') 
+    mask_pil.save(f'./mask_pil.png')     
         
     image, alpha_channel = load_img(image)
     # Origin image shape: (512, 512, 3)
@@ -185,9 +186,8 @@ def model_process(image, mask):
     # logger.info(f"Resized image shape: {image.shape}")
     print(f"Resized image shape: {image.shape} / {image[250][250]}")
     
-    #mask, _ = load_img(input["mask"].read(), gray=True)
-    mask_pil.save(f'./mask_pil.png')
-    mask = np.array(mask_pil)
+    mask, _ = load_img(mask), gray=True)
+    # mask = np.array(mask_pil)
     mask = resize_max_size(mask, size_limit=size_limit, interpolation=interpolation)
     print(f"mask image shape: {mask.shape} / {type(mask)} / {mask[250][250]}")