The detection contract
A detection is four uint8 values — normalized
xyxy, one byte each. That is the entire on-the-wire output, and it is the
fixed boundary between training, the codec, and the edge.
The arithmetic that makes one byte per coordinate enough:
- 256 levels across
[0, 1]→ at 480 px the quantization step is480 / 255 ≈ 1.88 px. - Worst-case round-trip error is half a step,
≈ 0.94 px— comfortably inside the 3 px budget. - All post-processing stays in 8-bit math: no float boxes to carry around the edge.
This lives in tensor_factory.codec; the canonical box is
tensor_factory.geometry.BBox. The model's ONNX output is a float
(1, 4) tensor; the runtime decodes it to a BBox and the four
bytes. The ~1.9 px localization figure on the landing page is the mock
(exact-geometry) number — on real photoreal data, box localization is the open quality
ceiling (~25 px median), tracked in TODO.md.
The model
TinyDetector is a single-object box regressor with a soft-argmax coordinate
head. The backbone is four stride-2 conv blocks (conv → batchnorm → ReLU) that take the
480 px input down a spatial pyramid:
480 -> 240 -> 120 -> 60 -> 30 # 4 stride-2 blocks, NO global pool channels: 3 -> c -> 2c -> 4c -> 4c # c = --width (default 16)
There is no global pooling on the box path, so spatial information survives to a
30 × 30 feature map. A 1×1 conv makes one heatmap per
box edge (channels x1, y1, x2, y2). Each heatmap is spatially
soft-maxed, and a marginal soft-argmax reads off that edge:
x1from channel 0's x-marginal,x2from channel 2's x-marginal,y1from channel 1's y-marginal,y2from channel 3's y-marginal.
Every coordinate is localized sub-pixel by the same mechanism — there is no width/height
regression, which was the precision bottleneck. The output is normalized
xyxy, tiny enough to int8-quantize and run well above 10 fps (the helicoils
model: ~204 fps @ 480 px, 81 KB).
The soft-argmax gain
A plain softmax of a diffuse or multimodal heatmap has its marginal expectation pulled toward 0.5 — the image centre. On clean synthetic data the peaks are already sharp, so it's invisible; on ambiguous real-world features it is the dominant localization error (the box regresses to the middle).
The fix is a learnable gain — an inverse softmax temperature, stored in
log-space as TinyDetector.log_gain so gain = exp(·) stays
positive. It scales the logits before the spatial softmax: gain > 1 sharpens
the distribution so the expectation tracks the true peak instead of drifting centre-ward.
- Init 0 → gain 1.0 → plain softmax, so existing checkpoints are unaffected and the optimizer is free to anneal it up.
- It is one scalar parameter and exports cleanly to ONNX — no change to the box/presence output contract.
--freeze-gainregisters it as a buffer fixed at 1 (an ablation).
A/B on the combined real dataset (frozen vs learnable, same seed/split): median
validation centre-error 43.4 → 27.4 px, with the gain self-learning to
1.35. Whether it moves the ceiling further is still being validated
(TODO.md).
The presence head
With presence=True, a single YOLO-style objectness head rides alongside the
box head: concatenated global mean and max pooling → a linear layer → one
logit. Mean alone washes out texture (a coil and a smooth bore have near-identical channel
averages); max captures whether a discriminative texture fires anywhere. forward
then returns (box, presence), where sigmoid(presence) is
P(target present) — the runtime thresholds it to decide between one box and no box. There
is no class label and no background class.
Negatives (no-object images) turn the presence head on and train its
objectness toward absent (target 0). Their box loss is masked (a no-object frame
has no box to fit), so they teach the model to report absent rather than emit a
spurious box. That is how detect can return present: false.
Training & checkpoint selection
- Losses: Smooth-L1 on the box, binary cross-entropy on the objectness logit, balanced by
--box-weight/--presence-weight(the box loss is far smaller, so it's weighted up to keep the presence gradient from drowning it). - Determinism: the split is seeded and
torch.manual_seed(seed)fixes weight init, so two runs differ only by hyperparameters — the A/B above is apples-to-apples. - Best-checkpoint export: with
--val-frac, each epoch is scoredaccuracy − err/size(accuracy leads, but boxes count ~50× more than the old scaling), and the best state — not the last, possibly-overfit epoch — is what gets exported. - Box error on positives only: negatives have no box, so a collapsed box can't hide behind presence accuracy.
Quantization & export
The path from a trained graph to the file the edge loads:
- Export on CPU (ONNX export from an MPS graph is unreliable). Legacy TorchScript exporter, opset 17 — a static tiny CNN needs nothing fancier, and it avoids the
onnxscriptdependency the dynamo exporter requires. - Dynamic quantization via onnxruntime: weights to
QUInt8. This is what yields the ~81 KB int8 file. - Fixed output names are the whole contract — the runtime reads each by name regardless of graph order, and the presence of a
presenceoutput is what tells it the model can say absent.
Output names are fixed: image in, box (and presence
for a presence-head model) out, named so the runtime reads each by name regardless of graph order.
A dynamic batch axis is declared on every IO.
Compute & device resolution
The core is dependency-free and CPU-only — geometry, codec, formats, inference. Everything heavy is a sibling package behind an extra:
- Generation is a hosted Gemini API call (no GPU) behind the
geminiextra. - Auto-labeling (GroundingDINO) and training are GPU-heavy and resolve their device cuda → mps → cpu.
So the dev loop runs on this Mac Studio's MPS and scales out to a CUDA box or AWS g5/g6
for bulk runs. The auto-labeler also sets PYTORCH_ENABLE_MPS_FALLBACK=1 so
unimplemented MPS ops fall back to CPU instead of erroring. Inference itself defaults to
the CPUExecutionProvider.
The review gate
AI-generated labels are never trainable until a human has validated them — and this is
enforced in one place, tensor_factory.review, so the rule is identical across
generation, labeling, and training.
Every COCO annotation (and image) carries two extra keys:
| Key | Values | Trainable? |
|---|---|---|
review | pending | no (the default; a missing key normalizes here) |
approved | yes | |
rejected | no | |
source | groundingdino | a guess → starts pending |
human | corrected in Label Studio → approved | |
synthetic_gt | mock generator's exact box → approved on creation |
Untrusted-by-default is the safe failure mode: an unmarked label doesn't sneak into
training. The loaders enforce it (require_review=True); fit
refuses an all-pending dataset with guidance to triage it first. --allow-unreviewed
is the deliberate override — the fast path that trains straight on GroundingDINO labels
after a contact-sheet QC. Negatives are exempt: a known absence is trainable by
construction.
Two serving surfaces, one core
The MCP server and the HTTP endpoint are both thin wrappers over
tensor_factory_mcp.core, so they return byte-identical JSON. The core caches
Detector instances by (model, input_size) (an LRU), so repeated
calls reuse the ONNX session.
| tensor-factory-mcp | tensor-factory-http | |
|---|---|---|
| Transport | stdio (MCP) | HTTP (stdlib http.server) |
| Deps | mcp, pydantic | none beyond the core |
| Default bind | — | 127.0.0.1 |
| Surface | 3 read-only tools | /health, /model_info, /detect |
The HTTP layer maps undecodable bytes to a 400 (not a 500 trace), caps the body at --max-mb (413), and keeps the server version out of its headers — small defensive choices for a thing that takes raw bytes over the wire.
Package layout
A uv workspace of five packages, one lockfile, src layout, each shipping py.typed:
| Package | Deps | Role |
|---|---|---|
tensor-factory | none (infer extra: onnxruntime, numpy, pillow) | geometry, codec, formats, review, inference + CLI |
tensor-factory-synth | pillow (gemini, gpu extras) | generation + GroundingDINO labeling + export |
tensor-factory-train | numpy, onnxruntime (train extra: torch) | the tiny detector → int8 ONNX |
tensor-factory-mcp | tensor-factory[infer], mcp, pydantic | MCP + HTTP serving (bundled models) |
tensor-factory-label | httpx | Label Studio push/pull |
Shared dev tooling and lint/test config live at the root, not per-package: ruff
(line length 100), ty for types, pytest with unit /
integration markers. New libraries drop under packages/<name>/ and
[tool.uv.workspace] members = ["packages/*"] picks them up automatically.