Newer
Older
SmTIAS-LightSim / src / sm_tias_lightsim / scene / schema.py
"""YAML シーン定義スキーマ(Pydantic モデル).

SPEC_02_シーン定義仕様 に準拠したモデル群と,YAML ファイルローダを提供する.
"""

from __future__ import annotations

from pathlib import Path
from typing import Annotated, Literal

import yaml
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator

# ---------------------------------------------------------------------------
# 共通基底・型エイリアス
# ---------------------------------------------------------------------------


# 全モデルで frozen=True を共有するため基底クラスに集約する
class _FrozenBase(BaseModel):
    model_config = ConfigDict(frozen=True)


# 3D ベクトル(mm 座標・方向ベクトル・RGB 輝度)の共通型
Vec3 = tuple[float, float, float]


# ---------------------------------------------------------------------------
# Film / Sensor
# ---------------------------------------------------------------------------


class FilmConfig(_FrozenBase):
    """レンダリングフィルム(出力解像度・ピクセルフィルタ)設定."""

    width: int = Field(..., ge=1, description="フィルム幅 [px]")
    height: int = Field(..., ge=1, description="フィルム高さ [px]")
    rfilter: Literal["gaussian", "box", "tent"] = Field(
        default="gaussian", description="ピクセル再構成フィルタ"
    )


class SensorConfig(_FrozenBase):
    """カメラセンサ設定.perspective または orthographic に対応.

    - perspective: 通常のカメラ.fov_deg を必須.
    - orthographic: 平行投影.ortho_scale (mm) を必須(撮影範囲の半幅 (sx, sy)).
      画像座標は X 軸が ±sx,Y 軸が ±sy にマップされる.
    """

    type: Literal["perspective", "orthographic"]
    position: Vec3 = Field(..., description="カメラ位置 [mm]")
    target: Vec3 = Field(..., description="注視点 [mm]")
    up: Vec3 = Field(..., description="上方向ベクトル")
    fov_deg: float | None = Field(
        default=None, description="水平方向 FOV [deg].perspective のみ必須"
    )
    ortho_scale: tuple[float, float] | None = Field(
        default=None,
        description="撮影範囲の半幅 (sx, sy) [mm].orthographic のみ必須",
    )
    film: FilmConfig

    @field_validator("fov_deg")
    @classmethod
    def _validate_fov(cls, v: float | None) -> float | None:
        """FOV が 0 < fov_deg < 180 の範囲内であることを検証する."""
        if v is None:
            return v
        if not (0.0 < v < 180.0):
            msg = f"fov_deg must be in range (0, 180), got {v}"
            raise ValueError(msg)
        return v

    @model_validator(mode="after")
    def _validate_type_specific_fields(self) -> SensorConfig:
        """type に応じた必須フィールドの有無を検証する."""
        if self.type == "perspective" and self.fov_deg is None:
            msg = "perspective sensor requires fov_deg"
            raise ValueError(msg)
        if self.type == "orthographic" and self.ortho_scale is None:
            msg = "orthographic sensor requires ortho_scale (sx, sy) in mm"
            raise ValueError(msg)
        return self


# ---------------------------------------------------------------------------
# Render
# ---------------------------------------------------------------------------


class RenderConfig(_FrozenBase):
    """レンダリングパラメータ設定."""

    variant: str = Field(default="llvm_ad_rgb", description="Mitsuba variant 文字列")
    integrator: Literal["path", "direct", "volpath", "ptracer"] = Field(
        default="path", description="積分器の種類.ptracer はフォワード光線追跡(光源→センサ)"
    )
    max_depth: int = Field(default=8, ge=1, description="最大光線追跡深度")
    spp: int = Field(default=256, ge=1, description="ピクセルあたりのサンプル数")


# ---------------------------------------------------------------------------
# Materials(discriminated union)
# ---------------------------------------------------------------------------


class DiffuseMaterial(_FrozenBase):
    """Lambertian 拡散反射材料."""

    type: Literal["diffuse"]
    reflectance: float | Vec3 = Field(..., description="反射率(スカラまたは RGB タプル)")


class _IorBase(_FrozenBase):
    """屈折率フィールドを持つ誘電体材料の共通基底."""

    int_ior: float = Field(..., description="内側の屈折率")
    ext_ior: float = Field(..., description="外側の屈折率")


class DielectricMaterial(_IorBase):
    """誘電体材料(厚みあり)."""

    type: Literal["dielectric"]


class ThinDielectricMaterial(_IorBase):
    """薄膜誘電体材料(ディフューザー薄板向け)."""

    type: Literal["thindielectric"]


class RoughConductorBSDFConfig(_FrozenBase):
    """マイクロファセット鏡面反射材料(完全導体近似).

    Mitsuba 3 の ``roughconductor`` BSDF に対応する.
    ``material="none"`` で eta/k なし(完全反射体)として扱い,
    内壁の光沢成分を表現する感度分析に使用する.
    """

    type: Literal["roughconductor"]
    alpha: float = Field(
        ..., ge=0.0, le=1.0, description="マイクロファセット粗さ(0=鏡面,1=マット)"
    )
    distribution: Literal["ggx", "beckmann"] = Field(
        default="ggx", description="マイクロファセット法線分布(ggx または beckmann)"
    )


class BlendBSDFConfig(_FrozenBase):
    """2 つの BSDF を線形補間するブレンド材料.

    Mitsuba 3 の ``blendbsdf`` に対応する.
    ``weight=0`` で component[0] のみ,``weight=1`` で component[1] のみ有効.
    内壁の diffuse(拡散)と roughconductor(光沢)を weight でブレンドする
    感度分析に使用する.

    再帰的な入れ子も可能だが,感度分析では 2 要素のフラット構造を想定する.
    """

    type: Literal["blendbsdf"]
    weight: float = Field(
        ...,
        ge=0.0,
        le=1.0,
        description="component[1] の混合比(0=全 component[0],1=全 component[1])",
    )
    components: list[
        Annotated[
            DiffuseMaterial
            | DielectricMaterial
            | ThinDielectricMaterial
            | RoughConductorBSDFConfig
            | BlendBSDFConfig,
            Field(discriminator="type"),
        ]
    ] = Field(
        ...,
        min_length=2,
        max_length=2,
        description="ブレンドする 2 つの BSDF([component0, component1])",
    )


# 再帰型の forward reference を解決する
BlendBSDFConfig.model_rebuild()


# type フィールドで識別される discriminated union
MaterialConfig = Annotated[
    DiffuseMaterial
    | DielectricMaterial
    | ThinDielectricMaterial
    | RoughConductorBSDFConfig
    | BlendBSDFConfig,
    Field(discriminator="type"),
]


# ---------------------------------------------------------------------------
# Lights(discriminated union)
# ---------------------------------------------------------------------------


class PointLight(_FrozenBase):
    """点光源(LED 単体モデル等)."""

    id: str = Field(..., description="シーン内ユニーク ID")
    type: Literal["point"]
    position: Vec3 = Field(..., description="位置 [mm]")
    intensity: float | Vec3 = Field(..., description="強度(スカラまたは RGB タプル)")
    profile: Literal["lambertian", "isotropic"] = Field(..., description="配光プロファイル")


class AreaLight(_FrozenBase):
    """面光源(矩形)."""

    id: str = Field(..., description="シーン内ユニーク ID")
    type: Literal["area"]
    shape: Literal["rectangle"] = Field(..., description="面形状")
    center: Vec3 = Field(..., description="中心位置 [mm]")
    size: tuple[float, float] = Field(..., description="サイズ [mm, mm]")
    direction: Vec3 = Field(..., description="放射方向ベクトル")
    radiance: float | Vec3 = Field(..., description="放射輝度(スカラまたは RGB タプル)")


class DirectionalLight(_FrozenBase):
    """平行光源(無限遠光源)."""

    id: str = Field(..., description="シーン内ユニーク ID")
    type: Literal["directional"]
    direction: Vec3 = Field(..., description="光の進行方向ベクトル")
    irradiance: float | Vec3 = Field(..., description="照度(スカラまたは RGB タプル)")


class SpotLight(_FrozenBase):
    """スポット光源(円錐配光・点光源ベース).

    Mitsuba 3 の `spot` emitter に対応.LED の指向性をデータシートの配光角に
    合わせて表現するために使用する.`cutoff_angle`(光が 0 になる角度)と
    `beam_width`(中央フラット領域)の組み合わせで smoothstep 補間の配光を
    定義する.半値全幅(FWHM) = `(cutoff_angle + beam_width) / 2 × 2`.
    """

    id: str = Field(..., description="シーン内ユニーク ID")
    type: Literal["spot"]
    position: Vec3 = Field(..., description="位置 [mm]")
    direction: Vec3 = Field(..., description="光軸方向ベクトル")
    cutoff_angle: float = Field(..., description="光が完全に 0 になる角度(度,0 < cutoff ≤ 90)")
    beam_width: float = Field(
        ..., description="中央フラット領域の半角(度,0 ≤ beam_width ≤ cutoff_angle)"
    )
    intensity: float | Vec3 = Field(
        ..., description="軸上強度 I_0 [W/sr](スカラまたは RGB タプル)"
    )


# type フィールドで識別される discriminated union
LightConfig = Annotated[
    PointLight | AreaLight | DirectionalLight | SpotLight,
    Field(discriminator="type"),
]


# ---------------------------------------------------------------------------
# Geometry
# ---------------------------------------------------------------------------


class MeshConfig(_FrozenBase):
    """メッシュ(STL)定義."""

    id: str = Field(..., description="シーン内ユニーク ID")
    path: str = Field(..., description="プロジェクトルートからの相対パス(存在チェックは行わない)")
    unit: Literal["mm", "m"] = Field(..., description="モデルの単位系")
    material: str = Field(..., description="materials の dict キー参照")


class GeometryConfig(_FrozenBase):
    """ジオメトリ設定(メッシュ一覧)."""

    meshes: list[MeshConfig] = Field(..., min_length=1, description="メッシュ定義リスト")


# ---------------------------------------------------------------------------
# OETF(光電変換関数)設定
# ---------------------------------------------------------------------------


class OetfExposureConfig(_FrozenBase):
    """OETF 適用時の露出(線形ゲイン)設定."""

    mode: Literal["anchor_mean"] = Field(
        default="anchor_mean",
        description=(
            "露出決定モード.anchor_mean: OETF 後 ROI mean が target_mean に"
            "一致するよう線形ゲインを対数二分で決定する(実測 AE のアンカー)."
        ),
    )
    target_mean: float = Field(
        ...,
        description=(
            "目標 ROI mean(uint8 スケール [0, 255]).実測 ROI mean(例: 146.7146)を指定する."
        ),
    )


class OetfConfig(_FrozenBase):
    """OETF(光電変換関数)順方向適用設定.

    enabled=True のとき sm-evaluate は ROI 抽出後・指標算出前に
    sRGB OETF を適用し,実測 sRGB PNG と同じドメインで評価する.

    後方互換: EvaluationConfig.oetf を省略すると enabled=False(無効)扱いで
    既存の linear 評価が継続される.
    """

    enabled: bool = Field(default=False, description="OETF 順方向適用の有効/無効.")
    transfer: Literal["srgb"] = Field(
        default="srgb",
        description="変換カーブ種別.現状 srgb(IEC 61966-2-1)のみサポート.",
    )
    exposure: OetfExposureConfig | None = Field(
        default=None,
        description="露出設定.enabled=True かつ anchor_mean モードを使う場合は必須.",
    )
    quantize_uint8: bool = Field(
        default=True,
        description=(
            "True のとき OETF 後を np.clip(np.round(x*255),0,255) で"
            "uint8 量子化する(実測 PNG 同等)."
        ),
    )


# ---------------------------------------------------------------------------
# Normalization(linear mean anchor 正規化)設定
# ---------------------------------------------------------------------------


class NormalizationConfig(_FrozenBase):
    """linear ドメイン mean anchor 正規化設定.

    enabled=True のとき sm-evaluate は ROI 抽出後・指標算出前に
    ``apply_mean_anchor`` で mean を target_mean に揃える線形ゲインを適用する.

    線形ゲインは形状指標(CoV / 勾配 / C/P 比 / max-min 比)に不変なので,
    絶対値スケールだけを実測と合わせた上で形状の公平比較が可能になる.

    OETF との併用禁止: oetf.enabled=True かつ normalization.enabled=True は
    EvaluationConfig の model_validator でバリデーションエラーになる.
    (OETF が内部で露出アンカーを持つため,二重アンカーになるのを防ぐ)

    後方互換: EvaluationConfig.normalization を省略すると無効扱いで
    既存の linear 評価が継続される.
    """

    enabled: bool = Field(default=False, description="mean anchor 正規化の有効/無効.")
    type: Literal["mean_anchor"] = Field(
        default="mean_anchor",
        description=(
            "正規化種別.mean_anchor: ROI mean を target_mean に揃える線形ゲイン適用."
            "将来の拡張(max_anchor / range_anchor 等)に備えた Literal 設計."
        ),
    )
    target_mean: float = Field(
        ...,
        description=(
            "目標 ROI mean(uint8 スケール [0, 255] 想定だが任意 float 可)."
            "実測 ROI mean(例: 146.7146)を指定する."
        ),
    )


# ---------------------------------------------------------------------------
# Evaluation
# ---------------------------------------------------------------------------


class RoiConfig(_FrozenBase):
    """評価 ROI 設定(image_pixel 型のみ対応).

    scene_face 型は将来実装予定.
    """

    type: Literal["image_pixel"]
    x: int = Field(..., ge=0, description="ROI 左上 X 座標 [px]")
    y: int = Field(..., ge=0, description="ROI 左上 Y 座標 [px]")
    width: int = Field(..., ge=1, description="ROI 幅 [px]")
    height: int = Field(..., ge=1, description="ROI 高さ [px]")


class ZonesConfig(_FrozenBase):
    """3 ゾーン分析設定(中心 / 中間 / 周辺)."""

    enabled: bool = Field(default=True, description="ゾーン分析の有効/無効")
    boundaries: tuple[float, float] = Field(
        default=(0.33, 0.66),
        description="境界値 (b1, b2).0 < b1 < b2 < 1 を満たす正規化距離",
    )

    @field_validator("boundaries")
    @classmethod
    def _validate_boundaries(cls, v: tuple[float, float]) -> tuple[float, float]:
        """境界値が 0 < b1 < b2 < 1 を満たすことを検証する."""
        b1, b2 = v
        if not (0.0 < b1 < b2 < 1.0):
            msg = f"boundaries must satisfy 0 < b1 < b2 < 1, got ({b1}, {b2})"
            raise ValueError(msg)
        return v


class MetricsConfig(_FrozenBase):
    """算出する均一性指標フラグ設定."""

    cov: bool = Field(default=True, description="変動係数 (CoV)")
    std: bool = Field(default=True, description="標準偏差")
    max_min: bool = Field(default=True, description="最大値/最小値比")
    radial_profile: bool = Field(default=True, description="放射状プロファイル")
    radial_min_max: bool = Field(
        default=False,
        description=(
            "動径 min/max 比(radial_min_max_ratio / radial_min_distance / "
            "radial_min / radial_max の 4 列を metrics.csv に追加)."
            "True のとき radial_profile の有無に関わらず内部で radial_profile を算出する."
        ),
    )
    shape_metrics: bool = Field(
        default=False,
        description=(
            "形状指標を算出する(輝度重心 / 角度-放射分解 / PCA 主軸)."
            "SSIM / NCC は shape_reference が指定されたときのみ算出する."
            "True のとき metrics.csv に 10〜12 列を末尾に追加する."
        ),
    )
    shape_reference: str | None = Field(
        default=None,
        description=(
            "SSIM / NCC 算出用の参照画像パス(.npy または .png)."
            "shape_metrics=True のときのみ有効."
            "省略時(null)は SSIM / NCC の列は空欄になる."
        ),
    )


class EvaluationConfig(_FrozenBase):
    """評価設定(ROI・ゾーン・指標・OETF・正規化).

    oetf フィールドを省略した既存 YAML は後方互換(OETF 無効扱い)で動作する.
    normalization フィールドを省略した既存 YAML も後方互換(正規化無効扱い)で動作する.
    """

    roi: RoiConfig
    zones: ZonesConfig | None = None
    metrics: MetricsConfig | None = None
    oetf: OetfConfig | None = Field(
        default=None,
        description=(
            "OETF 順方向適用設定.省略時は OETF 無効(linear のまま評価)."
            "eval_real.yaml(実測 PNG)には設定しないこと"
            "(実測は既に sRGB 符号化済みのため二重適用になる)."
        ),
    )
    normalization: NormalizationConfig | None = Field(
        default=None,
        description=(
            "linear ドメイン mean anchor 正規化設定.省略時は正規化無効."
            "oetf.enabled=True との併用は禁止(バリデーションエラー)."
        ),
    )

    @model_validator(mode="after")
    def _validate_oetf_normalization_exclusive(self) -> EvaluationConfig:
        """OETF と normalization の同時有効化を禁止する.

        OETF は内部で露出アンカー(mean ベースの線形ゲイン)を持つため,
        normalization との併用は二重アンカーになり意味をなさない.

        Returns:
            検証済みの EvaluationConfig 自身.

        Raises:
            ValueError: oetf.enabled=True かつ normalization.enabled=True の場合.
        """
        oetf_on = self.oetf is not None and self.oetf.enabled
        norm_on = self.normalization is not None and self.normalization.enabled
        if oetf_on and norm_on:
            msg = (
                "oetf.enabled=True と normalization.enabled=True は同時に設定できません."
                "OETF が内部で露出アンカーを持つため,linear mean anchor との併用は"
                "二重アンカーになります.どちらか一方のみ有効にしてください."
            )
            raise ValueError(msg)
        return self


# ---------------------------------------------------------------------------
# Top-level SceneConfig
# ---------------------------------------------------------------------------


class SceneConfig(_FrozenBase):
    """シーン定義全体を表すトップレベルモデル.

    SPEC_02_シーン定義仕様 に準拠した構造を持つ.
    model_validator でメッシュ・光源 ID の一意性,および材料参照整合性を検証する.
    """

    name: str = Field(..., description="設計識別名(出力ディレクトリ名に使用)")
    description: str | None = Field(default=None, description="自由記述")
    geometry: GeometryConfig
    materials: dict[str, MaterialConfig]
    lights: list[LightConfig] = Field(..., min_length=1, description="光源リスト(1 個以上)")
    sensor: SensorConfig
    render: RenderConfig = Field(default_factory=RenderConfig)
    evaluation: EvaluationConfig | None = None

    @model_validator(mode="after")
    def _validate_cross_references(self) -> SceneConfig:
        """メッシュ/光源 ID の一意性と材料参照整合性を検証する.

        検証内容:
            1. geometry.meshes の各 id が重複していないこと.
            2. lights の各 id が重複していないこと.
            3. 各 mesh.material が materials の dict キーに存在すること.

        Returns:
            検証済みの SceneConfig 自身.

        Raises:
            ValueError: 重複 ID または未定義材料参照が見つかった場合.
        """
        # 1. mesh.id ユニーク性チェック
        seen_mesh_ids: set[str] = set()
        for mesh in self.geometry.meshes:
            if mesh.id in seen_mesh_ids:
                msg = f"Duplicate mesh id: '{mesh.id}'"
                raise ValueError(msg)
            seen_mesh_ids.add(mesh.id)

        # 2. light.id ユニーク性チェック
        # LightConfig は discriminated union なので各エントリは必ず id を持つ
        seen_light_ids: set[str] = set()
        for light in self.lights:
            if light.id in seen_light_ids:
                msg = f"Duplicate light id: '{light.id}'"
                raise ValueError(msg)
            seen_light_ids.add(light.id)

        # 3. mesh.material 参照整合性チェック
        material_keys = set(self.materials.keys())
        for mesh in self.geometry.meshes:
            if mesh.material not in material_keys:
                msg = f"mesh '{mesh.id}' references unknown material '{mesh.material}'"
                raise ValueError(msg)

        return self


# ---------------------------------------------------------------------------
# YAML Loader
# ---------------------------------------------------------------------------


def load_scene_yaml(path: str | Path) -> SceneConfig:
    """YAML ファイルを読み込み,検証済み SceneConfig を返す.

    Args:
        path: YAML ファイルのパス.

    Returns:
        検証済みの SceneConfig インスタンス.

    Raises:
        FileNotFoundError: ファイルが存在しない場合.
        pydantic.ValidationError: スキーマ違反.
        yaml.YAMLError: YAML 構文エラー.
    """
    resolved = Path(path)
    if not resolved.exists():
        msg = f"Scene YAML file not found: {resolved}"
        raise FileNotFoundError(msg)

    with resolved.open(encoding="utf-8") as f:
        data = yaml.safe_load(f)

    return SceneConfig.model_validate(data)