Update app.py

app.py

@@ -120,13 +120,14 @@ def load_img_1_(nparr, gray: bool = False):
     return np_img, alpha_channel
     
 model = None
-def model_process_1(image, mask):
+def model_process(input):
     global model
     
     # input = request.files
     # RGB
     # origin_image_bytes = input["image"].read()
-    
+    image = input['image']
+    mask = input['mask']
     print(f'liuyz_2_here_', type(image), image.shape)
         
     image_pil = Image.fromarray(image)
@@ -138,7 +139,7 @@ def model_process_1(image, mask):
     #image, alpha_channel = load_img(image)
     # Origin image shape: (512, 512, 3)
     
-    alpha_channel = None
+    alpha_channel = np.ones((image.shape[0],image.shape[1]))*255
     original_shape = image.shape
     interpolation = cv2.INTER_CUBIC    
     
@@ -188,7 +189,7 @@ def model_process_1(image, mask):
     print(f"Resized image shape: {image.shape} / {image[250][250]}")
     
     # mask, _ = load_img(mask, gray=True)
-    mask = np.array(mask_pil)
+    # 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]}")
 
@@ -201,10 +202,7 @@ def model_process_1(image, mask):
     print(f"process time: {(time.time() - start) * 1000}ms, {res_np_img.shape} / {res_np_img[250][250]}")
 
     torch.cuda.empty_cache()
-    image = Image.fromarray(res_np_img)
-    image.save(f'./result_image.png')
-    return image
-    '''
+
     if alpha_channel is not None:
         if alpha_channel.shape[:2] != res_np_img.shape[:2]:
             alpha_channel = cv2.resize(
@@ -213,12 +211,16 @@ def model_process_1(image, mask):
         res_np_img = np.concatenate(
             (res_np_img, alpha_channel[:, :, np.newaxis]), axis=-1
         )
-
+    image = Image.fromarray(res_np_img)
+    image.save(f'./result_image.png')
+    return image
+    
+    '''
     ext = get_image_ext(origin_image_bytes)
     return ext
     '''
     
-def model_process(input): #image, mask):
+def model_process_2(input): #image, mask):
     global model
     # {'image': '/tmp/tmp8mn9xw93.png', 'mask': '/tmp/tmpn5ars4te.png'}
     # input = request.files
@@ -355,7 +357,7 @@ def predict(input):
     print(f'liuyz_3_', image.convert("RGB").resize((512, 512)).shape)
     # mask = dict["mask"] # .convert("RGB") #.resize((512, 512))
     '''   
-    # output = model_process(input) # dict["image"], dict["mask"])
+    output = model_process(input) # dict["image"], dict["mask"])
     
     # output = mask #output.images[0]
     # output = pipe(prompt = prompt, image=init_image, mask_image=mask,guidance_scale=7.5)