93 lines
3.2 KiB
Python
93 lines
3.2 KiB
Python
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
|