Initial commit

utils/ocr_google.py

@@ -0,0 +1,92 @@
+import concurrent.futures
+import json
+from pathlib import Path
+
+from google.cloud import vision_v1
+from google.oauth2 import service_account
+from loguru import logger
+from PIL import Image
+
+import config
+
+
+def get_vision_client() -> vision_v1.ImageAnnotatorClient:
+    credentials = service_account.Credentials.from_service_account_file(
+        config.SERVICE_ACCOUNT_FILE,
+        scopes=["https://www.googleapis.com/auth/cloud-platform"],
+    )
+    return vision_v1.ImageAnnotatorClient(credentials=credentials)
+
+
+def ocr_single(image_path: str, client: vision_v1.ImageAnnotatorClient) -> dict:
+    """
+    이미지 1장 → words + 정규화 boxes (0~1000) 반환
+    """
+    image_pil = Image.open(image_path).convert("RGB")
+    w, h      = image_pil.size
+
+    with open(image_path, "rb") as f:
+        content = f.read()
+
+    response = client.document_text_detection(
+        image=vision_v1.Image(content=content)
+    )
+
+    words, boxes = [], []
+    for page in response.full_text_annotation.pages:
+        for block in page.blocks:
+            for para in block.paragraphs:
+                for word in para.words:
+                    text = "".join([s.text for s in word.symbols])
+                    if not text.strip():
+                        continue
+                    v = word.bounding_box.vertices
+                    boxes.append([
+                        max(0,    int(1000 * v[0].x / w)),
+                        max(0,    int(1000 * v[0].y / h)),
+                        min(1000, int(1000 * v[2].x / w)),
+                        min(1000, int(1000 * v[2].y / h)),
+                    ])
+                    words.append(text)
+
+    result = {"image_path": str(image_path), "words": words, "boxes": boxes, "box_type": "normalized"}
+
+    # OCR 결과 저장
+    # json_path = Path(image_path).with_name(Path(image_path).stem + "_google.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"OCR 완료: {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:
+    """
+    ThreadPoolExecutor 병렬 OCR
+    Vision API는 네트워크 I/O 기반 → 병렬 효과 큼
+    """
+    # client  = get_vision_client()
+    if client is None:
+        client = get_vision_client()
+    results = [None] * len(image_paths)
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
+        future_to_idx = {
+            executor.submit(ocr_single, str(p), client): 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"OCR 실패: {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