430 lines
17 KiB
Python
430 lines
17 KiB
Python
import asyncio
|
|
from pathlib import Path
|
|
|
|
from fastapi import APIRouter, File, Form, HTTPException, UploadFile
|
|
from fastapi.responses import HTMLResponse
|
|
from sse_starlette.sse import EventSourceResponse
|
|
from loguru import logger
|
|
import json
|
|
|
|
from utils.ocr import get_ocr_functions
|
|
import config
|
|
from train.schemas import DatasetStatus
|
|
from utils.cache import (
|
|
save_cache_by_file,
|
|
get_cached_paths, cache_status, load_label2id, save_label2id,
|
|
)
|
|
from utils.ocr import ocr_batch
|
|
from utils.pdf import is_pdf, pdf_to_images
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
# ── OCR 진행 상태 전역 객체 ──────────────────────────
|
|
class OcrBuildStatus:
|
|
def __init__(self):
|
|
self.running: bool = False
|
|
self.total: int = 0
|
|
self.done: int = 0
|
|
self.failed: int = 0
|
|
self.message: str = "대기 중"
|
|
self.current: str = "" # 현재 처리 파일명
|
|
|
|
def reset(self, total: int):
|
|
self.running = True
|
|
self.total = total
|
|
self.done = 0
|
|
self.failed = 0
|
|
self.message = f"OCR 시작 ({total}건)"
|
|
self.current = ""
|
|
|
|
@property
|
|
def progress(self) -> float:
|
|
return round(self.done / self.total * 100, 1) if self.total > 0 else 0
|
|
|
|
ocr_status = OcrBuildStatus()
|
|
|
|
|
|
# ── 업로드 ────────────────────────────────────────────
|
|
@router.post("/upload", summary="학습 이미지/PDF 업로드")
|
|
async def upload_images(
|
|
label: str = Form(..., description="양식 종류 (예: TSMC_TypeA)"),
|
|
files: list[UploadFile] = File(..., description="이미지 또는 PDF 파일"),
|
|
):
|
|
save_dir = config.DATA_DIR / label
|
|
save_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
saved = []
|
|
for file in files:
|
|
content = await file.read()
|
|
|
|
if is_pdf(file.filename):
|
|
tmp_pdf = save_dir / file.filename
|
|
with open(tmp_pdf, "wb") as f:
|
|
f.write(content)
|
|
img_paths = pdf_to_images(str(tmp_pdf), dpi=config.OCR_DPI)
|
|
saved.extend([str(p) for p in img_paths])
|
|
tmp_pdf.unlink()
|
|
logger.info(f"PDF 업로드: {file.filename} → {len(img_paths)}페이지 변환")
|
|
else:
|
|
# dest = save_dir / file.filename
|
|
# with open(dest, "wb") as f:
|
|
# f.write(content)
|
|
# saved.append(str(dest))
|
|
# logger.info(f"이미지 업로드: {dest}")
|
|
|
|
dest = save_dir / file.filename
|
|
with open(dest, "wb") as f:
|
|
f.write(content)
|
|
|
|
# 이미지 리사이즈 적용
|
|
if hasattr(config, 'PDF_MAX_WIDTH'):
|
|
from PIL import Image
|
|
img = Image.open(dest)
|
|
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 = dest.suffix.lower()
|
|
if ext in (".jpg", ".jpeg"):
|
|
img.save(str(dest), quality=config.PDF_JPEG_QUALITY)
|
|
else:
|
|
img.save(str(dest))
|
|
logger.info(f"이미지 리사이즈: {dest.name} → {new_size}")
|
|
|
|
saved.append(str(dest))
|
|
logger.info(f"이미지 업로드: {dest}")
|
|
|
|
return {"label": label, "uploaded": len(files), "saved": len(saved), "files": saved}
|
|
|
|
|
|
# ── 데이터셋 빌드 (백그라운드 + SSE) ──────────────────
|
|
async def _run_build():
|
|
"""OCR 빌드를 백그라운드에서 실행하며 ocr_status 갱신"""
|
|
if not config.DATA_DIR.exists():
|
|
ocr_status.running = False
|
|
ocr_status.message = "오류: data/ 폴더가 없습니다."
|
|
return
|
|
|
|
classes = sorted([d.name for d in config.DATA_DIR.iterdir() if d.is_dir()])
|
|
label2id = {c: i for i, c in enumerate(classes)}
|
|
|
|
cached_paths = get_cached_paths()
|
|
|
|
new_paths, new_labels = [], []
|
|
for cls in classes:
|
|
for ext in ["*.png", "*.jpg", "*.jpeg", "*.tiff"]:
|
|
for img_path in (config.DATA_DIR / cls).rglob(ext):
|
|
if str(img_path) not in cached_paths:
|
|
new_paths.append(img_path)
|
|
new_labels.append(label2id[cls])
|
|
|
|
# for cls in classes:
|
|
# for ext in ["*.png", "*.jpg", "*.jpeg", "*.tiff"]:
|
|
# for img_path in (config.DATA_DIR / cls).rglob(ext):
|
|
# s = str(img_path)
|
|
# if s in cached_paths:
|
|
# print(f"[캐시HIT] {s}")
|
|
# else:
|
|
# print(f"[캐시MISS] {s}")
|
|
# print(f" → cached 샘플: {next(iter(cached_paths))}")
|
|
|
|
if not new_paths:
|
|
status = cache_status()
|
|
ocr_status.running = False
|
|
ocr_status.message = "신규 이미지 없음. 캐시가 최신 상태입니다."
|
|
return
|
|
|
|
ocr_status.reset(len(new_paths))
|
|
|
|
import time
|
|
start = time.time()
|
|
|
|
loop = asyncio.get_event_loop()
|
|
total_new = 0
|
|
|
|
# 완료 콜백: ocr_batch가 1건 끝날 때마다 호출
|
|
# def on_done(idx: int, result: dict):
|
|
# ocr_status.current = Path(new_paths[idx]).name
|
|
# ocr_status.message = f"[{ocr_status.done+1}/{ocr_status.total}] {ocr_status.current}"
|
|
# if result.get("error"):
|
|
# ocr_status.failed += 1
|
|
# ocr_status.done += 1
|
|
|
|
def on_done(idx: int, result: dict):
|
|
nonlocal total_new
|
|
# ocr_status.current = Path(new_paths[idx]).name
|
|
ocr_status.current = f"{Path(new_paths[idx]).parent.name}/{Path(new_paths[idx]).name}"
|
|
ocr_status.message = f"[{ocr_status.done + 1}/{ocr_status.total}] {ocr_status.current}"
|
|
if result.get("error"):
|
|
ocr_status.failed += 1
|
|
else:
|
|
# OCR 완료 즉시 저장
|
|
if result.get("words"):
|
|
result["label"] = new_labels[idx]
|
|
save_cache_by_file(result)
|
|
total_new += 1
|
|
ocr_status.done += 1
|
|
|
|
_, batch_fn, client = get_ocr_functions(config.OCR_ENGINE)
|
|
|
|
results = await loop.run_in_executor(
|
|
None, lambda: batch_fn(new_paths, client=client, max_workers=config.OCR_MAX_WORKERS, on_done=on_done)
|
|
)
|
|
|
|
# results = await loop.run_in_executor(
|
|
# None, ocr_batch, new_paths, config.OCR_MAX_WORKERS, on_done
|
|
# )
|
|
|
|
# for ocr, label in zip(results, new_labels):
|
|
# if not ocr.get("words"):
|
|
# continue
|
|
# ocr["label"] = label
|
|
# save_cache_by_file(ocr)
|
|
# total_new += 1
|
|
|
|
elapsed = time.time() - start
|
|
mins, secs = divmod(int(elapsed), 60)
|
|
save_label2id(label2id)
|
|
|
|
st = cache_status()
|
|
ocr_status.running = False
|
|
ocr_status.message = (
|
|
f"완료: 신규 {total_new}건 추가 | "
|
|
f"실패 {ocr_status.failed}건 | "
|
|
f"소요 {mins}분 {secs}초"
|
|
)
|
|
logger.info(ocr_status.message)
|
|
|
|
|
|
@router.post("/build", summary="데이터셋 생성 (OCR 실행)")
|
|
async def build_dataset(background_tasks=None):
|
|
from fastapi import BackgroundTasks
|
|
if ocr_status.running:
|
|
raise HTTPException(status_code=409, detail="이미 OCR이 실행 중입니다.")
|
|
|
|
# 백그라운드 실행 (SSE로 상태 조회 가능)
|
|
asyncio.create_task(_run_build())
|
|
return {"message": "OCR 빌드 시작. /dataset/stream 에서 진행상태 확인"}
|
|
|
|
|
|
# ── SSE 스트리밍 ──────────────────────────────────────
|
|
@router.get("/stream", summary="OCR 진행상태 SSE 스트리밍")
|
|
async def stream_build():
|
|
async def generator():
|
|
prev = ""
|
|
while True:
|
|
data = json.dumps({
|
|
"running": ocr_status.running,
|
|
"total": ocr_status.total,
|
|
"done": ocr_status.done,
|
|
"failed": ocr_status.failed,
|
|
"progress": ocr_status.progress,
|
|
"message": ocr_status.message,
|
|
"current": ocr_status.current,
|
|
}, ensure_ascii=False)
|
|
|
|
if data != prev:
|
|
yield {"event": "status", "data": data}
|
|
prev = data
|
|
|
|
if not ocr_status.running and ocr_status.done > 0:
|
|
yield {"event": "done", "data": json.dumps({"message": ocr_status.message})}
|
|
break
|
|
|
|
await asyncio.sleep(0.5)
|
|
|
|
return EventSourceResponse(generator())
|
|
|
|
|
|
# ── 상태 조회 ─────────────────────────────────────────
|
|
@router.get("/status", summary="데이터셋 현황 조회", response_model=DatasetStatus)
|
|
def dataset_status():
|
|
label2id = load_label2id()
|
|
by_label = cache_status()
|
|
total = sum(by_label.values())
|
|
return DatasetStatus(total_samples=total, by_label=by_label, label2id=label2id)
|
|
|
|
|
|
# ── OCR 모니터 HTML ────────────────────────────────────
|
|
@router.get("/monitor", summary="OCR 빌드 모니터", response_class=HTMLResponse)
|
|
def dataset_monitor():
|
|
html = """
|
|
<!DOCTYPE html>
|
|
<html>
|
|
<head>
|
|
<meta charset="utf-8">
|
|
<title>OCR 빌드 모니터</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: #00796b; 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: 120px; 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: #00796b; }
|
|
.card .value.fail { color: #c62828; }
|
|
|
|
.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: #00796b; border-radius: 9px;
|
|
transition: width 0.4s ease; width: 0%; }
|
|
|
|
.current-file { 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: #e0f2f1; color: #00695c; }
|
|
.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: #00796b; color: #fff; }
|
|
.btn-start:hover { background: #00695c; }
|
|
.btn-start:disabled { background: #80cbc4; cursor: not-allowed; }
|
|
.actions { display: flex; gap: 10px; align-items: center; margin-bottom: 20px; }
|
|
</style>
|
|
</head>
|
|
<body>
|
|
<div class="header"><h2>🔍 OCR 빌드 모니터</h2></div>
|
|
<div class="container">
|
|
|
|
<div class="actions">
|
|
<button class="btn btn-start" id="btnStart" onclick="startBuild()">▶ 빌드 시작</button>
|
|
<span id="badge" class="badge idle">대기 중</span>
|
|
</div>
|
|
|
|
<div class="progress-wrap">
|
|
<div class="progress-label">
|
|
<span id="progressLabel">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="cardTotal">-</div>
|
|
</div>
|
|
<div class="card">
|
|
<div class="label">완료</div>
|
|
<div class="value" id="cardDone">-</div>
|
|
</div>
|
|
<div class="card">
|
|
<div class="label">실패</div>
|
|
<div class="value fail" id="cardFail">-</div>
|
|
</div>
|
|
<div class="card">
|
|
<div class="label">진행률</div>
|
|
<div class="value" id="cardPct">0%</div>
|
|
</div>
|
|
</div>
|
|
|
|
<div class="current-file" id="currentFile">대기 중...</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) {
|
|
document.getElementById("progressFill").style.width = d.progress + "%";
|
|
document.getElementById("progressLabel").textContent = `${d.done} / ${d.total}`;
|
|
document.getElementById("pctLabel").textContent = d.progress + "%";
|
|
document.getElementById("cardTotal").textContent = d.total || "-";
|
|
document.getElementById("cardDone").textContent = d.done || "-";
|
|
document.getElementById("cardFail").textContent = d.failed || "0";
|
|
document.getElementById("cardPct").textContent = d.progress + "%";
|
|
document.getElementById("currentFile").textContent = d.message || "";
|
|
|
|
const badge = document.getElementById("badge");
|
|
if (d.running) {
|
|
badge.textContent = "실행 중"; badge.className = "badge running";
|
|
}
|
|
}
|
|
|
|
function startBuild() {
|
|
fetch("/dataset/build", { 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("/dataset/stream");
|
|
|
|
es.addEventListener("status", e => {
|
|
const d = JSON.parse(e.data);
|
|
updateUI(d);
|
|
});
|
|
|
|
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("/dataset/status")
|
|
.then(r => r.json())
|
|
.then(d => {
|
|
// ocr_status는 별도이므로, 단순히 running 여부만 폴링으로 확인
|
|
});
|
|
</script>
|
|
</body>
|
|
</html>
|
|
"""
|
|
return HTMLResponse(content=html) |