"""セグメンテーションの評価を手元で通すための小さなダミーの 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())
