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 = """ Classifier 학습 모니터

🧠 Classifier 학습 모니터

대기 중
Epoch 0 / 0 0%
샘플 수
-
클래스 수
-
진행률
0%
대기 중...

📋 진행 로그

""" 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)