tensor·factory

Developer reference

How it works

The design decisions behind the factory: why a detection is four bytes, how a tiny soft-argmax CNN localizes sub-pixel, how it quantizes to int8, and why a human gate sits between auto-labels and training.

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 is 480 / 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:

  • x1 from channel 0's x-marginal, x2 from channel 2's x-marginal,
  • y1 from channel 1's y-marginal, y2 from 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-gain registers it as a buffer fixed at 1 (an ablation).
Measured

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 scored accuracy − 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:

  1. 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 onnxscript dependency the dynamo exporter requires.
  2. Dynamic quantization via onnxruntime: weights to QUInt8. This is what yields the ~81 KB int8 file.
  3. Fixed output names are the whole contract — the runtime reads each by name regardless of graph order, and the presence of a presence output 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 gemini extra.
  • 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:

KeyValuesTrainable?
reviewpendingno (the default; a missing key normalizes here)
approvedyes
rejectedno
sourcegroundingdinoa guess → starts pending
humancorrected in Label Studio → approved
synthetic_gtmock 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-mcptensor-factory-http
Transportstdio (MCP)HTTP (stdlib http.server)
Depsmcp, pydanticnone beyond the core
Default bind127.0.0.1
Surface3 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:

PackageDepsRole
tensor-factorynone (infer extra: onnxruntime, numpy, pillow)geometry, codec, formats, review, inference + CLI
tensor-factory-synthpillow (gemini, gpu extras)generation + GroundingDINO labeling + export
tensor-factory-trainnumpy, onnxruntime (train extra: torch)the tiny detector → int8 ONNX
tensor-factory-mcptensor-factory[infer], mcp, pydanticMCP + HTTP serving (bundled models)
tensor-factory-labelhttpxLabel 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.

Prev← Python API Back toGetting started →