108 lines
3.6 KiB
Python
108 lines
3.6 KiB
Python
import concurrent.futures
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import httpx
|
|
from loguru import logger
|
|
|
|
import config
|
|
|
|
_client = httpx.Client(
|
|
timeout=httpx.Timeout(
|
|
connect=3.0,
|
|
read=30.0, # 파일 전송이므로 기존 30초 유지
|
|
write=30.0,
|
|
pool=3.0
|
|
),
|
|
limits=httpx.Limits(
|
|
max_keepalive_connections=5,
|
|
keepalive_expiry=50 # A서버 --timeout-keep-alive 60 보다 약간 낮게
|
|
)
|
|
)
|
|
|
|
|
|
def ocr_single_pdf(pdf_path: str) -> dict:
|
|
with open(pdf_path, "rb") as f:
|
|
resp = httpx.post(
|
|
config.PADDLE_OCR_URL,
|
|
files={"file": (Path(pdf_path).name, f, "application/pdf")},
|
|
data={"group_id": "infer"},
|
|
timeout=120, # PDF는 이미지보다 오래 걸리므로 timeout 증가
|
|
)
|
|
resp.raise_for_status()
|
|
ocr_result = resp.json()
|
|
|
|
words, boxes = [], []
|
|
for page in ocr_result: # PDF는 페이지가 여러 장일 수 있으므로 전체 순회
|
|
for text, box in zip(page.get("words", []), page.get("boxes", [])):
|
|
text = text.strip()
|
|
if not text or not box:
|
|
continue
|
|
words.append(text)
|
|
boxes.append(box)
|
|
|
|
logger.debug(f"PaddleOCR PDF 완료: {pdf_path} | 단어수={len(words)}")
|
|
return {"image_path": str(pdf_path), "words": words, "boxes": boxes, "box_type": "pixel"}
|
|
|
|
|
|
def ocr_single(image_path: str, client=None) -> dict:
|
|
with open(image_path, "rb") as f:
|
|
data = {"group_id": "infer"}
|
|
if config.SERVICE == "ucar":
|
|
data["gubun"] = "ocr_with_boxes"
|
|
|
|
resp = _client.post(
|
|
config.PADDLE_OCR_URL,
|
|
files={"file": (Path(image_path).name, f, "image/jpeg")},
|
|
data=data,
|
|
)
|
|
|
|
resp.raise_for_status()
|
|
ocr_result = resp.json()
|
|
|
|
words, boxes = [], []
|
|
|
|
# 페이지가 1장이므로 첫 번째 항목만 사용
|
|
page = ocr_result[0] if ocr_result else {}
|
|
|
|
for text, box in zip(page.get("words", []), page.get("boxes", [])):
|
|
text = text.strip()
|
|
if not text or not box:
|
|
continue
|
|
words.append(text)
|
|
boxes.append(box)
|
|
|
|
result = {"image_path": str(image_path), "words": words, "boxes": boxes, "box_type": "pixel"}
|
|
|
|
# OCR 결과 저장
|
|
# json_path = Path(image_path).with_name(Path(image_path).stem + "_paddle.json")
|
|
# with open(json_path, "w", encoding="utf-8") as f:
|
|
# json.dump(result, f, ensure_ascii=False, indent=2)
|
|
# logger.debug(f"OCR 저장: {json_path}")
|
|
|
|
logger.debug(f"PaddleOCR 완료: {image_path} | 단어수={len(words)}")
|
|
return result
|
|
|
|
|
|
def ocr_batch(image_paths: list, client=None, max_workers: int = config.OCR_MAX_WORKERS,
|
|
on_done=None) -> list: # ← on_done 추가
|
|
results = [None] * len(image_paths)
|
|
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
|
|
future_to_idx = {
|
|
executor.submit(ocr_single, str(p)): i
|
|
# executor.submit(ocr_single_pdf, str(p)): i
|
|
for i, p in enumerate(image_paths)
|
|
}
|
|
for future in concurrent.futures.as_completed(future_to_idx):
|
|
idx = future_to_idx[future]
|
|
try:
|
|
results[idx] = future.result()
|
|
except Exception as e:
|
|
logger.error(f"PaddleOCR 실패: {image_paths[idx]} | {e}")
|
|
results[idx] = {"image_path": str(image_paths[idx]),
|
|
"words": [], "boxes": [], "error": str(e)}
|
|
if on_done: # ← 콜백 호출
|
|
on_done(idx, results[idx])
|
|
return results
|