Treemble 内部节点检测器 (ResNet-34 U-Net, ONNX)

本仓库托管了 Treemble 使用的 ONNX 导出热力图模型。Treemble 是一个跨平台桌面工具，用于数字化系统发育树图形。

Treemble 允许您加载树形图像、标记节点，并导出机器可读的树表示（newick 文件）。要了解更多关于 Treemble 的信息、下载安装程序或浏览文档，请访问：

• https://www.treemble.org

模型概述

• 架构： ResNet-34 编码器 + U-Net 风格解码器（全卷积）
• 格式： ONNX（从 PyTorch 导出）
• 输入通道： 1（灰度）
• 输入形状： [1, 1, H, W]，其中 H 和 W（高度/宽度）为动态值
• 输出： 热力图 logits [1, C, H, W]
  • 对于此检查点：C = 1（仅内部节点）
• 典型用途： 给定树形图像的裁剪区域，模型生成一个密集概率图，其中峰值对应于内部节点位置。

与 model.onnx 一起存储的 model.config.json 文件编码了关键的预处理和解码参数（例如 resize/max_side、padding multiple 和峰值检测阈值），以便下游应用程序能够一致地重现训练时的推理行为。

文件

模型仓库包含：

• model.onnx – 节点检测器的 ONNX 计算图。
• model.config.json – 包含以下字段的 JSON 配置：
  • backbone（例如 \"resnet34\"）
  • in_channels（1 表示灰度）
  • internal_only / no_root_pred
  • normalize_input
  • max_side、pad_multiple
  • decode：
    • thresh
    • window
    • per_channel_topk
    • max_peaks
    • fallback_topk
    • assign_root_leftmost_when_two_classes

Treemble 同时使用这两个文件：model.onnx 用于网络，model.config.json 用于预处理/峰值解码逻辑。

预期用途

主要用例

• 在 Treemble 桌面应用程序内自动检测树形图像中的内部节点（以及可选的根节点）。
• 用户选择一个覆盖树的矩形区域；Treemble：
  1. 裁剪并预处理该区域（灰度、调整大小、填充）。
  2. 通过 ONNX Runtime 在本地运行此 ONNX 模型。
  3. 将峰值解码为节点坐标。
  4. 将坐标映射回全图坐标系并将其作为可编辑节点插入。

非目标

• 该模型不是通用对象检测器。
• 它专门用于矩形系统发育树图（根在左侧），在无关领域可能表现不佳。

训练数据（概述）

该模型在带有手动标注节点位置（末端、内部、根）的树形图形集合上进行训练。每个训练样本包括：

• 栅格图像（树形图）。
• 包含节点坐标和类型的 CSV 文件。

对于此特定检查点：

• 模型在仅内部模式下训练，其中根节点和内部节点在训练期间被视为单一类别。
• 末端节点检测由 Treemble 内部单独处理，不由此 ONNX 模型提供。

确切的数据集和详细注释不随此模型分发；它主要用于 Treemble 的内部节点检测工作流程。

推理示例（Python + onnxruntime）

以下是一个最小示例，展示如何使用 onnxruntime 从 Python 调用此模型，从树形图像的灰度裁剪中获取近似的内部节点位置。

import json
import numpy as np
from PIL import Image
import onnxruntime as ort

# Load config
with open("model.config.json", "r") as f:
  cfg = json.load(f)

max_side = cfg["max_side"]
pad_multiple = cfg["pad_multiple"]

# Load and preprocess image crop
img = Image.open("tree_crop.png").convert("L")
w, h = img.size

# Resize with aspect ratio preserved
scale = min(1.0, max_side / max(w, h))
new_w = max(1, int(round(w * scale)))
new_h = max(1, int(round(h * scale)))
img_resized = img.resize((new_w, new_h), resample=Image.BILINEAR)

# Pad so dimensions are multiples of pad_multiple
pad_w = (pad_multiple - (new_w % pad_multiple)) % pad_multiple
pad_h = (pad_multiple - (new_h % pad_multiple)) % pad_multiple
canvas = Image.new("L", (new_w + pad_w, new_h + pad_h), 0)
canvas.paste(img_resized, (0, 0))

arr = np.array(canvas, dtype=np.float32) / 255.0
input_tensor = arr[None, None, :, :]  # [1,1,H,W]

# Run ONNX inference (CPU)
sess = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])
logits, = sess.run(None, {"input": input_tensor})

# logits: [1, C, H, W] where C=1 for internal heatmap
logits = logits[0, 0]
prob = 1.0 / (1.0 + np.exp(-logits))

# Simple thresholding (Treemble uses more advanced NMS/peak limiting)
thresh = cfg["decode"]["thresh"]
ys, xs = np.where(prob > thresh)

# Map back to original crop coordinates
sx = new_w / w
sy = new_h / h
coords = [((x / sx), (y / sy)) for x, y in zip(xs, ys)]

print("Detected internal-node candidates (sample):", coords[:20])

对于生产用途，应用程序应实现配置文件 decode 部分中编码的非极大值抑制和 top-k 限制。

许可证和使用

本模型以 MIT 许可证 发布。如果您在科学工作或下游工具中使用它，请在适当的地方同时致谢 Treemble 及其作者。

如果您在工作中使用 Treemble，请引用：

Treemble: A Graphical Tool to Generate Newick Strings from Phylogenetic Tree Images.
John B. Allard and Sudhir Kumar (2025).
arXiv: 2508.07081. DOI: 10.48550/arXiv.2508.07081

有关 Treemble 本身的更多信息，请参阅：

• https://www.treemble.org