s-canvas/_unused/test_gate_render.py

"""구조물 조감도 스탠드얼론 테스트 스크립트.

Gate_Sample DXF → 파라미터 추출 → 3D 모델 → 다각도 캡처 → AI 렌더링.

사용법:
    python test_gate_render.py
    python test_gate_render.py --interactive   # 인터랙티브 3D 뷰어
    python test_gate_render.py --ai            # AI 렌더링 포함 (GEMINI_API_KEY 필요)
"""

from __future__ import annotations

import argparse
import os
import sys
import math
from pathlib import Path

import numpy as np
import pyvista as pv
from PIL import Image, ImageDraw, ImageFilter

from spillway_parser import parse_spillway_dxf, SpillwayParams
from spillway_3d_builder import SpillwayBuilder, COLORS


# ---------------------------------------------------------------------------
# 렌더링 유틸
# ---------------------------------------------------------------------------

def add_all_meshes(plotter: pv.Plotter, meshes: list):
    """메쉬 리스트를 플로터에 추가."""
    for mesh, color, opacity in meshes:
        try:
            plotter.add_mesh(
                mesh,
                color=color,
                opacity=opacity,
                smooth_shading=True,
                specular=0.3,
                specular_power=10,
            )
        except Exception as e:
            print(f"  mesh 추가 실패: {e}")


def compute_camera_for_birdseye(params: SpillwayParams,
                                elevation_deg: float = 35.0,
                                azimuth_deg: float = 225.0,
                                zoom: float = 1.2) -> tuple:
    """여수로 구조물을 프레임에 담는 카메라 위치 계산.

    elevation_deg: 앙각 (수평=0°, 수직 아래로=90°)
    azimuth_deg: 방위각 (북=0°, 동=90°, 남=180°, 서=270°)
    zoom: 값이 클수록 카메라가 멀어짐 (전체를 더 많이 담음)
    """
    # 씬 전체 범위 (수면 + 구조물 + 하류 에이프런 포함)
    # 수면: Y ∈ [-40, 0.5], 구조물: Y ∈ [0, pier_length], 하류: Y ∈ [pier_length, pier_length+30]
    scene_y_min = -40.0
    scene_y_max = params.pier_length + 30.0
    scene_x_span = params.total_span
    scene_z_min = min(params.el_upstream_bed, params.el_downstream, params.el_gate_sill) - 1.0
    scene_z_max = params.el_bridge_top + 6.0  # 권양기 지붕까지

    # 씬의 중심 (focal point) — 구조물 위주
    cx = scene_x_span / 2
    cy = (0 + params.pier_length) / 2  # 구조물 중심
    cz = (params.el_gate_sill + params.el_bridge_top) / 2

    # 씬의 대각선 크기
    scene_dx = scene_x_span * 1.2  # X 방향 여유
    scene_dy = scene_y_max - scene_y_min
    scene_dz = scene_z_max - scene_z_min
    scene_diag = math.sqrt(scene_dx ** 2 + scene_dy ** 2 + scene_dz ** 2)

    # 카메라 거리: PyVista 기본 viewAngle ≈ 30° 가정, tan(15°) ≈ 0.268
    # 프레임 꽉 채우려면 dist = (scene_diag/2) / tan(15°) ≈ scene_diag * 1.87
    dist = scene_diag * 1.0 * zoom

    # 방위각/앙각 → 구면 좌표
    el_rad = math.radians(elevation_deg)
    az_rad = math.radians(azimuth_deg)

    dx = dist * math.cos(el_rad) * math.sin(az_rad)
    dy = -dist * math.cos(el_rad) * math.cos(az_rad)  # Y축 반전 (북→양)
    dz = dist * math.sin(el_rad)

    camera_pos = (cx + dx, cy + dy, cz + dz)
    focal_point = (cx, cy, cz)
    view_up = (0, 0, 1)

    return camera_pos, focal_point, view_up


def capture_view(params: SpillwayParams, meshes: list,
                 elevation_deg: float, azimuth_deg: float,
                 size: int = 1536, bg_color: str = "#C8D4E0",
                 zoom: float = 1.3) -> Image.Image:
    """지정한 각도에서 3D 씬을 캡처."""
    plotter = pv.Plotter(off_screen=True, window_size=(size, size))
    plotter.set_background(bg_color)

    add_all_meshes(plotter, meshes)

    # 카메라 설정
    cam_pos, focal, up = compute_camera_for_birdseye(
        params, elevation_deg, azimuth_deg, zoom
    )
    plotter.camera_position = [cam_pos, focal, up]

    # 조명 설정: 기본 헤드라이트 + 방향 조명
    plotter.enable_3_lights()

    img_arr = plotter.screenshot(return_img=True, window_size=(size, size))
    plotter.close()

    img = Image.fromarray(img_arr)
    return img


def capture_depth(params: SpillwayParams, meshes: list,
                  elevation_deg: float, azimuth_deg: float,
                  size: int = 1536, zoom: float = 1.3) -> Image.Image:
    """depth map 캡처 (제어맵용)."""
    plotter = pv.Plotter(off_screen=True, window_size=(size, size))
    plotter.set_background("black")

    for mesh, _, _ in meshes:
        plotter.add_mesh(mesh, color="white", smooth_shading=False)

    cam_pos, focal, up = compute_camera_for_birdseye(params, elevation_deg, azimuth_deg, zoom)
    plotter.camera_position = [cam_pos, focal, up]

    # show()로 렌더 파이프라인 초기화 후 depth 추출
    plotter.show(auto_close=False)
    try:
        z_img = plotter.get_image_depth()
    except Exception:
        z_img = None
    plotter.close()

    if z_img is None:
        return Image.new("L", (size, size), 0)

    # NaN 처리 + 정규화
    z_img = np.array(z_img, dtype=np.float32)
    z_finite = z_img[np.isfinite(z_img)]
    if len(z_finite) == 0:
        return Image.new("L", (size, size), 0)

    z_min, z_max = z_finite.min(), z_finite.max()
    if z_max - z_min < 1e-6:
        return Image.new("L", (size, size), 128)

    z_norm = (z_img - z_min) / (z_max - z_min)
    z_norm = np.where(np.isfinite(z_norm), z_norm, 1.0)
    # 가까울수록 밝게 (invert)
    z_norm = 1.0 - z_norm
    z_8bit = (z_norm * 255).astype(np.uint8)

    return Image.fromarray(z_8bit, "L")


def capture_lineart(params: SpillwayParams, meshes: list,
                    elevation_deg: float, azimuth_deg: float,
                    size: int = 1536, zoom: float = 1.3) -> Image.Image:
    """라인아트 캡처 (흰 배경 + 검은 엣지)."""
    plotter = pv.Plotter(off_screen=True, window_size=(size, size))
    plotter.set_background("white")

    for mesh, _, _ in meshes:
        plotter.add_mesh(mesh, color="white", show_edges=True, edge_color="black", line_width=1)

    cam_pos, focal, up = compute_camera_for_birdseye(params, elevation_deg, azimuth_deg, zoom)
    plotter.camera_position = [cam_pos, focal, up]

    img_arr = plotter.screenshot(return_img=True, window_size=(size, size))
    plotter.close()

    return Image.fromarray(img_arr)


def compose_guide_image(capture: Image.Image, depth: Image.Image, lineart: Image.Image) -> Image.Image:
    """캡처 + depth + lineart를 가이드 이미지로 합성."""
    # 모두 동일 크기로 맞춤
    base = capture.convert("RGB")
    d = depth.convert("RGB").resize(base.size)
    la = lineart.convert("RGB").resize(base.size)

    # 80% base + 20% depth, 그 위에 lineart 살짝
    arr_base = np.array(base, dtype=np.float32)
    arr_depth = np.array(d, dtype=np.float32)
    arr_line = np.array(la, dtype=np.float32)

    blend = arr_base * 0.80 + arr_depth * 0.20
    # 라인아트: 검은 픽셀만 선택해서 덧씌움
    line_mask = (arr_line.mean(axis=2, keepdims=True) < 100).astype(np.float32)
    final = blend * (1 - line_mask * 0.4) + arr_line * (line_mask * 0.4)
    final = np.clip(final, 0, 255).astype(np.uint8)

    return Image.fromarray(final)


# ---------------------------------------------------------------------------
# 인터랙티브 뷰어
# ---------------------------------------------------------------------------

def show_interactive(params: SpillwayParams, meshes: list):
    """PyVista 인터랙티브 뷰어. q로 종료."""
    plotter = pv.Plotter(title="EG-VIEW Gate: Interactive Preview")
    plotter.set_background("#2B3A4A")
    add_all_meshes(plotter, meshes)

    # 카메라 초기 위치: bird's eye
    cam_pos, focal, up = compute_camera_for_birdseye(params, 35, 225, 1.2)
    plotter.camera_position = [cam_pos, focal, up]

    plotter.enable_3_lights()
    plotter.show_grid(color="#555")
    plotter.add_axes()

    plotter.add_text(
        params.summary().replace("\n", "  "),
        font_size=10, color="white", position="upper_left",
    )

    plotter.show()


# ---------------------------------------------------------------------------
# AI 렌더링 (Gemini)
# ---------------------------------------------------------------------------

def render_with_gemini(guide_img: Image.Image, prompt: str, api_key: str) -> Image.Image | None:
    """Gemini API로 AI 렌더링. 실패 시 None."""
    try:
        from google import genai
        from google.genai import types
        import io as _io

        client = genai.Client(api_key=api_key)

        # 이미지를 PNG 바이트로 변환
        buf = _io.BytesIO()
        guide_img.save(buf, format="PNG")
        img_bytes = buf.getvalue()

        response = client.models.generate_content(
            model="gemini-2.5-flash-image",
            contents=[
                prompt,
                types.Part.from_bytes(data=img_bytes, mime_type="image/png"),
            ],
        )

        # 응답에서 이미지 추출
        for part in response.candidates[0].content.parts:
            if hasattr(part, "inline_data") and part.inline_data:
                img_data = part.inline_data.data
                return Image.open(_io.BytesIO(img_data))

        print("  Gemini 응답에 이미지 없음")
        return None

    except Exception as e:
        print(f"  Gemini 렌더링 오류: {e}")
        return None


def build_gate_prompt(params: SpillwayParams, time_of_day: str = "daytime") -> str:
    """수문 구조물용 AI 프롬프트."""
    return (
        f"Photorealistic bird's eye view of a dam spillway gate facility, "
        f"{params.n_gates} radial (Tainter) gates each {params.gate_width:.0f}m wide by {params.gate_height:.0f}m tall, "
        f"ogee-profile concrete weir, service bridge on top, "
        f"hoist houses above each gate, {time_of_day}, "
        f"crystal clear water upstream, concrete apron downstream, "
        f"maintain exact structural geometry, layout, and proportions from the input image, "
        f"preserve gate positions and pier locations precisely, "
        f"professional architectural rendering, "
        f"8K ultra sharp detail, high dynamic range, "
        f"realistic concrete texture, steel gate panels, "
        f"bright daylight with sharp shadows, clear blue sky"
    )


# ---------------------------------------------------------------------------
# 메인
# ---------------------------------------------------------------------------

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--plan", default=None, help="Plan DXF (1/2)")
    ap.add_argument("--section", default=None, help="Section DXF (2/2)")
    ap.add_argument("--interactive", action="store_true", help="인터랙티브 3D 뷰어")
    ap.add_argument("--ai", action="store_true", help="AI 렌더링 (Gemini)")
    ap.add_argument("--output", default="gate_render_output", help="출력 디렉토리")
    ap.add_argument("--time", default="daytime", choices=["daytime", "sunset", "overcast"])
    ap.add_argument("--size", type=int, default=1536, help="렌더 해상도")
    args = ap.parse_args()

    # 기본 샘플 경로
    if args.plan is None:
        base = Path("Gate_Sample")
        args.plan = str(base / "12995740-M40-001 여수로 수문 설치도(1／2).dxf")
        args.section = str(base / "12995740-M40-002 여수로 수문 설치도(2／2).dxf")

    # 출력 디렉토리
    out = Path(args.output)
    out.mkdir(exist_ok=True)

    # 1) 파라미터 추출
    print("=" * 60)
    print("Step 1: DXF 파싱")
    print("=" * 60)
    params = parse_spillway_dxf(args.plan, args.section)
    print(params.summary())

    # 2) 3D 모델 빌드
    print("\n" + "=" * 60)
    print("Step 2: 3D 모델 빌드")
    print("=" * 60)
    builder = SpillwayBuilder(params)
    meshes = builder.build_all()
    print(f"{len(meshes)}개 메쉬 컴포넌트 생성")

    # 3) 인터랙티브 모드?
    if args.interactive:
        print("\n인터랙티브 뷰어 실행 중...")
        show_interactive(params, meshes)
        return

    # 4) 다각도 캡처
    print("\n" + "=" * 60)
    print("Step 3: 다각도 캡처")
    print("=" * 60)

    views = [
        ("top_down",     75, 180, 1.0),    # 수직 상부
        ("bird_eye_1",   35, 225, 1.2),    # 조감도 (북동)
        ("bird_eye_2",   35, 135, 1.2),    # 조감도 (북서)
        ("bird_eye_3",   25, 180, 1.3),    # 조감도 (정면)
        ("elevation",    5, 180, 1.1),     # 정면 입면
    ]

    for name, elev, azim, zoom in views:
        print(f"  [{name}] elev={elev}°, azim={azim}°")
        img = capture_view(params, meshes, elev, azim, size=args.size, zoom=zoom)
        img_path = out / f"capture_{name}.png"
        img.save(img_path)
        print(f"    → {img_path}")

    # 5) 제어맵 추출 (bird_eye_1 기준)
    print("\n" + "=" * 60)
    print("Step 4: 제어맵 추출 (bird_eye_1)")
    print("=" * 60)
    main_elev, main_azim, main_zoom = 35, 225, 1.2

    capture = capture_view(params, meshes, main_elev, main_azim, args.size, zoom=main_zoom)
    capture.save(out / "capture_main.png")
    print(f"  capture_main.png")

    depth = capture_depth(params, meshes, main_elev, main_azim, args.size, zoom=main_zoom)
    depth.save(out / "depth_map.png")
    print(f"  depth_map.png")

    lineart = capture_lineart(params, meshes, main_elev, main_azim, args.size, zoom=main_zoom)
    lineart.save(out / "lineart_map.png")
    print(f"  lineart_map.png")

    guide = compose_guide_image(capture, depth, lineart)
    guide.save(out / "guide_composite.png")
    print(f"  guide_composite.png")

    # 6) AI 렌더링 (선택적)
    if args.ai:
        api_key = os.environ.get("GEMINI_API_KEY", "")
        if not api_key:
            print("\n경고: GEMINI_API_KEY 환경변수 필요 (AI 렌더링 건너뜀)")
        else:
            print("\n" + "=" * 60)
            print("Step 5: AI 렌더링 (Gemini)")
            print("=" * 60)
            prompt = build_gate_prompt(params, args.time)
            print(f"  프롬프트: {prompt[:100]}...")

            ai_img = render_with_gemini(guide, prompt, api_key)
            if ai_img:
                ai_img.save(out / "ai_rendered.png")
                print(f"  → ai_rendered.png")
            else:
                print(f"  AI 렌더링 실패")

    print("\n" + "=" * 60)
    print(f"완료. 출력 디렉토리: {out.absolute()}")
    print("=" * 60)


if __name__ == "__main__":
    main()