Update app.py

app.py

@@ -104,6 +104,21 @@ def preprocess_mask(mask):
     mask = torch.from_numpy(mask)
     return mask
     
+ def load_img(nparr, gray: bool = False):
+    # alpha_channel = None
+    # nparr = np.frombuffer(img_bytes, np.uint8)
+    if gray:
+        np_img = cv2.imdecode(nparr, cv2.IMREAD_GRAYSCALE)
+    else:
+        np_img = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED)
+        if len(np_img.shape) == 3 and np_img.shape[2] == 4:
+            alpha_channel = np_img[:, :, -1]
+            np_img = cv2.cvtColor(np_img, cv2.COLOR_BGRA2RGB)
+        else:
+            np_img = cv2.cvtColor(np_img, cv2.COLOR_BGR2RGB)
+
+    return np_img, alpha_channel
+    
 model = None
 def model_process(image, mask):
     global model
@@ -112,19 +127,23 @@ def model_process(image, mask):
     # RGB
     # origin_image_bytes = input["image"].read()
     
-    
     print(f'liuyz_2_here_')
-    
-    # image, alpha_channel = load_img(origin_image_bytes)
+    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}')  
+        
+    image, alpha_channel = load_img(image)
     # Origin image shape: (512, 512, 3)
-    alpha_channel = None
+    
+    # alpha_channel = None
     original_shape = image.shape
     interpolation = cv2.INTER_CUBIC    
     
     # form = request.form
-    print(f'liuyz_3_here_', original_shape)
+    print(f'liuyz_3_here_', original_shape, alpha_channel)
 
-    size_limit = 512 # : Union[int, str] = form.get("sizeLimit", "1080")
+    size_limit = image_pil.shape[1] # : Union[int, str] = form.get("sizeLimit", "1080")
     if size_limit == "Original":
         size_limit = max(image.shape)
     else:
@@ -167,9 +186,8 @@ def model_process(image, mask):
     print(f"Resized image shape: {image.shape} / {image[250][250]}")
     
     #mask, _ = load_img(input["mask"].read(), gray=True)
-    mask_image = Image.fromarray(mask).convert("L")
-    mask_image.save(f'./mask_image.png')
-    mask = np.array(mask_image)
+    mask_pil.save(f'./mask_pil.png')
+    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]}")
 
@@ -238,11 +256,7 @@ def predict(dict):
     print(f'liuyz_3_', image.convert("RGB").resize((512, 512)).shape)
     # mask = dict["mask"] # .convert("RGB") #.resize((512, 512))
     '''
-    print(f'size__', dict["image"].size)
-    image = Image.fromarray(dict["image"])
-    mask = np.array(Image.fromarray(dict["mask"]).convert("L"))
-    print(f'mask___1 = {mask.shape}')
-    
+   
     output = model_process(dict["image"], dict["mask"])
     # output = mask #output.images[0]
     # output = pipe(prompt = prompt, image=init_image, mask_image=mask,guidance_scale=7.5)