share ocr_reader to accelerate inferenec

app.py

@@ -17,7 +17,7 @@ from caption_anything.text_refiner import build_text_refiner
 from caption_anything.segmenter import build_segmenter
 from caption_anything.utils.chatbot import ConversationBot, build_chatbot_tools, get_new_image_name
 from segment_anything import sam_model_registry
-
+import easyocr
 
 args = parse_augment()
 args.segmenter = "huge"
@@ -30,6 +30,8 @@ else:
     
 shared_captioner = build_captioner(args.captioner, args.device, args)
 shared_sam_model = sam_model_registry[seg_model_map[args.segmenter]](checkpoint=segmenter_checkpoint).to(args.device)
+ocr_lang = ["ch_tra", "en"]
+shared_ocr_reader = easyocr.Reader(ocr_lang)
 tools_dict = {e.split('_')[0].strip(): e.split('_')[1].strip() for e in args.chat_tools_dict.split(',')}
 shared_chatbot_tools = build_chatbot_tools(tools_dict)
 
@@ -57,13 +59,13 @@ class ImageSketcher(gr.Image):
         return super().preprocess(x)
 
 
-def build_caption_anything_with_models(args, api_key="", captioner=None, sam_model=None, text_refiner=None,
+def build_caption_anything_with_models(args, api_key="", captioner=None, sam_model=None, ocr_reader=None, text_refiner=None,
                                        session_id=None):
     segmenter = build_segmenter(args.segmenter, args.device, args, model=sam_model)
     captioner = captioner
     if session_id is not None:
         print('Init caption anything for session {}'.format(session_id))
-    return CaptionAnything(args, api_key, captioner=captioner, segmenter=segmenter, text_refiner=text_refiner)
+    return CaptionAnything(args, api_key, captioner=captioner, segmenter=segmenter, ocr_reader=ocr_reader, text_refiner=text_refiner)
 
 
 def init_openai_api_key(api_key=""):
@@ -146,6 +148,7 @@ def upload_callback(image_input, state, visual_chatgpt=None):
         api_key="",
         captioner=shared_captioner,
         sam_model=shared_sam_model,
+        ocr_reader=shared_ocr_reader,
         session_id=iface.app_id
     )
     model.segmenter.set_image(image_input)
@@ -154,6 +157,7 @@ def upload_callback(image_input, state, visual_chatgpt=None):
     input_size = model.input_size
     
     if visual_chatgpt is not None:
+        print('upload_callback: add caption to chatGPT memory')
         new_image_path = get_new_image_name('chat_image', func_name='upload')
         image_input.save(new_image_path)
         visual_chatgpt.current_image = new_image_path
@@ -192,6 +196,7 @@ def inference_click(image_input, point_prompt, click_mode, enable_wiki, language
         api_key="",
         captioner=shared_captioner,
         sam_model=shared_sam_model,
+        ocr_reader=shared_ocr_reader,
         text_refiner=text_refiner,
         session_id=iface.app_id
     )
@@ -213,6 +218,7 @@ def inference_click(image_input, point_prompt, click_mode, enable_wiki, language
     x, y = input_points[-1]
     
     if visual_chatgpt is not None:
+        print('inference_click: add caption to chatGPT memory')
         new_crop_save_path = get_new_image_name('chat_image', func_name='crop')
         Image.open(out["crop_save_path"]).save(new_crop_save_path)
         point_prompt = f'You should primarly use tools on the selected regional image (description: {text}, path: {new_crop_save_path}), which is a part of the whole image (path: {visual_chatgpt.current_image}). If human mentioned some objects not in the selected region, you can use tools on the whole image.'
@@ -273,6 +279,7 @@ def inference_traject(sketcher_image, enable_wiki, language, sentiment, factuali
         api_key="",
         captioner=shared_captioner,
         sam_model=shared_sam_model,
+        ocr_reader=shared_ocr_reader,
         text_refiner=text_refiner,
         session_id=iface.app_id
     )
@@ -325,6 +332,7 @@ def cap_everything(image_input, visual_chatgpt, text_refiner):
         api_key="",
         captioner=shared_captioner,
         sam_model=shared_sam_model,
+        ocr_reader=shared_ocr_reader,
         text_refiner=text_refiner,
         session_id=iface.app_id
     )

caption_anything/captioner/blip2.py

@@ -6,6 +6,7 @@ from transformers import AutoProcessor, Blip2ForConditionalGeneration
 
 from caption_anything.utils.utils import is_platform_win, load_image
 from .base_captioner import BaseCaptioner
+import time
 
 class BLIP2Captioner(BaseCaptioner):
     def __init__(self, device, dialogue: bool = False, enable_filter: bool = False):
@@ -33,8 +34,7 @@ class BLIP2Captioner(BaseCaptioner):
         if not self.dialogue:
             inputs = self.processor(image, text = args['text_prompt'], return_tensors="pt").to(self.device, self.torch_dtype)
             out = self.model.generate(**inputs, return_dict_in_generate=True, output_scores=True, max_new_tokens=50)
-            captions = self.processor.batch_decode(out.sequences, skip_special_tokens=True)
-            caption = [caption.strip() for caption in captions][0]
+            caption = self.processor.decode(out.sequences[0], skip_special_tokens=True).strip()
             if self.enable_filter and filter:
                 print('reference caption: {}, caption: {}'.format(args['reference_caption'], caption))
                 clip_score = self.filter_caption(image, caption, args['reference_caption'])

caption_anything/model.py

@@ -8,6 +8,7 @@ import numpy as np
 from PIL import Image
 import easyocr
 import copy
+import time
 from caption_anything.captioner import build_captioner, BaseCaptioner
 from caption_anything.segmenter import build_segmenter, build_segmenter_densecap
 from caption_anything.text_refiner import build_text_refiner
@@ -16,14 +17,15 @@ from caption_anything.utils.utils import mask_painter_foreground_all, mask_paint
 from caption_anything.utils.densecap_painter import draw_bbox
             
 class CaptionAnything:
-    def __init__(self, args, api_key="", captioner=None, segmenter=None, text_refiner=None):
+    def __init__(self, args, api_key="", captioner=None, segmenter=None, ocr_reader=None, text_refiner=None):
         self.args = args
         self.captioner = build_captioner(args.captioner, args.device, args) if captioner is None else captioner
         self.segmenter = build_segmenter(args.segmenter, args.device, args) if segmenter is None else segmenter
         self.segmenter_densecap = build_segmenter_densecap(args.segmenter, args.device, args, model=self.segmenter.model)
+        self.ocr_lang = ["ch_tra", "en"]
+        self.ocr_reader = ocr_reader if ocr_reader is not None else easyocr.Reader(self.ocr_lang)
         
-        self.lang = ["ch_tra", "en"]
-        self.reader = easyocr.Reader(self.lang)
+
         self.text_refiner = None
         if not args.disable_gpt:
             if text_refiner is not None:
@@ -31,6 +33,7 @@ class CaptionAnything:
             elif api_key != "":
                 self.init_refiner(api_key)
         self.require_caption_prompt = args.captioner == 'blip2'
+        print('text_refiner init time: ', time.time() - t0)
         
     @property
     def image_embedding(self):
@@ -213,7 +216,7 @@ class CaptionAnything:
     def parse_ocr(self, image, thres=0.2):
         width, height = get_image_shape(image)
         image = load_image(image, return_type='numpy')
-        bounds = self.reader.readtext(image)
+        bounds = self.ocr_reader.readtext(image)
         bounds = [bound for bound in bounds if bound[2] > thres]
         print('Process OCR Text:\n', bounds)
         
@@ -257,7 +260,7 @@ class CaptionAnything:
 if __name__ == "__main__":
     from caption_anything.utils.parser import parse_augment
     args = parse_augment()
-    image_path = 'image/ocr/Untitled.png'
+    image_path = 'result/wt/memes/87226084.jpg'
     image = Image.open(image_path)
     prompts = [
         {