Files
layoutlmv3_service/routers/visualize.py
2026-06-15 09:54:01 +09:00

289 lines
11 KiB
Python

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'<tr onclick="focusBox({i})" id="row-{i}">'
f'<td>{i+1}</td><td>{w}</td><td>{b}</td></tr>'
for i, (w, b) in enumerate(zip(words, boxes))
])
html = f"""
<!DOCTYPE html>
<html>
<head>
<meta charset="utf-8">
<title>OCR 시각화 - {label}/{filename}</title>
<style>
* {{ box-sizing: border-box; }}
body {{ font-family: sans-serif; margin: 0; background: #f5f5f5;
height: 100vh; display: flex; flex-direction: column; overflow: hidden; }}
.header {{ padding: 12px 20px 4px; flex-shrink: 0; }}
h2 {{ color: #333; margin: 0 0 4px; font-size: 16px; }}
.meta {{ font-size: 13px; color: #555; margin-bottom: 8px; }}
/* 토글 고정 영역 */
.toolbar {{ padding: 6px 20px; background: #fff;
border-bottom: 1px solid #ddd; flex-shrink: 0;
display: flex; gap: 24px; align-items: center; font-size: 13px; }}
.switch {{ position: relative; width: 40px; height: 22px; }}
.switch input{{ opacity: 0; width: 0; height: 0; }}
.slider {{ position: absolute; inset: 0; background: #ccc;
border-radius: 22px; cursor: pointer; transition: .3s; }}
.slider:before{{ content:""; position: absolute; width: 16px; height: 16px;
left: 3px; bottom: 3px; background: white;
border-radius: 50%; transition: .3s; }}
input:checked + .slider {{ background: #1976d2; }}
input:checked + .slider:before {{ transform: translateX(18px); }}
.toggle-item{{ display: flex; align-items: center; gap: 8px; }}
/* 본문 컨테이너 - 남은 높이 전부 */
.container {{ display: flex; gap: 0; flex: 1; overflow: hidden; }}
/* 이미지 영역 - 독립 스크롤 */
.left {{ flex: 1; overflow: auto; padding: 12px 12px 12px 20px;
border-right: 1px solid #ddd; background: #fff; }}
.canvas-wrap{{ position: relative; display: inline-block; }}
canvas {{ position: absolute; top: 0; left: 0; pointer-events: none; }}
img {{ display: block; border: 1px solid #ccc; max-width: none; }}
/* 테이블 영역 - 독립 스크롤 */
.right {{ width: 400px; flex-shrink: 0; overflow: auto;
padding: 12px 20px 12px 12px; background: #fafafa; }}
.right h3 {{ margin: 0 0 8px; font-size: 14px; position: sticky;
top: 0; background: #fafafa; padding: 4px 0; z-index: 1; }}
table {{ border-collapse: collapse; width: 100%; font-size: 13px; }}
th, td {{ border: 1px solid #ddd; padding: 5px 8px; text-align: left; cursor: pointer; }}
th {{ background: #e8e8e8; position: sticky; top: 30px; z-index: 1; }}
tr:hover {{ background: #fff9c4; }}
tr.active {{ background: #ffe082 !important; font-weight: bold; }}
</style>
</head>
<body>
<!-- 헤더 고정 -->
<div class="header">
<h2>OCR 시각화 | {label} / {filename}</h2>
<div class="meta">총 단어 수: <b>{len(words)}</b>개</div>
</div>
<!-- 토글 고정 툴바 -->
<div class="toolbar">
<div class="toggle-item">
<label class="switch">
<input type="checkbox" id="boxToggle" checked onchange="drawAll(activeIdx)">
<span class="slider"></span>
</label>
<span>바운딩박스</span>
</div>
<div class="toggle-item">
<label class="switch">
<input type="checkbox" id="textToggle" onchange="drawAll(activeIdx)">
<span class="slider"></span>
</label>
<span>OCR 텍스트</span>
</div>
</div>
<!-- 이미지(좌) + 테이블(우) 각각 독립 스크롤 -->
<div class="container">
<div class="left">
<div class="canvas-wrap">
<img id="img" src="data:image/png;base64,{img_b64}"
onload="initCanvas()" />
<canvas id="overlay"></canvas>
</div>
</div>
<div class="right" id="tablePane">
<h3>OCR 텍스트 목록</h3>
<table>
<thead>
<tr><th>#</th><th>텍스트</th><th>박스 (정규화)</th></tr>
</thead>
<tbody id="tbody">{ocr_rows}</tbody>
</table>
</div>
</div>
<script>
const BOXES = {boxes_json};
const WORDS = {words_json};
const img = document.getElementById("img");
const canvas = document.getElementById("overlay");
const ctx = canvas.getContext("2d");
let scaleX = 1, scaleY = 1;
let activeIdx = -1;
function initCanvas() {{
canvas.width = img.clientWidth;
canvas.height = img.clientHeight;
scaleX = img.clientWidth / {img_w};
scaleY = img.clientHeight / {img_h};
drawAll(-1);
}}
function drawAll(highlight) {{
ctx.clearRect(0, 0, canvas.width, canvas.height);
const showBox = document.getElementById("boxToggle").checked;
const showText = document.getElementById("textToggle").checked;
BOXES.forEach((b, i) => {{
const x1 = b[0] * scaleX, y1 = b[1] * scaleY;
const x2 = b[2] * scaleX, y2 = b[3] * scaleY;
const bw = x2 - x1, bh = y2 - y1;
// 강조 배경
if (i === highlight) {{
ctx.fillStyle = "rgba(255,235,59,0.45)";
ctx.fillRect(x1, y1, bw, bh);
}}
// 바운딩박스
if (showBox) {{
ctx.strokeStyle = i === highlight ? "#e53935" : "rgba(229,57,53,0.55)";
ctx.lineWidth = i === highlight ? 3 : 1.5;
ctx.strokeRect(x1, y1, bw, bh);
}}
// OCR 텍스트
if (showText) {{
const fs = Math.max(11, Math.min(bh * 0.85, 15));
ctx.font = `bold ${{fs}}px sans-serif`;
ctx.lineWidth = 2.5;
ctx.strokeStyle = "rgba(255,255,255,0.95)";
ctx.fillStyle = i === highlight ? "#b71c1c" : "#1565c0";
ctx.strokeText(WORDS[i], x1 + 1, y1 + fs);
ctx.fillText(WORDS[i], x1 + 1, y1 + fs);
}}
}});
}}
function focusBox(i) {{
activeIdx = i;
drawAll(i);
// 테이블 행 강조 + 스크롤
document.querySelectorAll("#tbody tr").forEach(r => r.classList.remove("active"));
const row = document.getElementById("row-" + i);
if (row) {{
row.classList.add("active");
row.scrollIntoView({{ block: "center", behavior: "smooth" }});
}}
// 이미지 영역 해당 박스 위치로 스크롤
const b = BOXES[i];
const cx = (b[0] + b[2]) / 2 * scaleX;
const cy = (b[1] + b[3]) / 2 * scaleY;
const leftPane = document.querySelector(".left");
const scrollX = cx - leftPane.clientWidth / 2;
const scrollY = cy - leftPane.clientHeight / 2;
leftPane.scrollTo({{ left: scrollX, top: scrollY, behavior: "smooth" }});
}}
window.addEventListener("resize", initCanvas);
</script>
</body>
</html>
"""
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,
}