update app.py

api_client.py

@@ -52,18 +52,34 @@ def base64_to_PILImage(im_b64):
     pil_img = Image.open(io.BytesIO(im_bytes))
     return pil_img
 
+def cleaner_img(image_file, remove_texts, mask_extend=20, disp_debug=True):
+    data = {'remove_texts': remove_texts,
+        'mask_extend': mask_extend,
+        'img': imgFile_to_base64(image_file),
+        }
+    ret = request_post(url, data, timeout=600, headers = None)  
+    if ret['code'] == 0:
+        if disp_debug:
+            for img in ret['result']['imgs']:
+                pilImage = base64_to_PILImage(img)
+                plt.imshow(pilImage)
+                plt.show()
+                plt.clf()
+                plt.close('all')    
+        img_len = len(ret['result']['imgs'])
+        pilImage = base64_to_PILImage(ret['result']['imgs'][img_len-1])  
+    else:
+        pilImage = None
+    return pilImage, ret
+
 image_file = 'dog.png'
-data = {'remove_texts': "小狗 . 椅子",
-    'extend': 20,
-    'img': imgFile_to_base64(image_file),
-    }
+remove_texts = "小狗 . 椅子"
 
-ret = request_post(url, data, timeout=600, headers = None)
-print(len(ret['result']['imgs']))
+mask_extend = 20
+pil_image, ret = cleaner_img(image_file, remove_texts, mask_extend, disp_debug=False) 
 
-for img in ret['result']['imgs']:
-    pilImage = base64_to_PILImage(img)
-    plt.imshow(pilImage)
-    plt.show()
-    plt.clf()
+plt.imshow(pil_image)
+plt.show()
+plt.clf()
+plt.close()
 

app.py

@@ -3,7 +3,17 @@ import warnings
 warnings.filterwarnings('ignore')
 
 import subprocess, io, os, sys, time
-os.system("pip install gradio==3.40.1")
+
+run_gradio = False
+if os.environ.get('IS_MY_DEBUG') is None:
+    run_gradio = True    
+else:
+    run_gradio = False
+    # run_gradio = True
+
+if run_gradio:
+    os.system("pip install gradio==3.40.1")
+
 import gradio as gr
 
 from loguru import logger
@@ -35,7 +45,10 @@ from GroundingDINO.groundingdino.util.utils import clean_state_dict, get_phrases
 
 import cv2
 import numpy as np
-import matplotlib.pyplot as plt
+import matplotlib
+matplotlib.use('AGG')
+plt = matplotlib.pyplot
+# import matplotlib.pyplot as plt
 
 groundingdino_enable = True
 sam_enable = True
@@ -332,60 +345,63 @@ def load_lama_cleaner_model(device):
         )
 
 def lama_cleaner_process(image, mask, cleaner_size_limit=1080):
-    ori_image = image
-    if mask.shape[0] == image.shape[1] and mask.shape[1] == image.shape[0] and mask.shape[0] != mask.shape[1]:
-        # rotate image
-        ori_image = np.transpose(image[::-1, ...][:, ::-1], axes=(1, 0, 2))[::-1, ...]
-        image = ori_image
-    
-    original_shape = ori_image.shape
-    interpolation = cv2.INTER_CUBIC
-    
-    size_limit = cleaner_size_limit
-    if size_limit == -1:
-        size_limit = max(image.shape)
-    else:
-        size_limit = int(size_limit)
-
-    config = lama_Config(
-        ldm_steps=25,
-        ldm_sampler='plms',
-        zits_wireframe=True,
-        hd_strategy='Original',
-        hd_strategy_crop_margin=196,
-        hd_strategy_crop_trigger_size=1280,
-        hd_strategy_resize_limit=2048,
-        prompt='',
-        use_croper=False,
-        croper_x=0,
-        croper_y=0,
-        croper_height=512,
-        croper_width=512,
-        sd_mask_blur=5,
-        sd_strength=0.75,
-        sd_steps=50,
-        sd_guidance_scale=7.5,
-        sd_sampler='ddim',
-        sd_seed=42,
-        cv2_flag='INPAINT_NS',
-        cv2_radius=5,
-    )
-    
-    if config.sd_seed == -1:
-        config.sd_seed = random.randint(1, 999999999)
+    try:
+        ori_image = image
+        if mask.shape[0] == image.shape[1] and mask.shape[1] == image.shape[0] and mask.shape[0] != mask.shape[1]:
+            # rotate image
+            ori_image = np.transpose(image[::-1, ...][:, ::-1], axes=(1, 0, 2))[::-1, ...]
+            image = ori_image
+        
+        original_shape = ori_image.shape
+        interpolation = cv2.INTER_CUBIC
+        
+        size_limit = cleaner_size_limit
+        if size_limit == -1:
+            size_limit = max(image.shape)
+        else:
+            size_limit = int(size_limit)
+
+        config = lama_Config(
+            ldm_steps=25,
+            ldm_sampler='plms',
+            zits_wireframe=True,
+            hd_strategy='Original',
+            hd_strategy_crop_margin=196,
+            hd_strategy_crop_trigger_size=1280,
+            hd_strategy_resize_limit=2048,
+            prompt='',
+            use_croper=False,
+            croper_x=0,
+            croper_y=0,
+            croper_height=512,
+            croper_width=512,
+            sd_mask_blur=5,
+            sd_strength=0.75,
+            sd_steps=50,
+            sd_guidance_scale=7.5,
+            sd_sampler='ddim',
+            sd_seed=42,
+            cv2_flag='INPAINT_NS',
+            cv2_radius=5,
+        )
+        
+        if config.sd_seed == -1:
+            config.sd_seed = random.randint(1, 999999999)
 
-    # logger.info(f"Origin image shape_0_: {original_shape} / {size_limit}")
-    image = resize_max_size(image, size_limit=size_limit, interpolation=interpolation)
-    # logger.info(f"Resized image shape_1_: {image.shape}")
-    
-    # logger.info(f"mask image shape_0_: {mask.shape} / {type(mask)}")
-    mask = resize_max_size(mask, size_limit=size_limit, interpolation=interpolation)
-    # logger.info(f"mask image shape_1_: {mask.shape} / {type(mask)}")
+        # logger.info(f"Origin image shape_0_: {original_shape} / {size_limit}")
+        image = resize_max_size(image, size_limit=size_limit, interpolation=interpolation)
+        # logger.info(f"Resized image shape_1_: {image.shape}")
+        
+        # logger.info(f"mask image shape_0_: {mask.shape} / {type(mask)}")
+        mask = resize_max_size(mask, size_limit=size_limit, interpolation=interpolation)
+        # logger.info(f"mask image shape_1_: {mask.shape} / {type(mask)}")
 
-    res_np_img = lama_cleaner_model(image, mask, config)
-    torch.cuda.empty_cache()
-  
-    image = Image.open(io.BytesIO(numpy_to_bytes(res_np_img, 'png')))
+        res_np_img = lama_cleaner_model(image, mask, config)
+        torch.cuda.empty_cache()
+    
+        image = Image.open(io.BytesIO(numpy_to_bytes(res_np_img, 'png')))
+    except Exception as e:
+        image = None        
     return  image
 
 class Ram_Predictor(RamPredictor):
@@ -691,6 +707,8 @@ def run_anything_task(input_image, text_prompt, task_type, inpaint_prompt, box_t
         plt.axis('off')
         image_path = os.path.join(output_dir, f"grounding_seg_output_{file_temp}.jpg")
         plt.savefig(image_path, bbox_inches="tight")
+        plt.clf()
+        plt.close('all')
         segment_image_result = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB)
         os.remove(image_path)
         output_images.append(Image.fromarray(segment_image_result)) 
@@ -757,6 +775,10 @@ def run_anything_task(input_image, text_prompt, task_type, inpaint_prompt, box_t
 
             logger.info(f'run_anything_task_[{file_temp}]_{task_type}_6_')            
             image_inpainting = lama_cleaner_process(np.array(image_pil), np.array(mask_pil.convert("L")), cleaner_size_limit)
+            if image_inpainting is None:
+                logger.info(f'run_anything_task_failed_')
+                return None, None, None, None, None, None, None
+
             # output_images.append(image_inpainting)
             # run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
 
@@ -975,7 +997,10 @@ class API_Starter:
                 request_data = request.data.decode('utf-8')
                 data = json.loads(request_data)
                 result = self.handle_data(data)
-                ret_json = {'code': 0, 'result':result}
+                if result is None:
+                    ret_json = {'code': -2, 'reason':'handle error'}
+                else:
+                    ret_json = {'code': 0, 'result':result}
             return jsonify(ret_json)
         
         self.app = app
@@ -996,15 +1021,18 @@ class API_Starter:
                             inpaint_mode = "merge", 
                             mask_source_radio = "type what to detect below", 
                             remove_mode = "rectangle",   # ["segment", "rectangle"]
-                            remove_mask_extend = "10", 
+                            remove_mask_extend = f"{data['mask_extend']}", 
                             num_relation = 5,
                             kosmos_input = None,
                             cleaner_size_limit = -1,
                             )
         output_images = results[0]
+        if output_images is None:
+            return None
         ret_json_images = []
         file_temp = int(time.time())
         count = 0
+        output_images = output_images[-1:]
         for image_pil in output_images:
             try:
                 img_format = image_pil.format.lower()
@@ -1086,16 +1114,12 @@ if __name__ == "__main__":
     # print(f'ram_model__{get_model_device(ram_model)}')
     # print(f'kosmos_model__{get_model_device(kosmos_model)}')    
 
-    if os.environ.get('IS_MY_DEBUG') is None:
+    if run_gradio:
         # Provide gradio services
         main_gradio(args)
     else:        
-        if 0 == 0:
-            # Provide API services
-            main_api(args)
-        else:
-            # Provide gradio services
-            main_gradio(args)
+        # Provide API services
+        main_api(args)