treemble-1

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