Newer
Older
SmTIAS-LightSim / src / sm_tias_lightsim / visualize / projection.py
"""orthographic センサ投影モジュール.

sm-render が出力した 2D 輝度マップ (luminance.npy) を
3D メッシュの各頂点へ射影し、per-vertex scalar 配列を返す純関数群を提供する.

## 投影規約(orthographic センサとの整合)

SPEC_03 と scene/sensors.py の look_at + scale 変換を厳密に再現する.

センサ設定(scene_baseline.yaml):
  - position = (0, 0, -50) mm → m では (0, 0, -0.05)
  - target   = (0, 0, 0)
  - up       = (0, 1, 0)
  - ortho_scale = (50, 52) mm → m では sx=0.05, sy=0.052

look_at 基底の導出:
  - forward d = normalize(target - origin) = (0, 0, 1) (+Z 向き)
  - right   r = normalize(cross(up, d))    = normalize((0,1,0)×(0,0,1)) = (1, 0, 0) (+X 向き)
    ※ Mitsuba look_at は right = cross(up, forward) で計算する
  - 実 up   u = cross(d, r)               = (0,0,1)×(1,0,0) = (0,-1,0)
    ※ ただし Mitsuba は渡した up ベクトルを直接使うため (0,1,0) のまま

to_world = look_at @ scale(sx, sy, 1):
  [sx  0   0   0  ]
  [0   sy  0   0  ]
  [0   0   1  -0.05]
  [0   0   0   1  ]

clip_space → world_space(Z=-0.05 平面):
  world_x = clip_x * sx_m
  world_y = clip_y * sy_m
  world_z = -0.05 (センサ位置、投影面ではなくレイ起点)

画像座標 (row, col) と clip_space の対応(Mitsuba/標準慣例):
  clip_x = (2 * col / W) - 1  … col=0 → -1 (物理 X-), col=W → +1 (物理 X+)
  clip_y = 1 - (2 * row / H)  … row=0 → +1 (物理 Y+), row=H → -1 (物理 Y-)

実データ検証: luminance.npy (spp=1024) で確認済み
  - LED Y-方向 (y=-70.2 mm) → 下部 (row > H/2) が明るい
  - row 大 → clip_y 小 → world_y 小 → 物理 Y- = LED 方向 ✓
  - 左右対称 (LED は X+/-対称配置) ✓

逆変換(world → 画像ピクセル):
  clip_x = world_x / sx_m
  clip_y = world_y / sy_m
  col = (clip_x + 1) / 2 * W
  row = (1 - clip_y) / 2 * H

評価板メッシュの特定:
  scene_config 中で id="evaluation_board" のメッシュを優先的に使用する.
  そのような id がない場合は、全メッシュの Z 座標範囲中心が sensor.target に最も
  近いメッシュを選択する(docstring に記載した決定論的フォールバック).
"""

from __future__ import annotations

import logging

import numpy as np

from sm_tias_lightsim.scene.schema import SceneConfig, SensorConfig

_logger = logging.getLogger(__name__)

# 評価板メッシュのデフォルト id(scene_baseline.yaml に準拠)
_EVAL_BOARD_ID = "evaluation_board"

# YAML mm → m 変換係数
_MM_TO_M = 0.001


def _compute_camera_basis_and_col_row(
    vertices_m: np.ndarray,
    sensor: SensorConfig,
) -> tuple[np.ndarray, np.ndarray, int, int]:
    """orthographic センサの look_at 基底を計算し、頂点の浮動小数点ピクセル座標を返す内部ヘルパ.

    投影規約(本モジュール冒頭 docstring・SPEC_03「画像↔物理の向き」)に準拠し、
    `project_luminance_to_vertices` と `compute_orthographic_uv` が共有する
    カメラ基底計算〜(col_f, row_f) 算出を一か所に集約する(DRY).

    Args:
        vertices_m: 頂点配列 (N_v, 3),m 単位.
        sensor: SceneConfig.sensor(position / target / up / ortho_scale を使用).

    Returns:
        (col_f, row_f, n_cols, n_rows) のタプル.
            col_f: float 列座標 (N_v,),範囲外頂点もクランプ前の生値を返す.
            row_f: float 行座標 (N_v,),同上.
            n_cols: 輝度マップの列数(ortho_scale[0] の比率を表す整数代替ではなく
                この関数内部では不使用;呼び出し側で形状から求める).
            n_rows: 同上.
        実際には (col_f, row_f, clip_x, clip_y) を返す(クランプ前 clip 値は
        範囲外チェックに使うため).

    Raises:
        ValueError: sensor.type が 'orthographic' でない場合,ortho_scale が None の場合,
            または position=target / up ∥ forward のとき.

    Note:
        n_cols / n_rows はこの関数の返り値には含まれない(輝度マップのサイズは
        呼び出し側が持つため).代わりに clip_x, clip_y を返し、呼び出し側で
        画像サイズを掛けてピクセル座標に変換する.
    """
    if sensor.type != "orthographic":
        msg = f"投影は orthographic センサのみサポートします (type={sensor.type!r})"
        raise ValueError(msg)
    if sensor.ortho_scale is None:
        msg = "sensor.ortho_scale が未設定です"
        raise ValueError(msg)

    # ortho_scale mm → m
    sx_m = sensor.ortho_scale[0] * _MM_TO_M
    sy_m = sensor.ortho_scale[1] * _MM_TO_M

    # カメラ look_at 基底の計算(Mitsuba 規約)
    pos_m = np.array(sensor.position, dtype=np.float64) * _MM_TO_M
    target_m = np.array(sensor.target, dtype=np.float64) * _MM_TO_M
    up = np.array(sensor.up, dtype=np.float64)

    forward = target_m - pos_m
    forward_len = np.linalg.norm(forward)
    if forward_len < 1e-9:
        msg = "sensor position と target が一致しています"
        raise ValueError(msg)
    forward = forward / forward_len

    right = np.cross(up, forward)
    right_len = np.linalg.norm(right)
    if right_len < 1e-9:
        msg = "sensor up ベクトルが forward と平行です"
        raise ValueError(msg)
    right = right / right_len

    cam_up = np.cross(forward, right)

    rot_cam_to_world = np.stack([right, cam_up, forward], axis=1)  # (3, 3)

    v_rel = vertices_m - pos_m  # (N_v, 3)
    cam_coords = v_rel @ rot_cam_to_world  # (N_v, 3)

    clip_x = cam_coords[:, 0] / sx_m
    clip_y = cam_coords[:, 1] / sy_m

    return clip_x, clip_y


def find_eval_board_mesh_id(scene_config: SceneConfig) -> str:
    """評価板メッシュの id を scene_config から決定する.

    優先順位:
      1. id == "evaluation_board" のメッシュ
      2. 全メッシュ中で Z 座標範囲中心が sensor.target に最も近いもの
         (評価板は sensor の正面にある)

    Args:
        scene_config: 検証済み SceneConfig.

    Returns:
        評価板として使用するメッシュの id 文字列.

    Raises:
        ValueError: geometry.meshes が空の場合.
    """
    if not scene_config.geometry.meshes:
        msg = "geometry.meshes が空です"
        raise ValueError(msg)

    # 優先 1: id="evaluation_board"
    for mesh_cfg in scene_config.geometry.meshes:
        if mesh_cfg.id == _EVAL_BOARD_ID:
            _logger.debug("evaluation_board mesh found by id")
            return mesh_cfg.id

    # フォールバック: メッシュリストの先頭を使用(合理的デフォルト)
    # ジオメトリ情報なしで Z 中心を計算するには STL 読み込みが必要で、
    # ここでは project_root が不明なため、先頭メッシュを返す.
    # 呼び出し側が project_root を持つ場合は個別に最近傍メッシュを選択可能.
    _logger.warning(
        "id='evaluation_board' が見つかりません.先頭メッシュ '%s' を使用します.",
        scene_config.geometry.meshes[0].id,
    )
    return scene_config.geometry.meshes[0].id


def compute_face_camera_mask(
    vertices_m: np.ndarray,
    faces: np.ndarray,
    sensor: SensorConfig,
    *,
    tol_m: float = 0.001,
) -> np.ndarray:
    """各 face が評価板の「カメラ側(センサ側)平坦評価プレート面」かを判定する bool マスクを返す.

    ## 判定基準(幾何位置ベース・巻き順非依存)

    STL の巻き順は作者依存で信頼できないため、face 法線の内積による判定をやめ、
    face centroid の「カメラ向き主軸方向座標」が評価板メッシュのその端に集中しているか
    で判定する.

    具体的な手順:
      1. sensor の view_dir = normalize(target - position) を計算する.
      2. view_dir の各頂点への射影成分 proj = dot(vertex, view_dir) を求め、
         メッシュ全体での proj_min(センサ側端)と proj_max(反センサ側端)を得る.
         baseline: sensor が -Z から +Z を向く → view_dir ≈ (0,0,1)、
                   proj ≈ world_z → proj_min が評価板の最小 Z 端(Z≈-4.5mm).
      3. 各 face の centroid の proj 値が「proj_min + tol_m 以下」の face を
         「カメラ側平坦評価プレート面」(内側面)として True にする.
         - tol_m(既定 1mm): 平坦面の数値誤差とわずかな厚みを吸収する余裕幅.
         - 評価板の評価プレートは単一平面(SPEC_03: Z=-4.5mm)のため tol_m=1mm は十分大きく、
           かつ外側 AR マーカー等(Z≈-2.5mm 以上)を取り込まない適切な値.
      4. その他の face(外側の楕円窓・AR マーカー・+Z 構造)は False(ニュートラルグレー描画).

    view_dir から向きを決めるため sensor.position / sensor.target のみに依存し、
    STL の巻き順・法線符号には依存しない.

    Args:
        vertices_m: 頂点配列 (N_v, 3),m 単位.
        faces: 三角形面配列 (N_f, 3),頂点インデックス(0-based).
        sensor: SceneConfig.sensor(position / target を使用).
        tol_m: センサ側端から数えた許容幅(m 単位,既定 0.001m = 1mm).
            評価プレートが単一平面なら数値誤差を吸収するだけの小さい値で十分.
            外側構造が近すぎる場合は小さくする(0.0005m = 0.5mm 等).

    Returns:
        bool 配列 (N_f,).True の face が評価対象(カメラ側平坦面・投影着色の対象).
    """
    pos_m = np.array(sensor.position, dtype=np.float64) * _MM_TO_M
    target_m = np.array(sensor.target, dtype=np.float64) * _MM_TO_M

    # カメラ視線方向(センサが向いている方向)
    view_dir = target_m - pos_m
    view_len = np.linalg.norm(view_dir)
    if view_len < 1e-9:
        msg = "sensor position と target が一致しています"
        raise ValueError(msg)
    view_dir = view_dir / view_len  # (3,) 単位ベクトル

    # 全頂点を view_dir へ射影(スカラ値)
    # baseline: view_dir = (0,0,1) のとき proj ≈ world_z
    proj_v = vertices_m @ view_dir  # (N_v,) 各頂点の view_dir 方向成分

    # メッシュ全体の射影範囲
    proj_min = float(proj_v.min())  # センサ側端(baseline では最小 Z 端)

    # face centroid の射影値を計算
    # face centroid = (v0 + v1 + v2) / 3 の view_dir 方向成分 = 各頂点 proj の平均
    proj_centroids = proj_v[faces].mean(axis=1)  # (N_f,)

    # centroid の proj が「センサ側端 + tol_m 以内」= カメラ側平坦評価プレート面
    mask = proj_centroids <= proj_min + tol_m

    n_camera = int(mask.sum())
    n_total = len(faces)
    _logger.debug(
        "face カメラ側端判定 (幾何位置ベース): 評価面=%d / 全 face=%d "
        "(proj_min=%.5f m, tol_m=%.4f m, 閾値=%.5f m)",
        n_camera,
        n_total,
        proj_min,
        tol_m,
        proj_min + tol_m,
    )

    return mask


def project_luminance_to_vertices(
    vertices_m: np.ndarray,
    luminance_map: np.ndarray,
    sensor: SensorConfig,
    *,
    clip_low: float = 1.0,
    clip_high: float = 99.0,
) -> np.ndarray:
    """メッシュ頂点を orthographic センサで 2D 輝度マップに射影し per-vertex 輝度を返す.

    投影規約は本モジュール冒頭の docstring を参照(sensors.py の look_at + scale と
    厳密に整合している).

    Args:
        vertices_m: 頂点配列 (N_v, 3),m 単位.
        luminance_map: 2D 輝度マップ (H, W) または (H, W, C).
            C チャネルある場合は最初のチャネルを使用する.
        sensor: SceneConfig.sensor(orthographic が前提).
        clip_low: percentile クリップの下限(%).既定 1.
        clip_high: percentile クリップの上限(%).既定 99.

    Returns:
        per-vertex 輝度スカラ配列 (N_v,) float32.
        センサ外(範囲外)の頂点は最近傍クリップ(境界値で折り返す)で補完する.

    Raises:
        ValueError: sensor.type が 'orthographic' でない場合,または
            sensor.ortho_scale が None の場合.
    """
    # 2D マップの準備: チャネルがあれば最初のチャネルを使用
    if luminance_map.ndim == 3:
        lum2d = luminance_map[:, :, 0].astype(np.float64)
    else:
        lum2d = luminance_map.astype(np.float64)

    n_rows, n_cols = lum2d.shape

    # カメラ基底計算(_compute_camera_basis_and_col_row と同一ロジック、DRY)
    clip_x, clip_y = _compute_camera_basis_and_col_row(vertices_m, sensor)

    # clip → 画像ピクセル座標
    # clip_x=-1 → col=0 (左端), clip_x=+1 → col=n_cols (右端)
    # clip_y=+1 → row=0 (上端), clip_y=-1 → row=n_rows (下端)
    # ※ Mitsuba は Y-up のため、clip_y と画像 row が反転する
    col_f = (clip_x + 1.0) / 2.0 * n_cols  # float 列座標
    row_f = (1.0 - clip_y) / 2.0 * n_rows  # float 行座標(Y 軸反転)

    # bilinear サンプリング(範囲外は境界でクランプ)
    col_f_clamped = np.clip(col_f, 0.0, n_cols - 1)
    row_f_clamped = np.clip(row_f, 0.0, n_rows - 1)

    col0 = np.floor(col_f_clamped).astype(np.int32)
    row0 = np.floor(row_f_clamped).astype(np.int32)
    col1 = np.clip(col0 + 1, 0, n_cols - 1)
    row1 = np.clip(row0 + 1, 0, n_rows - 1)

    # bilinear 補間の重み
    wc = col_f_clamped - col0.astype(np.float64)
    wr = row_f_clamped - row0.astype(np.float64)

    # 2D → 1D インデックスでサンプリング(4 隣接ピクセル)
    lum_flat = lum2d.ravel()
    idx00 = row0 * n_cols + col0
    idx01 = row0 * n_cols + col1
    idx10 = row1 * n_cols + col0
    idx11 = row1 * n_cols + col1

    sampled = (
        lum_flat[idx00] * (1 - wr) * (1 - wc)
        + lum_flat[idx01] * (1 - wr) * wc
        + lum_flat[idx10] * wr * (1 - wc)
        + lum_flat[idx11] * wr * wc
    )

    # percentile クリップで高ダイナミックレンジ潰れを防止
    p_low = float(np.percentile(sampled, clip_low))
    p_high = float(np.percentile(sampled, clip_high))
    sampled_clipped = np.clip(sampled, p_low, p_high)

    # 0–1 正規化(scalar bar が物理単位でなく相対輝度を表示する場合に使用)
    # → 呼び出し側でカラーマップに渡す際はクリップ後の値そのまま使用する
    #   (scalar bar タイトルで単位を明示する)

    # 範囲外頂点の割合をログ出力
    out_of_range = np.sum((clip_x < -1) | (clip_x > 1) | (clip_y < -1) | (clip_y > 1))
    if out_of_range > 0:
        _logger.debug(
            "投影範囲外の頂点: %d / %d(境界クランプ済み)",
            int(out_of_range),
            len(clip_x),
        )

    return sampled_clipped.astype(np.float32)


def compute_orthographic_uv(
    vertices_m: np.ndarray,
    luminance_shape: tuple[int, int],
    sensor: SensorConfig,
) -> np.ndarray:
    """orthographic センサ投影で各頂点の UV テクスチャ座標 ∈ [0, 1] を返す純関数.

    `project_luminance_to_vertices` と**同一の投影規約**(look_at 基底 +
    ortho_scale mm→m + clip→col/row 変換)で各頂点の浮動小数点ピクセル座標を求め、
    VTK テクスチャ座標(原点: 左下)に変換する.

    ## UV 変換式

    ピクセル座標から UV への変換(_compute_camera_basis_and_col_row と共有):

        col_f = (clip_x + 1) / 2 * W    # 画像座標: 左端=0, 右端=W
        row_f = (1 - clip_y) / 2 * H    # 画像座標: 上端=0, 下端=H(Y 反転)

        u = col_f / W                    # ∈ [0, 1](画像左=u=0, 右=u=1)
        v = 1 - row_f / H               # ∈ [0, 1](VTK 原点は左下=v=0, 上=v=1)

    ## 向きの根拠(SPEC_03「画像↔物理の向き(検証済み)」)

    - 画像右(col 大)= 物理 X+  → u 大 = X+ ✓
    - 画像上(row 小)= 物理 Y+  → v 大(v=1-row/H, row 小で v 大)= Y+ ✓
    - VTK テクスチャ座標は Y 軸が画像とは逆(原点が左下)であるため、
      v = 1 - row_f / H の反転で「画像上 = Y+」と整合する.

    Args:
        vertices_m: 頂点配列 (N_v, 3),m 単位.
        luminance_shape: (H, W) 輝度マップの形状(UV 比率の計算に使用).
        sensor: SceneConfig.sensor(orthographic が前提).

    Returns:
        UV テクスチャ座標 (N_v, 2) float32,各行 (u, v) ∈ [0, 1].
        センサ外(範囲外)の頂点は [0, 1] にクランプされる.

    Raises:
        ValueError: sensor.type が 'orthographic' でない場合,または
            sensor.ortho_scale が None の場合.
    """
    h, w = luminance_shape

    # 共通ヘルパで clip_x, clip_y を計算(DRY: project_luminance_to_vertices と同一基底)
    clip_x, clip_y = _compute_camera_basis_and_col_row(vertices_m, sensor)

    # clip → 正規化ピクセル座標(0〜1)
    # col_f / W = (clip_x + 1) / 2
    u = (clip_x + 1.0) / 2.0
    # row_f / H = (1 - clip_y) / 2  →  v = 1 - row_f/H = (1 + clip_y) / 2
    # VTK テクスチャ座標は原点が左下(画像は左上原点)なので v を反転する
    # v = 1 - (row_f / H) = 1 - (1 - clip_y) / 2 = (1 + clip_y) / 2
    v = (1.0 + clip_y) / 2.0

    # 範囲外をクランプ
    u = np.clip(u, 0.0, 1.0)
    v = np.clip(v, 0.0, 1.0)

    out_of_range = np.sum((clip_x < -1) | (clip_x > 1) | (clip_y < -1) | (clip_y > 1))
    if out_of_range > 0:
        _logger.debug(
            "UV 算出: 投影範囲外の頂点 %d / %d ([0,1] クランプ済み)",
            int(out_of_range),
            len(clip_x),
        )

    uv = np.stack([u, v], axis=1).astype(np.float32)
    _logger.debug(
        "compute_orthographic_uv: shape=%s  u=[%.4f, %.4f]  v=[%.4f, %.4f]",
        uv.shape,
        float(u.min()),
        float(u.max()),
        float(v.min()),
        float(v.max()),
    )
    return uv