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}
"""
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,
}