mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-06 11:28:38 +08:00
fa71f8d0c7
**Summary** This PR tackles a significant memory bottleneck when processing image-heavy Word documents. Previously, our pipeline eagerly decoded DOCX images into `PIL.Image` objects, which caused high peak memory usage. To solve this, I've introduced a **lazy-loading approach**: images are now stored as raw blobs and only decoded exactly when and where they are consumed. This successfully reduces the memory footprint while keeping the parsing output completely identical to before. **What's Changed** Instead of a dry file-by-file list, here is the logical breakdown of the updates: * **The Core Abstraction (`lazy_image.py`)**: Introduced `LazyDocxImage` along with helper APIs to handle lazy decoding, image-type checks, and NumPy compatibility. It also supports `.close()` and detached PIL access to ensure safe lifecycle management and prevent memory leaks. * **Pipeline Integration (`naive.py`, `figure_parser.py`, etc.)**: Updated the general DOCX picture extraction to return these new lazy images. Downstream consumers (like the figure/VLM flow and base64 encoding paths) now decode images right at the use site using detached PIL instances, avoiding shared-instance side effects. * **Compatibility Hooks (`operators.py`, `book.py`, etc.)**: Added necessary compatibility conversions so these lazy images flow smoothly through existing merging, filtering, and presentation steps without breaking. **Scope & What is Intentionally Left Out** To keep this PR focused, I have restricted these changes strictly to the **general Word pipeline** and its downstream consumers. The `QA` and `manual` Word parsing pipelines are explicitly **not modified** in this PR. They can be safely migrated to this new lazy-load model in a subsequent, standalone PR. **Design Considerations** I briefly considered adding image compression during processing, but decided against it to avoid any potential quality degradation in the derived outputs. I also held off on a massive pipeline re-architecture to avoid overly invasive changes right now. **Validation & Testing** I've tested this to ensure no regressions: * Compared identical DOCX inputs before and after this branch: chunk counts, extracted text, table HTML, and image descriptions match perfectly. * **Confirmed a noticeable drop in peak memory usage when processing image-dense documents.** For a 30MB Word document containing 243 1080p screenshots, memory consumption is reduced by approximately 1.5GB. **Breaking Changes** None.
110 lines
3.4 KiB
Python
110 lines
3.4 KiB
Python
#
|
|
# Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
|
|
import base64
|
|
import logging
|
|
from functools import partial
|
|
from io import BytesIO
|
|
|
|
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)
|
|
|
|
|
|
async def image2id(d: dict, storage_put_func: partial, objname: str, bucket: str = "imagetemps"):
|
|
import logging
|
|
from io import BytesIO
|
|
from rag.svr.task_executor import minio_limiter
|
|
|
|
if "image" not in d:
|
|
return
|
|
if not d["image"]:
|
|
del d["image"]
|
|
return
|
|
|
|
def encode_image():
|
|
with BytesIO() as buf:
|
|
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")
|
|
except OSError as e:
|
|
logging.warning(f"Saving image exception: {e}")
|
|
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)
|
|
if jpeg_binary is None:
|
|
del d["image"]
|
|
return
|
|
|
|
async with minio_limiter:
|
|
await thread_pool_exec(
|
|
lambda: storage_put_func(bucket=bucket, fnm=objname, binary=jpeg_binary)
|
|
)
|
|
|
|
d["img_id"] = f"{bucket}-{objname}"
|
|
|
|
if not isinstance(d["image"], bytes):
|
|
d["image"].close()
|
|
del d["image"]
|
|
|
|
|
|
def id2image(image_id: str | None, storage_get_func: partial):
|
|
if not image_id:
|
|
return
|
|
arr = image_id.split("-")
|
|
if len(arr) != 2:
|
|
return
|
|
bkt, nm = image_id.split("-")
|
|
try:
|
|
blob = storage_get_func(bucket=bkt, fnm=nm)
|
|
if not blob:
|
|
return
|
|
return Image.open(BytesIO(blob))
|
|
except Exception as e:
|
|
logging.exception(e)
|