후처리 로직 추가

router/deepseek_router.py

@@ -27,7 +27,7 @@ async def perform_ocr(
             raise HTTPException(status_code=400, detail="파일 내용이 비어있습니다.")
 
         # 모든 처리 로직을 OCR 엔진에 위임
-        result = process_document(
+        result = await process_document(
             file_bytes=file_content,
             content_type=document.content_type,
             filename=document.filename,

services/ocr_engine.py

@@ -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 파일을 업로드해주세요."
         )
+

test/test.py

@@ -228,10 +228,10 @@ def process_pdf(llm, sampling_params, pdf_path):
         contents += content + page_num_separator
 
     # Save results
-    json_path = os.path.join(
+    result_json_path = os.path.join(
         f"{config.OUTPUT_PATH}/result", f"{file_name_without_ext}.json"
     )
-    pdf_out_path = os.path.join(
+    result_pdf_path = os.path.join(
         config.OUTPUT_PATH, f"{file_name_without_ext}_layouts.pdf"
     )
 
@@ -247,10 +247,10 @@ def process_pdf(llm, sampling_params, pdf_path):
         "parsed": contents,
     }
 
-    with open(json_path, "w", encoding="utf-8") as f:
+    with open(result_json_path, "w", encoding="utf-8") as f:
         json.dump(output_data, f, ensure_ascii=False, indent=4)
 
-    pil_to_pdf_img2pdf(draw_images, pdf_out_path)
+    pil_to_pdf_img2pdf(draw_images, result_pdf_path)
     print(
         f"{Colors.GREEN}Finished processing {pdf_path}. Results saved in {config.OUTPUT_PATH}{Colors.RESET}"
     )
@@ -283,7 +283,7 @@ def process_image(llm, sampling_params, image_path):
     print(result_out)
 
     # Save results
-    output_json_path = os.path.join(
+    result_json_path = os.path.join(
         f"{config.OUTPUT_PATH}/result", f"{file_name_without_ext}.json"
     )
     result_image_path = os.path.join(
@@ -316,7 +316,7 @@ def process_image(llm, sampling_params, image_path):
         "parsed": processed_text,
     }
 
-    with open(output_json_path, "w", encoding="utf-8") as f:
+    with open(result_json_path, "w", encoding="utf-8") as f:
         json.dump(output_data, f, ensure_ascii=False, indent=4)
 
     result_image.save(result_image_path)