Update app.py

app.py

@@ -120,7 +120,7 @@ def load_img_1_(nparr, gray: bool = False):
     return np_img, alpha_channel
     
 model = None
-def model_process(input):
+def model_process_pil(input):
     global model
     
     # input = request.files
@@ -225,7 +225,7 @@ def model_process(input):
     return ext
     '''
     
-def model_process_2(input): #image, mask):
+def model_process_filepath(input): #image, mask):
     global model
     # {'image': '/tmp/tmp8mn9xw93.png', 'mask': '/tmp/tmpn5ars4te.png'}
     # input = request.files
@@ -237,6 +237,9 @@ def model_process_2(input): #image, mask):
     original_shape = image.shape
     interpolation = cv2.INTER_CUBIC
 
+    image_pil = Image.fromarray(image)
+    # mask_pil = Image.fromarray(mask).convert("L")
+    
     # form = request.form
     # print(f'size_limit_1_ =  ', form["sizeLimit"], type(input["image"]))
     size_limit = "Original" #: Union[int, str] = form.get("sizeLimit", "1080")
@@ -285,26 +288,33 @@ def model_process_2(input): #image, mask):
     mask = resize_max_size(mask, size_limit=size_limit, interpolation=interpolation)
     print(f"mask image shape: {mask.shape} / {type(mask)} / {mask[250][250]} / {alpha_channel}")
 
+    if model is None:
+        return None
+        
     start = time.time()
     res_np_img = model(image, mask, config)
-    logger.info(f"process time: {(time.time() - start) * 1000}ms")
-    print(f"process time: {(time.time() - start) * 1000}ms, {res_np_img.shape} / {res_np_img[250][250]}")
+    logger.info(f"process time: {(time.time() - start) * 1000}ms, {res_np_img.shape}")
+    print(f"process time_1_: {(time.time() - start) * 1000}ms, {alpha_channel.shape}, {res_np_img.shape} / {res_np_img[250][250]} / {res_np_img.dtype}")
 
     torch.cuda.empty_cache()
 
     if alpha_channel is not None:
-        print(f"liuyz_here_1_: {alpha_channel}")
+        print(f"liuyz_here_10_: {alpha_channel.shape} / {alpha_channel.dtype} / {res_np_img.dtype}") 
         if alpha_channel.shape[:2] != res_np_img.shape[:2]:
+            print(f"liuyz_here_20_: {alpha_channel.shape} / {res_np_img.shape}")
             alpha_channel = cv2.resize(
                 alpha_channel, dsize=(res_np_img.shape[1], res_np_img.shape[0])
             )
+        print(f"liuyz_here_30_: {alpha_channel.shape} / {res_np_img.shape} / {alpha_channel.dtype} / {res_np_img.dtype}")
         res_np_img = np.concatenate(
             (res_np_img, alpha_channel[:, :, np.newaxis]), axis=-1
         )
-
+        print(f"liuyz_here_40_: {alpha_channel.shape} / {res_np_img.shape} / {alpha_channel.dtype} / {res_np_img.dtype}")
+    print(f"process time_2_: {(time.time() - start) * 1000}ms, {alpha_channel.shape}, {res_np_img.shape} / {res_np_img[250][250]} / {res_np_img.dtype}")
+  
     image = Image.fromarray(res_np_img)
     image.save(f'./result_image.png')
-    return image
+    return  image_pil # image
     '''
     ext = get_image_ext(origin_image_bytes)
 
@@ -346,6 +356,7 @@ def read_content(file_path):
 
     return content
 
+image_type = 'filepath' #'pil'
 def predict(input):
     print(f'liuyz_0_', input)  
     '''
@@ -362,12 +373,15 @@ 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"])
+    if image_type == 'filepath':
+        output = model_process_filepath(input) # dict["image"], dict["mask"])
+    elif image_type == 'pil':
+        output = model_process_pil(input)
     
     # output = mask #output.images[0]
     # output = pipe(prompt = prompt, image=init_image, mask_image=mask,guidance_scale=7.5)
     # output = input["mask"]
-    output = None
+    # output = None
     return output #, gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
 
 print(f'liuyz_500_here_')
@@ -460,7 +474,7 @@ with image_blocks as demo:
         with gr.Box():
             with gr.Row():
                 with gr.Column():
-                    image = gr.Image(source='upload', tool='sketch',type='filepath', label="Upload").style(height=512)
+                    image = gr.Image(source='upload', tool='sketch',type=f'{image_type}', label="Upload").style(height=512)
                     with gr.Row(elem_id="prompt-container").style(mobile_collapse=False, equal_height=True):
                         # prompt = gr.Textbox(placeholder = 'Your prompt (what you want in place of what is erased)', show_label=False, elem_id="input-text")
                         btn = gr.Button("Done!").style(