feat: add ecdetseg

[wep21]

• add new segmentation model from https://github.com/Intellindust-AI-Lab/EdgeCrafter
• onnx conversion note: I've replaced grid sample with gahter op due to trt/trt rtx degrade.

"""Dynamic-batch plugin-free fix: replace each GridSample (1:1) with an explicit
gather-bilinear subgraph of identical I/O. Works for any batch N. No padding needed
(per-GridSample native point count)."""
import sys
import numpy as np, onnx
from onnx import TensorProto as T
import onnx_graphsurgeon as gs

src = sys.argv[1] if len(sys.argv) > 1 else "onnx_export/ecdetseg-s-dyn.onnx"
dst = sys.argv[2] if len(sys.argv) > 2 else "onnx_export/ecdetseg-s-gather-dyn.onnx"
graph = gs.import_onnx(onnx.shape_inference.infer_shapes(onnx.load(src)))
uid = [0]
def nm(s): uid[0] += 1; return f"gd_{s}_{uid[0]}"
def N(op, ins, dtype=np.float32, **a):
    o = gs.Variable(nm(op.lower()), dtype=dtype)  # shape=None -> inferred
    graph.nodes.append(gs.Node(op, inputs=list(ins), outputs=[o], attrs=a)); return o
def K(arr): return gs.Constant(nm("c"), np.asarray(arr))

def gather_bilinear(X, grid, H, W):
    """X [NM,D,H,W], grid [NM,LQ,P,2] in [-1,1] -> [NM,D,LQ,P] (align_corners=False, zeros pad)."""
    vflat = N("Reshape", [X, K(np.array([0, 0, -1], np.int64))])         # [NM,D,HW]
    lx = N("Gather", [grid, K(np.int64(0))], axis=3)                      # [NM,LQ,P]
    ly = N("Gather", [grid, K(np.int64(1))], axis=3)
    # grid in [-1,1] -> pixel (align_corners=False): x=(g+1)*0.5*W-0.5
    x = N("Sub", [N("Mul", [N("Add", [lx, K(np.float32(1))]), K(np.float32(0.5 * W))]), K(np.float32(0.5))])
    y = N("Sub", [N("Mul", [N("Add", [ly, K(np.float32(1))]), K(np.float32(0.5 * H))]), K(np.float32(0.5))])
    x0 = N("Floor", [x]); y0 = N("Floor", [y])
    x1 = N("Add", [x0, K(np.float32(1))]); y1 = N("Add", [y0, K(np.float32(1))])
    wx = N("Sub", [x, x0]); wy = N("Sub", [y, y0])
    wx0 = N("Sub", [K(np.float32(1)), wx]); wy0 = N("Sub", [K(np.float32(1)), wy])
    # dynamic expand shape [NM, D, LQ*P]
    sh = N("Shape", [vflat], dtype=np.int64)                             # [NM,D,HW]
    nmd = N("Slice", [sh, K(np.array([0], np.int64)), K(np.array([2], np.int64)), K(np.array([0], np.int64))], dtype=np.int64)
    acc = None
    for cx, wcx in [(x0, wx0), (x1, wx)]:
        for cy, wcy in [(y0, wy0), (y1, wy)]:
            vx = N("And", [N("GreaterOrEqual", [cx, K(np.float32(0))], dtype=bool),
                           N("LessOrEqual", [cx, K(np.float32(W - 1))], dtype=bool)], dtype=bool)
            vy = N("And", [N("GreaterOrEqual", [cy, K(np.float32(0))], dtype=bool),
                           N("LessOrEqual", [cy, K(np.float32(H - 1))], dtype=bool)], dtype=bool)
            valid = N("Cast", [N("And", [vx, vy], dtype=bool)], to=int(T.FLOAT))
            cxc = N("Cast", [N("Clip", [cx, K(np.float32(0)), K(np.float32(W - 1))])], dtype=np.int64, to=int(T.INT64))
            cyc = N("Cast", [N("Clip", [cy, K(np.float32(0)), K(np.float32(H - 1))])], dtype=np.int64, to=int(T.INT64))
            lin = N("Add", [N("Mul", [cyc, K(np.int64(W))], dtype=np.int64), cxc], dtype=np.int64)  # [NM,LQ,P]
            linr = N("Reshape", [lin, K(np.array([0, 1, -1], np.int64))], dtype=np.int64)            # [NM,1,LQ*P]
            shl = N("Shape", [linr], dtype=np.int64)
            lqp = N("Slice", [shl, K(np.array([2], np.int64)), K(np.array([3], np.int64)), K(np.array([0], np.int64))], dtype=np.int64)
            eshape = N("Concat", [nmd, lqp], dtype=np.int64, axis=0)                                  # [NM,D,LQ*P]
            line = N("Expand", [linr, eshape], dtype=np.int64)
            gv = N("GatherElements", [vflat, line], axis=2)                                           # [NM,D,LQ*P]
            gvr = N("Reshape", [gv, K(np.array([0, 0, 300, -1], np.int64))])                          # [NM,D,LQ,P]
            wv = N("Reshape", [N("Mul", [N("Mul", [wcx, wcy]), valid]), K(np.array([0, 1, 300, -1], np.int64))])  # [NM,1,LQ,P]
            term = N("Mul", [gvr, wv])
            acc = term if acc is None else N("Add", [acc, term])
    return acc  # [NM,D,LQ,P]

# replace each GridSample 1:1
LEVELS = [(80, 80), (40, 40), (20, 20)]      # Split_output_{0,1,2}
gss = [n for n in graph.nodes if n.op == "GridSample"]
print("GridSample to replace:", len(gss))
for n in gss:
    X, grid = n.inputs[0], n.inputs[1]
    lvl = int(grid.name.rsplit("_", 1)[-1])  # .../Split_output_{0,1,2}
    H, W = LEVELS[lvl]
    out = gather_bilinear(X, grid, int(H), int(W))
    old = n.outputs[0]
    for c in graph.nodes:
        c.inputs = [out if i is old else i for i in c.inputs]

graph.cleanup().toposort()
m = gs.export_onnx(graph)
m = onnx.shape_inference.infer_shapes(m)
onnx.save(m, dst)
print("saved", dst, "| GridSample left:", sum(1 for n in onnx.load(dst).graph.node if n.op_type == "GridSample"))

[wep21]

LD_LIBRARY_PATH=/home/daisuke/Downloads/TensorRT-RTX-1.4.0.76/lib  cargo run -r -F nvrtx-full --example image-segmentation --features vision,annotator -- --source assets/bus.jpg --confs 0.5 ecdetseg --scale m --dtype fp32 --device nvrtx:0 --processor-device cuda:0

trt rtx fp16 result

• s

- m - l - x

[wep21]

@jamjamjon Could you review this PR?

[jamjamjon]

@jamjamjon Could you review this PR?

@wep21 This is awesome! Thank you for the contribution and for taking the time to optimize the ONNX conversion. The GridSample → Gather replacement for better TensorRT compatibility is particularly helpful.

I've merged the PR. Thanks again for your work!

One quick question: are you planning to add the corresponding pose model in the future? It would be great to have both the segmentation and pose variants available in the project.