Files
DarkRoom/tools/export-seg-model.sh
T
dtourolleandClaude Sonnet 5 e6b01226eb Let the segmentation export script pick its own input size
Comparing a larger model against a larger input size meant re-exporting at
resolutions other than the shipped 640, and the script only ever wrote that
one number. `IMGSZ` is now a second positional argument, defaulted to 640 so
every existing call is unchanged.

The experiment this was built for found bigger input a net loss on its own
merits — yolo26n-seg and yolo26s-seg at 1280 both lost track of large,
frame-filling subjects (a bus's box shrank and its score nearly halved)
in exchange for catching small or partially-occluded ones tiling already
handles. Nothing shipped from it, but the ability to re-run that comparison
is worth keeping.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-08-23 13:16:51 +02:00

65 lines
2.3 KiB
Bash
Executable File

#!/usr/bin/env bash
# Re-export the segmentation model that ships in core/dr-segment/models/.
#
# The .onnx is committed (D14), so this is not part of any build — it exists so
# the committed artefact is reproducible rather than a binary someone once
# produced and nobody can regenerate. Run it when bumping the model.
#
# ./tools/export-seg-model.sh
#
# Requires `uv`. Everything else is fetched into a throwaway venv.
#
# ## Why these export flags
#
# `dynamic=False` is not a default we failed to change: **tract cannot parse
# the dynamic-shape graph at all**, failing shape inference on the neck's
# Concat. A fixed input shape is a hard requirement of the pure-Rust backend
# (see the workspace manifest for why that backend was chosen), and it is what
# makes the tiling option in `semantic.rs` the only route to more resolution.
#
# `imgsz` square rather than a rectangle matched to 3:2: one graph has to
# serve portrait, landscape, square crops and panoramas. A landscape-shaped
# graph trades letterbox waste on 3:2 for worse waste on everything else.
set -euo pipefail
HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO="$(cd "${HERE}/.." && pwd)"
OUT="${REPO}/core/dr-segment/models"
MODEL="${1:-yolo26n-seg}"
IMGSZ="${2:-640}"
WORK="$(mktemp -d)"
trap 'rm -rf "${WORK}"' EXIT
echo "==> exporting ${MODEL} at imgsz=${IMGSZ} in ${WORK}"
cd "${WORK}"
uv venv --python 3.12 venv
VIRTUAL_ENV="${WORK}/venv" uv pip install ultralytics onnx onnxslim
VIRTUAL_ENV="${WORK}/venv" "${WORK}/venv/bin/python" - "${MODEL}" "${IMGSZ}" <<'PY'
import sys, json
from ultralytics import YOLO
name = sys.argv[1]
imgsz = int(sys.argv[2])
m = YOLO(f"{name}.pt")
path = m.export(format="onnx", opset=17, simplify=True, imgsz=imgsz, dynamic=False)
print("ONNX:", path)
# The class names travel with the model rather than being retyped into Rust —
# a hand-copied vocabulary is a silent mismatch waiting to happen when the
# model is bumped.
with open("classes.json", "w") as f:
json.dump([m.names[i] for i in range(len(m.names))], f, indent=1)
print("classes:", len(m.names))
PY
mkdir -p "${OUT}"
cp "${WORK}/${MODEL}.onnx" "${OUT}/${MODEL}.onnx"
cp "${WORK}/classes.json" "${OUT}/${MODEL}.classes.json"
echo "==> wrote:"
ls -la "${OUT}"
echo
echo "Remember: these weights are AGPL-3.0 (see ${OUT}/LICENCE.md)."