commit ce22ba962dd29bebaaaa1020504a9f24df672d41 Author: jskim Date: Mon Jun 15 09:54:01 2026 +0900 Initial commit diff --git a/classifier_service.py b/classifier_service.py new file mode 100644 index 0000000..6ccd5a4 --- /dev/null +++ b/classifier_service.py @@ -0,0 +1,113 @@ +import argparse +import gc +import os + +import psutil +import uvicorn +from fastapi import FastAPI, Request +from fastapi.middleware.cors import CORSMiddleware +from loguru import logger +from starlette.responses import JSONResponse + +from routers import router +from routers import classifier +from utils.classifier_model_store import get_model + +app = FastAPI(title="classifier_service") +app.state.request_counter = 0 + +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + +app.include_router(router) + +# app.include_router(classifier.router) + +# ── 로거 설정 ───────────────────────────────────── +logger.remove() +logger.add( + "logs/{time:YYYY-MM-DD}.log", + level="DEBUG", + rotation="00:00", + retention="30 days", + format="{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {name}:{function}:{line} - {message}", +) + + +# ── 메모리 유틸 ─────────────────────────────────── +def log_memory_usage() -> float: + process = psutil.Process(os.getpid()) + return process.memory_info().rss / 1024 / 1024 + + +# ── HTTP 미들웨어 ───────────────────────────────── +@app.middleware("http") +async def log_requests(request: Request, call_next): + try: + memory_before = log_memory_usage() + request_counter = getattr(app.state, "request_counter", 0) + 1 + app.state.request_counter = request_counter + + logger.info(f"request_counter: {request_counter}") + logger.info(f"Received: {request.method} {request.url}") + logger.info(f"Headers: {dict(request.headers)}") + + response = await call_next(request) + logger.info(f"Response status: {response.status_code}") + + if request_counter % 100 == 0: + logger.info("gc.collect 실행") + gc.collect() + + memory_after = log_memory_usage() + logger.info( + f"Memory: Before={memory_before:.2f}MB, " + f"After={memory_after:.2f}MB, " + f"Diff={memory_after - memory_before:.2f}MB" + ) + return response + + except Exception as e: + logger.exception(f"요청 처리 중 예외 발생: {e}") + return JSONResponse( + status_code=500, + content={"detail": "Internal Server Error", "error": str(e)}, + ) + + +@app.get("/", tags=["기타"]) +async def root(): + return {"message": "document classifier service"} + + +@app.get("/health", tags=["기타"]) +async def health(): + m = get_model() + return { + "status": "ok", + "model_loaded": m is not None, + "classes": list(m.classes_) if m else [], + } + + +# ── 진입점 ──────────────────────────────────────── +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("-host", default="0.0.0.0") + parser.add_argument("-port", default=10941) + args = parser.parse_args() + + logger.info(f"서버 시작: {args.host}:{args.port}") + uvicorn.run( + app, + host=args.host, + port=int(args.port), + limit_concurrency=1000, + timeout_keep_alive=120, + log_level="debug", + ) \ No newline at end of file diff --git a/layoutlmv3_service.py b/layoutlmv3_service.py new file mode 100644 index 0000000..49a10f8 --- /dev/null +++ b/layoutlmv3_service.py @@ -0,0 +1,113 @@ +import argparse +import gc +import os +import sys + +import psutil +import uvicorn +from fastapi import FastAPI, Request +from fastapi.middleware.cors import CORSMiddleware +from loguru import logger +from starlette.responses import JSONResponse + +from routers import router + +app = FastAPI(title="layoutlmv3_service") +app.state.request_counter = 0 + +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + +app.include_router(router) + +# ── 로거 설정 ───────────────────────────────────── +logger.remove() +logger.add( + "logs/{time:YYYY-MM-DD}.log", + level="DEBUG", + rotation="00:00", + retention="30 days", + format="{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {name}:{function}:{line} - {message}", +) + + +# ── 메모리 유틸 ─────────────────────────────────── +def log_memory_usage() -> float: + process = psutil.Process(os.getpid()) + return process.memory_info().rss / 1024 / 1024 # MB + + +# ── HTTP 미들웨어 ───────────────────────────────── +@app.middleware("http") +async def log_requests(request: Request, call_next): + try: + memory_before = log_memory_usage() + request_counter = getattr(app.state, "request_counter", 0) + 1 + app.state.request_counter = request_counter + + logger.info(f"request_counter: {request_counter}") + logger.info(f"Received: {request.method} {request.url}") + logger.info(f"Headers: {dict(request.headers)}") + + response = await call_next(request) + logger.info(f"Response status: {response.status_code}") + + if request_counter % 100 == 0: + logger.info("gc.collect 실행") + gc.collect() + + memory_after = log_memory_usage() + logger.info( + f"Memory: Before={memory_before:.2f}MB, " + f"After={memory_after:.2f}MB, " + f"Diff={memory_after - memory_before:.2f}MB" + ) + return response + + except Exception as e: + logger.exception(f"요청 처리 중 예외 발생: {e}") + return JSONResponse( + status_code=500, + content={"detail": "Internal Server Error", "error": str(e)}, + ) + + +@app.get("/", tags=["기타"]) +async def root(): + return {"message": "layoutlmv3 invoice classifier"} + + +@app.get("/health", tags=["기타"]) +async def health(): + import config + from utils.cache import load_cache + return { + "status": "ok", + "device": str(config.DEVICE), + "model_ready": (config.MODEL_SAVE_PATH / "config.json").exists(), + "dataset_size": len(load_cache()), + } + + +# ── 진입점 ──────────────────────────────────────── +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("-host", default="0.0.0.0") + parser.add_argument("-port", default=10941) + parser.add_argument("-pool_size", default="10") + args = parser.parse_args() + + logger.info(f"서버 시작: {args.host}:{args.port}") + uvicorn.run( + app, + host=args.host, + port=int(args.port), + limit_concurrency=1000, + timeout_keep_alive=120, + log_level="debug", + ) \ No newline at end of file diff --git a/ocrlogin-431508-8db9e68e7a47.json b/ocrlogin-431508-8db9e68e7a47.json new file mode 100644 index 0000000..1a862e3 --- /dev/null +++ b/ocrlogin-431508-8db9e68e7a47.json @@ -0,0 +1,13 @@ +{ + "type": "service_account", + "project_id": "ocrlogin-431508", + "private_key_id": "8db9e68e7a47fc0116656b68cc99b51eb070a6b4", + "private_key": "-----BEGIN PRIVATE KEY-----\nMIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQCpFAIMT0E2S0eH\nwXxApGev8kGqaCRXN8DuG8pSM/rRjs/TvYTu8FDlAWkgiMgymxYrpzhWvqAFWoD8\n837L+ntqF2Ll8nkZOmlluSg+1TgKvGWS00OA6pgpVYkMJKcD3S7mJHqlO4LIKRDy\nHhcc6fmjv67KxUROllJoMAZCO6qsyGVB0xJmV+rzXZavXNJ00Gmi0kYG9gSvCW38\nur98gmaa0cZLIbDi6Tml3r5Zfb/oNzX1ktEjpEIIbwoamFujjr1tWuXtQH7Sf7OJ\ne3N+F83/EurkSWrORoF3SziltXQ1Sutuw3Ns5R+Ky1Mg/gO1Wyk3Bw0F4Keefobt\nWxSfD+H5AgMBAAECggEAPtZJRpLj7Q5AONNvXsTbJjhWMENBEksNwFCCuldIJb66\nPXrHX1ff8KQ8ElPTd39M149vsElrRmIS4y+JlbxzRoQHhOc/G2GqjxwnuWZbzB2l\ncFJk2ZIWV/JKm0E58wUua2juTd9WpRYiDqGhPGU2mqVgDEsRLlXOrZr/kHkFXu4F\nSSGuDSwSRSEN02sgdhhK9/DcfuyHS8YwP9Cq+fb2W+Jdd6kMaA3eQjVZo+UAG9zy\nxQxq+6qPXPstocIjSk5Aa18Z7eOQr+bwd4oIjZpthlJa9V5lNuWYrYmAYk3lBh7r\nLJ+KCsdoS3lIN1gom9EK4p6oLLSAVuXIpxHBK3znYwKBgQC22s/DYBJH+/uMnyvV\nTdUzAcDnpzhJ2Bmup+cyLwLIhQDZPar72Bh/wU0uptsW1/QUFNqHYiZZLOeAeUC8\nVGdb/rUl9r+eMFr+8KqFGRpB4yLqRZniRtGxmGXd/wbx3MYA4GV7AdNgEoavQkip\nl0iKVn7K587wJjow1HQLBJnoMwKBgQDstmjwPfNIJzUw+TtJA9jemdagOuQTQvgi\nchq3PQd1JfU8eFBOZS1EpvcVkA86oB05hBwJvuk6J/jgeWjuwkve7ppCOFt24xzr\nJEEVPspBJ+T+tGAu4Fs/5ACNC8I1LV5oY0xCJXa7gNljeHU6SrRD7tN+sO+4SRFs\nN99uPT9RIwKBgGHQAI1hacYJ29CoIHl0rhQf3wHL6IdPysUr2bd1gEalJwQOQdWA\nDfLhAxludgntMQpA8Xi0HxFavOdzdRaJC9UhFeOd73h+I172fDDAcdRG3Rl2a8+n\n1GnsvKkYz603TM+ROZeoLVrZ7iP4EAhv/YTKqf5+K6s4t64BJ6XxKycTAoGAWNAT\nzVehCMhVJ7vLJ5j+7H4RzepqmmN9EAd5yJhoTObh/T8y+kbx1hlDCV8Up61dabAM\niQeNIBnRQf+rhDF4H/ur+v6EKrYJqpveo2b8obejLoFkuRHKis0z+7eWtTcBfe8L\ntKGzy6QLbEvMyAMxYW+hAJ7IQn9/vvezp/vo3rsCgYEAsjIb+H82uY2wwNJlf4ar\nuASzvE7debIvuNX/ThaX0lyBX0a1dXGxtLJp7U/E/eky/IXXCE/hm4qrOV8UZxwv\nYQak8whLeu3xgiu6DIlYsEwRCtV3sCdh1pSfIWjhDKGdaHn1BKCBd+Cfxqa6bon5\nphAyYfmLSEYotI4JPlbGCwU=\n-----END PRIVATE KEY-----\n", + "client_email": "vertaxai-ocr@ocrlogin-431508.iam.gserviceaccount.com", + "client_id": "111900626593542552794", + "auth_uri": "https://accounts.google.com/o/oauth2/auth", + "token_uri": "https://oauth2.googleapis.com/token", + "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs", + "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/vertaxai-ocr%40ocrlogin-431508.iam.gserviceaccount.com", + "universe_domain": "googleapis.com" +} diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..70052bd --- /dev/null +++ b/requirements.txt @@ -0,0 +1,87 @@ +annotated-doc==0.0.4 +annotated-types==0.7.0 +anyio==4.12.1 +cachetools==5.5.2 +certifi==2026.2.25 +charset-normalizer==3.4.4 +click==8.3.1 +colorama==0.4.6 +dnspython==2.8.0 +email-validator==2.3.0 +et_xmlfile==2.0.0 +fastapi==0.135.1 +fastapi-cli==0.0.24 +filelock==3.20.0 +fsspec==2025.12.0 +google-api-core==2.30.0 +google-auth==2.29.0 +google-cloud-vision==3.7.2 +googleapis-common-protos==1.72.0 +grpcio==1.78.0 +grpcio-status==1.62.3 +h11==0.16.0 +httpcore==1.0.9 +httptools==0.7.1 +httpx==0.28.1 +huggingface_hub==0.36.2 +idna==3.11 +Jinja2==3.1.6 +joblib==1.5.3 +loguru==0.7.2 +markdown-it-py==4.0.0 +MarkupSafe==3.0.2 +mdurl==0.1.2 +mpmath==1.3.0 +networkx==3.6.1 +numpy==1.26.4 +openpyxl==3.1.5 +orjson==3.11.7 +packaging @ file:///C:/miniconda3/conda-bld/packaging_1761049137378/work +pandas==3.0.1 +pillow==10.3.0 +proto-plus==1.27.1 +protobuf==4.25.8 +psutil==5.9.8 +pyasn1==0.6.2 +pyasn1_modules==0.4.2 +pydantic==2.7.1 +pydantic_core==2.18.2 +Pygments==2.19.2 +PyMuPDF==1.24.3 +PyMuPDFb==1.24.3 +python-dateutil==2.9.0.post0 +python-dotenv==1.2.2 +python-multipart==0.0.9 +PyYAML==6.0.3 +regex==2026.2.28 +requests==2.32.5 +rich==14.3.3 +rich-toolkit==0.19.7 +rsa==4.9.1 +safetensors==0.7.0 +scikit-learn==1.8.0 +scipy==1.17.1 +setuptools==80.10.2 +shellingham==1.5.4 +six==1.17.0 +sse-starlette==3.3.2 +starlette==0.52.1 +sympy==1.13.1 +threadpoolctl==3.6.0 +timm==0.9.16 +tokenizers==0.19.1 +torch==2.6.0+cu124 +torchvision==0.21.0+cu124 +tqdm==4.66.4 +transformers==4.41.2 +typer==0.24.1 +typing-inspection==0.4.2 +typing_extensions==4.15.0 +tzdata==2025.3 +ujson==5.11.0 +urllib3==2.6.3 +uvicorn==0.30.1 +watchfiles==1.1.1 +websockets==16.0 +wheel==0.46.3 +win32_setctime==1.2.0 diff --git a/routers/__init__.py b/routers/__init__.py new file mode 100644 index 0000000..d66e319 --- /dev/null +++ b/routers/__init__.py @@ -0,0 +1,10 @@ +from fastapi import APIRouter +from routers import dataset, train, evaluate, predict, visualize, classifier + +router = APIRouter() +router.include_router(dataset.router, prefix="/dataset", tags=["1. 데이터셋"]) +router.include_router(train.router, prefix="/train", tags=["2. 학습"]) +router.include_router(evaluate.router, prefix="/evaluate", tags=["3. 평가"]) +router.include_router(predict.router, prefix="/predict", tags=["4. 추론"]) +router.include_router(visualize.router, prefix="/visualize", tags=["5. 시각화"]) +router.include_router(classifier.router, prefix="/classifier", tags=["6. 문서분류"]) # 추가 \ No newline at end of file diff --git a/routers/classifier.py b/routers/classifier.py new file mode 100644 index 0000000..8090054 --- /dev/null +++ b/routers/classifier.py @@ -0,0 +1,520 @@ +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) + diff --git a/routers/dataset.py b/routers/dataset.py new file mode 100644 index 0000000..dd7161f --- /dev/null +++ b/routers/dataset.py @@ -0,0 +1,430 @@ +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 빌드 모니터 + + + +

    🔍 OCR 빌드 모니터

    +
    + +
    + + 대기 중 +
    + +
    +
    + 0 / 0 + 0% +
    +
    +
    +
    +
    + +
    +
    +
    전체
    +
    -
    +
    +
    +
    완료
    +
    -
    +
    +
    +
    실패
    +
    -
    +
    +
    +
    진행률
    +
    0%
    +
    +
    + +
    대기 중...
    + +
    +

    📋 진행 로그

    +
      +
      +
      + + + + + """ + return HTMLResponse(content=html) \ No newline at end of file diff --git a/routers/evaluate.py b/routers/evaluate.py new file mode 100644 index 0000000..392edea --- /dev/null +++ b/routers/evaluate.py @@ -0,0 +1,17 @@ +from fastapi import APIRouter, HTTPException +from loguru import logger + +import config +from train.trainer import run_evaluate, _get_latest_model_path + +router = APIRouter() + + +@router.post("", summary="모델 평가") +def evaluate(test_size: float = 0.2): + if not _get_latest_model_path(): + raise HTTPException( + status_code=400, detail="저장된 모델 없음. 학습 먼저 실행하세요." + ) + logger.info(f"평가 요청 | test_size={test_size}") + return run_evaluate(test_size) \ No newline at end of file diff --git a/routers/predict.py b/routers/predict.py new file mode 100644 index 0000000..f4c06de --- /dev/null +++ b/routers/predict.py @@ -0,0 +1,272 @@ +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) \ No newline at end of file diff --git a/routers/train.py b/routers/train.py new file mode 100644 index 0000000..3e7fc9d --- /dev/null +++ b/routers/train.py @@ -0,0 +1,305 @@ +from fastapi import APIRouter, BackgroundTasks, HTTPException +from fastapi.responses import HTMLResponse +from sse_starlette.sse import EventSourceResponse +from loguru import logger +import asyncio +import json + +import config +import train.trainer as trainer +from train.schemas import TrainRequest, TrainStatus +from utils.cache import load_cache + +router = APIRouter() + + +@router.post("/start", summary="학습 실행") +async def start_train( + background_tasks: BackgroundTasks, + req: TrainRequest = None, +): + if req is None: + req = TrainRequest() + + if trainer.train_status.running: + raise HTTPException(status_code=409, detail="이미 학습이 실행 중입니다.") + + samples = load_cache() + if not samples: + raise HTTPException( + status_code=400, + detail="학습 데이터가 없습니다. /dataset/build 먼저 실행하세요.", + ) + + logger.info(f"학습 요청: {req.dict()}") + background_tasks.add_task(trainer.run_train, req) + return {"message": "학습 시작", "config": req.dict()} + + + +@router.post("/stop", summary="학습 중단") +def stop_train(): + if not trainer.train_status.running: + raise HTTPException(status_code=409, detail="학습이 실행 중이 아닙니다.") + trainer.stop_requested = True + trainer.train_status.message = "중단 요청됨..." + return {"message": "중단 요청 완료"} + + +@router.get("/status", summary="학습 진행 상태 조회", response_model=TrainStatus) +def get_train_status(): + return trainer.train_status + + +@router.get("/stream", summary="학습 진행 상태 SSE 스트리밍") +async def stream_status(): + """ + SSE(Server-Sent Events)로 학습 상태 실시간 전송 + 학습 완료 시 자동 종료 + """ + async def event_generator(): + prev_msg = "" + while True: + status = trainer.train_status + data = json.dumps({ + "running": status.running, + "epoch": status.epoch, + "total": status.total, + "loss": status.loss, + "val_acc": status.val_acc, + "message": status.message, + # "progress": round(status.epoch / status.total * 100, 1) if status.total > 0 else 0, + "progress": round(status.batch / status.total_batches * 100, 1) if status.total_batches > 0 else 0, # ← 배치 기준으로 변경 + "batch_message": status.message, # 배치 단위 메시지 + }, ensure_ascii=False) + + if data != prev_msg: + yield {"event": "status", "data": data} + prev_msg = data + + if not status.running and status.epoch > 0: + yield {"event": "done", "data": json.dumps({"message": status.message})} + break + + await asyncio.sleep(0.5) # 0.5초마다 polling → 배치 진행 반영 + + return EventSourceResponse(event_generator()) + + +@router.get("/monitor", summary="학습 모니터 HTML", response_class=HTMLResponse) +def train_monitor(): + """브라우저에서 학습 진행 상태 실시간 확인""" + html = """ + + + + + 학습 모니터 + + + +
      +

      🤖 LayoutLMv3 학습 모니터

      +
      +
      +
      + + + 대기 중 +
      + + +
      +
      + Epoch 0 / 0 + 0% +
      +
      +
      +
      +
      + + +
      +
      +
      현재 Loss
      +
      -
      +
      +
      +
      Val Accuracy
      +
      -
      +
      +
      +
      진행률
      +
      0%
      +
      +
      + + +
      대기 중...
      + + +
      +

      📋 진행 로그

      +
        +
        +
        + + + + + """ + return HTMLResponse(content=html) \ No newline at end of file diff --git a/routers/visualize.py b/routers/visualize.py new file mode 100644 index 0000000..dd4fca4 --- /dev/null +++ b/routers/visualize.py @@ -0,0 +1,289 @@ +import json +from pathlib import Path + +from fastapi import APIRouter, HTTPException +from fastapi.responses import HTMLResponse +from PIL import Image, ImageDraw, ImageFont +import base64 +import io + +import config +from utils.cache import _cache_path + +router = APIRouter() + + +def _bbox_denormalize(box: list, w: int, h: int) -> list: + """0~1000 정규화 좌표 → 픽셀 좌표 복원""" + return [ + int(box[0] * w / 1000), + int(box[1] * h / 1000), + int(box[2] * w / 1000), + int(box[3] * h / 1000), + ] + + +def _draw_ocr(image_path: str, words: list, boxes: list) -> tuple[str, int, int]: + """ + 원본 이미지를 base64로 변환만 함 + 바운딩박스/텍스트는 canvas에서 처리 + """ + image = Image.open(image_path).convert("RGB") + w, h = image.size + buf = io.BytesIO() + image.save(buf, format="PNG") + return base64.b64encode(buf.getvalue()).decode("utf-8"), w, h + + +@router.get( + "/{label}/{filename:path}", + summary="OCR 결과 시각화", + description="양식명과 파일명으로 OCR 바운딩박스를 이미지 위에 표시합니다.", + response_class=HTMLResponse, +) +def visualize(label: str, filename: str): + """ + 사용법: + GET /visualize/TSMC_TypeA/invoice_001 + → 이미지 + 바운딩박스 + 텍스트를 HTML로 반환 + """ + # json 캐시 로드 + json_path = config.OCR_CACHE_DIR / label / f"{filename}.json" + if not json_path.exists(): + raise HTTPException( + status_code=404, + detail=f"캐시 없음: {label}/{filename}.json | /dataset/build 먼저 실행하세요." + ) + + with open(json_path, encoding="utf-8") as f: + ocr = json.load(f) + + image_path = ocr.get("image_path") + if not image_path or not Path(image_path).exists(): + raise HTTPException(status_code=404, detail=f"이미지 파일 없음: {image_path}") + + words = ocr.get("words", []) + boxes = ocr.get("boxes", []) + + # 시각화 이미지 생성 + img_b64, img_w, img_h = _draw_ocr(image_path, words, boxes) + + # boxes를 픽셀 좌표로 변환 (JS에서 사용) + # pixel_boxes = [_bbox_denormalize(b, img_w, img_h) for b in boxes] + + box_type = ocr.get("box_type", "normalized") + + if box_type == "pixel": + pixel_boxes = boxes # 픽셀 좌표 그대로 사용 + else: + pixel_boxes = [_bbox_denormalize(b, img_w, img_h) for b in boxes] # 역변환 + + boxes_json = json.dumps(pixel_boxes) + words_json = json.dumps(words) + + # OCR 텍스트 목록 + ocr_rows = "".join([ + f'' + f'{i+1}{w}{b}' + for i, (w, b) in enumerate(zip(words, boxes)) + ]) + + html = f""" + + + + + OCR 시각화 - {label}/{filename} + + + + +
        +

        OCR 시각화 | {label} / {filename}

        +
        총 단어 수: {len(words)}개
        +
        + + +
        +
        + + 바운딩박스 +
        +
        + + OCR 텍스트 +
        +
        + + +
        +
        +
        + + +
        +
        +
        +

        OCR 텍스트 목록

        + + + + + {ocr_rows} +
        #텍스트박스 (정규화)
        +
        +
        + + + + + """ + return HTMLResponse(content=html) + + +@router.get( + "/{label}", + summary="양식별 파일 목록 조회", + description="캐시된 파일 목록을 반환합니다.", +) +def list_files(label: str): + cls_dir = config.OCR_CACHE_DIR / label + if not cls_dir.exists(): + raise HTTPException(status_code=404, detail=f"양식 없음: {label}") + + # 서브폴더 포함 전체 json 탐색 + files = [ + str(p.relative_to(cls_dir).with_suffix("")) # 서브폴더/파일명 형태로 반환 + for p in sorted(cls_dir.rglob("*.json")) + ] + return { + "label": label, + "count": len(files), + "files": files, + } \ No newline at end of file diff --git a/train/classifier_dataset.py b/train/classifier_dataset.py new file mode 100644 index 0000000..af38d28 --- /dev/null +++ b/train/classifier_dataset.py @@ -0,0 +1,183 @@ +""" +files/json/{LABEL}/{파일명}/page_xxx.json 구조를 읽어 +페이지 단위로 샘플을 생성하여 캐시 저장 +문서 1개 = 샘플 N개 (페이지 수만큼) +""" +import json +import os +import pickle +from pathlib import Path + +from loguru import logger + +import config + +# JSON_DIR = os.getenv("CLASSIFIER_JSON_DIR", "./files/json") +# CACHE_PATH = os.getenv("CLASSIFIER_CACHE_PATH", "./files/classifier_dataset.pkl") +# MIN_WORDS = int(os.getenv("CLASSIFIER_MIN_WORDS", "20")) # 페이지당 최소 단어 수 + +JSON_DIR = config.CLASSIFIER_JSON_DIR +CACHE_PATH = config.CLASSIFIER_CACHE_PATH +MIN_WORDS = config.CLASSIFIER_MIN_WORDS + +_status: dict = {"status": "idle", "message": "", "total": 0} + + +def get_status() -> dict: + return _status + + +# def build() -> None: +# global _status +# try: +# _status = {"status": "building", "message": "JSON 파일 스캔 중...", "total": 0} +# json_path = Path(JSON_DIR) +# +# if not json_path.exists(): +# raise FileNotFoundError(f"JSON 디렉토리 없음: {JSON_DIR}") +# +# data = [] +# skip_count = 0 +# +# for label_dir in sorted(json_path.iterdir()): +# if not label_dir.is_dir(): +# continue +# label = label_dir.name +# label_count = 0 +# +# for doc_dir in sorted(label_dir.iterdir()): +# if not doc_dir.is_dir(): +# continue +# +# for page_file in sorted(doc_dir.glob("page_*.json")): +# try: +# with open(page_file, encoding="utf-8") as f: +# d = json.load(f) +# words = d.get("words", []) +# except Exception as e: +# logger.warning(f"파일 읽기 실패 {page_file}: {e}") +# continue +# +# # 단어 수 미달 페이지 제외 +# if len(words) < MIN_WORDS: +# print(f"skip file : {len(words)} : {page_file}") +# skip_count += 1 +# continue +# +# data.append({ +# "text": " ".join(words), +# "label": label, +# "src": str(page_file), # 디버깅용 +# }) +# label_count += 1 +# +# _status["message"] = f"{label}: {label_count}페이지 샘플 추가 (누적 {len(data)}개)" +# logger.info(_status["message"]) +# +# if not data: +# raise ValueError("수집된 데이터가 없습니다.") +# +# os.makedirs(os.path.dirname(CACHE_PATH), exist_ok=True) +# with open(CACHE_PATH, "wb") as f: +# pickle.dump(data, f) +# +# _status = { +# "status": "done", +# "message": f"데이터셋 생성 완료 (단어 {MIN_WORDS}개 미만 {skip_count}페이지 제외)", +# "total": len(data), +# } +# logger.info(f"데이터셋 저장: {CACHE_PATH} ({len(data)}개, 제외 {skip_count}개)") +# +# except Exception as e: +# _status = {"status": "error", "message": str(e), "total": 0} +# logger.exception(f"데이터셋 생성 실패: {e}") + + +def _process_page_file(page_file: Path, label: str, data: list, skip_count: int) -> tuple[int, int]: + """단일 page_*.json 파일 처리. (label_count 증가분, skip_count) 반환""" + try: + with open(page_file, encoding="utf-8") as f: + d = json.load(f) + words = d.get("words", []) + except Exception as e: + logger.warning(f"파일 읽기 실패 {page_file}: {e}") + return 0, skip_count + + if len(words) < MIN_WORDS: + print(f"skip file : {len(words)} : {page_file}") + return 0, skip_count + 1 + + data.append({"text": " ".join(words), "label": label, "src": str(page_file)}) + return 1, skip_count + + +def _process_label_dir(label_dir: Path, data: list, skip_count: int) -> tuple[int, int]: + """label 폴더 처리. label_count, skip_count 반환""" + label = label_dir.name + label_count = 0 + + # label 폴더 직속 page_*.json + for page_file in sorted(label_dir.glob("*.json")): + added, skip_count = _process_page_file(page_file, label, data, skip_count) + label_count += added + + # 하위 doc_dir 폴더 + for doc_dir in sorted(label_dir.iterdir()): + if not doc_dir.is_dir(): + continue + for page_file in sorted(doc_dir.glob("*.json")): + added, skip_count = _process_page_file(page_file, label, data, skip_count) + label_count += added + + return label_count, skip_count + + +def build() -> None: + global _status + try: + _status = {"status": "building", "message": "JSON 파일 스캔 중...", "total": 0} + json_path = Path(JSON_DIR) + + if not json_path.exists(): + raise FileNotFoundError(f"JSON 디렉토리 없음: {JSON_DIR}") + + data = [] + skip_count = 0 + + for label_dir in sorted(json_path.iterdir()): + if not label_dir.is_dir(): + continue + + label_count, skip_count = _process_label_dir(label_dir, data, skip_count) + + _status["message"] = f"{label_dir.name}: {label_count}페이지 샘플 추가 (누적 {len(data)}개)" + logger.info(_status["message"]) + + if not data: + raise ValueError("수집된 데이터가 없습니다.") + + os.makedirs(os.path.dirname(CACHE_PATH), exist_ok=True) + with open(CACHE_PATH, "wb") as f: + pickle.dump(data, f) + + _status = { + "status": "done", + "message": f"데이터셋 생성 완료 (단어 {MIN_WORDS}개 미만 {skip_count}페이지 제외)", + "total": len(data), + } + logger.info(f"데이터셋 저장: {CACHE_PATH} ({len(data)}개, 제외 {skip_count}개)") + + except Exception as e: + _status = {"status": "error", "message": str(e), "total": 0} + logger.exception(f"데이터셋 생성 실패: {e}") + + +def load() -> list[dict]: + if not os.path.exists(CACHE_PATH): + return [] + try: + with open(CACHE_PATH, "rb") as f: + return pickle.load(f) + except Exception as e: + logger.error(f"데이터셋 로드 실패: {e}") + return [] \ No newline at end of file diff --git a/train/classifier_trainer.py b/train/classifier_trainer.py new file mode 100644 index 0000000..1dc8e92 --- /dev/null +++ b/train/classifier_trainer.py @@ -0,0 +1,215 @@ +""" +classifier_dataset.load() 로 데이터 로드 후 TF-IDF + SVM 학습 +""" +import random + +from collections import Counter + +from loguru import logger +from sklearn.calibration import CalibratedClassifierCV +from sklearn.feature_extraction.text import TfidfVectorizer +from sklearn.metrics import classification_report +from sklearn.model_selection import train_test_split +from sklearn.pipeline import Pipeline +from sklearn.svm import LinearSVC + +from config import SERVICE +from train import classifier_dataset +from utils import classifier_model_store + +_status: dict = {"status": "idle", "message": ""} + + +def get_status() -> dict: + return _status + + +def augment_text(text: str, n: int = 25) -> list: + """텍스트 증강 - 단어 순서 변경, 일부 단어 제거""" + words = text.split() + augmented = [] + + for _ in range(n): + aug = words.copy() + + # 1. 랜덤 단어 10% 제거 + drop_count = max(1, int(len(aug) * 0.1)) + for _ in range(drop_count): + if aug: + aug.pop(random.randint(0, len(aug) - 1)) + + # 2. 랜덤 단어 일부 순서 섞기 + if len(aug) > 5: + idx = random.randint(0, len(aug) - 3) + aug[idx], aug[idx + 1] = aug[idx + 1], aug[idx] + + augmented.append(" ".join(aug)) + + return augmented + + +def run_train() -> None: + global _status + try: + _status = {"status": "training", "message": "데이터셋 로딩 중..."} + + data = classifier_dataset.load() + if not data: + raise ValueError("데이터셋이 없습니다. /classifier/dataset/build 먼저 실행하세요.") + + texts = [d["text"] for d in data] + labels = [d["label"] for d in data] + + # ── 특정 클래스 데이터 증강 ────────────────────── + if SERVICE == "amko": + AUGMENT_LABELS = ["AMKOR TECHNOLOGY TAIWAN_B"] # 증강할 클래스 + AUGMENT_TARGET = 30 # 목표 샘플 수 + + augmented_texts, augmented_labels = [], [] + label_groups = {} + for t, l in zip(texts, labels): + label_groups.setdefault(l, []).append(t) + + for label in AUGMENT_LABELS: + samples = label_groups.get(label, []) + current_count = len(samples) + + if current_count < AUGMENT_TARGET: + need = AUGMENT_TARGET - current_count + logger.info(f"{label}: {current_count}개 → {AUGMENT_TARGET}개 증강 ({need}개 추가)") + + for i in range(need): + base_text = samples[i % current_count] # 기존 샘플 순환 + aug = augment_text(base_text, n=1)[0] + augmented_texts.append(aug) + augmented_labels.append(label) + + texts = texts + augmented_texts + labels = labels + augmented_labels + logger.info(f"증강 후 클래스별 샘플: {Counter(labels)}") + # ──────────────────────────────────────────────── + + # # ── 클래스 언더샘플링 ───────────────────────── + # MAX_SAMPLES_PER_CLASS = 20 # B타입 8개의 2~3배 수준으로 제한 + # + # balanced_data = [] + # class_counts = Counter(labels) + # + # for label in set(labels): + # class_texts = [t for t, l in zip(texts, labels) if l == label] + # # MAX_SAMPLES_PER_CLASS 초과 시 랜덤 샘플링 + # if len(class_texts) > MAX_SAMPLES_PER_CLASS: + # class_texts = random.sample(class_texts, MAX_SAMPLES_PER_CLASS) + # balanced_data.extend([(t, label) for t in class_texts]) + # + # texts, labels = zip(*balanced_data) + # texts, labels = list(texts), list(labels) + # logger.info(f"언더샘플링 후: {Counter(labels)}") + + # # ── 핵심 키워드 가중치 부여 (특정 클래스만) ────────────────────── + # KEYWORD_BOOST_BY_LABEL = { + # "AMKOR TECHNOLOGY TAIWAN_A": { # A타입에만 + # "Process Fee": 5, + # "Total Process Fee": 5, + # "Consigned Value": 5, + # "Total Process Value": 5, + # }, + # "AMKOR TECHNOLOGY TAIWAN_B": { # B타입에만 + # "PACKING LIST SUMMARY": 5, + # "FOC Y": 3, + # }, + # } + # + # def boost_keywords(text: str, label: str) -> str: + # boost_map = KEYWORD_BOOST_BY_LABEL.get(label, {}) # 해당 label 없으면 빈 dict + # boosted = text + # for keyword, repeat in boost_map.items(): + # if keyword.lower() in text.lower(): + # boosted += f" {keyword}" * repeat + # return boosted + # + # texts = [boost_keywords(t, l) for t, l in zip(texts, labels)] + # logger.info("핵심 키워드 가중치 부여 완료") + # # ──────────────────────────────────────────────── + + # # ── 샘플 부족 클래스 필터링 ────────────────────── + # counts = Counter(labels) + # MIN_SAMPLES = 5 + # before = len(texts) + # filtered = [(t, l) for t, l in zip(texts, labels) if counts[l] >= MIN_SAMPLES] + # if not filtered: + # raise ValueError("모든 클래스가 샘플 부족으로 제외되었습니다.") + # texts, labels = map(list, zip(*filtered)) + # excluded = {k: v for k, v in counts.items() if v < MIN_SAMPLES} + # if excluded: + # logger.warning(f"샘플 부족 제외: {excluded}") + # logger.info(f"필터링: {before}개 → {len(texts)}개") + # # ──────────────────────────────────────────────── + # + # counts = Counter(labels) + # min_count = min(counts.values()) + # cv = 5 + # stratify = labels if min_count >= cv else None + + MIN_SAMPLES = 5 + + # ── 필터링 ─────────────────────────────────────── + counts = Counter(labels) + before = len(texts) + filtered = [(t, l) for t, l in zip(texts, labels) if counts[l] >= MIN_SAMPLES] + if not filtered: + raise ValueError("모든 클래스가 샘플 부족으로 제외되었습니다.") + texts, labels = map(list, zip(*filtered)) + excluded = {k: v for k, v in counts.items() if v < MIN_SAMPLES} + if excluded: + logger.warning(f"샘플 부족 제외: {excluded}") + logger.info(f"필터링: {before}개 → {len(texts)}개") + + # ── MIN_SAMPLES 기준으로 cv/stratify 자동 결정 ── + counts = Counter(labels) + min_count = min(counts.values()) + cv = max(2, min(5, int(MIN_SAMPLES * 0.8))) # MIN_SAMPLES의 80% (train 비율) + stratify = labels if min_count >= cv else None + + logger.info(f"MIN_SAMPLES={MIN_SAMPLES} → cv={cv}, stratify={'적용' if stratify else '비적용'}") + + _status["message"] = f"{len(texts)}개 샘플 / {len(counts)}개 클래스 학습 중..." + _status["samples"] = len(texts) + _status["classes"] = len(counts) + logger.info(_status["message"]) + + X_train, X_test, y_train, y_test = train_test_split( + texts, labels, + test_size=0.2, + stratify=stratify, + random_state=42, + ) + + # train_counts = Counter(y_train) + # min_train_count = min(train_counts.values()) + # cv = max(2, min(5, min_train_count)) + + pipeline = Pipeline([ + ("tfidf", TfidfVectorizer(max_features=50000, ngram_range=(1, 2))), + ("clf", CalibratedClassifierCV(LinearSVC(max_iter=2000, class_weight="balanced"), cv=cv)), + ]) + pipeline.fit(X_train, y_train) + + report = classification_report(y_test, pipeline.predict(X_test), zero_division=0) + logger.info(f"\n{report}") + + classifier_model_store.set_model(pipeline) + + _status = { + "status": "done", + "message": f"학습 완료 | 샘플={len(texts)}, 클래스={len(counts)}", + "samples": len(texts), + "classes": len(counts), + "report": report, # 필요 시 별도 API로 노출 + } + + logger.info("classifier 학습 완료") + + except Exception as e: + _status = {"status": "error", "message": str(e)} + logger.exception(f"classifier 학습 실패: {e}") \ No newline at end of file diff --git a/train/schemas.py b/train/schemas.py new file mode 100644 index 0000000..8e87fc2 --- /dev/null +++ b/train/schemas.py @@ -0,0 +1,37 @@ +from pydantic import BaseModel, Field + + +class TrainRequest(BaseModel): + epochs: int = Field(10, description="학습 에폭 수") + lr: float = Field(1e-5, description="학습률") + batch_size: int = Field(4, description="배치 사이즈") + ewc_lambda: float = Field(1000., description="EWC 패널티 강도 (클수록 기존 지식 보호)") + strategy: str = Field("full",description="auto | full | ewc | replay") + + +class TrainStatus(BaseModel): + running: bool + epoch: int + total: int + loss: float + val_acc: float = 0.0 + message: str + batch: int = 0 # 현재 배치 + total_batches: int = 0 # 전체 배치 수 + +class PredictResult(BaseModel): + filename: str + label: str | None = None + confidence: float | None = None + all_probs: dict | None = None + key_tokens: list[str] | None = None + ocr_word_count: int | None = None + total_pages: int | None = None + per_page: list | None = None + error: str | None = None + + +class DatasetStatus(BaseModel): + total_samples: int + by_label: dict + label2id: dict \ No newline at end of file diff --git a/train/trainer.py b/train/trainer.py new file mode 100644 index 0000000..b81ff0f --- /dev/null +++ b/train/trainer.py @@ -0,0 +1,348 @@ +import random +import traceback +from pathlib import Path + +import torch +import torch.nn as nn +from PIL import Image +from sklearn.model_selection import train_test_split +from sklearn.metrics import classification_report +from torch.utils.data import Dataset, DataLoader +from torch.optim import AdamW +from transformers import ( + LayoutLMv3Processor, + LayoutLMv3ForSequenceClassification, + get_scheduler, +) +from tqdm import tqdm +from loguru import logger + +import config +from train.schemas import TrainRequest, TrainStatus +from utils.cache import ( + load_cache, save_label2id, + load_replay_buffer, update_replay_buffer, load_label2id +) + +# 학습 상태 전역 +train_status = TrainStatus( + running=False, epoch=0, total=0, loss=0.0, message="대기 중" +) +stop_requested = False + + +# ── Dataset ─────────────────────────────────────── +class InvoiceDataset(Dataset): + def __init__(self, samples: list, processor): + self.samples = samples + self.processor = processor + + def __len__(self): + return len(self.samples) + + def __getitem__(self, idx): + s = self.samples[idx] + image = Image.open(s["image_path"]).convert("RGB") + encoding = self.processor( + image, text=s["words"], boxes=s["boxes"], + return_tensors="pt", truncation=True, + padding="max_length", max_length=config.MAX_LEN, + ) + return { + "input_ids": encoding["input_ids"].squeeze(), + "attention_mask": encoding["attention_mask"].squeeze(), + "bbox": encoding["bbox"].squeeze(), + "pixel_values": encoding["pixel_values"].squeeze(), + "labels": torch.tensor(s["label"], dtype=torch.long), + } + + +# ── Fisher (EWC) ────────────────────────────────── +def compute_fisher(model, loader) -> dict: + fisher = {n: torch.zeros_like(p) + for n, p in model.named_parameters() if p.requires_grad} + model.eval() + for batch in loader: + model.zero_grad() + out = model( + input_ids = batch["input_ids"].to(config.DEVICE), + attention_mask = batch["attention_mask"].to(config.DEVICE), + bbox = batch["bbox"].to(config.DEVICE), + pixel_values = batch["pixel_values"].to(config.DEVICE), + labels = batch["labels"].to(config.DEVICE), + ) + out.loss.backward() + for n, p in model.named_parameters(): + if p.requires_grad and p.grad is not None: + fisher[n] += p.grad.pow(2) + return {n: f / max(len(loader), 1) for n, f in fisher.items()} + + +# ── 검증 ────────────────────────────────────────── +def evaluate_loader(model, loader) -> float: + model.eval() + correct, total = 0, 0 + with torch.no_grad(): + for batch in loader: + out = model( + input_ids = batch["input_ids"].to(config.DEVICE), + attention_mask = batch["attention_mask"].to(config.DEVICE), + bbox = batch["bbox"].to(config.DEVICE), + pixel_values = batch["pixel_values"].to(config.DEVICE), + ) + preds = out.logits.argmax(dim=-1).cpu() + correct += (preds == batch["labels"]).sum().item() + total += len(batch["labels"]) + return correct / total if total else 0.0 + + +def _get_latest_model_path(): + """ + 날짜/차수 폴더에서 가장 최신 모델 경로 반환 + files/models/20260305/0002/ ← 최신 + """ + base = config.MODEL_BASE_PATH + all_models = sorted([ + d for d in base.rglob("config.json") + ]) + if not all_models: + return None + + print(all_models[-1].parent) + return all_models[-1].parent + + +# ── 학습 루프 ───────────────────────────────────── +def run_train_loop(model, processor, train_s, val_s, req: TrainRequest, + label2id: dict, strategy: str, save_path: Path): + global train_status + + train_loader = DataLoader( + InvoiceDataset(train_s, processor), + batch_size=req.batch_size, shuffle=True + ) + val_loader = DataLoader( + InvoiceDataset(val_s, processor), + batch_size=req.batch_size + ) + optimizer = AdamW(model.parameters(), lr=req.lr) + scheduler = get_scheduler( + "cosine", optimizer=optimizer, + num_warmup_steps=len(train_loader), + num_training_steps=len(train_loader) * req.epochs, + ) + + # EWC 사전 계산 + ewc_params, ewc_fisher = {}, {} + if strategy == "ewc": + ewc_params = {n: p.clone().detach() + for n, p in model.named_parameters() if p.requires_grad} + ewc_fisher = compute_fisher(model, val_loader) + logger.info("EWC Fisher 계산 완료") + + best_acc = 0.0 + total_batches = len(train_loader) * req.epochs + + for epoch in range(req.epochs): + model.train() + total_loss = 0 + + pbar = tqdm( + train_loader, + desc=f"Epoch {epoch+1}/{req.epochs}", + ncols=100, + leave=True, + ) + for batch_idx, batch in enumerate(pbar): + if stop_requested: + logger.info("학습 중단 요청 감지") + return + + out = model( + input_ids = batch["input_ids"].to(config.DEVICE), + attention_mask = batch["attention_mask"].to(config.DEVICE), + bbox = batch["bbox"].to(config.DEVICE), + pixel_values = batch["pixel_values"].to(config.DEVICE), + labels = batch["labels"].to(config.DEVICE), + ) + loss = out.loss + if strategy == "ewc" and ewc_fisher: + for n, p in model.named_parameters(): + if n in ewc_fisher: + loss += req.ewc_lambda * ( + ewc_fisher[n] * (p - ewc_params[n]).pow(2) + ).sum() + + loss.backward() + nn.utils.clip_grad_norm_(model.parameters(), 1.0) + optimizer.step() + scheduler.step() + optimizer.zero_grad() + total_loss += loss.item() + + pbar.set_postfix(loss=f"{loss.item():.4f}") + + # 배치 단위 상태 업데이트 + done_batches = epoch * len(train_loader) + (batch_idx + 1) + train_status.epoch = epoch + 1 + train_status.total = req.epochs + train_status.loss = round(loss.item(), 4) + train_status.batch = done_batches # 완료된 전체 배치 + train_status.total_batches = total_batches # 전체 배치 수 + train_status.message = ( + f"Epoch {epoch+1}/{req.epochs} | " + f"batch {batch_idx+1}/{len(train_loader)} | " + f"loss={loss.item():.4f}" + ) + + avg_loss = total_loss / len(train_loader) + val_acc = evaluate_loader(model, val_loader) + + msg = ( + f"Epoch {epoch+1}/{req.epochs} | " + f"loss={avg_loss:.4f} | val_acc={val_acc:.4f}" + ) + tqdm.write(f"[결과] {msg}") + logger.info(msg) + + train_status.epoch = epoch + 1 + train_status.loss = round(avg_loss, 4) + train_status.val_acc = round(val_acc, 4) + train_status.message = msg + + if val_acc > best_acc: + best_acc = val_acc + model.save_pretrained(str(save_path)) + processor.save_pretrained(str(save_path)) + save_label2id(label2id) + logger.info(f"모델 저장 (val_acc={val_acc:.4f}) → {save_path}") + + +# ── 학습 진입점 ─────────────────────────────────── +def run_train(req: TrainRequest): + global train_status + stop_requested = False + train_status = TrainStatus( + running=True, epoch=0, total=req.epochs, + loss=0.0, message="학습 준비 중" + ) + try: + from utils.cache import load_label2id + + # 날짜/차수 폴더 생성 + save_path = config.get_model_save_path() + logger.info(f"모델 저장 경로: {save_path}") + train_status.message = f"저장 경로: {save_path}" + + samples = load_cache() + label2id = load_label2id() + id2label = {v: k for k, v in label2id.items()} + + processor_path = ( + str(config.MODEL_BASE_PATH / "preprocessor_config.json") + if (config.MODEL_BASE_PATH / "preprocessor_config.json").exists() + else config.MODEL_NAME + ) + # 최신 모델 폴더 찾기 (이어 학습) + latest_model = _get_latest_model_path() + model_path = str(latest_model) if latest_model else config.MODEL_NAME + logger.info(f"기반 모델: {model_path}") + + processor = LayoutLMv3Processor.from_pretrained( + latest_model and str(latest_model) or config.MODEL_NAME, + apply_ocr=False + ) + model = LayoutLMv3ForSequenceClassification.from_pretrained( + model_path, + num_labels=len(label2id), + id2label=id2label, + label2id=label2id, + ignore_mismatched_sizes=True, + ).to(config.DEVICE) + + # 전략 자동 선택 + strategy = req.strategy + if strategy == "auto": + buf = load_replay_buffer() + existing_types = set(buf.keys()) + new_types = {id2label[s["label"]] for s in samples} + strategy = "replay" if (new_types - existing_types) else "ewc" + train_status.message = f"전략: {strategy}" + logger.info(f"학습 전략: {strategy}") + + train_s, val_s = train_test_split( + samples, test_size=0.2, random_state=42, + stratify=[s["label"] for s in samples], + ) + + if strategy == "replay": + replay_s = [s for v in load_replay_buffer().values() for s in v] + combined = train_s + replay_s + random.shuffle(combined) + run_train_loop(model, processor, combined, val_s, req, label2id, strategy, save_path) + update_replay_buffer(train_s, id2label) + else: + run_train_loop(model, processor, train_s, val_s, req, label2id, strategy, save_path) + update_replay_buffer(train_s, id2label) + + train_status.message = f"학습 완료 → {save_path}" + logger.info(f"학습 완료: {save_path}") + + except Exception: + train_status.message = f"오류: {traceback.format_exc()}" + logger.exception("학습 중 예외 발생") + finally: + stop_requested = False + train_status.running = False + + +# ── 평가 진입점 ─────────────────────────────────── +def run_evaluate(test_size: float = 0.2) -> dict: + samples = load_cache() + label2id = load_label2id() + id2label = {v: k for k, v in label2id.items()} + + _, test_s = train_test_split( + samples, test_size=test_size, random_state=42, + stratify=[s["label"] for s in samples], + ) + + model_path = _get_latest_model_path() + if not model_path: + raise RuntimeError("저장된 모델이 없습니다.") + + processor = LayoutLMv3Processor.from_pretrained( + str(model_path), apply_ocr=False + ) + model = LayoutLMv3ForSequenceClassification.from_pretrained( + str(model_path) + ).to(config.DEVICE) + model.eval() + + loader = DataLoader(InvoiceDataset(test_s, processor), batch_size=config.BATCH_SIZE) + all_preds, all_labels = [], [] + with torch.no_grad(): + for batch in loader: + out = model( + input_ids = batch["input_ids"].to(config.DEVICE), + attention_mask = batch["attention_mask"].to(config.DEVICE), + bbox = batch["bbox"].to(config.DEVICE), + pixel_values = batch["pixel_values"].to(config.DEVICE), + ) + all_preds.extend(out.logits.argmax(dim=-1).cpu().tolist()) + all_labels.extend(batch["labels"].tolist()) + + target_names = [id2label[i] for i in range(len(id2label))] + report_dict = classification_report( + all_labels, all_preds, target_names=target_names, output_dict=True + ) + report_str = classification_report( + all_labels, all_preds, target_names=target_names + ) + logger.info(f"평가 완료 | accuracy={report_dict['accuracy']:.4f}") + return { + "test_samples": len(test_s), + "accuracy": round(report_dict["accuracy"], 4), + "report": report_dict, + "report_text": report_str, + } \ No newline at end of file diff --git a/utils/cache.py b/utils/cache.py new file mode 100644 index 0000000..635e6dc --- /dev/null +++ b/utils/cache.py @@ -0,0 +1,123 @@ +import json +import random +from pathlib import Path + +from loguru import logger + +import config + + +# ── OCR 캐시 (파일별 json 분리) ─────────────────── + +def _cache_path(image_path: str) -> Path: + """ + 이미지 경로 → json 저장 경로 + files/data/TSMC_A/1000114581_INV/page_001.png + → files/json/TSMC_A/1000114581_INV/page_001.json + """ + img = Path(image_path) + try: + rel = img.relative_to(config.DATA_DIR) + except ValueError: + rel = Path(img.parent.name) / img.name + return config.OCR_CACHE_DIR / rel.with_suffix(".json") + + +def save_cache_by_file(ocr_result: dict): + """이미지 1장의 OCR 결과를 개별 json으로 저장""" + path = _cache_path(ocr_result["image_path"]) + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "w", encoding="utf-8") as f: + json.dump(ocr_result, f, ensure_ascii=False, indent=2) + logger.debug(f"캐시 저장: {path}") + + +def load_cache() -> list: + """전체 캐시 통합 로드 (학습/평가 시 사용)""" + samples = [] + for p in config.OCR_CACHE_DIR.rglob("*.json"): + try: + with open(p, encoding="utf-8") as f: + samples.append(json.load(f)) + except (json.JSONDecodeError, UnicodeDecodeError) as e: + logger.warning(f"캐시 로드 실패 (건너뜀): {p} | {e}") + return samples + + +def get_cached_paths() -> set: + """캐시된 이미지 경로 전체 반환 (중복 방지용)""" + cached = [] + for p in config.OCR_CACHE_DIR.rglob("*.json"): + try: + with open(p, encoding="utf-8") as f: + data = json.load(f) + cached.append(data.get("image_path", "")) + # image_path = data.get("image_path", "") + # if image_path: + # cached.append(str(Path(image_path))) # 정규화 + # print(str(Path(image_path))) + except (json.JSONDecodeError, UnicodeDecodeError): + pass + return set(cached) + + +def cache_status() -> dict: + """양식별 캐시 현황 (폴더별 json 파일 수 집계)""" + if not config.OCR_CACHE_DIR.exists(): + return {} + status = {} + for cls_dir in sorted(config.OCR_CACHE_DIR.iterdir()): + if cls_dir.is_dir(): + status[cls_dir.name] = len(list(cls_dir.rglob("*.json"))) + return status + + +# ── label2id ────────────────────────────────────── + +def load_label2id() -> dict: + path = config.MODEL_SAVE_PATH / "label2id.json" + if path.exists(): + with open(path, encoding="utf-8") as f: + return json.load(f) + if config.DATA_DIR.exists(): + classes = sorted([d.name for d in config.DATA_DIR.iterdir() if d.is_dir()]) + return {c: i for i, c in enumerate(classes)} + return {} + + +def save_label2id(label2id: dict): + config.MODEL_SAVE_PATH.mkdir(exist_ok=True) + with open(config.MODEL_SAVE_PATH / "label2id.json", "w", encoding="utf-8") as f: + json.dump(label2id, f, ensure_ascii=False, indent=2) + logger.info(f"label2id 저장: {label2id}") + + +# ── 리플레이 버퍼 ────────────────────────────────── + +def load_replay_buffer() -> dict: + if not config.REPLAY_BUFFER_FILE.exists(): + return {} + try: + with open(config.REPLAY_BUFFER_FILE, encoding="utf-8") as f: + return json.load(f) + except (json.JSONDecodeError, UnicodeDecodeError) as e: + backup = config.REPLAY_BUFFER_FILE.with_suffix(".json.bak") + config.REPLAY_BUFFER_FILE.rename(backup) + logger.warning(f"replay_buffer.json 손상 → 백업 후 초기화: {e}") + return {} + + +def update_replay_buffer(new_samples: list, id2label: dict, per_class: int = 10): + buf = load_replay_buffer() + by_label = {} + for s in new_samples: + name = id2label[s["label"]] + by_label.setdefault(name, []).append(s) + + for name, samples in by_label.items(): + combined = buf.get(name, []) + samples + buf[name] = random.sample(combined, min(per_class, len(combined))) + + with open(config.REPLAY_BUFFER_FILE, "w", encoding="utf-8") as f: + json.dump(buf, f, ensure_ascii=False, indent=2) + logger.info(f"리플레이 버퍼 갱신: { {k: len(v) for k, v in buf.items()} }") \ No newline at end of file diff --git a/utils/classifier_model_store.py b/utils/classifier_model_store.py new file mode 100644 index 0000000..b79eec7 --- /dev/null +++ b/utils/classifier_model_store.py @@ -0,0 +1,68 @@ +""" +모델 로드/저장/예측 담당 +""" +from pathlib import Path +from typing import Optional + +import joblib +import numpy as np +from loguru import logger +from sklearn.pipeline import Pipeline + +import config + +MODEL_FILENAME = "classifier.pkl" + + +def _get_latest_path() -> Optional[Path]: + candidates = sorted(config.CLASSIFIER_MODEL_BASE_PATH.rglob(MODEL_FILENAME)) + return candidates[-1] if candidates else None + + +def _get_save_path() -> Path: + from datetime import datetime + date_dir = config.CLASSIFIER_MODEL_BASE_PATH / datetime.now().strftime("%Y%m%d") + date_dir.mkdir(exist_ok=True) + existing = sorted([d for d in date_dir.iterdir() if d.is_dir() and d.name.isdigit()]) + next_idx = int(existing[-1].name) + 1 if existing else 1 + save_dir = date_dir / f"{next_idx:04d}" + save_dir.mkdir(exist_ok=True) + return save_dir / MODEL_FILENAME + + +_model: Optional[Pipeline] = None + +# 서버 시작 시 최신 모델 자동 로드 +_latest = _get_latest_path() +if _latest: + try: + _model = joblib.load(_latest) + logger.info(f"classifier 모델 로드: {_latest} | classes={list(_model.classes_)}") + except Exception as e: + logger.warning(f"모델 로드 실패: {e}") + + +def get_model() -> Optional[Pipeline]: + return _model + + +def set_model(pipeline: Pipeline) -> None: + global _model + path = _get_save_path() + joblib.dump(pipeline, path) + _model = pipeline + logger.info(f"classifier 모델 저장: {path}") + + +def predict( + pipeline: Pipeline, + text: str, + threshold: float = 0.6, +) -> tuple[str, float, dict[str, float]]: + probs = pipeline.predict_proba([text])[0] + classes = pipeline.classes_ + max_idx = int(np.argmax(probs)) + confidence = float(probs[max_idx]) + label = "OTHER" if confidence < threshold else classes[max_idx] + all_probs = {c: round(float(p), 4) for c, p in zip(classes, probs)} + return label, round(confidence, 4), all_probs \ No newline at end of file diff --git a/utils/ocr.py b/utils/ocr.py new file mode 100644 index 0000000..1e43599 --- /dev/null +++ b/utils/ocr.py @@ -0,0 +1,38 @@ +# import config +# +# if config.OCR_ENGINE == "google": +# from utils.ocr_google import ocr_single, ocr_batch, get_vision_client +# else: +# from utils.ocr_paddle import ocr_single, ocr_batch +# def get_vision_client(): +# return None # paddle은 client 불필요, 호환성 유지용 + + +from utils.ocr_paddle import ocr_single as paddle_ocr_single, ocr_batch as paddle_ocr_batch +from utils.ocr_google import ocr_single as google_ocr_single, ocr_batch as google_ocr_batch, get_vision_client +import config + +_google_client = None + + +def get_vision_client_cached(): + global _google_client + if _google_client is None: + _google_client = get_vision_client() + return _google_client + + +def get_ocr_functions(engine: str): + if engine == "google": + return google_ocr_single, google_ocr_batch, get_vision_client_cached() # 캐싱된 클라이언트 사용 + return paddle_ocr_single, paddle_ocr_batch, None + + +def ocr_single(image_path, client=None): + fn, _, c = get_ocr_functions(config.OCR_ENGINE) + return fn(image_path, c) + + +def ocr_batch(img_paths, client=None): + _, fn, c = get_ocr_functions(config.OCR_ENGINE) + return fn(img_paths, c) \ No newline at end of file diff --git a/utils/ocr_google.py b/utils/ocr_google.py new file mode 100644 index 0000000..6fe8fb8 --- /dev/null +++ b/utils/ocr_google.py @@ -0,0 +1,92 @@ +import concurrent.futures +import json +from pathlib import Path + +from google.cloud import vision_v1 +from google.oauth2 import service_account +from loguru import logger +from PIL import Image + +import config + + +def get_vision_client() -> vision_v1.ImageAnnotatorClient: + credentials = service_account.Credentials.from_service_account_file( + config.SERVICE_ACCOUNT_FILE, + scopes=["https://www.googleapis.com/auth/cloud-platform"], + ) + return vision_v1.ImageAnnotatorClient(credentials=credentials) + + +def ocr_single(image_path: str, client: vision_v1.ImageAnnotatorClient) -> dict: + """ + 이미지 1장 → words + 정규화 boxes (0~1000) 반환 + """ + image_pil = Image.open(image_path).convert("RGB") + w, h = image_pil.size + + with open(image_path, "rb") as f: + content = f.read() + + response = client.document_text_detection( + image=vision_v1.Image(content=content) + ) + + words, boxes = [], [] + for page in response.full_text_annotation.pages: + for block in page.blocks: + for para in block.paragraphs: + for word in para.words: + text = "".join([s.text for s in word.symbols]) + if not text.strip(): + continue + v = word.bounding_box.vertices + boxes.append([ + max(0, int(1000 * v[0].x / w)), + max(0, int(1000 * v[0].y / h)), + min(1000, int(1000 * v[2].x / w)), + min(1000, int(1000 * v[2].y / h)), + ]) + words.append(text) + + result = {"image_path": str(image_path), "words": words, "boxes": boxes, "box_type": "normalized"} + + # OCR 결과 저장 + # json_path = Path(image_path).with_name(Path(image_path).stem + "_google.json") + # with open(json_path, "w", encoding="utf-8") as f: + # json.dump(result, f, ensure_ascii=False, indent=2) + # logger.debug(f"OCR 저장: {json_path}") + + logger.debug(f"OCR 완료: {image_path} | 단어수={len(words)}") + return result + + +def ocr_batch(image_paths: list, client=None, max_workers: int = config.OCR_MAX_WORKERS, on_done=None) -> list: + """ + ThreadPoolExecutor 병렬 OCR + Vision API는 네트워크 I/O 기반 → 병렬 효과 큼 + """ + # client = get_vision_client() + if client is None: + client = get_vision_client() + results = [None] * len(image_paths) + + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: + future_to_idx = { + executor.submit(ocr_single, str(p), client): i + for i, p in enumerate(image_paths) + } + for future in concurrent.futures.as_completed(future_to_idx): + idx = future_to_idx[future] + try: + results[idx] = future.result() + except Exception as e: + logger.error(f"OCR 실패: {image_paths[idx]} | {e}") + results[idx] = { + "image_path": str(image_paths[idx]), + "words": [], "boxes": [], "error": str(e), + } + if on_done: + on_done(idx, results[idx]) + + return results diff --git a/utils/ocr_paddle.py b/utils/ocr_paddle.py new file mode 100644 index 0000000..d1d6ea7 --- /dev/null +++ b/utils/ocr_paddle.py @@ -0,0 +1,107 @@ +import concurrent.futures +import json +from pathlib import Path + +import httpx +from loguru import logger + +import config + +_client = httpx.Client( + timeout=httpx.Timeout( + connect=3.0, + read=30.0, # 파일 전송이므로 기존 30초 유지 + write=30.0, + pool=3.0 + ), + limits=httpx.Limits( + max_keepalive_connections=5, + keepalive_expiry=50 # A서버 --timeout-keep-alive 60 보다 약간 낮게 + ) +) + + +def ocr_single_pdf(pdf_path: str) -> dict: + with open(pdf_path, "rb") as f: + resp = httpx.post( + config.PADDLE_OCR_URL, + files={"file": (Path(pdf_path).name, f, "application/pdf")}, + data={"group_id": "infer"}, + timeout=120, # PDF는 이미지보다 오래 걸리므로 timeout 증가 + ) + resp.raise_for_status() + ocr_result = resp.json() + + words, boxes = [], [] + for page in ocr_result: # PDF는 페이지가 여러 장일 수 있으므로 전체 순회 + for text, box in zip(page.get("words", []), page.get("boxes", [])): + text = text.strip() + if not text or not box: + continue + words.append(text) + boxes.append(box) + + logger.debug(f"PaddleOCR PDF 완료: {pdf_path} | 단어수={len(words)}") + return {"image_path": str(pdf_path), "words": words, "boxes": boxes, "box_type": "pixel"} + + +def ocr_single(image_path: str, client=None) -> dict: + with open(image_path, "rb") as f: + data = {"group_id": "infer"} + if config.SERVICE == "ucar": + data["gubun"] = "ocr_with_boxes" + + resp = _client.post( + config.PADDLE_OCR_URL, + files={"file": (Path(image_path).name, f, "image/jpeg")}, + data=data, + ) + + resp.raise_for_status() + ocr_result = resp.json() + + words, boxes = [], [] + + # 페이지가 1장이므로 첫 번째 항목만 사용 + page = ocr_result[0] if ocr_result else {} + + for text, box in zip(page.get("words", []), page.get("boxes", [])): + text = text.strip() + if not text or not box: + continue + words.append(text) + boxes.append(box) + + result = {"image_path": str(image_path), "words": words, "boxes": boxes, "box_type": "pixel"} + + # OCR 결과 저장 + # json_path = Path(image_path).with_name(Path(image_path).stem + "_paddle.json") + # with open(json_path, "w", encoding="utf-8") as f: + # json.dump(result, f, ensure_ascii=False, indent=2) + # logger.debug(f"OCR 저장: {json_path}") + + logger.debug(f"PaddleOCR 완료: {image_path} | 단어수={len(words)}") + return result + + +def ocr_batch(image_paths: list, client=None, max_workers: int = config.OCR_MAX_WORKERS, + on_done=None) -> list: # ← on_done 추가 + results = [None] * len(image_paths) + + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: + future_to_idx = { + executor.submit(ocr_single, str(p)): i + # executor.submit(ocr_single_pdf, str(p)): i + for i, p in enumerate(image_paths) + } + for future in concurrent.futures.as_completed(future_to_idx): + idx = future_to_idx[future] + try: + results[idx] = future.result() + except Exception as e: + logger.error(f"PaddleOCR 실패: {image_paths[idx]} | {e}") + results[idx] = {"image_path": str(image_paths[idx]), + "words": [], "boxes": [], "error": str(e)} + if on_done: # ← 콜백 호출 + on_done(idx, results[idx]) + return results diff --git a/utils/pdf.py b/utils/pdf.py new file mode 100644 index 0000000..c7749f5 --- /dev/null +++ b/utils/pdf.py @@ -0,0 +1,85 @@ +from pathlib import Path + +import fitz # PyMuPDF +from loguru import logger + +import config +import concurrent.futures + + +def is_pdf(filename: str) -> bool: + return Path(filename).suffix.lower() == ".pdf" + + +def pdf_to_images(pdf_path: str, dpi: int = config.OCR_DPI, + max_pages: int = config.PDF_MAX_PAGES) -> list[Path]: + pdf_path = Path(pdf_path) + output_dir = pdf_path.parent / pdf_path.stem + output_dir.mkdir(exist_ok=True) + + ext = config.PDF_IMAGE_FORMAT # "jpeg" + page_indices = list(range(max_pages)) + + # 이미 변환된 이미지가 모두 있으면 스킵 + cached = [output_dir / f"page_{i+1:03d}.{ext}" for i in page_indices] + if all(p.exists() for p in cached): + logger.debug(f"PDF 변환 캐시 사용: {pdf_path.name}") + return cached + + pdf_path = Path(pdf_path) + print(pdf_path.exists()) + print(pdf_path.stat().st_size) + if not pdf_path.exists() or pdf_path.stat().st_size == 0: + raise ValueError(f"유효하지 않은 PDF 파일입니다: {pdf_path.name}") + + matrix = fitz.Matrix(dpi / 72, dpi / 72) + doc = fitz.open(str(pdf_path)) + + img_paths = [] + for i in range(min(max_pages, len(doc))): + img_path = output_dir / f"page_{i+1:03d}.{ext}" + if img_path.exists(): + logger.debug(f"페이지 캐시 사용: {img_path.name}") + else: + pixmap = doc[i].get_pixmap(matrix=matrix, alpha=False) + + # 최대 너비 초과 시 비율 유지하며 축소 + if hasattr(config, 'PDF_MAX_WIDTH') and pixmap.width > config.PDF_MAX_WIDTH: + scale = config.PDF_MAX_WIDTH / pixmap.width + pixmap = doc[i].get_pixmap(matrix=fitz.Matrix(dpi / 72 * scale, dpi / 72 * scale), alpha=False) + + if ext in ("jpeg", "jpg"): + pixmap.save(str(img_path), jpg_quality=config.PDF_JPEG_QUALITY) # JPEG 저장 + else: + pixmap.save(str(img_path)) + logger.debug(f"PDF 변환: {pdf_path.name} | 페이지 {i+1}") + img_paths.append(img_path) + + doc.close() + logger.info(f"PDF 변환 완료: {pdf_path.name} | {len(img_paths)}페이지 | {dpi}dpi {ext}") + return img_paths + + +def pdf_to_images_batch(pdf_paths: list[str], + max_workers: int = config.OCR_MAX_WORKERS) -> dict[str, list[Path]]: + """ + 여러 PDF를 ThreadPoolExecutor로 병렬 변환 + 반환: {pdf_path: [img_path, ...]} + """ + results = {} + + def _convert(pdf_path: str): + return pdf_path, pdf_to_images(pdf_path) + + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = {executor.submit(_convert, p): p for p in pdf_paths} + for future in concurrent.futures.as_completed(futures): + try: + path, imgs = future.result() + results[path] = imgs + except Exception as e: + path = futures[future] + logger.error(f"PDF 변환 실패: {path} | {e}") + results[path] = [] + + return results \ No newline at end of file