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)