fix bugs of example images and api keys

app.py

@@ -40,16 +40,16 @@ description = """Gradio demo for Caption Anything, image to dense captioning gen
 """
 
 examples = [
+    ["test_img/img35.webp"],
     ["test_img/img2.jpg"],
     ["test_img/img5.jpg"],
     ["test_img/img12.jpg"],
     ["test_img/img14.jpg"],
+    ["test_img/img0.png"],
+    ["test_img/img1.jpg"],
 ]
 
 args = parse_augment()
-args.captioner = 'blip2'
-args.seg_crop_mode = 'wo_bg'
-args.regular_box = True
 # args.device = 'cuda:5'
 # args.disable_gpt = False
 # args.enable_reduce_tokens = True
@@ -57,9 +57,10 @@ args.regular_box = True
 model = CaptionAnything(args)
 
 def init_openai_api_key(api_key):
-    os.environ['OPENAI_API_KEY'] = api_key
-    model.init_refiner()
-
+    # os.environ['OPENAI_API_KEY'] = api_key
+    model.init_refiner(api_key)
+    openai_available = model.text_refiner is not None
+    return gr.update(visible = openai_available), gr.update(visible = openai_available), gr.update(visible = openai_available), gr.update(visible = True), gr.update(visible = True)
 
 def get_prompt(chat_input, click_state):    
     points = click_state[0]
@@ -78,7 +79,7 @@ def get_prompt(chat_input, click_state):
     return prompt
 
 def chat_with_points(chat_input, click_state, state):
-    if not hasattr(model, "text_refiner"):
+    if model.text_refiner is None:
         response = "Text refiner is not initilzed, please input openai api key."
         state = state + [(chat_input, response)]
         return state, state
@@ -132,7 +133,7 @@ def inference_seg_cap(image_input, point_prompt, language, sentiment, factuality
     image_input = create_bubble_frame(image_input, text, (evt.index[0], evt.index[1]))
 
     yield state, state, click_state, chat_input, image_input, wiki
-    if not args.disable_gpt and hasattr(model, "text_refiner"):
+    if not args.disable_gpt and model.text_refiner:
         refined_caption = model.text_refiner.inference(query=text, controls=controls, context=out['context_captions'])
         # new_cap = 'Original: ' + text + '. Refined: ' + refined_caption['caption']
         new_cap = refined_caption['caption']
@@ -143,10 +144,16 @@ def inference_seg_cap(image_input, point_prompt, language, sentiment, factuality
 def upload_callback(image_input, state):
     state = [] + [('Image size: ' + str(image_input.size), None)]
     click_state = [[], [], []]
+    res = 1024
+    width, height = image_input.size
+    ratio = min(1.0 * res / max(width, height), 1.0)
+    if ratio < 1.0:
+        image_input = image_input.resize((int(width * ratio), int(height * ratio)))
+        print('Scaling input image to {}'.format(image_input.size))
     model.segmenter.image = None
     model.segmenter.image_embedding = None
     model.segmenter.set_image(image_input)
-    return state, image_input, click_state
+    return state, image_input, click_state, image_input
 
 with gr.Blocks(
     css='''
@@ -163,55 +170,62 @@ with gr.Blocks(
 
     with gr.Row():
         with gr.Column(scale=1.0):
-            image_input = gr.Image(type="pil", interactive=True, elem_id="image_upload")
-            with gr.Row(scale=1.0):
-                point_prompt = gr.Radio(
-                    choices=["Positive",  "Negative"],
-                    value="Positive",
-                    label="Point Prompt",
-                    interactive=True)
-                clear_button_clike = gr.Button(value="Clear Clicks", interactive=True)
-                clear_button_image = gr.Button(value="Clear Image", interactive=True)
-            with gr.Row(scale=1.0):
-                language = gr.Dropdown(['English', 'Chinese', 'French', "Spanish", "Arabic", "Portuguese", "Cantonese"], value="English", label="Language", interactive=True)
-               
-                sentiment = gr.Radio(
-                    choices=["Positive", "Natural", "Negative"],
-                    value="Natural",
-                    label="Sentiment",
-                    interactive=True,
-                )
-            with gr.Row(scale=1.0):
-                factuality = gr.Radio(
-                    choices=["Factual", "Imagination"],
-                    value="Factual",
-                    label="Factuality",
-                    interactive=True,
-                )
-                length = gr.Slider(
-                    minimum=10,
-                    maximum=80,
-                    value=10,
-                    step=1,
-                    interactive=True,
-                    label="Length",
-                )
+            with gr.Column(visible=False) as modules_not_need_gpt:
+                image_input = gr.Image(type="pil", interactive=True, elem_id="image_upload")
+                example_image = gr.Image(type="pil", interactive=False, visible=False)
+                with gr.Row(scale=1.0):
+                    point_prompt = gr.Radio(
+                        choices=["Positive",  "Negative"],
+                        value="Positive",
+                        label="Point Prompt",
+                        interactive=True)
+                    clear_button_clike = gr.Button(value="Clear Clicks", interactive=True)
+                    clear_button_image = gr.Button(value="Clear Image", interactive=True)
+            with gr.Column(visible=False) as modules_need_gpt:
+                with gr.Row(scale=1.0):
+                    language = gr.Dropdown(['English', 'Chinese', 'French', "Spanish", "Arabic", "Portuguese", "Cantonese"], value="English", label="Language", interactive=True)
+                
+                    sentiment = gr.Radio(
+                        choices=["Positive", "Natural", "Negative"],
+                        value="Natural",
+                        label="Sentiment",
+                        interactive=True,
+                    )
+                with gr.Row(scale=1.0):
+                    factuality = gr.Radio(
+                        choices=["Factual", "Imagination"],
+                        value="Factual",
+                        label="Factuality",
+                        interactive=True,
+                    )
+                    length = gr.Slider(
+                        minimum=10,
+                        maximum=80,
+                        value=10,
+                        step=1,
+                        interactive=True,
+                        label="Length",
+                    )
         
         with gr.Column(scale=0.5):
             openai_api_key = gr.Textbox(
-                placeholder="Input your openAI API key and press Enter",
+                placeholder="Input openAI API key and press Enter (Input blank will disable GPT)",
                 show_label=False,
                 label = "OpenAI API Key",
                 lines=1,
                 type="password"
                 )
-            openai_api_key.submit(init_openai_api_key, inputs=[openai_api_key])
-            wiki_output = gr.Textbox(lines=6, label="Wiki")
-            chatbot = gr.Chatbot(label="Chat about Selected Object",).style(height=450,scale=0.5)
-            chat_input = gr.Textbox(lines=1, label="Chat Input")
-            with gr.Row():
-                clear_button_text = gr.Button(value="Clear Text", interactive=True)
-                submit_button_text = gr.Button(value="Submit", interactive=True, variant="primary")
+            with gr.Column(visible=False) as modules_need_gpt2:
+                wiki_output = gr.Textbox(lines=6, label="Wiki")
+            with gr.Column(visible=False) as modules_not_need_gpt2:
+                chatbot = gr.Chatbot(label="Chat about Selected Object",).style(height=450,scale=0.5)
+                with gr.Column(visible=False) as modules_need_gpt3:
+                    chat_input = gr.Textbox(lines=1, label="Chat Input")
+                    with gr.Row():
+                        clear_button_text = gr.Button(value="Clear Text", interactive=True)
+                        submit_button_text = gr.Button(value="Submit", interactive=True, variant="primary")
+                    
+    openai_api_key.submit(init_openai_api_key, inputs=[openai_api_key], outputs=[modules_need_gpt,modules_need_gpt2, modules_need_gpt3, modules_not_need_gpt, modules_not_need_gpt2])
     clear_button_clike.click(
         lambda x: ([[], [], []], x, ""),
         [origin_image],
@@ -220,9 +234,9 @@ with gr.Blocks(
         show_progress=False
     )
     clear_button_image.click(
-        lambda: (None, [], [], [[], [], []], ""),
+        lambda: (None, [], [], [[], [], []], "", ""),
         [],
-        [image_input, chatbot, state, click_state, wiki_output],
+        [image_input, chatbot, state, click_state, wiki_output, origin_image],
         queue=False,
         show_progress=False
     )
@@ -234,20 +248,25 @@ with gr.Blocks(
         show_progress=False
     )
     image_input.clear(
-        lambda: (None, [], [], [[], [], []], ""),
+        lambda: (None, [], [], [[], [], []], "", ""),
         [],
-        [image_input, chatbot, state, click_state, wiki_output],
+        [image_input, chatbot, state, click_state, wiki_output, origin_image],
         queue=False,
         show_progress=False
     )
 
-    examples = gr.Examples(
+    def example_callback(x):
+        model.image_embedding = None
+        return x
+        
+    gr.Examples(
         examples=examples,
-        inputs=[image_input],
+        inputs=[example_image],
     )
 
-    image_input.upload(upload_callback,[image_input, state], [state, origin_image, click_state])
+    image_input.upload(upload_callback,[image_input, state], [state, origin_image, click_state, image_input])
     chat_input.submit(chat_with_points, [chat_input, click_state, state], [chatbot, state])
+    example_image.change(upload_callback,[example_image, state], [state, origin_image, click_state, image_input])
 
     # select coordinate
     image_input.select(inference_seg_cap, 
@@ -264,5 +283,5 @@ with gr.Blocks(
         outputs=[chatbot, state, click_state, chat_input, image_input, wiki_output],
         show_progress=False, queue=True)
     
-iface.queue(concurrency_count=5, api_open=False, max_size=10)
-iface.launch(server_name="0.0.0.0", enable_queue=True)
+iface.queue(concurrency_count=1, api_open=False, max_size=10)
+iface.launch(server_name="0.0.0.0", enable_queue=True, server_port=args.port, share=args.gradio_share)  

app_old.py

@@ -98,9 +98,9 @@ def chat_with_points(chat_input, click_state, state):
     return state, state
 
 def init_openai_api_key(api_key):
-    os.environ['OPENAI_API_KEY'] = api_key
+    # os.environ['OPENAI_API_KEY'] = api_key
     global model
-    model = CaptionAnything(args)
+    model = CaptionAnything(args, api_key)
 
 css='''
 #image_upload{min-height:200px}

caption_anything.py

@@ -8,18 +8,22 @@ import time
 from PIL import Image
 
 class CaptionAnything():
-    def __init__(self, args):
+    def __init__(self, args, api_key=""):
         self.args = args
         self.captioner = build_captioner(args.captioner, args.device, args)
         self.segmenter = build_segmenter(args.segmenter, args.device, args)
+        self.text_refiner = None
         if not args.disable_gpt:
-            self.init_refiner()
+            self.init_refiner(api_key)
 
-
-    def init_refiner(self):
-        if os.environ.get('OPENAI_API_KEY', None):
-            self.text_refiner = build_text_refiner(self.args.text_refiner, self.args.device, self.args)
-            
+    def init_refiner(self, api_key):
+        try:
+            self.text_refiner = build_text_refiner(self.args.text_refiner, self.args.device, self.args, api_key)
+            self.text_refiner.llm('hi') # test
+        except:
+            self.text_refiner = None
+            print('Openai api key is NOT given')
+    
     def inference(self, image, prompt, controls, disable_gpt=False):
         #  segment with prompt
         print("CA prompt: ", prompt, "CA controls",controls)
@@ -35,14 +39,14 @@ class CaptionAnything():
         print("seg_mask.shape: ", seg_mask.shape)
         #  captioning with mask
         if self.args.enable_reduce_tokens:
-            caption, crop_save_path = self.captioner.inference_with_reduced_tokens(image, seg_mask, crop_mode=self.args.seg_crop_mode, filter=self.args.clip_filter, regular_box = self.args.regular_box)
+            caption, crop_save_path = self.captioner.inference_with_reduced_tokens(image, seg_mask, crop_mode=self.args.seg_crop_mode, filter=self.args.clip_filter, disable_regular_box = self.args.disable_regular_box)
         else:
-            caption, crop_save_path = self.captioner.inference_seg(image, seg_mask, crop_mode=self.args.seg_crop_mode, filter=self.args.clip_filter, regular_box = self.args.regular_box)
+            caption, crop_save_path = self.captioner.inference_seg(image, seg_mask, crop_mode=self.args.seg_crop_mode, filter=self.args.clip_filter, disable_regular_box = self.args.disable_regular_box)
         #  refining with TextRefiner
         context_captions = []
         if self.args.context_captions:
             context_captions.append(self.captioner.inference(image))
-        if not disable_gpt and hasattr(self, "text_refiner"):
+        if not disable_gpt and self.text_refiner is not None:
             refined_caption = self.text_refiner.inference(query=caption, controls=controls, context=context_captions)
         else:
             refined_caption = {'raw_caption': caption}                
@@ -54,14 +58,14 @@ class CaptionAnything():
     
 def parse_augment():
     parser = argparse.ArgumentParser()
-    parser.add_argument('--captioner', type=str, default="blip")
+    parser.add_argument('--captioner', type=str, default="blip2")
     parser.add_argument('--segmenter', type=str, default="base")
     parser.add_argument('--text_refiner', type=str, default="base")
     parser.add_argument('--segmenter_checkpoint', type=str, default="segmenter/sam_vit_h_4b8939.pth")
-    parser.add_argument('--seg_crop_mode', type=str, default="w_bg", choices=['wo_bg', 'w_bg'], help="whether to add or remove background of the image when captioning")
+    parser.add_argument('--seg_crop_mode', type=str, default="wo_bg", choices=['wo_bg', 'w_bg'], help="whether to add or remove background of the image when captioning")
     parser.add_argument('--clip_filter', action="store_true", help="use clip to filter bad captions")
     parser.add_argument('--context_captions', action="store_true", help="use surrounding captions to enhance current caption (TODO)")
-    parser.add_argument('--regular_box', action="store_true", default = False, help="crop image with a regular box")
+    parser.add_argument('--disable_regular_box', action="store_true", default = False, help="crop image with a regular box")
     parser.add_argument('--device', type=str, default="cuda:0")
     parser.add_argument('--port', type=int, default=6086, help="only useful when running gradio applications")  
     parser.add_argument('--debug', action="store_true")
@@ -101,7 +105,7 @@ if __name__ == "__main__":
             "language": "English",
         }
     
-    model = CaptionAnything(args)
+    model = CaptionAnything(args, os.environ['OPENAI_API_KEY'])
     for prompt in prompts:
         print('*'*30)
         print('Image path: ', image_path)

captioner/base_captioner.py

@@ -135,7 +135,7 @@ class BaseCaptioner:
         return caption, crop_save_path
         
 
-    def inference_seg(self, image: Union[np.ndarray, str], seg_mask: Union[np.ndarray, Image.Image, str], crop_mode="w_bg", filter=False, regular_box = False):
+    def inference_seg(self, image: Union[np.ndarray, str], seg_mask: Union[np.ndarray, Image.Image, str], crop_mode="w_bg", filter=False, disable_regular_box = False):
         if type(image) == str:
             image = Image.open(image)
         if type(seg_mask) == str:
@@ -151,14 +151,14 @@ class BaseCaptioner:
         else:
             image = np.array(image)
 
-        if regular_box:
-            min_area_box = new_seg_to_box(seg_mask)
-        else:
+        if disable_regular_box:
             min_area_box = seg_to_box(seg_mask)
+        else:
+            min_area_box = new_seg_to_box(seg_mask)
         return self.inference_box(image, min_area_box, filter)
     
 
-    def generate_seg_cropped_image(self, image: Union[np.ndarray, str], seg_mask: Union[np.ndarray, Image.Image, str], crop_mode="w_bg", regular_box = False):
+    def generate_seg_cropped_image(self, image: Union[np.ndarray, str], seg_mask: Union[np.ndarray, Image.Image, str], crop_mode="w_bg", disable_regular_box = False):
         if type(image) == str:
             image = Image.open(image)
         if type(seg_mask) == str:
@@ -173,10 +173,10 @@ class BaseCaptioner:
         else:
             image = np.array(image)
 
-        if regular_box:
-            box = new_seg_to_box(seg_mask)
-        else:
+        if disable_regular_box:
             box = seg_to_box(seg_mask)
+        else:
+            box = new_seg_to_box(seg_mask)
         
         if np.array(box).size == 4: # [x0, y0, x1, y1], where (x0, y0), (x1, y1) represent top-left and bottom-right corners
             size = max(image.shape[0], image.shape[1])

captioner/blip.py

@@ -25,15 +25,15 @@ class BLIPCaptioner(BaseCaptioner):
                 image = Image.open(image)
         inputs = self.processor(image, return_tensors="pt").to(self.device, self.torch_dtype)
         out = self.model.generate(**inputs, max_new_tokens=50)
-        captions = self.processor.decode(out[0], skip_special_tokens=True)
+        captions = self.processor.decode(out[0], skip_special_tokens=True).strip()
         if self.enable_filter and filter:
             captions = self.filter_caption(image, captions)
         print(f"\nProcessed ImageCaptioning by BLIPCaptioner, Output Text: {captions}")
         return captions
     
     @torch.no_grad()
-    def inference_with_reduced_tokens(self, image: Union[np.ndarray, Image.Image, str], seg_mask, crop_mode="w_bg", filter=False, regular_box = False):
-        crop_save_path = self.generate_seg_cropped_image(image=image, seg_mask=seg_mask, crop_mode=crop_mode, regular_box=regular_box)
+    def inference_with_reduced_tokens(self, image: Union[np.ndarray, Image.Image, str], seg_mask, crop_mode="w_bg", filter=False, disable_regular_box = False):
+        crop_save_path = self.generate_seg_cropped_image(image=image, seg_mask=seg_mask, crop_mode=crop_mode, disable_regular_box=disable_regular_box)
         if type(image) == str: # input path
             image = Image.open(image)
         inputs = self.processor(image, return_tensors="pt")
@@ -45,7 +45,7 @@ class BLIPCaptioner(BaseCaptioner):
         seg_mask = seg_mask.float()
         pixel_masks = seg_mask.unsqueeze(0).to(self.device)
         out = self.model.generate(pixel_values=pixel_values, pixel_masks=pixel_masks, max_new_tokens=50)
-        captions = self.processor.decode(out[0], skip_special_tokens=True)
+        captions = self.processor.decode(out[0], skip_special_tokens=True).strip()
         if self.enable_filter and filter:
             captions = self.filter_caption(image, captions)
         print(f"\nProcessed ImageCaptioning by BLIPCaptioner, Output Text: {captions}")

captioner/blip2.py

@@ -22,9 +22,10 @@ class BLIP2Captioner(BaseCaptioner):
                 image = Image.open(image)
 
         if not self.dialogue:
-            inputs = self.processor(image, text = 'Ignore the black background! This is a photo of ', return_tensors="pt").to(self.device, self.torch_dtype)
+            text_prompt = 'Context: ignore the white part in this image. Question: describe this image. Answer:'
+            inputs = self.processor(image, text = text_prompt, return_tensors="pt").to(self.device, self.torch_dtype)
             out = self.model.generate(**inputs, max_new_tokens=50)
-            captions = self.processor.decode(out[0], skip_special_tokens=True)
+            captions = self.processor.decode(out[0], skip_special_tokens=True).strip()
             if self.enable_filter and filter:
                 captions = self.filter_caption(image, captions)
             print(f"\nProcessed ImageCaptioning by BLIP2Captioner, Output Text: {captions}")

captioner/git.py

@@ -22,15 +22,15 @@ class GITCaptioner(BaseCaptioner):
             image = Image.open(image)
         pixel_values = self.processor(images=image, return_tensors="pt").pixel_values.to(self.device, self.torch_dtype)
         generated_ids = self.model.generate(pixel_values=pixel_values, max_new_tokens=50)
-        generated_caption = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        generated_caption = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
         if self.enable_filter and filter:
             captions = self.filter_caption(image, captions)
         print(f"\nProcessed ImageCaptioning by GITCaptioner, Output Text: {generated_caption}")
         return generated_caption
 
     @torch.no_grad()
-    def inference_with_reduced_tokens(self, image: Union[np.ndarray, Image.Image, str], seg_mask, crop_mode="w_bg", filter=False, regular_box = False):
-        crop_save_path = self.generate_seg_cropped_image(image=image, seg_mask=seg_mask, crop_mode=crop_mode, regular_box=regular_box)
+    def inference_with_reduced_tokens(self, image: Union[np.ndarray, Image.Image, str], seg_mask, crop_mode="w_bg", filter=False, disable_regular_box = False):
+        crop_save_path = self.generate_seg_cropped_image(image=image, seg_mask=seg_mask, crop_mode=crop_mode, disable_regular_box=disable_regular_box)
         if type(image) == str: # input path
             image = Image.open(image)
         inputs = self.processor(images=image, return_tensors="pt")
@@ -42,7 +42,7 @@ class GITCaptioner(BaseCaptioner):
         seg_mask = seg_mask.float()
         pixel_masks = seg_mask.unsqueeze(0).to(self.device)
         out = self.model.generate(pixel_values=pixel_values, pixel_masks=pixel_masks, max_new_tokens=50)
-        captions = self.processor.decode(out[0], skip_special_tokens=True)
+        captions = self.processor.decode(out[0], skip_special_tokens=True).strip()
         if self.enable_filter and filter:
             captions = self.filter_caption(image, captions)
         print(f"\nProcessed ImageCaptioning by BLIPCaptioner, Output Text: {captions}")

image_editing_utils.py

@@ -35,7 +35,8 @@ def create_bubble_frame(image, text, point, font_path='DejaVuSansCondensed-Bold.
 
     # Wrap the text to fit within the max_text_width
     lines = wrap_text(text, font, max_text_width)
-    text_width, text_height = font.getsize(lines[0])
+    text_width = max([font.getsize(line)[0] for line in lines])
+    _, text_height = font.getsize(lines[0])
     text_height = text_height * len(lines)
 
     # Define bubble frame dimensions
@@ -48,7 +49,7 @@ def create_bubble_frame(image, text, point, font_path='DejaVuSansCondensed-Bold.
 
     # Draw the bubble frame on the new image
     draw = ImageDraw.Draw(bubble)
-    draw.rectangle([(0, 0), (bubble_width - 1, bubble_height - 1)], fill=(255, 255, 255, 0), outline=(255, 255, 255, 0), width=2)
+    # draw.rectangle([(0, 0), (bubble_width - 1, bubble_height - 1)], fill=(255, 255, 255, 0), outline=(255, 255, 255, 0), width=2)
 
     # Draw the wrapped text line by line
     y_text = padding

segmenter/base_segmenter.py

@@ -46,7 +46,7 @@ class BaseSegmenter:
             new_masks = np.concatenate([mask["segmentation"][np.newaxis,:] for mask in masks])
             return new_masks
         else:
-            if not self.reuse_feature:
+            if not self.reuse_feature or self.image_embedding is None:
                 self.set_image(image)
                 self.predictor.set_image(self.image)
             else:

text_refiner/__init__.py

@@ -1,6 +1,6 @@
 from text_refiner.text_refiner import TextRefiner
 
 
-def build_text_refiner(type, device, args=None):
+def build_text_refiner(type, device, args=None, api_key=""):
     if type == 'base':
-        return TextRefiner(device)
+        return TextRefiner(device, api_key)

text_refiner/text_refiner.py

@@ -5,12 +5,9 @@ from transformers import pipeline, BlipProcessor, BlipForConditionalGeneration,
 import pdb
 
 class TextRefiner:
-    def __init__(self, device):
+    def __init__(self, device, api_key=""):
         print(f"Initializing TextRefiner to {device}")
-        try:
-            self.llm = OpenAI(model_name="gpt-3.5-turbo", temperature=0)
-        except:
-            print('Openai api key is NOT given')
+        self.llm = OpenAI(model_name="gpt-3.5-turbo", temperature=0, openai_api_key=api_key)
         self.prompt_tag = {
             "imagination": {"True": "could",
                             "False": "could not"}

tools.py

@@ -1,7 +1,9 @@
 import cv2
+import torch
 import numpy as np
 from PIL import Image
 import copy
+import time
 
 
 def colormap(rgb=True):
@@ -100,16 +102,6 @@ color_list = colormap()
 color_list = color_list.astype('uint8').tolist()
 
 
-def gauss_filter(kernel_size, sigma):
-	max_idx = kernel_size // 2
-	idx = np.linspace(-max_idx, max_idx, kernel_size)
-	Y, X = np.meshgrid(idx, idx)
-	gauss_filter = np.exp(-(X**2 + Y**2) / (2*sigma**2))
-	gauss_filter /= np.sum(np.sum(gauss_filter))
-	
-	return gauss_filter
-
-
 def vis_add_mask(image, mask, color, alpha, kernel_size):
 	color = np.array(color)
 	mask = mask.astype('float').copy()
@@ -129,6 +121,23 @@ def vis_add_mask_wo_blur(image, mask, color, alpha):
 	return image
 
 
+def vis_add_mask_wo_gaussian(image, background_mask, contour_mask, background_color, contour_color, background_alpha, contour_alpha):
+	background_color = np.array(background_color)
+	contour_color = np.array(contour_color)
+
+	# background_mask = 1 - background_mask
+	# contour_mask = 1 - contour_mask
+
+	for i in range(3):
+		image[:, :, i] = image[:, :, i] * (1-background_alpha+background_mask*background_alpha) \
+			+ background_color[i] * (background_alpha-background_mask*background_alpha)
+		
+		image[:, :, i] = image[:, :, i] * (1-contour_alpha+contour_mask*contour_alpha) \
+			+ contour_color[i] * (contour_alpha-contour_mask*contour_alpha)
+
+	return image.astype('uint8')
+
+
 def mask_painter(input_image, input_mask, background_alpha=0.7, background_blur_radius=7, contour_width=3, contour_color=3, contour_alpha=1):
 	"""
 	Input:
@@ -146,11 +155,7 @@ def mask_painter(input_image, input_mask, background_alpha=0.7, background_blur_
 	assert input_image.shape[:2] == input_mask.shape, 'different shape'
 	assert background_blur_radius % 2 * contour_width % 2 > 0, 'background_blur_radius and contour_width must be ODD'
 
-	width, height = input_image.shape[0], input_image.shape[1]
-	res = 1024
-	ratio = min(1.0 * res / max(width, height), 1.0)  
-	input_image = cv2.resize(input_image, (int(height*ratio), int(width*ratio)))
-	input_mask = cv2.resize(input_mask, (int(height*ratio), int(width*ratio)))
+	
 	# 0: background, 1: foreground
 	input_mask[input_mask>0] = 255
 
@@ -163,15 +168,120 @@ def mask_painter(input_image, input_mask, background_alpha=0.7, background_blur_
 	kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (contour_width, contour_width))
 	contour_mask = cv2.dilate(contour_mask, kernel)
 	painted_image = vis_add_mask(painted_image, 255-contour_mask, color_list[contour_color], contour_alpha, contour_width)
-	painted_image = cv2.resize(painted_image, (height, width))
+
+	# painted_image = background_dist_map
+	
+	return painted_image
+
+
+def mask_generator_00(mask, background_radius, contour_radius):
+	# no background width when '00'
+	# distance map
+	dist_transform_fore = cv2.distanceTransform(mask, cv2.DIST_L2, 3)
+	dist_transform_back = cv2.distanceTransform(1-mask, cv2.DIST_L2, 3)
+	dist_map = dist_transform_fore - dist_transform_back
+	# ...:::!!!:::...
+	contour_radius += 2
+	contour_mask = np.abs(np.clip(dist_map, -contour_radius, contour_radius))
+	contour_mask = contour_mask / np.max(contour_mask)
+	contour_mask[contour_mask>0.5] = 1.
+
+	return mask, contour_mask
+
+
+def mask_generator_01(mask, background_radius, contour_radius):
+	# no background width when '00'
+	# distance map
+	dist_transform_fore = cv2.distanceTransform(mask, cv2.DIST_L2, 3)
+	dist_transform_back = cv2.distanceTransform(1-mask, cv2.DIST_L2, 3)
+	dist_map = dist_transform_fore - dist_transform_back
+	# ...:::!!!:::...
+	contour_radius += 2
+	contour_mask = np.abs(np.clip(dist_map, -contour_radius, contour_radius))
+	contour_mask = contour_mask / np.max(contour_mask)
+	return mask, contour_mask
+
+
+def mask_generator_10(mask, background_radius, contour_radius):
+	# distance map
+	dist_transform_fore = cv2.distanceTransform(mask, cv2.DIST_L2, 3)
+	dist_transform_back = cv2.distanceTransform(1-mask, cv2.DIST_L2, 3)
+	dist_map = dist_transform_fore - dist_transform_back
+	# .....:::::!!!!!
+	background_mask = np.clip(dist_map, -background_radius, background_radius)
+	background_mask = (background_mask - np.min(background_mask))
+	background_mask = background_mask / np.max(background_mask)
+	# ...:::!!!:::...
+	contour_radius += 2
+	contour_mask = np.abs(np.clip(dist_map, -contour_radius, contour_radius))
+	contour_mask = contour_mask / np.max(contour_mask)
+	contour_mask[contour_mask>0.5] = 1.
+	return background_mask, contour_mask
+
+
+def mask_generator_11(mask, background_radius, contour_radius):
+	# distance map
+	dist_transform_fore = cv2.distanceTransform(mask, cv2.DIST_L2, 3)
+	dist_transform_back = cv2.distanceTransform(1-mask, cv2.DIST_L2, 3)
+	dist_map = dist_transform_fore - dist_transform_back
+	# .....:::::!!!!!
+	background_mask = np.clip(dist_map, -background_radius, background_radius)
+	background_mask = (background_mask - np.min(background_mask))
+	background_mask = background_mask / np.max(background_mask)
+	# ...:::!!!:::...
+	contour_radius += 2
+	contour_mask = np.abs(np.clip(dist_map, -contour_radius, contour_radius))
+	contour_mask = contour_mask / np.max(contour_mask)
+	return background_mask, contour_mask
+
+
+def mask_painter_wo_gaussian(input_image, input_mask, background_alpha=0.5, background_blur_radius=7, contour_width=3, contour_color=3, contour_alpha=1, mode='11'):
+	"""
+	Input:
+	input_image: numpy array
+	input_mask: numpy array
+	background_alpha: transparency of background, [0, 1], 1: all black, 0: do nothing
+	background_blur_radius: radius of background blur, must be odd number
+	contour_width: width of mask contour, must be odd number
+	contour_color: color index (in color map) of mask contour, 0: black, 1: white, >1: others
+	contour_alpha: transparency of mask contour, [0, 1], if 0: no contour highlighted
+	mode: painting mode, '00', no blur, '01' only blur contour, '10' only blur background, '11' blur both
+
+	Output:
+	painted_image: numpy array
+	"""
+	assert input_image.shape[:2] == input_mask.shape, 'different shape'
+	assert background_blur_radius % 2 * contour_width % 2 > 0, 'background_blur_radius and contour_width must be ODD'
+	assert mode in ['00', '01', '10', '11'], 'mode should be 00, 01, 10, or 11'
+
+	# downsample input image and mask
+	width, height = input_image.shape[0], input_image.shape[1]
+	res = 1024
+	ratio = min(1.0 * res / max(width, height), 1.0)  
+	input_image = cv2.resize(input_image, (int(height*ratio), int(width*ratio)))
+	input_mask = cv2.resize(input_mask, (int(height*ratio), int(width*ratio)))
+	
+	# 0: background, 1: foreground
+	msk = np.clip(input_mask, 0, 1)
+
+	# generate masks for background and contour pixels
+	background_radius = (background_blur_radius - 1) // 2
+	contour_radius = (contour_width - 1) // 2
+	generator_dict = {'00':mask_generator_00, '01':mask_generator_01, '10':mask_generator_10, '11':mask_generator_11}
+	background_mask, contour_mask = generator_dict[mode](msk, background_radius, contour_radius)
+
+	# paint
+	painted_image = vis_add_mask_wo_gaussian\
+		(input_image, background_mask, contour_mask, color_list[0], color_list[contour_color], background_alpha, contour_alpha)	# black for background
+
 	return painted_image
 
 
 if __name__ == '__main__':
 	
 	background_alpha = 0.7  	# transparency of background 1: all black, 0: do nothing
-	background_blur_radius = 35	# radius of background blur, must be odd number
-	contour_width = 7       	# contour width, must be odd number
+	background_blur_radius = 31	# radius of background blur, must be odd number
+	contour_width = 11       	# contour width, must be odd number
 	contour_color = 3      		# id in color map, 0: black, 1: white, >1: others
 	contour_alpha = 1       	# transparency of background, 0: no contour highlighted
 
@@ -180,8 +290,54 @@ if __name__ == '__main__':
 	input_mask = np.array(Image.open('./test_img/painter_input_mask.jpg').convert('P'))
 	
 	# paint
-	painted_image = mask_painter(input_image, input_mask, background_alpha, background_blur_radius, contour_width, contour_color, contour_alpha)
+	overall_time_1 = 0
+	overall_time_2 = 0
+	overall_time_3 = 0
+	overall_time_4 = 0
+	overall_time_5 = 0
+	
+	for i in range(50):
+		t2 = time.time()
+		painted_image_00 = mask_painter_wo_gaussian(input_image, input_mask, background_alpha, background_blur_radius, contour_width, contour_color, contour_alpha, mode='00')
+		e2 = time.time()
+
+		t3 = time.time()
+		painted_image_10 = mask_painter_wo_gaussian(input_image, input_mask, background_alpha, background_blur_radius, contour_width, contour_color, contour_alpha, mode='10')
+		e3 = time.time()
+
+		t1 = time.time()
+		painted_image = mask_painter(input_image, input_mask, background_alpha, background_blur_radius, contour_width, contour_color, contour_alpha)
+		e1 = time.time()
+
+		t4 = time.time()
+		painted_image_01 = mask_painter_wo_gaussian(input_image, input_mask, background_alpha, background_blur_radius, contour_width, contour_color, contour_alpha, mode='01')
+		e4 = time.time()
+
+		t5 = time.time()
+		painted_image_11 = mask_painter_wo_gaussian(input_image, input_mask, background_alpha, background_blur_radius, contour_width, contour_color, contour_alpha, mode='11')
+		e5 = time.time()
+
+		overall_time_1 += (e1 - t1)
+		overall_time_2 += (e2 - t2)
+		overall_time_3 += (e3 - t3)
+		overall_time_4 += (e4 - t4)
+		overall_time_5 += (e5 - t5)
+
+	print(f'average time w gaussian: {overall_time_1/50}')
+	print(f'average time w/o gaussian00: {overall_time_2/50}')
+	print(f'average time w/o gaussian10: {overall_time_3/50}')
+	print(f'average time w/o gaussian01: {overall_time_4/50}')
+	print(f'average time w/o gaussian11: {overall_time_5/50}')
 
 	# save
-	painted_image = Image.fromarray(painted_image)
-	painted_image.save('./test_img/painter_output_image.png')
+	painted_image_00 = Image.fromarray(painted_image_00)
+	painted_image_00.save('./test_img/painter_output_image_00.png')
+
+	painted_image_10 = Image.fromarray(painted_image_10)
+	painted_image_10.save('./test_img/painter_output_image_10.png')
+
+	painted_image_01 = Image.fromarray(painted_image_01)
+	painted_image_01.save('./test_img/painter_output_image_01.png')
+
+	painted_image_11 = Image.fromarray(painted_image_11)
+	painted_image_11.save('./test_img/painter_output_image_11.png')