Initial commit

routers/predict.py

@@ -0,0 +1,272 @@
+import json
+from pathlib import Path
+
+import torch
+from fastapi import APIRouter, File, HTTPException, UploadFile, Form
+from PIL import Image
+from transformers import LayoutLMv3ForSequenceClassification, LayoutLMv3Processor
+from loguru import logger
+
+import config
+from train.schemas import PredictResult
+from utils.cache import load_label2id
+from utils.ocr import get_vision_client, ocr_single, ocr_batch
+from utils.pdf import is_pdf, pdf_to_images, pdf_to_images_batch
+
+from train.trainer import _get_latest_model_path
+from collections import defaultdict
+from datetime import datetime
+
+router = APIRouter()
+
+
+_model_cache: dict = {
+    "model":     None,
+    "processor": None,
+    "id2label":  None,
+    "path":      None,   # 로드된 모델 경로 (재학습 후 갱신 감지용)
+}
+
+
+def _infer(img_path: str, filename: str,
+           model, processor, id2label: dict,
+           threshold: float = 0.6) -> dict:
+    """이미지 1장 → 추론 결과"""
+    # client = get_vision_client()
+    # ocr    = ocr_single(img_path, client)
+
+    # 이미지와 같은 폴더에 json 캐시 저장/로드
+    cache_file = Path(img_path).with_suffix(".json")
+    if cache_file.exists():
+        with open(cache_file, encoding="utf-8") as f:
+            ocr = json.load(f)
+        logger.debug(f"OCR 캐시 사용: {cache_file}")
+    else:
+        # 캐시 없으면 OCR 실행 후 같은 폴더에 저장
+        client = get_vision_client()
+        ocr = ocr_single(img_path, client)
+        if ocr["words"]:
+            with open(cache_file, "w", encoding="utf-8") as f:
+                json.dump(ocr, f, ensure_ascii=False, indent=2)
+        logger.debug(f"OCR 신규 실행 → 저장: {cache_file}")
+
+    if not ocr["words"]:
+        return {"filename": filename, "error": "OCR 결과 없음. 이미지 품질 확인 필요."}
+
+    image    = Image.open(img_path).convert("RGB")
+    encoding = processor(
+        image, text=ocr["words"], boxes=ocr["boxes"],
+        return_tensors="pt", truncation=True,
+        padding="max_length", max_length=config.MAX_LEN,
+    )
+
+    with torch.no_grad():
+        outputs = model(
+            input_ids        = encoding["input_ids"].to(config.DEVICE),
+            attention_mask   = encoding["attention_mask"].to(config.DEVICE),
+            bbox             = encoding["bbox"].to(config.DEVICE),
+            pixel_values     = encoding["pixel_values"].to(config.DEVICE),
+            output_attentions= True,
+        )
+
+    probs   = torch.softmax(outputs.logits, dim=-1).squeeze().cpu()
+    pred_id = int(probs.argmax())
+
+    confidence = round(float(probs[pred_id]), 4)
+    label = id2label[pred_id] if confidence >= threshold else "OTHER"
+
+    attn     = outputs.attentions[-1][0].mean(0)[0]
+    n_words  = min(len(ocr["words"]), attn.shape[0] - 1)
+    top5_idx = attn[1:n_words+1].topk(min(5, n_words)).indices.tolist()
+
+    return {
+        "filename": filename,
+        "label": label,
+        "confidence": confidence,
+        "all_probs": {id2label[i]: round(float(p), 4) for i, p in enumerate(probs)},
+        "key_tokens": [ocr["words"][i] for i in top5_idx],
+        "ocr_word_count": len(ocr["words"]),
+    }
+
+    # return {
+    #     "filename":       filename,
+    #     "label":          id2label[pred_id],
+    #     "confidence":     round(float(probs[pred_id]), 4),
+    #     "all_probs":      {id2label[i]: round(float(p), 4) for i, p in enumerate(probs)},
+    #     "key_tokens":     [ocr["words"][i] for i in top5_idx],
+    #     "ocr_word_count": len(ocr["words"]),
+    # }
+
+
+def _load_model_and_processor():
+    model_path = _get_latest_model_path()
+    if not model_path:
+        raise HTTPException(status_code=400, detail="저장된 모델 없음. 학습 먼저 실행하세요.")
+
+    # 이미 로드된 모델이고 경로가 같으면 캐시 반환
+    if (_model_cache["model"] is not None and
+            _model_cache["path"] == str(model_path)):
+        logger.debug("모델 캐시 사용")
+        return _model_cache["model"], _model_cache["processor"], _model_cache["id2label"]
+
+    # 최초 로드 or 재학습 후 경로 변경 시
+    logger.info(f"모델 로드: {model_path}")
+    label2id  = load_label2id()
+    id2label  = {v: k for k, v in label2id.items()}
+    processor = LayoutLMv3Processor.from_pretrained(str(model_path), apply_ocr=False)
+    model     = LayoutLMv3ForSequenceClassification.from_pretrained(
+        str(model_path)
+    ).to(config.DEVICE)
+    model.eval()
+
+    # 캐시 갱신
+    _model_cache["model"]     = model
+    _model_cache["processor"] = processor
+    _model_cache["id2label"]  = id2label
+    _model_cache["path"]      = str(model_path)
+
+    return model, processor, id2label
+
+
+def _check_model():
+    if not _get_latest_model_path():
+        raise HTTPException(status_code=400, detail="저장된 모델 없음. 학습 먼저 실행하세요.")
+
+
+@router.post("", summary="단일 이미지/PDF 추론", response_model=PredictResult)
+async def predict(file: UploadFile = File(...), group_id: str = Form(...)):
+    _check_model()
+
+    # tmp_path  = Path(f"tmp_{file.filename}")
+    # img_paths = []
+
+    date_str = datetime.now().strftime("%Y%m%d")
+    tmp_dir = config.TMP_DIR / group_id / date_str  # files/tmp/{group_id}/년월일
+    tmp_dir.mkdir(parents=True, exist_ok=True)
+
+    tmp_path = tmp_dir / file.filename  # ← 수정
+    img_paths = []
+
+    with open(tmp_path, "wb") as f:
+        f.write(await file.read())
+
+    # try:
+        model, processor, id2label = _load_model_and_processor()
+
+        if is_pdf(file.filename):
+            img_paths = pdf_to_images(str(tmp_path), dpi=config.OCR_DPI)
+
+            if not img_paths:
+                raise HTTPException(status_code=422, detail="PDF 변환 실패.")
+
+            page_results = [
+                _infer(str(p), file.filename, model, processor, id2label)
+                for p in img_paths
+            ]
+
+            valid = [r for r in page_results
+                     if "label" in r and r.get("ocr_word_count", 0) >= config.OCR_MIN_WORDS]
+
+            if not valid:
+                best = {"error": "전체 페이지 OCR 실패"}
+            elif len(valid) == 1:
+                best = valid[0]
+            else:
+                # 2페이지 confidence 합산 → 라벨별 점수가 높은 쪽 선택
+                score = defaultdict(float)
+                for r in valid:
+                    for label, prob in r["all_probs"].items():
+                        score[label] += prob
+                best_label = max(score, key=score.get)
+                best = max(valid, key=lambda r: r["all_probs"].get(best_label, 0))
+                best = {**best, "label": best_label,
+                        "confidence": round(score[best_label] / len(valid), 4)}  # 평균 confidence
+
+            logger.info(f"PDF 추론 완료: {file.filename} | {best.get('label')}")
+            return PredictResult(**{**best,
+                                    "total_pages": len(img_paths),
+                                    "per_page": page_results})
+        else:
+            result = _infer(str(tmp_path), file.filename, model, processor, id2label)
+            logger.info(f"이미지 추론 완료: {file.filename} | {result.get('label')}")
+            return PredictResult(**result)
+
+    # finally:
+    #     tmp_path.unlink(missing_ok=True)
+    #     for p in img_paths:
+    #         Path(p).unlink(missing_ok=True)
+
+
+@router.post("/batch", summary="다중 이미지/PDF 배치 추론")
+async def predict_batch(files: list[UploadFile] = File(...), group_id: str = Form(...)):
+    _check_model()
+    model, processor, id2label = _load_model_and_processor()
+
+    date_str = datetime.now().strftime("%Y%m%d")
+    tmp_dir = config.TMP_DIR / group_id / date_str
+    tmp_dir.mkdir(parents=True, exist_ok=True)
+
+    # 1. 파일 저장
+    tmp_paths, pdf_tmp_paths, img_meta = [], [], []
+    for file in files:
+        tmp = Path(f"tmp_{file.filename}")
+        with open(tmp, "wb") as f:
+            f.write(await file.read())
+        tmp_paths.append(tmp)
+
+        if is_pdf(file.filename):
+            pdf_tmp_paths.append((file.filename, str(tmp)))
+        else:
+            img_meta.append((file.filename, [tmp]))
+
+    # 2. PDF 병렬 변환
+    if pdf_tmp_paths:
+        paths_only = [p for _, p in pdf_tmp_paths]
+        converted  = pdf_to_images_batch(paths_only)          # ← 병렬 변환
+        for filename, tmp_path in pdf_tmp_paths:
+            img_meta.append((filename, converted.get(tmp_path, [])))
+
+    try:
+        results = []
+        for filename, img_list in img_meta:
+            if not img_list:
+                results.append(PredictResult(filename=filename, error="PDF 변환 실패"))
+                continue
+
+            page_results = [
+                _infer(str(p), filename, model, processor, id2label)
+                for p in img_list
+            ]
+            valid = [r for r in page_results if "label" in r]
+
+            if not valid:
+                best = {"filename": filename, "error": "OCR 실패"}
+            elif len(valid) == 1:
+                best = valid[0]
+            else:
+                score = defaultdict(float)
+                for r in valid:
+                    for label, prob in r["all_probs"].items():
+                        score[label] += prob
+                best_label = max(score, key=score.get)
+                best = max(valid, key=lambda r: r["all_probs"].get(best_label, 0))
+                best = {**best, "label": best_label,
+                        "confidence": round(score[best_label] / len(valid), 4)}
+
+            is_multi = len(img_list) > 1
+            results.append(PredictResult(**{
+                **best,
+                "total_pages": len(img_list) if is_multi else None,
+                "per_page":    page_results if is_multi else None,
+            }))
+            logger.info(f"배치 추론: {filename} | {best.get('label')}")
+
+        return {"total": len(files), "results": [r.dict() for r in results]}
+
+    finally:
+        for p in tmp_paths:
+            p.unlink(missing_ok=True)
+        for _, img_list in img_meta:
+            if len(img_list) > 1:
+                for p in img_list:
+                    Path(p).unlink(missing_ok=True)