521 lines
22 KiB
Python
521 lines
22 KiB
Python
import os
|
|
from datetime import datetime
|
|
import time
|
|
from typing import List
|
|
|
|
from fastapi import APIRouter, BackgroundTasks, HTTPException
|
|
from fastapi.responses import HTMLResponse
|
|
from sse_starlette.sse import EventSourceResponse
|
|
from loguru import logger
|
|
from pydantic import BaseModel
|
|
import asyncio
|
|
import json
|
|
|
|
from utils import classifier_model_store
|
|
from train import classifier_trainer
|
|
from train import classifier_dataset
|
|
|
|
from fastapi import APIRouter, BackgroundTasks, File, Form, HTTPException, UploadFile
|
|
from pathlib import Path
|
|
import asyncio
|
|
|
|
import config
|
|
from utils.ocr import ocr_batch, get_ocr_functions
|
|
from utils.pdf import is_pdf, pdf_to_images
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
# ── 스키마 ────────────────────────────────────────
|
|
class PredictRequest(BaseModel):
|
|
words: list[str]
|
|
threshold: float = 0.6
|
|
|
|
|
|
class PredictResponse(BaseModel):
|
|
label: str
|
|
file_name: str
|
|
confidence: float
|
|
all_probs: dict[str, float]
|
|
|
|
|
|
class PredictListResponse(BaseModel):
|
|
results: list[PredictResponse]
|
|
success: bool
|
|
|
|
|
|
class TrainResponse(BaseModel):
|
|
status: str
|
|
message: str
|
|
|
|
class DatasetResponse(BaseModel):
|
|
status: str
|
|
message: str
|
|
total: int = 0
|
|
|
|
|
|
# ── Dataset ───────────────────────────────────────
|
|
@router.post("/dataset/build", response_model=DatasetResponse, summary="데이터셋 생성")
|
|
async def build_dataset(background_tasks: BackgroundTasks):
|
|
background_tasks.add_task(classifier_dataset.build)
|
|
return DatasetResponse(status="building", message="데이터셋 생성 시작. /classifier/dataset/status 확인하세요.")
|
|
|
|
|
|
@router.get("/dataset/status", response_model=DatasetResponse, summary="데이터셋 생성 상태")
|
|
async def dataset_status():
|
|
s = classifier_dataset.get_status()
|
|
return DatasetResponse(**s)
|
|
|
|
|
|
@router.get("/dataset/info", summary="데이터셋 정보 조회")
|
|
async def dataset_info():
|
|
data = classifier_dataset.load()
|
|
if not data:
|
|
raise HTTPException(status_code=404, detail="데이터셋이 없습니다. /classifier/dataset/build 먼저 실행하세요.")
|
|
|
|
from collections import Counter
|
|
counts = Counter(item["label"] for item in data)
|
|
return {
|
|
"total": len(data),
|
|
"label_count": len(counts),
|
|
"labels": dict(sorted(counts.items(), key=lambda x: -x[1])),
|
|
}
|
|
|
|
|
|
# ── Train ─────────────────────────────────────────
|
|
@router.post("/train", response_model=TrainResponse, summary="모델 학습 시작")
|
|
async def train(background_tasks: BackgroundTasks):
|
|
data = classifier_dataset.load()
|
|
if not data:
|
|
raise HTTPException(status_code=400, detail="데이터셋이 없습니다. /classifier/dataset/build 먼저 실행하세요.")
|
|
|
|
s = classifier_trainer.get_status()
|
|
if s["status"] == "training":
|
|
raise HTTPException(status_code=409, detail="이미 학습 중입니다.")
|
|
|
|
background_tasks.add_task(classifier_trainer.run_train)
|
|
logger.info("classifier 학습 백그라운드 태스크 등록")
|
|
return TrainResponse(status="training", message="학습을 시작했습니다. /classifier/train/status 로 확인하세요.")
|
|
|
|
|
|
@router.get("/train/status", response_model=TrainResponse, summary="학습 상태 확인")
|
|
async def train_status():
|
|
return TrainResponse(**classifier_trainer.get_status())
|
|
|
|
|
|
# ── Train SSE 스트리밍 ────────────────────────────
|
|
@router.get("/train/stream", summary="학습 진행상태 SSE 스트리밍")
|
|
async def train_stream():
|
|
async def generator():
|
|
prev = ""
|
|
while True:
|
|
s = classifier_trainer.get_status()
|
|
data = json.dumps(s, ensure_ascii=False)
|
|
|
|
if data != prev:
|
|
yield {"event": "status", "data": data}
|
|
prev = data
|
|
|
|
if s.get("status") not in ("training", "building"):
|
|
yield {"event": "done", "data": json.dumps({"message": s.get("message", "완료")})}
|
|
break
|
|
|
|
await asyncio.sleep(0.5)
|
|
|
|
return EventSourceResponse(generator())
|
|
|
|
|
|
# ── Train 모니터 HTML ─────────────────────────────
|
|
@router.get("/train/monitor", summary="학습 모니터 HTML", response_class=HTMLResponse)
|
|
def train_monitor():
|
|
html = """
|
|
<!DOCTYPE html>
|
|
<html>
|
|
<head>
|
|
<meta charset="utf-8">
|
|
<title>Classifier 학습 모니터</title>
|
|
<style>
|
|
* { box-sizing: border-box; }
|
|
body { font-family: sans-serif; margin: 0; background: #f5f5f5;
|
|
display: flex; flex-direction: column; height: 100vh; }
|
|
.header { padding: 16px 24px; background: #6a1b9a; color: #fff; flex-shrink: 0; }
|
|
.header h2 { margin: 0; font-size: 18px; }
|
|
.container { padding: 20px 24px; flex: 1; overflow-y: auto; }
|
|
|
|
.cards { display: flex; gap: 16px; margin-bottom: 20px; flex-wrap: wrap; }
|
|
.card { background: #fff; border-radius: 8px; padding: 16px 20px;
|
|
flex: 1; min-width: 130px; box-shadow: 0 1px 4px rgba(0,0,0,0.1); }
|
|
.card .label { font-size: 12px; color: #888; margin-bottom: 6px; }
|
|
.card .value { font-size: 26px; font-weight: bold; color: #6a1b9a; }
|
|
|
|
.progress-wrap { background: #fff; border-radius: 8px; padding: 16px 20px;
|
|
margin-bottom: 20px; box-shadow: 0 1px 4px rgba(0,0,0,0.1); }
|
|
.progress-label { display: flex; justify-content: space-between;
|
|
font-size: 13px; color: #555; margin-bottom: 8px; }
|
|
.progress-bar { height: 18px; background: #e0e0e0; border-radius: 9px; overflow: hidden; }
|
|
.progress-fill { height: 100%; background: #6a1b9a; border-radius: 9px;
|
|
transition: width 0.4s ease; width: 0%; }
|
|
|
|
.msg-box { background: #fff; border-radius: 8px; padding: 12px 20px;
|
|
margin-bottom: 20px; box-shadow: 0 1px 4px rgba(0,0,0,0.1);
|
|
font-size: 13px; color: #555; font-family: monospace; word-break: break-all; }
|
|
|
|
.log-wrap { background: #fff; border-radius: 8px; padding: 16px 20px;
|
|
box-shadow: 0 1px 4px rgba(0,0,0,0.1); }
|
|
.log-wrap h3 { margin: 0 0 10px; font-size: 14px; color: #333; }
|
|
.log-list { list-style: none; margin: 0; padding: 0;
|
|
max-height: 400px; overflow-y: auto; }
|
|
.log-list li { padding: 5px 8px; font-size: 13px; border-bottom: 1px solid #f0f0f0;
|
|
display: flex; gap: 10px; }
|
|
.log-list li .time { color: #aaa; flex-shrink: 0; }
|
|
.log-list li.done { color: #2e7d32; font-weight: bold; }
|
|
.log-list li.error { color: #c62828; }
|
|
|
|
.badge { display: inline-block; padding: 3px 10px; border-radius: 12px;
|
|
font-size: 12px; font-weight: bold; }
|
|
.badge.running { background: #f3e5f5; color: #6a1b9a; }
|
|
.badge.done { background: #e8f5e9; color: #2e7d32; }
|
|
.badge.idle { background: #f5f5f5; color: #757575; }
|
|
|
|
.btn { padding: 8px 20px; border: none; border-radius: 6px; cursor: pointer;
|
|
font-size: 13px; font-weight: bold; }
|
|
.btn-start { background: #6a1b9a; color: #fff; }
|
|
.btn-start:hover { background: #4a148c; }
|
|
.btn-start:disabled { background: #ce93d8; cursor: not-allowed; }
|
|
.actions { display: flex; gap: 10px; align-items: center; margin-bottom: 20px; }
|
|
</style>
|
|
</head>
|
|
<body>
|
|
<div class="header"><h2>🧠 Classifier 학습 모니터</h2></div>
|
|
<div class="container">
|
|
|
|
<div class="actions">
|
|
<button class="btn btn-start" id="btnStart" onclick="startTrain()">▶ 학습 시작</button>
|
|
<span id="badge" class="badge idle">대기 중</span>
|
|
</div>
|
|
|
|
<div class="progress-wrap">
|
|
<div class="progress-label">
|
|
<span id="epochLabel">Epoch 0 / 0</span>
|
|
<span id="pctLabel">0%</span>
|
|
</div>
|
|
<div class="progress-bar">
|
|
<div class="progress-fill" id="progressFill"></div>
|
|
</div>
|
|
</div>
|
|
|
|
<div class="cards">
|
|
<div class="card">
|
|
<div class="label">샘플 수</div>
|
|
<div class="value" id="cardSamples">-</div>
|
|
</div>
|
|
<div class="card">
|
|
<div class="label">클래스 수</div>
|
|
<div class="value" id="cardClasses">-</div>
|
|
</div>
|
|
<div class="card">
|
|
<div class="label">진행률</div>
|
|
<div class="value" id="cardPct">0%</div>
|
|
</div>
|
|
</div>
|
|
|
|
<div class="msg-box" id="msgBox">대기 중...</div>
|
|
|
|
<div class="log-wrap">
|
|
<h3>📋 진행 로그</h3>
|
|
<ul class="log-list" id="logList"></ul>
|
|
</div>
|
|
</div>
|
|
|
|
<script>
|
|
let es = null;
|
|
|
|
function now() { return new Date().toLocaleTimeString("ko-KR"); }
|
|
|
|
function addLog(msg, cls = "") {
|
|
const li = document.createElement("li");
|
|
li.className = cls;
|
|
li.innerHTML = `<span class="time">${now()}</span><span>${msg}</span>`;
|
|
document.getElementById("logList").prepend(li);
|
|
}
|
|
/*
|
|
function updateUI(d) {
|
|
const progress = d.status === "training" ? 50 : d.status === "done" ? 100 : 0;
|
|
|
|
document.getElementById("progressFill").style.width = progress + "%";
|
|
document.getElementById("epochLabel").textContent = d.status === "training" ? "학습 진행 중..." : d.status === "done" ? "완료" : "대기 중";
|
|
document.getElementById("pctLabel").textContent = progress + "%";
|
|
document.getElementById("cardSamples").textContent = d.samples ?? "-";
|
|
document.getElementById("cardClasses").textContent = d.classes ?? "-";
|
|
document.getElementById("cardPct").textContent = progress + "%";
|
|
document.getElementById("msgBox").textContent = d.message ?? "";
|
|
|
|
const badge = document.getElementById("badge");
|
|
if (d.status === "training") {
|
|
badge.textContent = "학습 중"; badge.className = "badge running";
|
|
}
|
|
}
|
|
*/
|
|
function updateUI(d) {
|
|
let progress = 0;
|
|
let stepLabel = "대기 중";
|
|
|
|
if (d.status === "training") {
|
|
if (d.message.includes("로딩")) {
|
|
progress = 10; stepLabel = "데이터 로딩 중...";
|
|
} else if (d.message.includes("학습 중")) {
|
|
progress = 40; stepLabel = "TF-IDF 변환 중...";
|
|
} else if (d.message.includes("fit") || d.message.includes("학습")) {
|
|
progress = 70; stepLabel = "SVM 학습 중...";
|
|
} else {
|
|
progress = 50; stepLabel = "학습 중...";
|
|
}
|
|
} else if (d.status === "done") {
|
|
progress = 100; stepLabel = "완료";
|
|
} else if (d.status === "error") {
|
|
progress = 0; stepLabel = "오류 발생";
|
|
}
|
|
|
|
document.getElementById("progressFill").style.width = progress + "%";
|
|
document.getElementById("epochLabel").textContent = stepLabel;
|
|
document.getElementById("pctLabel").textContent = progress + "%";
|
|
document.getElementById("cardSamples").textContent = d.samples ?? "-";
|
|
document.getElementById("cardClasses").textContent = d.classes ?? "-";
|
|
document.getElementById("cardPct").textContent = progress + "%";
|
|
document.getElementById("msgBox").textContent = d.message ?? "";
|
|
|
|
const badge = document.getElementById("badge");
|
|
if (d.status === "training") {
|
|
badge.textContent = "학습 중"; badge.className = "badge running";
|
|
} else if (d.status === "error") {
|
|
badge.textContent = "오류"; badge.className = "badge idle";
|
|
}
|
|
}
|
|
|
|
function startTrain() {
|
|
fetch("/classifier/train", { method: "POST" })
|
|
.then(r => r.json())
|
|
.then(data => {
|
|
addLog("학습 시작: " + data.message);
|
|
document.getElementById("btnStart").disabled = true;
|
|
subscribeSSE();
|
|
})
|
|
.catch(e => addLog("오류: " + e, "error"));
|
|
}
|
|
|
|
function subscribeSSE() {
|
|
if (es) es.close();
|
|
es = new EventSource("/classifier/train/stream");
|
|
|
|
es.addEventListener("status", e => {
|
|
const d = JSON.parse(e.data);
|
|
console.log(d);
|
|
updateUI(d);
|
|
if (d.message) addLog(d.message);
|
|
});
|
|
|
|
es.addEventListener("done", e => {
|
|
const d = JSON.parse(e.data);
|
|
addLog("✅ " + d.message, "done");
|
|
const badge = document.getElementById("badge");
|
|
badge.textContent = "완료"; badge.className = "badge done";
|
|
document.getElementById("btnStart").disabled = false;
|
|
es.close();
|
|
});
|
|
|
|
es.onerror = () => {
|
|
addLog("SSE 연결 끊김", "error");
|
|
es.close();
|
|
document.getElementById("btnStart").disabled = false;
|
|
};
|
|
}
|
|
|
|
// 페이지 로드 시 이미 학습 중이면 자동 연결
|
|
fetch("/classifier/train/status")
|
|
.then(r => r.json())
|
|
.then(d => {
|
|
updateUI(d);
|
|
if (d.status === "training") {
|
|
document.getElementById("btnStart").disabled = true;
|
|
addLog("학습 진행 중 - 자동 연결");
|
|
subscribeSSE();
|
|
}
|
|
});
|
|
</script>
|
|
</body>
|
|
</html>
|
|
"""
|
|
return HTMLResponse(content=html)
|
|
|
|
|
|
# ── Predict ───────────────────────────────────────
|
|
@router.post("/predict", response_model=PredictResponse, summary="문서 분류 예측 (PDF/이미지 파일)")
|
|
async def predict(
|
|
file: UploadFile = File(..., description="PDF 또는 이미지 파일"),
|
|
threshold: float = Form(0.4, description="신뢰도 임계값"),
|
|
ocr_engine: str = Form(config.OCR_ENGINE, description="ocr 엔진 선택: paddle | google"),
|
|
):
|
|
m = classifier_model_store.get_model()
|
|
if m is None:
|
|
raise HTTPException(status_code=503, detail="모델이 없습니다. /classifier/train 먼저 실행하세요.")
|
|
|
|
content = await file.read()
|
|
tmp_path = config.TMP_DIR / file.filename
|
|
with open(tmp_path, "wb") as f:
|
|
f.write(content)
|
|
f.flush()
|
|
os.fsync(f.fileno())
|
|
|
|
# try:
|
|
t_total = time.time()
|
|
|
|
# 1. 이미지 변환
|
|
t0 = time.time()
|
|
if is_pdf(file.filename):
|
|
img_paths = pdf_to_images(str(tmp_path), dpi=config.OCR_DPI, max_pages=1)
|
|
else:
|
|
# 이미지 리사이즈
|
|
if hasattr(config, 'PDF_MAX_WIDTH'):
|
|
from PIL import Image
|
|
img = Image.open(tmp_path)
|
|
if img.width > config.PDF_MAX_WIDTH:
|
|
scale = config.PDF_MAX_WIDTH / img.width
|
|
new_size = (int(img.width * scale), int(img.height * scale))
|
|
img = img.resize(new_size, Image.LANCZOS)
|
|
ext = tmp_path.suffix.lower()
|
|
if ext in (".jpg", ".jpeg"):
|
|
img.save(str(tmp_path), quality=config.PDF_JPEG_QUALITY)
|
|
else:
|
|
img.save(str(tmp_path))
|
|
img_paths = [tmp_path]
|
|
|
|
t_convert = time.time() - t0
|
|
logger.info(f"[{file.filename}] 이미지변환: {t_convert:.3f}s | {len(img_paths)}페이지")
|
|
|
|
if not img_paths:
|
|
raise HTTPException(status_code=400, detail="이미지 변환 실패")
|
|
|
|
# 2. OCR
|
|
t0 = time.time()
|
|
loop = asyncio.get_event_loop()
|
|
# results = await loop.run_in_executor(None, ocr_batch, img_paths)
|
|
_, batch_fn, client = get_ocr_functions(ocr_engine)
|
|
results = await loop.run_in_executor(None, batch_fn, img_paths, client)
|
|
|
|
t_ocr = time.time() - t0
|
|
logger.info(f"[{file.filename}] OCR: {t_ocr:.3f}s | {len(results)}페이지")
|
|
|
|
ocr_errors = [r for r in results if r.get("error")]
|
|
if ocr_errors:
|
|
raise HTTPException(status_code=422, detail=f"OCR 실패: {ocr_errors[0]['error']}")
|
|
|
|
# 3. 추론
|
|
t0 = time.time()
|
|
best_label, best_conf, best_probs = "OTHER", 0.0, {}
|
|
for r in results:
|
|
words = r.get("words", [])
|
|
if len(words) < config.OCR_MIN_WORDS:
|
|
continue
|
|
lbl, conf, probs = classifier_model_store.predict(m, " ".join(words), threshold)
|
|
if conf > best_conf:
|
|
best_label, best_conf, best_probs = lbl, conf, probs
|
|
t_infer = time.time() - t0
|
|
logger.info(f"[{file.filename}] 추론: {t_infer:.3f}s | label={best_label}, confidence={best_conf:.4f}")
|
|
|
|
# 4. 최종
|
|
t_total = time.time() - t_total
|
|
logger.info(f"[{file.filename}] 최종: {t_total:.3f}s (변환={t_convert:.3f}s / OCR={t_ocr:.3f}s / 추론={t_infer:.3f}s)")
|
|
|
|
if best_conf == 0.0:
|
|
raise HTTPException(status_code=422, detail="OCR 결과가 없거나 단어 수 미달입니다.")
|
|
|
|
return PredictResponse(label=best_label, file_name=file.filename, confidence=best_conf, all_probs=best_probs)
|
|
# finally:
|
|
# # ── 임시 파일 정리 ──────────────────────────
|
|
# tmp_path.unlink(missing_ok=True)
|
|
# if is_pdf(file.filename):
|
|
# for p in img_paths:
|
|
# Path(p).unlink(missing_ok=True)
|
|
|
|
|
|
@router.post("/predictMulti", response_model=PredictListResponse, summary="문서 분류 예측 (PDF/이미지 파일)")
|
|
async def predictMulti(
|
|
files: List[UploadFile] = File(..., description="PDF 또는 이미지 파일"),
|
|
threshold: float = Form(0.4, description="신뢰도 임계값"),
|
|
ocr_engine: str = Form(config.OCR_ENGINE, description="ocr 엔진 선택: paddle | google"),
|
|
):
|
|
m = classifier_model_store.get_model()
|
|
if m is None:
|
|
raise HTTPException(status_code=503, detail="모델이 없습니다. /classifier/train 먼저 실행하세요.")
|
|
|
|
responses = []
|
|
for file in files:
|
|
try:
|
|
content = await file.read()
|
|
# tmp_path = config.TMP_DIR / file.filename
|
|
today = datetime.now().strftime("%Y%m%d")
|
|
tmp_path = config.TMP_DIR / today / file.filename
|
|
tmp_path.parent.mkdir(parents=True, exist_ok=True)
|
|
with open(tmp_path, "wb") as f:
|
|
f.write(content)
|
|
|
|
# 1. 이미지 변환
|
|
if is_pdf(file.filename):
|
|
img_paths = pdf_to_images(str(tmp_path), dpi=config.OCR_DPI, max_pages=1)
|
|
else:
|
|
# 이미지 리사이즈
|
|
if hasattr(config, 'PDF_MAX_WIDTH'):
|
|
from PIL import Image
|
|
img = Image.open(tmp_path)
|
|
if img.width > config.PDF_MAX_WIDTH:
|
|
scale = config.PDF_MAX_WIDTH / img.width
|
|
new_size = (int(img.width * scale), int(img.height * scale))
|
|
img = img.resize(new_size, Image.LANCZOS)
|
|
ext = tmp_path.suffix.lower()
|
|
if ext in (".jpg", ".jpeg"):
|
|
img.save(str(tmp_path), quality=config.PDF_JPEG_QUALITY)
|
|
else:
|
|
img.save(str(tmp_path))
|
|
img_paths = [tmp_path]
|
|
|
|
if not img_paths:
|
|
raise ValueError("이미지 변환 실패")
|
|
|
|
# 2. OCR
|
|
def on_done(idx: int, result: dict):
|
|
if result.get("words"):
|
|
image_path = result.get("image_path")
|
|
json_path = Path(image_path).with_name(Path(image_path).stem + "_google.json")
|
|
with open(json_path, "w", encoding="utf-8") as f2:
|
|
json.dump(result, f2, ensure_ascii=False, indent=2)
|
|
logger.debug(f"OCR 저장: {json_path}")
|
|
|
|
loop = asyncio.get_event_loop()
|
|
# results = await loop.run_in_executor(None, ocr_batch, img_paths)
|
|
|
|
_, batch_fn, client = get_ocr_functions(ocr_engine)
|
|
results = await loop.run_in_executor(None, batch_fn, img_paths, client)
|
|
# print(results)
|
|
|
|
# 3. 추론
|
|
best_label, best_conf, best_probs = "OTHER", 0.0, {}
|
|
for r in results:
|
|
words = r.get("words", [])
|
|
if len(words) < config.OCR_MIN_WORDS:
|
|
continue
|
|
lbl, conf, probs = classifier_model_store.predict(m, " ".join(words), threshold)
|
|
if conf > best_conf:
|
|
best_label, best_conf, best_probs = lbl, conf, probs
|
|
|
|
if best_conf == 0.0:
|
|
raise ValueError("OCR 결과가 없거나 단어 수 미달입니다.")
|
|
|
|
responses.append(PredictResponse(label=best_label, file_name=file.filename, confidence=best_conf, all_probs=best_probs))
|
|
|
|
except Exception as e:
|
|
logger.warning(f"[{file.filename}] 처리 실패: {e}")
|
|
responses.append(PredictResponse(label="ERROR", confidence=0.0, all_probs={}))
|
|
|
|
return PredictListResponse(results=responses, success=True)
|
|
|