kyy/deepseek_ocr
1
0
Fork 0
Code Issues 1 Pull Requests Actions Packages Projects Releases Wiki Activity

후처리 로직 추가

Browse Source
This commit is contained in:
kyy
2025-11-06 14:38:54 +09:00
parent 8a6f2ae2d8
commit 7749492ae7
1 changed files with 92 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

109
services/ocr_engine.py
View File

@@ -1,10 +1,11 @@
import asyncio
import io
import logging
import re
import fitz
from config.model_settings import CROP_MODE, MODEL_PATH, PROMPT
from PIL import Image
from PIL import Image, ImageOps
from vllm import AsyncLLMEngine, SamplingParams
from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.model_executor.models.registry import ModelRegistry
@@ -52,12 +53,66 @@ sampling_params = SamplingParams(
processor = DeepseekOCRProcessor()
# --------------------------------------------------------------------------
# 2. 핵심 처리 함수
# 2. 후처리 함수 (공식 코드 기반)
# --------------------------------------------------------------------------
async def _process_single_image(image: Image.Image) -> str:
"""단일 PIL 이미지를 받아 OCR을 수행하고 텍스트를 반환합니다."""
def _postprocess_text(text: str, page_num: int = 0) -> str:
"""
모델의 원본 출력에서 태그를 제거/변경하고 텍스트를 정리합니다.
(test/test.py의 후처리 로직 기반으로 수정)
"""
if not text:
return ""
# 1. 정규식으로 모든 ref/det 태그 블록을 찾음
# 패턴은 (전체 매치, ref 내용, det 내용)을 캡처
pattern = r"(<\|ref\|>(.*?)<\|/ref\|><\|det\|>(.*?)<\|/det\|>)"
matches = re.findall(pattern, text, re.DOTALL)
# 2. 전체 매치된 문자열을 이미지 태그와 기타 태그로 분리
matches_images = []
matches_other = []
for match_tuple in matches:
full_match_str = match_tuple[0] # 전체 매치된 부분
ref_content = match_tuple[1] # <|ref|> 안의 내용
if "image" in ref_content:
matches_images.append(full_match_str)
else:
matches_other.append(full_match_str)
processed_text = text
# 3. 이미지 태그는 마크다운 링크로 대체
for idx, img_tag in enumerate(matches_images):
img_link = f"![](images/{page_num}_{idx}.jpg)\n"
processed_text = processed_text.replace(img_tag, img_link)
# 4. 이미지가 아닌 다른 모든 태그는 제거
for other_tag in matches_other:
processed_text = processed_text.replace(other_tag, "")
# 5. 특수 문자, 불필요한 토큰 및 추가 공백 정리
processed_text = (
processed_text.replace("<｜end of sentence｜>", "")
.replace("\\coloneqq", ":=")
.replace("\\eqqcolon", "=:")
.replace("\n\n\n", "\n\n")
.strip()
)
return processed_text
# --------------------------------------------------------------------------
# 3. 핵심 처리 함수
# --------------------------------------------------------------------------
async def _process_single_image(image: Image.Image, page_num: int = 0) -> str:
"""단일 PIL 이미지를 받아 OCR을 수행하고 후처리된 텍스트를 반환합니다."""
if _engine is None:
raise RuntimeError("vLLM engine not initialized yet")
if "<image>" not in PROMPT:
@@ -70,16 +125,18 @@ async def _process_single_image(image: Image.Image) -> str:
request = {"prompt": PROMPT, "multi_modal_data": {"image": image_features}}
request_id = f"request-{asyncio.get_running_loop().time()}"
final_output = ""
raw_output = ""
async for request_output in _engine.generate(request, sampling_params, request_id):
if request_output.outputs:
final_output = request_output.outputs[0].text
raw_output = request_output.outputs[0].text
return final_output
# 후처리 적용 (페이지 번호 전달)
clean_text = _postprocess_text(raw_output, page_num)
return clean_text
def _pdf_to_images(pdf_bytes: bytes, dpi=144) -> list[Image.Image]:
"""PDF 바이트를 받아 페이지별 PIL 이미지 리스트를 반환합니다."""
def _pdf_to_images_high_quality(pdf_bytes: bytes, dpi=144) -> list[Image.Image]:
"""PDF 바이트를 받아 페이지별 고품질 PIL 이미지 리스트를 반환합니다."""
images = []
pdf_document = fitz.open(stream=pdf_bytes, filetype="pdf")
zoom = dpi / 72.0
@@ -88,8 +145,14 @@ def _pdf_to_images(pdf_bytes: bytes, dpi=144) -> list[Image.Image]:
for page_num in range(pdf_document.page_count):
page = pdf_document[page_num]
pixmap = page.get_pixmap(matrix=matrix, alpha=False)
Image.MAX_IMAGE_PIXELS = None
img_data = pixmap.tobytes("png")
img = Image.open(io.BytesIO(img_data))
if img.mode in ("RGBA", "LA"):
background = Image.new("RGB", img.size, (255, 255, 255))
background.paste(img, mask=img.split()[-1] if img.mode == "RGBA" else None)
img = background
images.append(img)
pdf_document.close()
@@ -100,26 +163,37 @@ async def process_document(file_bytes: bytes, content_type: str, filename: str)
"""
업로드된 파일(이미지 또는 PDF)을 처리하여 OCR 결과를 반환합니다.
"""
if content_type.startswith("image/"):
# Content-Type이 generic할 경우, 파일 확장자로 타입을 유추
inferred_content_type = content_type
if content_type == "application/octet-stream":
if filename.lower().endswith(".pdf"):
inferred_content_type = "application/pdf"
elif any(filename.lower().endswith(ext) for ext in [".jpg", ".jpeg", ".png", ".bmp", ".gif", ".webp"]):
inferred_content_type = "image/jpeg" # 구체적인 타입은 중요하지 않음
if inferred_content_type.startswith("image/"):
try:
image = Image.open(io.BytesIO(file_bytes)).convert("RGB")
image = Image.open(io.BytesIO(file_bytes))
image = ImageOps.exif_transpose(image).convert("RGB")
except Exception as e:
raise ValueError(f"이미지 파일을 여는 데 실패했습니다: {e}")
result_text = await _process_single_image(image)
# 단일 이미지는 페이지 번호를 0으로 간주
result_text = await _process_single_image(image, page_num=0)
return {"filename": filename, "text": result_text}
elif content_type == "application/pdf":
elif inferred_content_type == "application/pdf":
try:
images = _pdf_to_images(file_bytes)
images = _pdf_to_images_high_quality(file_bytes)
except Exception as e:
raise ValueError(f"PDF 파일을 처리하는 데 실패했습니다: {e}")
# 각 페이지를 비동기적으로 처리
tasks = [_process_single_image(img) for img in images]
# 각 페이지를 비동기적으로 처리 (페이지 번호 전달)
tasks = [
_process_single_image(img, page_num=i) for i, img in enumerate(images)
]
page_results = await asyncio.gather(*tasks)
# 페이지 구분자를 넣어 전체 텍스트 합치기
full_text = "\n<--- Page Split --->\n".join(page_results)
return {"filename": filename, "text": full_text, "page_count": len(images)}
@@ -128,3 +202,4 @@ async def process_document(file_bytes: bytes, content_type: str, filename: str)
f"지원하지 않는 파일 형식입니다: {content_type}. "
"이미지(JPEG, PNG 등) 또는 PDF 파일을 업로드해주세요."
)