diff --git a/deepdoc/parser/figure_parser.py b/deepdoc/parser/figure_parser.py index ec5e333de2..d10795a431 100644 --- a/deepdoc/parser/figure_parser.py +++ b/deepdoc/parser/figure_parser.py @@ -24,19 +24,24 @@ from common.connection_utils import timeout from rag.app.picture import vision_llm_chunk as picture_vision_llm_chunk from rag.prompts.generator import vision_llm_figure_describe_prompt, vision_llm_figure_describe_prompt_with_context from rag.nlp import append_context2table_image4pdf +from rag.utils.lazy_image import ensure_pil_image, open_image_for_processing, is_image_like # need to delete before pr def vision_figure_parser_figure_data_wrapper(figures_data_without_positions): if not figures_data_without_positions: return [] - return [ - ( - (figure_data[1], [figure_data[0]]), - [(0, 0, 0, 0, 0)], + res = [] + for figure_data in figures_data_without_positions: + img = ensure_pil_image(figure_data[1]) + if not isinstance(img, Image.Image): + continue + res.append( + ( + (img, [figure_data[0]]), + [(0, 0, 0, 0, 0)], + ) ) - for figure_data in figures_data_without_positions - if isinstance(figure_data[1], Image.Image) - ] + return res def vision_figure_parser_docx_wrapper(sections, tbls, callback=None,**kwargs): if not sections: @@ -96,7 +101,7 @@ def vision_figure_parser_pdf_wrapper(tbls, callback=None, **kwargs): if vision_model: def is_figure_item(item): - return isinstance(item[0][0], Image.Image) and isinstance(item[0][1], list) + return is_image_like(item[0][0]) and isinstance(item[0][1], list) figures_data = [item for item in tbls if is_figure_item(item)] figure_contexts = [] @@ -134,6 +139,9 @@ def vision_figure_parser_docx_wrapper_naive(chunks, idx_lst, callback=None, **kw if vision_model: @timeout(30, 3) def worker(idx, ck): + img, close_after = open_image_for_processing(ck.get("image"), allow_bytes=True) + if not isinstance(img, Image.Image): + return idx, "" context_above = ck.get("context_above", "") context_below = ck.get("context_below", "") if context_above or context_below: @@ -149,13 +157,20 @@ def vision_figure_parser_docx_wrapper_naive(chunks, idx_lst, callback=None, **kw prompt = vision_llm_figure_describe_prompt() logging.info(f"[VisionFigureParser] figure={idx} context_len=0 prompt=default") - description_text = picture_vision_llm_chunk( - binary=ck.get("image"), - vision_model=vision_model, - prompt=prompt, - callback=callback, - ) - return idx, description_text + try: + description_text = picture_vision_llm_chunk( + binary=img, + vision_model=vision_model, + prompt=prompt, + callback=callback, + ) + return idx, description_text + finally: + if close_after and isinstance(img, Image.Image): + try: + img.close() + except Exception: + pass with ThreadPoolExecutor(max_workers=10) as executor: futures = [ @@ -187,13 +202,15 @@ class VisionFigureParser: # position if len(item) == 2 and isinstance(item[0], tuple) and len(item[0]) == 2 and isinstance(item[1], list) and isinstance(item[1][0], tuple) and len(item[1][0]) == 5: img_desc = item[0] - assert len(img_desc) == 2 and isinstance(img_desc[0], Image.Image) and isinstance(img_desc[1], list), "Should be (figure, [description])" - self.figures.append(img_desc[0]) + img = ensure_pil_image(img_desc[0]) + assert len(img_desc) == 2 and isinstance(img, Image.Image) and isinstance(img_desc[1], list), "Should be (figure, [description])" + self.figures.append(img) self.descriptions.append(img_desc[1]) self.positions.append(item[1]) else: - assert len(item) == 2 and isinstance(item[0], Image.Image) and isinstance(item[1], list), f"Unexpected form of figure data: get {len(item)=}, {item=}" - self.figures.append(item[0]) + img = ensure_pil_image(item[0]) + assert len(item) == 2 and isinstance(img, Image.Image) and isinstance(item[1], list), f"Unexpected form of figure data: get {len(item)=}, {item=}" + self.figures.append(img) self.descriptions.append(item[1]) def _assemble(self): diff --git a/deepdoc/vision/operators.py b/deepdoc/vision/operators.py index 65d2efa4cb..43b55ccd3a 100644 --- a/deepdoc/vision/operators.py +++ b/deepdoc/vision/operators.py @@ -22,6 +22,7 @@ import cv2 import numpy as np import math from PIL import Image +from rag.utils.lazy_image import ensure_pil_image class DecodeImage: @@ -128,8 +129,9 @@ class NormalizeImage: def __call__(self, data): img = data['image'] from PIL import Image - if isinstance(img, Image.Image): - img = np.array(img) + pil = ensure_pil_image(img) + if isinstance(pil, Image.Image): + img = np.array(pil) assert isinstance(img, np.ndarray), "invalid input 'img' in NormalizeImage" data['image'] = ( @@ -147,8 +149,9 @@ class ToCHWImage: def __call__(self, data): img = data['image'] from PIL import Image - if isinstance(img, Image.Image): - img = np.array(img) + pil = ensure_pil_image(img) + if isinstance(pil, Image.Image): + img = np.array(pil) data['image'] = img.transpose((2, 0, 1)) return data diff --git a/rag/app/book.py b/rag/app/book.py index d3c45b4251..847aec50a3 100644 --- a/rag/app/book.py +++ b/rag/app/book.py @@ -27,6 +27,7 @@ from rag.nlp import rag_tokenizer from deepdoc.parser import PdfParser, HtmlParser from deepdoc.parser.figure_parser import vision_figure_parser_docx_wrapper from PIL import Image +from rag.utils.lazy_image import LazyDocxImage class Pdf(PdfParser): @@ -85,7 +86,11 @@ def chunk(filename, binary=None, from_page=0, to_page=100000, lang="Chinese", ca tbls = vision_figure_parser_docx_wrapper(sections=sections, tbls=tbls, callback=callback, **kwargs) # tbls = [((None, lns), None) for lns in tbls] - sections = [(item[0], item[1] if item[1] is not None else "") for item in sections if not isinstance(item[1], Image.Image)] + sections = [ + (item[0], item[1] if item[1] is not None else "") + for item in sections + if not isinstance(item[1], (Image.Image, LazyDocxImage)) + ] callback(0.8, "Finish parsing.") elif re.search(r"\.pdf$", filename, re.IGNORECASE): diff --git a/rag/app/naive.py b/rag/app/naive.py index 6c49d53bfb..fcece22aba 100644 --- a/rag/app/naive.py +++ b/rag/app/naive.py @@ -33,6 +33,7 @@ from common.token_utils import num_tokens_from_string from common.constants import LLMType from api.db.services.llm_service import LLMBundle from rag.utils.file_utils import extract_embed_file, extract_links_from_pdf, extract_links_from_docx, extract_html +from rag.utils.lazy_image import LazyDocxImage from deepdoc.parser import DocxParser, ExcelParser, HtmlParser, JsonParser, MarkdownElementExtractor, MarkdownParser, PdfParser, TxtParser from deepdoc.parser.figure_parser import VisionFigureParser, vision_figure_parser_docx_wrapper_naive, vision_figure_parser_pdf_wrapper from deepdoc.parser.pdf_parser import PlainParser, VisionParser @@ -237,7 +238,7 @@ class Docx(DocxParser): imgs = paragraph._element.xpath(".//pic:pic") if not imgs: return None - res_img = None + image_blobs = [] for img in imgs: embed = img.xpath(".//a:blip/@r:embed") if not embed: @@ -261,17 +262,11 @@ class Docx(DocxParser): except Exception as e: logging.warning(f"The recognized image stream appears to be corrupted. Skipping image, exception: {e}") continue - try: - image = Image.open(BytesIO(image_blob)).convert("RGB") - if res_img is None: - res_img = image - else: - res_img = concat_img(res_img, image) - except Exception as e: - logging.warning(f"Fail to open or concat images, exception: {e}") - continue + image_blobs.append(image_blob) - return res_img + if not image_blobs: + return None + return LazyDocxImage(image_blobs) def __clean(self, line): line = re.sub(r"\u3000", " ", line).strip() diff --git a/rag/app/presentation.py b/rag/app/presentation.py index c6f922bf78..909fd61a30 100644 --- a/rag/app/presentation.py +++ b/rag/app/presentation.py @@ -20,7 +20,6 @@ import re from collections import defaultdict from io import BytesIO -from PIL import Image from PyPDF2 import PdfReader as pdf2_read from deepdoc.parser import PdfParser, PlainParser @@ -29,6 +28,7 @@ from rag.app.naive import by_plaintext, PARSERS from common.parser_config_utils import normalize_layout_recognizer from rag.nlp import rag_tokenizer from rag.nlp import tokenize +from rag.utils.lazy_image import ensure_pil_image, is_image_like class Pdf(PdfParser): @@ -228,8 +228,10 @@ def chunk(filename, binary=None, from_page=0, to_page=100000, lang="Chinese", ca for pn, (txt, img) in enumerate(sections): d = copy.deepcopy(doc) pn += from_page - if not isinstance(img, Image.Image): + if not is_image_like(img): img = None + else: + img = ensure_pil_image(img) d["image"] = img d["page_num_int"] = [pn + 1] d["top_int"] = [0] diff --git a/rag/nlp/__init__.py b/rag/nlp/__init__.py index d94d6301e6..364e953881 100644 --- a/rag/nlp/__init__.py +++ b/rag/nlp/__init__.py @@ -1212,6 +1212,10 @@ def docx_question_level(p, bull=-1): def concat_img(img1, img2): + from rag.utils.lazy_image import ensure_pil_image + + img1 = ensure_pil_image(img1) or img1 + img2 = ensure_pil_image(img2) or img2 if img1 and not img2: return img1 if not img1 and img2: diff --git a/rag/utils/base64_image.py b/rag/utils/base64_image.py index 7493834924..5a328dd3ea 100644 --- a/rag/utils/base64_image.py +++ b/rag/utils/base64_image.py @@ -24,6 +24,7 @@ from PIL import Image from common.misc_utils import thread_pool_exec +from rag.utils.lazy_image import open_image_for_processing test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAGQAAABkCAIAAAD/gAIDAAAA6ElEQVR4nO3QwQ3AIBDAsIP9d25XIC+EZE8QZc18w5l9O+AlZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBWYFZgVmBT+IYAHHLHkdEgAAAABJRU5ErkJggg==" test_image = base64.b64decode(test_image_base64) @@ -42,15 +43,24 @@ async def image2id(d: dict, storage_put_func: partial, objname: str, bucket: str def encode_image(): with BytesIO() as buf: - img = d["image"] + img, close_after = open_image_for_processing(d["image"], allow_bytes=False) if isinstance(img, bytes): buf.write(img) buf.seek(0) return buf.getvalue() + if not isinstance(img, Image.Image): + return None + if img.mode in ("RGBA", "P"): + orig_img = img img = img.convert("RGB") + if close_after: + try: + orig_img.close() + except Exception: + pass try: img.save(buf, format="JPEG") @@ -59,6 +69,11 @@ async def image2id(d: dict, storage_put_func: partial, objname: str, bucket: str return None buf.seek(0) + if close_after: + try: + img.close() + except Exception: + pass return buf.getvalue() jpeg_binary = await thread_pool_exec(encode_image) diff --git a/rag/utils/lazy_image.py b/rag/utils/lazy_image.py new file mode 100644 index 0000000000..7de2bfd5ce --- /dev/null +++ b/rag/utils/lazy_image.py @@ -0,0 +1,116 @@ +import logging +from io import BytesIO + +from PIL import Image + +from rag.nlp import concat_img + + +class LazyDocxImage: + def __init__(self, blobs, source=None): + self._blobs = [b for b in (blobs or []) if b] + self.source = source + self._pil = None + + def __bool__(self): + return bool(self._blobs) + + def to_pil(self): + if self._pil is not None: + try: + self._pil.load() + return self._pil + except Exception: + try: + self._pil.close() + except Exception: + pass + self._pil = None + res_img = None + for blob in self._blobs: + try: + image = Image.open(BytesIO(blob)).convert("RGB") + except Exception as e: + logging.info(f"LazyDocxImage: skip bad image blob: {e}") + continue + + if res_img is None: + res_img = image + continue + + new_img = concat_img(res_img, image) + if new_img is not res_img: + try: + res_img.close() + except Exception: + pass + try: + image.close() + except Exception: + pass + res_img = new_img + + self._pil = res_img + return self._pil + + def to_pil_detached(self): + pil = self.to_pil() + self._pil = None + return pil + + def close(self): + if self._pil is not None: + try: + self._pil.close() + except Exception: + pass + self._pil = None + return None + + def __getattr__(self, name): + pil = self.to_pil() + if pil is None: + raise AttributeError(name) + return getattr(pil, name) + + def __array__(self, dtype=None): + import numpy as np + + pil = self.to_pil() + if pil is None: + return np.array([], dtype=dtype) + return np.array(pil, dtype=dtype) + + def __enter__(self): + return self.to_pil() + + def __exit__(self, exc_type, exc, tb): + self.close() + return False + + +def ensure_pil_image(img): + if isinstance(img, Image.Image): + return img + if isinstance(img, LazyDocxImage): + return img.to_pil() + return None + + +def is_image_like(img): + return isinstance(img, Image.Image) or isinstance(img, LazyDocxImage) + + +def open_image_for_processing(img, *, allow_bytes=False): + if isinstance(img, Image.Image): + return img, False + if isinstance(img, LazyDocxImage): + return img.to_pil_detached(), True + if allow_bytes and isinstance(img, (bytes, bytearray)): + try: + pil = Image.open(BytesIO(img)).convert("RGB") + return pil, True + except Exception as e: + logging.info(f"open_image_for_processing: bad bytes: {e}") + return None, False + return img, False