Initial commit

utils/ocr_paddle.py

@@ -0,0 +1,107 @@
+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