Initial commit

utils/cache.py

@@ -0,0 +1,123 @@
+import json
+import random
+from pathlib import Path
+
+from loguru import logger
+
+import config
+
+
+# ── OCR 캐시 (파일별 json 분리) ───────────────────
+
+def _cache_path(image_path: str) -> Path:
+    """
+    이미지 경로 → json 저장 경로
+    files/data/TSMC_A/1000114581_INV/page_001.png
+    → files/json/TSMC_A/1000114581_INV/page_001.json
+    """
+    img = Path(image_path)
+    try:
+        rel = img.relative_to(config.DATA_DIR)
+    except ValueError:
+        rel = Path(img.parent.name) / img.name
+    return config.OCR_CACHE_DIR / rel.with_suffix(".json")
+
+
+def save_cache_by_file(ocr_result: dict):
+    """이미지 1장의 OCR 결과를 개별 json으로 저장"""
+    path = _cache_path(ocr_result["image_path"])
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(ocr_result, f, ensure_ascii=False, indent=2)
+    logger.debug(f"캐시 저장: {path}")
+
+
+def load_cache() -> list:
+    """전체 캐시 통합 로드 (학습/평가 시 사용)"""
+    samples = []
+    for p in config.OCR_CACHE_DIR.rglob("*.json"):
+        try:
+            with open(p, encoding="utf-8") as f:
+                samples.append(json.load(f))
+        except (json.JSONDecodeError, UnicodeDecodeError) as e:
+            logger.warning(f"캐시 로드 실패 (건너뜀): {p} | {e}")
+    return samples
+
+
+def get_cached_paths() -> set:
+    """캐시된 이미지 경로 전체 반환 (중복 방지용)"""
+    cached = []
+    for p in config.OCR_CACHE_DIR.rglob("*.json"):
+        try:
+            with open(p, encoding="utf-8") as f:
+                data = json.load(f)
+            cached.append(data.get("image_path", ""))
+            # image_path = data.get("image_path", "")
+            # if image_path:
+            #     cached.append(str(Path(image_path)))  # 정규화
+            #     print(str(Path(image_path)))
+        except (json.JSONDecodeError, UnicodeDecodeError):
+            pass
+    return set(cached)
+
+
+def cache_status() -> dict:
+    """양식별 캐시 현황 (폴더별 json 파일 수 집계)"""
+    if not config.OCR_CACHE_DIR.exists():
+        return {}
+    status = {}
+    for cls_dir in sorted(config.OCR_CACHE_DIR.iterdir()):
+        if cls_dir.is_dir():
+            status[cls_dir.name] = len(list(cls_dir.rglob("*.json")))
+    return status
+
+
+# ── label2id ──────────────────────────────────────
+
+def load_label2id() -> dict:
+    path = config.MODEL_SAVE_PATH / "label2id.json"
+    if path.exists():
+        with open(path, encoding="utf-8") as f:
+            return json.load(f)
+    if config.DATA_DIR.exists():
+        classes = sorted([d.name for d in config.DATA_DIR.iterdir() if d.is_dir()])
+        return {c: i for i, c in enumerate(classes)}
+    return {}
+
+
+def save_label2id(label2id: dict):
+    config.MODEL_SAVE_PATH.mkdir(exist_ok=True)
+    with open(config.MODEL_SAVE_PATH / "label2id.json", "w", encoding="utf-8") as f:
+        json.dump(label2id, f, ensure_ascii=False, indent=2)
+    logger.info(f"label2id 저장: {label2id}")
+
+
+# ── 리플레이 버퍼 ──────────────────────────────────
+
+def load_replay_buffer() -> dict:
+    if not config.REPLAY_BUFFER_FILE.exists():
+        return {}
+    try:
+        with open(config.REPLAY_BUFFER_FILE, encoding="utf-8") as f:
+            return json.load(f)
+    except (json.JSONDecodeError, UnicodeDecodeError) as e:
+        backup = config.REPLAY_BUFFER_FILE.with_suffix(".json.bak")
+        config.REPLAY_BUFFER_FILE.rename(backup)
+        logger.warning(f"replay_buffer.json 손상 → 백업 후 초기화: {e}")
+        return {}
+
+
+def update_replay_buffer(new_samples: list, id2label: dict, per_class: int = 10):
+    buf      = load_replay_buffer()
+    by_label = {}
+    for s in new_samples:
+        name = id2label[s["label"]]
+        by_label.setdefault(name, []).append(s)
+
+    for name, samples in by_label.items():
+        combined  = buf.get(name, []) + samples
+        buf[name] = random.sample(combined, min(per_class, len(combined)))
+
+    with open(config.REPLAY_BUFFER_FILE, "w", encoding="utf-8") as f:
+        json.dump(buf, f, ensure_ascii=False, indent=2)
+    logger.info(f"리플레이 버퍼 갱신: { {k: len(v) for k, v in buf.items()} }")

utils/classifier_model_store.py

@@ -0,0 +1,68 @@
+"""
+모델 로드/저장/예측 담당
+"""
+from pathlib import Path
+from typing import Optional
+
+import joblib
+import numpy as np
+from loguru import logger
+from sklearn.pipeline import Pipeline
+
+import config
+
+MODEL_FILENAME = "classifier.pkl"
+
+
+def _get_latest_path() -> Optional[Path]:
+    candidates = sorted(config.CLASSIFIER_MODEL_BASE_PATH.rglob(MODEL_FILENAME))
+    return candidates[-1] if candidates else None
+
+
+def _get_save_path() -> Path:
+    from datetime import datetime
+    date_dir = config.CLASSIFIER_MODEL_BASE_PATH / datetime.now().strftime("%Y%m%d")
+    date_dir.mkdir(exist_ok=True)
+    existing  = sorted([d for d in date_dir.iterdir() if d.is_dir() and d.name.isdigit()])
+    next_idx  = int(existing[-1].name) + 1 if existing else 1
+    save_dir  = date_dir / f"{next_idx:04d}"
+    save_dir.mkdir(exist_ok=True)
+    return save_dir / MODEL_FILENAME
+
+
+_model: Optional[Pipeline] = None
+
+# 서버 시작 시 최신 모델 자동 로드
+_latest = _get_latest_path()
+if _latest:
+    try:
+        _model = joblib.load(_latest)
+        logger.info(f"classifier 모델 로드: {_latest} | classes={list(_model.classes_)}")
+    except Exception as e:
+        logger.warning(f"모델 로드 실패: {e}")
+
+
+def get_model() -> Optional[Pipeline]:
+    return _model
+
+
+def set_model(pipeline: Pipeline) -> None:
+    global _model
+    path = _get_save_path()
+    joblib.dump(pipeline, path)
+    _model = pipeline
+    logger.info(f"classifier 모델 저장: {path}")
+
+
+def predict(
+    pipeline: Pipeline,
+    text: str,
+    threshold: float = 0.6,
+) -> tuple[str, float, dict[str, float]]:
+    probs      = pipeline.predict_proba([text])[0]
+    classes    = pipeline.classes_
+    max_idx    = int(np.argmax(probs))
+    confidence = float(probs[max_idx])
+    label      = "OTHER" if confidence < threshold else classes[max_idx]
+    all_probs  = {c: round(float(p), 4) for c, p in zip(classes, probs)}
+    return label, round(confidence, 4), all_probs

utils/ocr.py

@@ -0,0 +1,38 @@
+# import config
+#
+# if config.OCR_ENGINE == "google":
+#     from utils.ocr_google import ocr_single, ocr_batch, get_vision_client
+# else:
+#     from utils.ocr_paddle import ocr_single, ocr_batch
+#     def get_vision_client():
+#         return None   # paddle은 client 불필요, 호환성 유지용
+
+
+from utils.ocr_paddle import ocr_single as paddle_ocr_single, ocr_batch as paddle_ocr_batch
+from utils.ocr_google import ocr_single as google_ocr_single, ocr_batch as google_ocr_batch, get_vision_client
+import config
+
+_google_client = None
+
+
+def get_vision_client_cached():
+    global _google_client
+    if _google_client is None:
+        _google_client = get_vision_client()
+    return _google_client
+
+
+def get_ocr_functions(engine: str):
+    if engine == "google":
+        return google_ocr_single, google_ocr_batch, get_vision_client_cached()  # 캐싱된 클라이언트 사용
+    return paddle_ocr_single, paddle_ocr_batch, None
+
+
+def ocr_single(image_path, client=None):
+    fn, _, c = get_ocr_functions(config.OCR_ENGINE)
+    return fn(image_path, c)
+
+
+def ocr_batch(img_paths, client=None):
+    _, fn, c = get_ocr_functions(config.OCR_ENGINE)
+    return fn(img_paths, c)

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

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

utils/pdf.py

@@ -0,0 +1,85 @@
+from pathlib import Path
+
+import fitz  # PyMuPDF
+from loguru import logger
+
+import config
+import concurrent.futures
+
+
+def is_pdf(filename: str) -> bool:
+    return Path(filename).suffix.lower() == ".pdf"
+
+
+def pdf_to_images(pdf_path: str, dpi: int = config.OCR_DPI,
+                  max_pages: int = config.PDF_MAX_PAGES) -> list[Path]:
+    pdf_path   = Path(pdf_path)
+    output_dir = pdf_path.parent / pdf_path.stem
+    output_dir.mkdir(exist_ok=True)
+
+    ext          = config.PDF_IMAGE_FORMAT                     # "jpeg"
+    page_indices = list(range(max_pages))
+
+    # 이미 변환된 이미지가 모두 있으면 스킵
+    cached = [output_dir / f"page_{i+1:03d}.{ext}" for i in page_indices]
+    if all(p.exists() for p in cached):
+        logger.debug(f"PDF 변환 캐시 사용: {pdf_path.name}")
+        return cached
+
+    pdf_path = Path(pdf_path)
+    print(pdf_path.exists())
+    print(pdf_path.stat().st_size)
+    if not pdf_path.exists() or pdf_path.stat().st_size == 0:
+        raise ValueError(f"유효하지 않은 PDF 파일입니다: {pdf_path.name}")
+
+    matrix = fitz.Matrix(dpi / 72, dpi / 72)
+    doc = fitz.open(str(pdf_path))
+
+    img_paths = []
+    for i in range(min(max_pages, len(doc))):
+        img_path = output_dir / f"page_{i+1:03d}.{ext}"
+        if img_path.exists():
+            logger.debug(f"페이지 캐시 사용: {img_path.name}")
+        else:
+            pixmap = doc[i].get_pixmap(matrix=matrix, alpha=False)
+
+            # 최대 너비 초과 시 비율 유지하며 축소
+            if hasattr(config, 'PDF_MAX_WIDTH') and pixmap.width > config.PDF_MAX_WIDTH:
+                scale = config.PDF_MAX_WIDTH / pixmap.width
+                pixmap = doc[i].get_pixmap(matrix=fitz.Matrix(dpi / 72 * scale, dpi / 72 * scale), alpha=False)
+
+            if ext in ("jpeg", "jpg"):
+                pixmap.save(str(img_path), jpg_quality=config.PDF_JPEG_QUALITY)  # JPEG 저장
+            else:
+                pixmap.save(str(img_path))
+            logger.debug(f"PDF 변환: {pdf_path.name} | 페이지 {i+1}")
+        img_paths.append(img_path)
+
+    doc.close()
+    logger.info(f"PDF 변환 완료: {pdf_path.name} | {len(img_paths)}페이지 | {dpi}dpi {ext}")
+    return img_paths
+
+
+def pdf_to_images_batch(pdf_paths: list[str],
+                        max_workers: int = config.OCR_MAX_WORKERS) -> dict[str, list[Path]]:
+    """
+    여러 PDF를 ThreadPoolExecutor로 병렬 변환
+    반환: {pdf_path: [img_path, ...]}
+    """
+    results = {}
+
+    def _convert(pdf_path: str):
+        return pdf_path, pdf_to_images(pdf_path)
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
+        futures = {executor.submit(_convert, p): p for p in pdf_paths}
+        for future in concurrent.futures.as_completed(futures):
+            try:
+                path, imgs = future.result()
+                results[path] = imgs
+            except Exception as e:
+                path = futures[future]
+                logger.error(f"PDF 변환 실패: {path} | {e}")
+                results[path] = []
+
+    return results