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 = """
OCR 빌드 모니터
"""
return HTMLResponse(content=html)