143 lines
6.0 KiB
Python
143 lines
6.0 KiB
Python
"""① detect — 포스터에서 카메라가 겨냥할 영역을 검출한다.
|
|
|
|
격자 위에서 VLM이 대략 읽고, 그 주변에서만 잉크 경계로 다듬는다.
|
|
"""
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
from scipy import ndimage
|
|
|
|
from answers.detect_answer import DetectAnswer, GridBox
|
|
from models.detect import Box, DetectResult, Focal, Regions
|
|
from settings import settings
|
|
from utils.common_llm import StructuredLLM
|
|
from utils.image import to_data_uri
|
|
from utils.prompt import load_prompt
|
|
|
|
Image.MAX_IMAGE_PIXELS = None
|
|
|
|
WORK_WIDTH_PX = 1100
|
|
REGION_KINDS = ("title", "datetime", "place", "logo")
|
|
REGION_COLORS = {"title": (255, 60, 60), "datetime": (60, 220, 255),
|
|
"place": (140, 255, 140), "logo": (255, 200, 60)}
|
|
|
|
DETECT_PROMPT = load_prompt("detect")
|
|
|
|
detect_llm = StructuredLLM("gpt-4o", settings.chatgpt_api_key)
|
|
|
|
|
|
def grid_overlay(poster: Image.Image) -> Image.Image:
|
|
overlay = poster.convert("RGB").copy()
|
|
width, height = overlay.size
|
|
draw = ImageDraw.Draw(overlay, "RGBA")
|
|
font = ImageFont.load_default(max(14, width // 55))
|
|
for row in range(1, 20):
|
|
y = int(height * row * 0.05)
|
|
draw.line([(0, y), (width, y)], fill=(255, 0, 80, 150), width=max(1, width // 700))
|
|
if row % 2 == 0:
|
|
draw.text((4, y + 2), f"{row * 5}", font=font, fill=(255, 255, 0, 255))
|
|
for col in range(1, 10):
|
|
x = int(width * col * 0.10)
|
|
draw.line([(x, 0), (x, height)], fill=(0, 220, 255, 130), width=max(1, width // 700))
|
|
draw.text((x + 3, 3), f"{col * 10}", font=font, fill=(0, 255, 255, 255))
|
|
return overlay
|
|
|
|
|
|
def check_overlay(poster: Image.Image, regions: dict[str, Box],
|
|
focal_points: list[Focal]) -> Image.Image:
|
|
overlay = poster.convert("RGB").copy()
|
|
width, height = overlay.size
|
|
draw = ImageDraw.Draw(overlay)
|
|
font = ImageFont.load_default(max(16, width // 45))
|
|
for kind, box in regions.items():
|
|
color = REGION_COLORS.get(kind, (255, 255, 255))
|
|
draw.rectangle([box.x0 * width, box.y0 * height, box.x1 * width, box.y1 * height],
|
|
outline=color, width=max(3, width // 300))
|
|
draw.text((box.x0 * width + 6, box.y0 * height + 4), kind, font=font, fill=color)
|
|
for focal in focal_points:
|
|
x, y = focal.cx * width, focal.cy * height
|
|
radius = max(10, width // 60)
|
|
draw.ellipse([x - radius, y - radius, x + radius, y + radius],
|
|
outline=(255, 255, 0), width=max(3, width // 300))
|
|
draw.text((x + radius + 4, y - radius), focal.name, font=font, fill=(255, 255, 0))
|
|
return overlay
|
|
|
|
|
|
def ink_mask(poster: Image.Image) -> np.ndarray:
|
|
"""다중 스케일 국소 대비. 창이 획 두께보다 커야 굵은 글씨가 안 빠진다."""
|
|
brightness = np.asarray(poster).astype(np.float32).max(2)
|
|
height, width = brightness.shape
|
|
mask = np.zeros((height, width), bool)
|
|
for window in (width // 40, width // 12, width // 5):
|
|
local_mean = ndimage.uniform_filter(brightness, max(11, window) | 1)
|
|
mask |= brightness > local_mean + 18
|
|
mask |= brightness < local_mean - 36
|
|
return mask
|
|
|
|
|
|
def tighten_box(answer_box: GridBox, ink: np.ndarray,
|
|
pad_ratio: float = 0.03) -> dict[str, float] | None:
|
|
"""VLM이 준 대략 박스 주변에서만 잉크 경계를 찾아 정규화 좌표로 조인다."""
|
|
height, width = ink.shape
|
|
left = int(max(0, (answer_box.x0 / 100 - pad_ratio) * width))
|
|
right = int(min(width, (answer_box.x1 / 100 + pad_ratio) * width))
|
|
top = int(max(0, (answer_box.y0 / 100 - pad_ratio) * height))
|
|
bottom = int(min(height, (answer_box.y1 / 100 + pad_ratio) * height))
|
|
if right <= left or bottom <= top:
|
|
return None
|
|
patch = ink[top:bottom, left:right]
|
|
if patch.sum() < 50:
|
|
return None
|
|
row_density, col_density = patch.mean(1), patch.mean(0)
|
|
ink_rows = np.where(row_density > max(0.02, row_density.max() * 0.12))[0]
|
|
ink_cols = np.where(col_density > max(0.02, col_density.max() * 0.12))[0]
|
|
if len(ink_rows) == 0 or len(ink_cols) == 0:
|
|
return None
|
|
return {"x0": round((left + ink_cols.min()) / width, 4),
|
|
"x1": round((left + ink_cols.max() + 1) / width, 4),
|
|
"y0": round((top + ink_rows.min()) / height, 4),
|
|
"y1": round((top + ink_rows.max() + 1) / height, 4)}
|
|
|
|
|
|
async def detect(poster: Path | Image.Image) -> DetectResult:
|
|
source = (Image.open(poster) if isinstance(poster, Path) else poster).convert("RGB")
|
|
work_image = source.copy()
|
|
work_image.thumbnail((WORK_WIDTH_PX, WORK_WIDTH_PX * 3), Image.LANCZOS)
|
|
|
|
grid_image = grid_overlay(work_image)
|
|
answer = await detect_llm.ask_with_images(
|
|
DetectAnswer, DETECT_PROMPT, [("격자를 씌운 포스터:", to_data_uri(grid_image))])
|
|
ink = ink_mask(work_image)
|
|
|
|
regions: dict[str, Box] = {}
|
|
vlm_raw: dict = {}
|
|
for kind in REGION_KINDS:
|
|
answer_box: GridBox | None = getattr(answer, kind)
|
|
if answer_box is None:
|
|
continue
|
|
vlm_raw[kind] = answer_box.model_dump()
|
|
# 다듬기 실패 시 VLM 값을 그대로 쓴다
|
|
box = tighten_box(answer_box, ink) or {
|
|
corner: round(getattr(answer_box, corner) / 100, 4)
|
|
for corner in ("x0", "y0", "x1", "y1")}
|
|
regions[kind] = Box(**box, cx=round((box["x0"] + box["x1"]) / 2, 4),
|
|
cy=round((box["y0"] + box["y1"]) / 2, 4), text=answer_box.text)
|
|
|
|
focal_points = [Focal(name=focal.name, cx=round(focal.x / 100, 4), cy=round(focal.y / 100, 4))
|
|
for focal in answer.focal_points]
|
|
|
|
return DetectResult(
|
|
data=Regions(
|
|
festival_name=answer.festival_name,
|
|
title_style=answer.title_style,
|
|
recommended_mode="layered" if answer.title_style == "calligraphy" else "camera",
|
|
poster_size=source.size,
|
|
regions=regions,
|
|
focal_points=focal_points,
|
|
vlm_raw=vlm_raw,
|
|
),
|
|
grid=grid_image,
|
|
check=check_overlay(work_image, regions, focal_points),
|
|
)
|