import onnx

def check_onnx_num_classes(onnx_path):
    model = onnx.load(onnx_path)
    graph = model.graph

    print(f"📂 模型路徑: {onnx_path}")
    print(f"📦 輸出節點總數: {len(graph.output)}")

    for output in graph.output:
        name = output.name
        shape = []
        for dim in output.type.tensor_type.shape.dim:
            if dim.dim_param:
                shape.append(dim.dim_param)
            else:
                shape.append(dim.dim_value)
        print(f"🔎 輸出節點名稱: {name}")
        print(f"   輸出形狀: {shape}")
        if len(shape) == 4:
            num_classes = shape[1]
            print(f"✅ 偵測到類別數: {num_classes}")
            if num_classes == 19:
                print("⚠️ 這是 Cityscapes 預設模型 (19 類)")
            elif num_classes == 4:
                print("✅ 這是你訓練的 GolfDataset 模型 (4 類)")
            else:
                print("❓ 類別數未知，請確認是否正確訓練/轉換模型")

if __name__ == '__main__':
    onnx_path = r'C:\Users\rd_de\kneronstdc\work_dirs\meconfig\latest.onnx'
    check_onnx_num_classes(onnx_path)