272 lines
10 KiB
Python
272 lines
10 KiB
Python
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) |