John-Allard/treemble-1

Treemble Internal Node Detector (ResNet-34 U-Net, ONNX)

This repository hosts an ONNX-exported heatmap model used by Treemble, a cross-platform desktop tool for digitizing phylogenetic tree figures.

Treemble lets you load a tree image, mark or edit nodes, and export clean tree representations. To learn more about Treemble, download installers, or browse documentation, visit:

• https://www.treemble.org

Model overview

• Architecture: ResNet-34 encoder + U-Net style decoder (fully convolutional)
• Format: ONNX (exported from PyTorch)
• Input channels: 1 (grayscale)
• Input shape: [1, 1, H, W] with dynamic H and W (height/width)
• Output: Heatmap logits [1, C, H, W]
  • For this checkpoint: C = 1 (internal nodes only)
• Typical usage: Given a crop of a tree image, the model produces a dense probability map where peaks correspond to internal node locations.

The file model.config.json stored alongside model.onnx encodes key preprocessing and decoding parameters (e.g. resize/max_side, padding multiple, and peak-detection thresholds) so downstream applications can reproduce the training-time inference behavior consistently.

Files

The model repository contains:

• model.onnx – ONNX graph for the node detector.
• model.config.json – JSON configuration with fields such as:
  • backbone (e.g. "resnet34")
  • in_channels (1 for grayscale)
  • 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 uses both files together: model.onnx for the network and model.config.json for the preprocessing / peak-decoding logic.

Intended use

Primary use case

• Automatic detection of internal nodes (and optionally root) in tree images inside the Treemble desktop app.
• The user selects a rectangular region covering the tree; Treemble:
  1. Crops and preprocesses that region (grayscale, resize, pad).
  2. Runs this ONNX model locally via ONNX Runtime.
  3. Decodes peaks into node coordinates.
  4. Maps coordinates back to the full-image coordinate system and inserts them as editable nodes.

Non-goals

• The model is not a general-purpose object detector.
• It is specialized for tree-figure imagery (e.g. phylogenetic-style plots) and may not perform well on unrelated domains or very different visual styles.

Training data (high-level)

The model was trained on a collection of tree figures with manually annotated node locations (tip, internal, root). Each training sample consists of:

• A raster image (tree figure).
• A CSV with node coordinates and types.

For this particular checkpoint:

• The model was trained in internal-only mode, where roots and internals are treated as a single class during training.
• Tip node detection is handled separately inside Treemble and is not provided by this ONNX model.

Exact datasets and detailed annotations are not distributed with this model; it is intended primarily for Treemble’s internal node detection workflow.

Inference example (Python + onnxruntime)

Below is a minimal example showing how you might call this model from Python with onnxruntime to obtain approximate internal node locations from a grayscale crop of a tree image.

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])

For production use, applications should implement non-maximum suppression and top-k limiting as encoded in the decode section of the config.

License and usage

This model is released under the MIT License. If you use it in scientific work or downstream tools, please also acknowledge Treemble and its authors where appropriate.

If you use Treemble in your work, please cite:

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

For more information about Treemble itself, see:

• https://www.treemble.org