Newer
Older
SmTIAS-LightSim / scripts / show_scene_cutaway.py
"""シーン YAML の筐体を半割りにした 3D 断面図を出力する(幾何の目視確認用).

sm-visualize は閉じたキャビティを外から見るため内部の位置関係が読み取りにくい.
本スクリプトは筐体メッシュを y >= 0 で切り落とし,評価板・LED(実寸の円板 + 向き矢印)・
センサ位置を色分けして描くことで,「LED が壁のどこを向いているか」「板がどこに嵌まるか」を
断面で直接確認できるようにする.

実行方法:
    uv run python scripts/show_scene_cutaway.py \
        --scene configs/_tmp_sphere_radius/scene_r170.yaml --output cutaway_r170.png

    # ウィンドウを開いて回して見る場合は --output を省略する
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

import numpy as np
import pyvista as pv
import trimesh
import yaml

ROOT = Path(__file__).resolve().parents[1]

HOUSING_COLOR = "lightsteelblue"
BOARD_COLOR = "crimson"
LED_COLOR = "gold"
BAFFLE_COLOR = "darkorange"
# 半割りにして内部を見せるのは筐体だけ.評価板とバッフルは丸ごと残す
CLIPPED_MESH_IDS = ("sphere_housing", "cube_housing")


def _load_mesh(path: Path, unit: str) -> pv.PolyData:
    """STL を読み込み mm 単位の PolyData として返す."""
    mesh = trimesh.load(path, force="mesh")
    vertices = np.asarray(mesh.vertices, dtype=np.float64)
    if unit == "m":
        vertices = vertices * 1000.0
    faces = np.asarray(mesh.faces, dtype=np.int64)
    return pv.PolyData(vertices, np.hstack([np.full((len(faces), 1), 3), faces]).ravel())


def _add_led(plotter: pv.Plotter, light: dict, scale: float) -> None:
    """LED を実寸の円板 + 向き矢印として追加する."""
    center = np.asarray(light["center"], dtype=np.float64)
    direction = np.asarray(light.get("direction", [0.0, 0.0, -1.0]), dtype=np.float64)
    direction = direction / np.linalg.norm(direction)
    radius = float(light.get("radius", 2.0))
    disc = pv.Disc(
        center=center, normal=direction, inner=0.0, outer=max(radius, scale * 0.01), c_res=48
    )
    plotter.add_mesh(disc, color=LED_COLOR)
    plotter.add_mesh(
        pv.Arrow(start=center, direction=direction, scale=scale * 0.18), color=LED_COLOR
    )
    plotter.add_point_labels(
        [center], [f"LED z={center[2]:.1f}"], font_size=12, text_color="black",
        shape_opacity=0.4, always_visible=True,
    )


def build_plotter(scene_yaml: Path, off_screen: bool) -> pv.Plotter:
    """シーン YAML から半割り断面のプロッタを組み立てて返す."""
    config = yaml.safe_load(scene_yaml.read_text(encoding="utf-8"))
    plotter = pv.Plotter(off_screen=off_screen, window_size=(1400, 1000))
    plotter.set_background("white")

    extent = 100.0
    for entry in config["geometry"]["meshes"]:
        path = Path(entry["path"])
        if not path.is_absolute():
            path = (scene_yaml.parent / path).resolve()
        mesh = _load_mesh(path, entry.get("unit", "mm"))
        extent = max(extent, float(np.abs(mesh.bounds).max()))
        clipped = entry["id"] in CLIPPED_MESH_IDS
        shown = mesh.clip(normal=(0, 1, 0), origin=(0, 0, 0)) if clipped else mesh
        if clipped:
            color, opacity = HOUSING_COLOR, 1.0
        elif entry["id"] == "baffle":
            color, opacity = BAFFLE_COLOR, 1.0
        else:
            color, opacity = BOARD_COLOR, 0.55
        plotter.add_mesh(shown, color=color, opacity=opacity, show_edges=False,
                         label=entry["id"])

    for light in config.get("lights", []):
        _add_led(plotter, light, extent)

    sensor = config.get("sensor")
    if sensor:
        position = np.asarray(sensor["position"], dtype=np.float64)
        plotter.add_mesh(pv.Sphere(radius=extent * 0.012, center=position), color="royalblue")
        plotter.add_point_labels(
            [position], [f"センサ z={position[2]:.1f}"], font_size=12, text_color="black",
            shape_opacity=0.4, always_visible=True,
        )

    plotter.add_axes()
    plotter.show_grid(color="gray")
    # 筐体は y >= 0 を切り落としているので,+Y 側から覗き込むと内面が見える
    plotter.camera_position = [
        (extent * 1.6, extent * 2.6, extent * 0.7),
        (0.0, 0.0, -extent * 0.55),
        (0.0, 0.0, 1.0),
    ]
    return plotter


def main() -> int:
    """断面図を表示または PNG 保存する."""
    parser = argparse.ArgumentParser(description="シーン筐体の半割り断面を 3D 表示する")
    parser.add_argument("--scene", type=Path, required=True, help="シーン定義 YAML")
    parser.add_argument("--output", type=Path, default=None,
                        help="PNG 保存先(省略時はウィンドウを開く)")
    args = parser.parse_args()

    plotter = build_plotter(args.scene, off_screen=args.output is not None)
    if args.output is None:
        plotter.show()
        return 0
    args.output.parent.mkdir(parents=True, exist_ok=True)
    plotter.screenshot(str(args.output))
    print(f"出力: {args.output}")
    return 0


if __name__ == "__main__":
    sys.exit(main())