"""セグメンテーションの評価を手元で通すための小さなダミーの ONNX モデルを作る. 本物の ResUNet・検出モデルと**入出力の名前・形状・型だけ**を合わせたもので，合成データ （`scripts/make_synthetic_segmentation.py`）の赤い楕円を前景として返す．精度の評価には使えない． - `ResUNet.onnx`: 入力 `input.1`（`(1, 3, 256, 256)` float32．BGR・/255），出力 `output` （`(1, 2, 256, 256)` float32 のロジット）．1×1 畳み込みで ch0 = k・(R − B − 0.25)，ch1 = −ch0 （赤みの強い画素が前景．ch1 は出力チャンネルの取り違えを検出できるかの確認用） - `Detection_Normal.onnx`: 入力 `image_tensor:0`（`(1, H, W, 3)` uint8．RGB），出力 `detection_boxes:0`（`(1, 2, 4)`．正規化 `[ymin, xmin, ymax, xmax]`）・`detection_scores:0` （`(1, 2)`）．1 つ目は R − B が大きい画素の外接矩形（スコア 0.9），2 つ目は画像全体（スコア 0.3） ONNX のグラフは `onnx`（開発用の依存）で直接組み立てる． 使い方: uv run python scripts/make_dummy_seg_models.py --out data/synthetic_seg/models """ from __future__ import annotations import argparse import logging import sys from pathlib import Path import numpy as np import onnx from onnx import TensorProto, helper, numpy_helper logger = logging.getLogger(__name__) OPSET = 17 IR_VERSION = 9 RESUNET_SIZE = 256 RESUNET_GAIN = 30.0 # ロジットの傾き（シグモイド後に 0/1 に近くなるよう大きめ） RESUNET_OFFSET = 0.25 # R − B がこれを超える画素を前景とする（値域 0〜1） DETECTION_RED_MINUS_BLUE = 60.0 # 検出で前景とみなす R − B（値域 0〜255） DETECTION_SCORES = (0.9, 0.3) def build_resunet() -> onnx.ModelProto: """ダミーの ResUNet を作る. Returns: ONNX モデル（`input.1` `(1, 3, 256, 256)` → `output` `(1, 2, 256, 256)`）． """ # 入力のチャンネルは B・G・R の順（DeepTIAS は OpenCV の BGR のまま渡す）． fg = np.array([-RESUNET_GAIN, 0.0, RESUNET_GAIN], dtype=np.float32) weight = np.stack([fg, -fg]).reshape(2, 3, 1, 1) bias = np.array([-RESUNET_GAIN * RESUNET_OFFSET, RESUNET_GAIN * RESUNET_OFFSET], np.float32) graph = helper.make_graph( [helper.make_node("Conv", ["input.1", "W", "B"], ["output"])], "dummy_resunet", [ helper.make_tensor_value_info( "input.1", TensorProto.FLOAT, [1, 3, RESUNET_SIZE, RESUNET_SIZE] ) ], [ helper.make_tensor_value_info( "output", TensorProto.FLOAT, [1, 2, RESUNET_SIZE, RESUNET_SIZE] ) ], initializer=[numpy_helper.from_array(weight, "W"), numpy_helper.from_array(bias, "B")], ) return _finish(graph) def build_detector() -> onnx.ModelProto: """ダミーの検出モデルを作る. Returns: ONNX モデル（`image_tensor:0` `(1, H, W, 3)` uint8 → `detection_boxes:0` `(1, 2, 4)`・ `detection_scores:0` `(1, 2)`）． """ def const(name: str, value: np.ndarray) -> onnx.NodeProto: return helper.make_node( "Constant", [], [name], value=numpy_helper.from_array(value, name + "_v") ) nodes = [ helper.make_node("Cast", ["image_tensor:0"], ["xf"], to=TensorProto.FLOAT), const("c0", np.array([0], np.int64)), const("c1", np.array([1], np.int64)), const("c2", np.array([2], np.int64)), const("c3", np.array([3], np.int64)), const("ax3", np.array([3], np.int64)), helper.make_node("Slice", ["xf", "c0", "c1", "ax3"], ["red"]), helper.make_node("Slice", ["xf", "c2", "c3", "ax3"], ["blue"]), helper.make_node("Sub", ["red", "blue"], ["rb"]), const("thr", np.array(DETECTION_RED_MINUS_BLUE, np.float32)), helper.make_node("Greater", ["rb", "thr"], ["fg_bool"]), helper.make_node("Cast", ["fg_bool"], ["fg"], to=TensorProto.FLOAT), # 行・列ごとに前景があるか（(1, H)・(1, W)）． helper.make_node("ReduceMax", ["fg"], ["rows"], axes=[2, 3], keepdims=0), helper.make_node("ReduceMax", ["fg"], ["cols"], axes=[1, 3], keepdims=0), helper.make_node("ArgMax", ["rows"], ["ymin_i"], axis=1, keepdims=0), helper.make_node("ArgMax", ["rows"], ["ymax_i"], axis=1, keepdims=0, select_last_index=1), helper.make_node("ArgMax", ["cols"], ["xmin_i"], axis=1, keepdims=0), helper.make_node("ArgMax", ["cols"], ["xmax_i"], axis=1, keepdims=0, select_last_index=1), helper.make_node("Shape", ["image_tensor:0"], ["shape"]), helper.make_node("Gather", ["shape", "c1"], ["h_i"], axis=0), helper.make_node("Gather", ["shape", "c2"], ["w_i"], axis=0), helper.make_node("Cast", ["h_i"], ["h"], to=TensorProto.FLOAT), helper.make_node("Cast", ["w_i"], ["w"], to=TensorProto.FLOAT), const("one", np.array(1.0, np.float32)), ] for name in ("ymin", "ymax", "xmin", "xmax"): nodes.append(helper.make_node("Cast", [f"{name}_i"], [f"{name}_f"], to=TensorProto.FLOAT)) nodes += [ helper.make_node("Add", ["ymax_f", "one"], ["ymax_e"]), helper.make_node("Add", ["xmax_f", "one"], ["xmax_e"]), helper.make_node("Div", ["ymin_f", "h"], ["ymin_n"]), helper.make_node("Div", ["ymax_e", "h"], ["ymax_n"]), helper.make_node("Div", ["xmin_f", "w"], ["xmin_n"]), helper.make_node("Div", ["xmax_e", "w"], ["xmax_n"]), helper.make_node("Concat", ["ymin_n", "xmin_n", "ymax_n", "xmax_n"], ["box1"], axis=0), const("box_shape", np.array([1, 1, 4], np.int64)), helper.make_node("Reshape", ["box1", "box_shape"], ["box1_r"]), const("box2", np.array([[[0.0, 0.0, 1.0, 1.0]]], np.float32)), helper.make_node("Concat", ["box1_r", "box2"], ["detection_boxes:0"], axis=1), const("detection_scores:0", np.array([DETECTION_SCORES], np.float32)), ] graph = helper.make_graph( nodes, "dummy_detection", [helper.make_tensor_value_info("image_tensor:0", TensorProto.UINT8, [1, None, None, 3])], [ helper.make_tensor_value_info("detection_boxes:0", TensorProto.FLOAT, [1, 2, 4]), helper.make_tensor_value_info("detection_scores:0", TensorProto.FLOAT, [1, 2]), ], ) return _finish(graph) def _finish(graph: onnx.GraphProto) -> onnx.ModelProto: model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", OPSET)]) model.ir_version = IR_VERSION onnx.checker.check_model(model) return model def _parse_args(argv: list[str] | None) -> argparse.Namespace: parser = argparse.ArgumentParser(description="ダミーの ResUNet・検出モデル（ONNX）を作る") parser.add_argument( "--out", type=Path, default=Path("data/synthetic_seg/models"), help="出力先（既定: data/synthetic_seg/models）", ) return parser.parse_args(argv) def main(argv: list[str] | None = None) -> int: """エントリポイント. Args: argv: コマンドライン引数（None なら `sys.argv[1:]`）． Returns: 終了コード（成功 0）． """ logging.basicConfig( level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s" ) args = _parse_args(argv) args.out.mkdir(parents=True, exist_ok=True) onnx.save(build_resunet(), args.out / "ResUNet.onnx") onnx.save(build_detector(), args.out / "Detection_Normal.onnx") logger.info("ダミーのモデルを書き出した: %s", args.out) return 0 if __name__ == "__main__": sys.exit(main())