Weknora/Enterprise-Knowledge-Base
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: 新增文档模型类，调整配置与解析逻辑，优化日志及导入

Browse Source 
移除日志设置与冗余代码，优化导入、类型提示及OCR后端管理
统一调整各文件模块导入路径为绝对导入
调整导入路径，移除部分导入，优化日志及注释
升级文档解析器为 Docx2Parser，优化超时与图片处理逻辑
This commit is contained in:
begoniezhao
2025-11-07 10:30:02 +08:00
committed by lyingbug
parent af620806e0
commit 2d66abedf0
39 changed files with 2676 additions and 1570 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docreader/parser/base_parser.py
+171 -474
View File
@@ -1,65 +1,28 @@
# -*- coding: utf-8 -*-
import re
import os
import asyncio
from typing import List, Dict, Any, Optional, Tuple, Union
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
import logging
import sys
import traceback
import numpy as np
import time
import io
import json
from .ocr_engine import OCREngine
from .image_utils import image_to_base64
from .config import ChunkingConfig
from .storage import create_storage
import logging
import os
import re
import time
from abc import ABC, abstractmethod
from typing import Dict, List, Optional, Tuple
import requests
from PIL import Image
# Add parent directory to Python path for src imports
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
if parent_dir not in sys.path:
sys.path.insert(0, parent_dir)
try:
from services.docreader.src.parser.caption import Caption
except ImportError:
# Fallback: try relative import
try:
from .caption import Caption
except ImportError:
# If both imports fail, set to None
Caption = None
logging.warning(
"Failed to import Caption, image captioning will be unavailable"
)
from docreader.models.document import Chunk, Document
from docreader.models.read_config import ChunkingConfig
from docreader.parser.caption import Caption
from docreader.parser.ocr_engine import OCREngine
from docreader.parser.storage import create_storage
from docreader.splitter.splitter import TextSplitter
from docreader.utils import endecode
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
@dataclass
class Chunk:
"""Chunk result"""
content: str # Chunk content
seq: int # Chunk sequence number
start: int # Chunk start position
end: int # Chunk end position
images: List[Dict[str, Any]] = field(default_factory=list) # Images in the chunk
@dataclass
class ParseResult:
"""Parse result"""
text: str # Extracted text content
chunks: Optional[List[Chunk]] = None # Chunk results
class BaseParser(ABC):
"""Base parser interface"""
@@ -97,17 +60,17 @@ class BaseParser(ABC):
def __init__(
self,
file_name: str = "",
file_type: str = None,
file_type: Optional[str] = None,
enable_multimodal: bool = True,
chunk_size: int = 1000,
chunk_overlap: int = 200,
separators: list = ["\n\n", "\n", "。"],
separators: list[str] = ["\n\n", "\n", "。"],
ocr_backend: str = "paddle",
ocr_config: dict = None,
ocr_config: dict = {},
max_image_size: int = 1920, # Maximum image size
max_concurrent_tasks: int = 5, # Max concurrent tasks
max_chunks: int = 1000, # Max number of returned chunks
chunking_config: ChunkingConfig = None, # Chunking configuration object
chunking_config: Optional[ChunkingConfig] = None,
):
"""Initialize parser
@@ -125,7 +88,6 @@ class BaseParser(ABC):
max_chunks: Max number of returned chunks
"""
# Storage client instance
self._storage = None
self.file_name = file_name
self.file_type = file_type or os.path.splitext(file_name)[1]
self.enable_multimodal = enable_multimodal
@@ -133,15 +95,16 @@ class BaseParser(ABC):
self.chunk_overlap = chunk_overlap
self.separators = separators
self.ocr_backend = os.getenv("OCR_BACKEND", ocr_backend)
self.ocr_config = ocr_config or {}
self.ocr_config = ocr_config
self.max_image_size = max_image_size
self.max_concurrent_tasks = max_concurrent_tasks
self.max_chunks = max_chunks
self.chunking_config = chunking_config
logger.info(
f"Initializing {self.__class__.__name__} for file: {file_name}, type: {self.file_type}"
self.storage = create_storage(
self.chunking_config.storage_config if self.chunking_config else None
)
logger.info(f"Initializing parser for file: {file_name}, type: {file_type}")
logger.info(
f"Parser config: chunk_size={chunk_size}, "
f"overlap={chunk_overlap}, "
@@ -150,16 +113,24 @@ class BaseParser(ABC):
f"max_chunks={max_chunks}"
)
# Only initialize Caption service if multimodal is enabled
if self.enable_multimodal:
try:
self.caption_parser = Caption(self.chunking_config.vlm_config)
except Exception as e:
logger.warning(f"Failed to initialize Caption service: {str(e)}")
self.caption_parser = None
else:
self.caption_parser = None
vlm_config = self.chunking_config.vlm_config if self.chunking_config else None
self.caption_parser = (
Caption(vlm_config=vlm_config) if self.enable_multimodal else None
)
def perform_ocr(self, image):
@abstractmethod
def parse_into_text(self, content: bytes) -> Document:
"""Parse document content
Args:
content: Document content
Returns:
Either a string containing the parsed text, or a tuple of (text, image_map)
where image_map is a dict mapping image URLs to Image objects
"""
def perform_ocr(self, image: Image.Image):
"""Execute OCR recognition on the image
Args:
@@ -170,53 +141,23 @@ class BaseParser(ABC):
"""
start_time = time.time()
logger.info("Starting OCR recognition")
resized_image = None
try:
# Resize image to avoid processing large images
resized_image = self._resize_image_if_needed(image)
# Resize image to avoid processing large images
resized_image = self._resize_image_if_needed(image)
# Get OCR engine
ocr_engine = self.get_ocr_engine(
backend_type=self.ocr_backend, **self.ocr_config
)
if ocr_engine is None:
logger.error(
f"OCR engine ({self.ocr_backend}) initialization failed or unavailable, "
"skipping OCR recognition"
)
return ""
# Get OCR engine
ocr_engine = OCREngine.get_instance(self.ocr_backend)
# Execute OCR prediction
logger.info(f"Executing OCR prediction (using {self.ocr_backend} engine)")
# Add extra exception handling
try:
ocr_result = ocr_engine.predict(resized_image)
except RuntimeError as e:
# Handle common CUDA memory issues or other runtime errors
logger.error(f"OCR prediction runtime error: {str(e)}")
return ""
except Exception as e:
# Handle other prediction errors
logger.error(f"Unexpected OCR prediction error: {str(e)}")
return ""
# Execute OCR prediction
logger.info(f"Executing OCR prediction (using {self.ocr_backend} engine)")
ocr_result = ocr_engine.predict(resized_image)
process_time = time.time() - start_time
logger.info(f"OCR recognition completed, time: {process_time:.2f} seconds")
return ocr_result
except Exception as e:
process_time = time.time() - start_time
logger.error(
f"OCR recognition error: {str(e)}, time: {process_time:.2f} seconds"
)
return ""
finally:
# Release image resources
if resized_image is not image and hasattr(resized_image, "close"):
# Only close the new image we created, not the original image
resized_image.close()
process_time = time.time() - start_time
logger.info(f"OCR recognition completed, time: {process_time:.2f} seconds")
def _resize_image_if_needed(self, image):
return ocr_result
def _resize_image_if_needed(self, image: Image.Image) -> Image.Image:
"""Resize image if it exceeds maximum size limit
Args:
@@ -225,102 +166,21 @@ class BaseParser(ABC):
Returns:
Resized image object
"""
try:
# If it's a PIL Image
if hasattr(image, "size"):
width, height = image.size
if width > self.max_image_size or height > self.max_image_size:
logger.info(f"Resizing PIL image, original size: {width}x{height}")
scale = min(
self.max_image_size / width, self.max_image_size / height
)
new_width = int(width * scale)
new_height = int(height * scale)
resized_image = image.resize((new_width, new_height))
logger.info(f"Resized to: {new_width}x{new_height}")
return resized_image
else:
logger.info(
f"PIL image size {width}x{height} is within limits, no resizing needed"
)
return image
# If it's a numpy array
elif hasattr(image, "shape"):
height, width = image.shape[:2]
if width > self.max_image_size or height > self.max_image_size:
logger.info(
f"Resizing numpy image, original size: {width}x{height}"
)
scale = min(
self.max_image_size / width, self.max_image_size / height
)
new_width = int(width * scale)
new_height = int(height * scale)
# Use PIL for resizing numpy arrays
pil_image = Image.fromarray(image)
resized_pil = pil_image.resize((new_width, new_height))
resized_image = np.array(resized_pil)
logger.info(f"Resized to: {new_width}x{new_height}")
return resized_image
else:
logger.info(
f"Numpy image size {width}x{height} is within limits, no resizing needed"
)
return image
else:
logger.warning(f"Unknown image type: {type(image)}, cannot resize")
return image
except Exception as e:
logger.error(f"Error resizing image: {str(e)}")
return image
width, height = image.size
if width > self.max_image_size or height > self.max_image_size:
logger.info(f"Resizing PIL image, original size: {width}x{height}")
scale = min(self.max_image_size / width, self.max_image_size / height)
new_width = int(width * scale)
new_height = int(height * scale)
resized_image = image.resize((new_width, new_height))
logger.info(f"Resized to: {new_width}x{new_height}")
return resized_image
def process_image(self, image, image_url=None):
"""Process image: first perform OCR, then get caption if text is available
logger.info(f"PIL image size is {width}x{height}, no resizing needed")
return image
Args:
image: Image object (PIL.Image or numpy array)
image_url: Image URL (if uploaded)
Returns:
tuple: (ocr_text, caption, image_url)
- ocr_text: OCR extracted text
- caption: Image description (if OCR has text) or empty string
- image_url: Image URL (if provided)
"""
logger.info("Starting image processing (OCR + optional caption)")
# Resize image
image = self._resize_image_if_needed(image)
# Perform OCR recognition
ocr_text = self.perform_ocr(image)
caption = ""
if self.caption_parser:
logger.info(
f"OCR successfully extracted {len(ocr_text)} characters, continuing to get caption"
)
# Convert image to base64 for caption generation
img_base64 = image_to_base64(image)
if img_base64:
caption = self.get_image_caption(img_base64)
if caption:
logger.info(f"Successfully obtained image caption: {caption}")
else:
logger.warning("Failed to get caption")
else:
logger.warning("Failed to convert image to base64")
caption = ""
else:
logger.info("Caption service not initialized, skipping caption retrieval")
# Release image resources
del image
return ocr_text, caption, image_url
async def process_image_async(self, image, image_url=None):
"""Asynchronously process image: first perform OCR, then get caption if text is available
async def process_image_async(self, image: Image.Image, image_url: str):
"""Asynchronously process image: first perform OCR, then get caption
Args:
image: Image object (PIL.Image or numpy array)
@@ -333,84 +193,47 @@ class BaseParser(ABC):
- image_url: Image URL (if provided)
"""
logger.info("Starting asynchronous image processing (OCR + optional caption)")
resized_image = None
# Resize image
resized_image = self._resize_image_if_needed(image)
try:
# Resize image
resized_image = self._resize_image_if_needed(image)
# Perform OCR recognition (using run_in_executor to execute synchronous operations in the event loop)
# Perform OCR recognition
loop = asyncio.get_event_loop()
try:
# Add timeout mechanism to avoid infinite blocking (30 seconds timeout)
ocr_task = loop.run_in_executor(None, self.perform_ocr, resized_image)
ocr_text = await asyncio.wait_for(ocr_task, timeout=30.0)
except asyncio.TimeoutError:
logger.error(
"OCR processing timed out (30 seconds), skipping this image"
)
ocr_text = ""
except Exception as e:
logger.error(f"OCR processing error: {str(e)}")
logger.error(f"OCR processing error, skipping this image: {str(e)}")
ocr_text = ""
logger.info(
f"OCR successfully extracted {len(ocr_text)} characters, continuing to get caption"
)
caption = ""
if self.caption_parser:
try:
# Convert image to base64 for caption generation
img_base64 = image_to_base64(resized_image)
if img_base64:
# Add timeout to avoid blocking caption retrieval (30 seconds timeout)
caption_task = self.get_image_caption_async(img_base64)
image_data, caption = await asyncio.wait_for(
caption_task, timeout=30.0
)
if caption:
logger.info(
f"Successfully obtained image caption: {caption}"
)
else:
logger.warning("Failed to get caption")
else:
logger.warning("Failed to convert image to base64")
caption = ""
except asyncio.TimeoutError:
logger.warning("Caption retrieval timed out, skipping")
except Exception as e:
logger.error(f"Failed to get caption: {str(e)}")
else:
logger.info(
"Caption service not initialized, skipping caption retrieval"
)
logger.info(f"Successfully obtained image ocr: {ocr_text}")
img_base64 = endecode.decode_image(resized_image)
caption = self.get_image_caption(img_base64)
logger.info(f"Successfully obtained image caption: {caption}")
return ocr_text, caption, image_url
finally:
# Release image resources
if resized_image is not image and hasattr(resized_image, "close"):
# Only close the new image we created, not the original image
resized_image.close()
resized_image.close()
async def process_with_limit(self, idx, image, url, semaphore):
async def process_with_limit(
self, idx: int, image: Image.Image, url: str, semaphore: asyncio.Semaphore
):
"""Function to process a single image using a semaphore"""
try:
logger.info(f"Waiting to process image {idx+1}")
logger.info(f"Waiting to process image {idx + 1}")
async with semaphore: # Use semaphore to control concurrency
logger.info(f"Starting to process image {idx+1}")
logger.info(f"Starting to process image {idx + 1}")
result = await self.process_image_async(image, url)
logger.info(f"Completed processing image {idx+1}")
logger.info(f"Completed processing image {idx + 1}")
return result
except Exception as e:
logger.error(f"Error processing image {idx+1}: {str(e)}")
logger.error(f"Error processing image {idx + 1}: {str(e)}")
return ("", "", url) # Return empty result to avoid overall failure
finally:
# Manually release image resources
if hasattr(image, "close"):
image.close()
image.close()
async def process_multiple_images(self, images_data):
async def process_multiple_images(self, images_data: List[Tuple[Image.Image, str]]):
"""Process multiple images concurrently
Args:
@@ -450,7 +273,7 @@ class BaseParser(ABC):
for i, result in enumerate(completed_results):
if isinstance(result, Exception):
logger.error(
f"Image {i+1} processing returned an exception: {str(result)}"
f"Image {i + 1} processing returned an exception: {str(result)}"
)
# For exceptions, add empty results
if i < len(images_data):
@@ -467,47 +290,10 @@ class BaseParser(ABC):
logger.info("Image processing resource cleanup complete")
logger.info(
f"Completed concurrent processing of {len(results)}/{len(images_data)} images"
f"Concurrent processing of {len(results)}/{len(images_data)} images"
)
return results
def decode_bytes(self, content: bytes) -> str:
"""Intelligently decode byte stream, supports multiple encodings
Tries to decode in common encodings, if all fail, uses latin-1 as fallback
Args:
content: Byte stream to decode
Returns:
Decoded string
"""
logger.info(f"Attempting to decode bytes of length: {len(content)}")
# Common encodings, sorted by priority
encodings = ["utf-8", "gb18030", "gb2312", "gbk", "big5", "ascii", "latin-1"]
text = None
# Try decoding with each encoding format
for encoding in encodings:
try:
text = content.decode(encoding)
logger.info(f"Successfully decoded content using {encoding} encoding")
break
except UnicodeDecodeError:
logger.info(f"Failed to decode using {encoding} encoding")
continue
# If all encodings fail, use latin-1 as fallback
if text is None:
text = content.decode("latin-1")
logger.warning(
f"Unable to determine correct encoding, using latin-1 as fallback. "
f"This may cause character issues."
)
logger.info(f"Decoded text length: {len(text)} characters")
return text
def get_image_caption(self, image_data: str) -> str:
"""Get image description
@@ -517,6 +303,9 @@ class BaseParser(ABC):
Returns:
Image description
"""
if not self.caption_parser:
logger.warning("Caption parser not initialized")
return ""
start_time = time.time()
logger.info(
f"Getting caption for image: {image_data[:250]}..."
@@ -533,80 +322,7 @@ class BaseParser(ABC):
logger.warning("Failed to get caption for image")
return caption
async def get_image_caption_async(self, image_data: str) -> Tuple[str, str]:
"""Asynchronously get image description
Args:
image_data: Image data (base64 encoded string or URL)
Returns:
Tuple[str, str]: Image data and corresponding description
"""
caption = self.get_image_caption(image_data)
return image_data, caption
def __init_storage(self):
"""Initialize storage client based on configuration"""
if self._storage is None:
storage_config = (
self.chunking_config.storage_config if self.chunking_config else None
)
self._storage = create_storage(storage_config)
logger.info(
f"Initialized storage client: {self._storage.__class__.__name__}"
)
return self._storage
def upload_file(self, file_path: str) -> str:
"""Upload file to object storage
Args:
file_path: File path
Returns:
File URL
"""
logger.info(f"Uploading file: {file_path}")
try:
storage = self.__init_storage()
return storage.upload_file(file_path)
except Exception as e:
logger.error(f"Failed to upload file: {str(e)}")
return ""
def upload_bytes(self, content: bytes, file_ext: str = ".png") -> str:
"""Upload bytes to object storage
Args:
content: Byte content to upload
file_ext: File extension
Returns:
File URL
"""
logger.info(f"Uploading bytes content, size: {len(content)} bytes")
try:
storage = self.__init_storage()
return storage.upload_bytes(content, file_ext)
except Exception as e:
logger.error(f"Failed to upload bytes to storage: {str(e)}")
traceback.print_exc()
return ""
@abstractmethod
def parse_into_text(self, content: bytes) -> Union[str, Tuple[str, Dict[str, Any]]]:
"""Parse document content
Args:
content: Document content
Returns:
Either a string containing the parsed text, or a tuple of (text, image_map)
where image_map is a dict mapping image URLs to Image objects
"""
pass
def parse(self, content: bytes) -> ParseResult:
def parse(self, content: bytes) -> Document:
"""Parse document content
Args:
@@ -616,17 +332,19 @@ class BaseParser(ABC):
Parse result
"""
logger.info(
f"Parsing document with {self.__class__.__name__}, content size: {len(content)} bytes"
f"Parsing document with {self.__class__.__name__}, bytes: {len(content)}"
)
parse_result = self.parse_into_text(content)
if isinstance(parse_result, tuple):
text, image_map = parse_result
else:
text = parse_result
image_map = {}
logger.info(f"Extracted {len(text)} characters of text from {self.file_name}")
logger.info(f"Beginning chunking process for text")
chunks = self.chunk_text(text)
document = self.parse_into_text(content)
logger.info(
f"Extracted {len(document.content)} characters from {self.file_name}"
)
splitter = TextSplitter(
chunk_size=self.chunk_size,
chunk_overlap=self.chunk_overlap,
separators=self.separators,
)
chunk_str = splitter.split_text(document.content)
chunks = self._str_to_chunk(chunk_str)
logger.info(f"Created {len(chunks)} chunks from document")
# Limit the number of returned chunks
@@ -636,7 +354,7 @@ class BaseParser(ABC):
)
chunks = chunks[: self.max_chunks]
# If multimodal is enabled and file type is supported, process images in each chunk
# If multimodal is enabled and file type is supported, process images
if self.enable_multimodal:
# Get file extension and convert to lowercase
file_ext = (
@@ -647,11 +365,12 @@ class BaseParser(ABC):
# Define allowed file types for image processing
allowed_types = [
".pdf", # PDF files
# Text files
".pdf",
".md",
".markdown", # Markdown files
".markdown",
".doc",
".docx", # Word documents
".docx",
# Image files
".jpg",
".jpeg",
@@ -666,13 +385,21 @@ class BaseParser(ABC):
logger.info(
f"Processing images in each chunk for file type: {file_ext}"
)
chunks = self.process_chunks_images(chunks, image_map)
chunks = self.process_chunks_images(chunks, document.images)
else:
logger.info(
f"Skipping image processing for unsupported file type: {file_ext}"
)
return ParseResult(text=text, chunks=chunks)
document.chunks = chunks
return document
def _str_to_chunk(self, text: List[Tuple[int, int, str]]) -> List[Chunk]:
"""Convert string to Chunk object"""
return [
Chunk(seq=i, content=t, start=start, end=end)
for i, (start, end, t) in enumerate(text)
]
def _split_into_units(self, text: str) -> List[str]:
"""
@@ -682,9 +409,7 @@ class BaseParser(ABC):
Returns:
基本单元的列表
"""
logger.info(
f"Splitting text into basic units with robust structure protection, text length: {len(text)}"
)
logger.info(f"Splitting text into basic units, text length: {len(text)}")
# 定义所有需要作为整体保护的结构模式 ---
table_pattern = r"(?m)(^\|.*\|[ \t]*\r?\n(?:[ \t]*\r?\n)?^\|\s*:?--+.*\r?\n(?:^\|.*\|\r?\n?)*)"
@@ -710,7 +435,8 @@ class BaseParser(ABC):
# 按起始位置排序
protected_ranges.sort(key=lambda x: x[0])
logger.info(
f"Found {len(protected_ranges)} protected structures (tables, code, formulas, images, links)."
f"Found {len(protected_ranges)} protected structures "
"(tables, code, formulas, images, links)."
)
# 合并可能重叠的保护范围 ---
@@ -731,7 +457,7 @@ class BaseParser(ABC):
merged_ranges.append((current_start, current_end))
protected_ranges = merged_ranges
logger.info(
f"After merging overlaps, {len(protected_ranges)} protected ranges remain."
f"After overlaps, {len(protected_ranges)} protected ranges remain."
)
# 根据保护范围和分隔符来分割文本 ---
@@ -749,7 +475,7 @@ class BaseParser(ABC):
segments = re.split(separator_pattern, pre_text)
units.extend([s for s in segments if s]) # 添加所有非空部分
# b. 将整个受保护的块（例如，一个完整的表格）作为一个单独的、不可分割的单元添加
# b. 将整个受保护的块（例如，一个完整的表格）作为一个不可分割的单元添加
protected_text = text[start:end]
units.append(protected_text)
@@ -764,38 +490,6 @@ class BaseParser(ABC):
logger.info(f"Text splitting complete, created {len(units)} final basic units.")
return units
def _find_complete_units(self, units: List[str], target_size: int) -> List[str]:
"""Find a list of complete units that do not exceed the target size
Args:
units: List of units
target_size: Target size
Returns:
List of complete units
"""
logger.info(f"Finding complete units with target size: {target_size}")
result = []
current_size = 0
for unit in units:
unit_size = len(unit)
if current_size + unit_size > target_size and result:
logger.info(
f"Reached target size limit at {current_size} characters, stopping"
)
break
result.append(unit)
current_size += unit_size
logger.info(
f"Added unit of size {unit_size}, current total: {current_size}/{target_size}"
)
logger.info(
f"Found {len(result)} complete units totaling {current_size} characters"
)
return result
def chunk_text(self, text: str) -> List[Chunk]:
"""Chunk text, preserving Markdown structure
@@ -825,7 +519,7 @@ class BaseParser(ABC):
for i, unit in enumerate(units):
unit_size = len(unit)
logger.info(f"Processing unit {i+1}/{len(units)}, size: {unit_size}")
logger.info(f"Processing unit {i + 1}/{len(units)}, size: {unit_size}")
# If current chunk plus new unit exceeds size limit, create new chunk
if current_size + unit_size > self.chunk_size and current_chunk:
@@ -855,14 +549,12 @@ class BaseParser(ABC):
for u in reversed(current_chunk):
if overlap_size + len(u) > overlap_target:
logger.info(
f"Reached overlap target ({overlap_size}/{overlap_target})"
f"Overlap target ({overlap_size}/{overlap_target})"
)
break
overlap_units.insert(0, u)
overlap_size += len(u)
logger.info(
f"Added unit to overlap, current overlap size: {overlap_size}"
)
logger.info(f"Added unit to overlap, size: {overlap_size}")
# Remove elements from overlap that are included in separators
start_index = 0
@@ -883,7 +575,7 @@ class BaseParser(ABC):
overlap_units = overlap_units[start_index:]
logger.info(
f"Final overlap: {len(overlap_units)} units, {overlap_size} characters"
f"Overlap: {len(overlap_units)} units, {overlap_size} size"
)
current_chunk = overlap_units
@@ -899,7 +591,7 @@ class BaseParser(ABC):
current_chunk.append(unit)
current_size += unit_size
logger.info(
f"Added unit to current chunk, now at {current_size}/{self.chunk_size} characters"
f"Added unit to current chunk, at {current_size}/{self.chunk_size}"
)
# Add the last chunk
@@ -925,12 +617,13 @@ class BaseParser(ABC):
chunk: Document chunk
Returns:
List of image information, each element contains image URL and match position
List of image information
"""
logger.info(f"Extracting image information from Chunk #{chunk.seq}")
text = chunk.content
# Regex to extract image information from text, supporting Markdown images and HTML images
# Regex to extract image information from text,
# support: Markdown images, HTML images
img_pattern = r'!\[([^\]]*)\]\(([^)]+)\)|<img [^>]*src="([^"]+)" [^>]*>'
# Extract image information
@@ -954,28 +647,28 @@ class BaseParser(ABC):
images_info.append(image_info)
logger.info(
f"Image in Chunk #{chunk.seq} {match_idx+1}: " f"URL={img_url[:50]}..."
f"Image in Chunk #{chunk.seq} {match_idx + 1}: URL={img_url[:50]}..."
if len(img_url) > 50
else f"Image in Chunk #{chunk.seq} {match_idx+1}: URL={img_url}"
else f"Image in Chunk #{chunk.seq} {match_idx + 1}: URL={img_url}"
)
return images_info
async def download_and_upload_image(self, img_url: str):
"""Download image and upload to object storage, if it's already an object storage path or local path, use directly
async def download_and_upload_image(
self, img_url: str
) -> Tuple[str, str, Image.Image | None]:
"""Download image and upload to object storage,
if it's already an object storage path or local path, use directly
Args:
img_url: Image URL or local path
Returns:
tuple: (original URL, storage URL, image object), if failed returns (original URL, None, None)
tuple: (original URL, storage URL, image object),
if failed returns (original URL, None, None)
"""
try:
import requests
from PIL import Image
import io
# Check if it's already a storage URL (COS or MinIO)
is_storage_url = any(
pattern in img_url
@@ -997,12 +690,7 @@ class BaseParser(ABC):
response = requests.get(img_url, timeout=5, proxies=proxies)
if response.status_code == 200:
image = Image.open(io.BytesIO(response.content))
try:
return img_url, img_url, image
finally:
# Ensure image resources are also released after the function returns
# Image will be closed by the caller
pass
return img_url, img_url, image
else:
logger.warning(
f"Failed to get storage image: {response.status_code}"
@@ -1022,7 +710,7 @@ class BaseParser(ABC):
# Upload to storage
with open(img_url, "rb") as f:
content = f.read()
storage_url = self.upload_bytes(content)
storage_url = self.storage.upload_bytes(content)
logger.info(
f"Successfully uploaded local image to storage: {storage_url}"
)
@@ -1031,7 +719,7 @@ class BaseParser(ABC):
logger.error(f"Error processing local image: {str(e)}")
if image and hasattr(image, "close"):
image.close()
return img_url, None, None
return img_url, img_url, None
# Normal remote URL download handling
else:
@@ -1044,9 +732,7 @@ class BaseParser(ABC):
if https_proxy:
proxies["https"] = https_proxy
logger.info(
f"Downloading image {img_url}, using proxy: {proxies if proxies else 'None'}"
)
logger.info(f"Downloading image {img_url}, using proxy: {proxies}")
response = requests.get(img_url, timeout=5, proxies=proxies)
if response.status_code == 200:
@@ -1054,7 +740,7 @@ class BaseParser(ABC):
image = Image.open(io.BytesIO(response.content))
try:
# Upload to storage using the method in BaseParser
storage_url = self.upload_bytes(response.content)
storage_url = self.storage.upload_bytes(response.content)
logger.info(
f"Successfully uploaded image to storage: {storage_url}"
)
@@ -1064,11 +750,11 @@ class BaseParser(ABC):
pass
else:
logger.warning(f"Failed to download image: {response.status_code}")
return img_url, None, None
return img_url, img_url, None
except Exception as e:
logger.error(f"Error downloading or processing image: {str(e)}")
return img_url, None, None
return img_url, img_url, None
async def process_chunk_images_async(
self, chunk, chunk_idx, total_chunks, image_map=None
@@ -1086,18 +772,19 @@ class BaseParser(ABC):
"""
logger.info(
f"Starting to process images in Chunk #{chunk_idx+1}/{total_chunks}"
f"Starting to process images in Chunk #{chunk_idx + 1}/{total_chunks}"
)
# Extract image information from the Chunk
images_info = self.extract_images_from_chunk(chunk)
if not images_info:
logger.info(f"Chunk #{chunk_idx+1} found no images")
logger.info(f"Chunk #{chunk_idx + 1} found no images")
return chunk
# Prepare images that need to be downloaded and processed
images_to_process = []
url_to_info_map = {} # Map URL to image information
# Map URL to image information
url_to_info_map = {}
# Record all image URLs that need to be processed
for img_info in images_info:
@@ -1106,14 +793,21 @@ class BaseParser(ABC):
results = []
download_tasks = []
for img_url in url_to_info_map.keys(): # Check if image is already in the image_map
# Check if image is already in the image_map
for img_url in url_to_info_map.keys():
if image_map and img_url in image_map:
logger.info(f"Image already in image_map: {img_url}, using cached object")
results.append((img_url, img_url, image_map[img_url]))
logger.info(
f"Image already in image_map: {img_url}, using cached object"
)
image = Image.open(
io.BytesIO(endecode.encode_image(image_map[img_url]))
)
results.append((img_url, img_url, image))
else:
download_task = self.download_and_upload_image(img_url)
download_tasks.append(download_task)
# Concurrent download and upload of images, ignore images that are already in the image_map
# Concurrent download and upload of images,
# ignore images that are already in the image_map
results.extend(await asyncio.gather(*download_tasks))
# Process download results, prepare for OCR processing
@@ -1123,16 +817,17 @@ class BaseParser(ABC):
img_info["cos_url"] = cos_url
images_to_process.append((image, cos_url))
# If no images were successfully downloaded and uploaded, return the original Chunk
# If no images were successfully downloaded and uploaded,
# return the original Chunk
if not images_to_process:
logger.info(
f"Chunk #{chunk_idx+1} found no successfully downloaded and uploaded images"
f"Chunk #{chunk_idx + 1} not found downloaded and uploaded images"
)
return chunk
# Concurrent processing of all images (OCR + caption)
logger.info(
f"Processing {len(images_to_process)} images in Chunk #{chunk_idx+1}"
f"Processing {len(images_to_process)} images in Chunk #{chunk_idx + 1}"
)
# Concurrent processing of all images
@@ -1163,10 +858,12 @@ class BaseParser(ABC):
# Update image information in the Chunk
chunk.images = processed_images
logger.info(f"Completed image processing in Chunk #{chunk_idx+1}")
logger.info(f"Completed image processing in Chunk #{chunk_idx + 1}")
return chunk
def process_chunks_images(self, chunks: List[Chunk], image_map=None) -> List[Chunk]:
def process_chunks_images(
self, chunks: List[Chunk], image_map: Dict[str, str] = {}
) -> List[Chunk]:
"""Concurrent processing of images in all Chunks
Args:
@@ -1210,7 +907,7 @@ class BaseParser(ABC):
processed_chunks = []
for i, result in enumerate(results):
if isinstance(result, Exception):
logger.error(f"Error processing Chunk {i+1}: {str(result)}")
logger.error(f"Error processing Chunk {i + 1}: {str(result)}")
# Keep original Chunk
if i < len(chunks):
processed_chunks.append(chunks[i])
@@ -1235,7 +932,7 @@ class BaseParser(ABC):
# Execute processing for all Chunks
processed_chunks = loop.run_until_complete(process_all_chunks())
logger.info(
f"Successfully completed concurrent processing of {len(processed_chunks)}/{len(chunks)} chunks"
f"Completed processing of {len(processed_chunks)}/{len(chunks)} chunks"
)
return processed_chunks