Auspex — from-scratch multi-task vision model

A single PyTorch model — implemented entirely from scratch, random init, no pretrained weights, no model libraries — covering all five annotation types used in labeling work:

Backbone: a multi-resolution trunk that keeps parallel branches at strides 4/8/16/32 (channels 18/36/72/144) and repeatedly fuses them, so high-resolution detail survives to the last stage — ideal for masks, keypoints, and thin polylines. Every task is independently toggleable in config; instance tasks ride on the detection branch. Trains on COCO JSON + CVAT-for-images 1.1 XML + Pascal VOC XML, mixed in one run, with per-source provides: task masking.

Install & use

Distributed on PyPI as compiled binaries (see License) under the name auspex-ai — the import name is auspex:

pip install auspex-ai

Current wheel coverage: Python 3.11 on Windows x64 (more platforms via the CI wheel matrix as needed). Wheels can also be installed directly from a file: pip install auspex_ai-<version>-<platform>.whl.

from auspex import Auspex

model = Auspex(labelspace="labelspace.yaml")
model.train(data="data.yaml", epochs=150)          # train on your own data
results = model.predict("images/", save=True)      # detect

model = Auspex(weights="best.pt")                  # checkpoints are self-contained
model.val(data="data.yaml")

Same via the console command: auspex train --data data.yaml --labelspace labelspace.yaml, auspex predict --weights best.pt --source images/, auspex val ....

Owner-only: python build_protected.py produces the compiled distribution wheel (requires cython + a C compiler; wheels are per-Python/per-platform).

Quickstart (synthetic sanity run, from the repo)

pip install -r requirements.txt
python tools/make_synth_data.py --out runs/synth --n 8   # images + CVAT/COCO annotations
python tools/check_dataset.py --data runs/synth/data_synth.yaml --labelspace configs/labelspace.yaml
python tools/visualize_batch.py --data runs/synth/data_synth.yaml --labelspace configs/labelspace.yaml
python tools/measure_vram.py                              # pick micro-batch for your GPU
python tools/overfit8.py --tasks all                      # THE gate: must print GATE PASSED
python tools/train.py --data runs/synth/data_synth.yaml --epochs 50 --run-dir runs/exp1
python tools/val.py --weights runs/exp1/best.pt --data runs/synth/data_synth.yaml
python tools/predict.py --weights runs/exp1/best.pt --source runs/synth/images --out runs/predict

Training on your own data

1. Write configs/labelspace.yaml: canonical categories (order = class id), aliases, keypoint names + flip_pairs, optional kpt_sigmas, is_polyline: true for polyline classes, and image-level tags.
2. Write a data.yaml:

sources:
  - name: batch1
    path: annotations/batch1.xml        # CVAT 1.1 XML | COCO .json | VOC xml dir
    images_root: images/
    split: train
    provides: {rect: true, polygon: true, keypoint: false, polyline: true, tag: true}

provides is explicit and required — it declares which tasks that source actually labels. An image with zero instances of a provided task is a valid negative; a task not provided contributes nothing to that task's loss, not even background.

3. python tools/check_dataset.py --data data.yaml --labelspace configs/labelspace.yaml
4. Run overfit8.py against a small slice, then train.py.

Training from random init: plan for 150–300 epochs on small/medium datasets; accuracy scales with data since nothing is pretrained.

Layout

configs/    labelspace.yaml (label space), train.yaml overrides
auspex/
  data/     schema, labelspace, parsers (coco/cvat/voc), letterbox, augment,
            rasterize, dataset, collate, sources
  models/   backbone, neck, heads/ (detect, segment, instance, classify), model
  losses/   assigner (TAL), losses (multi-task criterion)
  engine/   trainer, ema, optim, schedule, checkpoint, validator
  metrics/  det_map, mask_map, keypoints, polyline, tags
  infer/    decode, predict (postprocess), polygonize
tools/      train, val, predict, overfit8, visualize_batch, measure_vram,
            check_dataset, make_synth_data, crosscheck_map
tests/      pytest suites incl. the trap-checklist invariants

The overfit-8 gate

tools/overfit8.py trains on 8 images (aug/EMA off) and demands near-perfect reproduction per task. Any task that cannot overfit 8 images has a target-building or loss bug — real training must not start until the gate passes. It is kept forever as a regression test.

Growing the model (new classes, new datasets, scale)

Neural networks forget what they stop seeing ("catastrophic forgetting") — so growing Auspex always means training on the combined data, old + new. Three built-in mechanisms make that scale:

• --transfer <ckpt> — start a new run from an old checkpoint whose labelspace has grown or been reordered. Weights carry over; class/tag head rows are mapped by name (names are stored in every checkpoint); brand-new classes start fresh. Use --resume only for continuing an identical run.
• --source-balance 0.5 — when sources are very unequal (500 apples vs 100k bananas), balance how often each source appears per batch (0 = natural frequency, 1 = every source equally likely).
• --cache-records — serve annotations from a packed on-disk cache (<data.yaml dir>/.auspex_cache/). RAM stays tiny regardless of dataset size, warm starts skip parsing entirely, and Windows workers stop paying N× memory. Turn on past ~100k images; caches auto-rebuild when an annotation file changes.

Scale guide:

Tip: list every class you expect to need in labelspace.yaml from day one — unused classes cost almost nothing and --resume then works with no surgery.

Where it runs

Plain PyTorch — nothing platform- or GPU-specific in the model. The trainer auto-picks the precision the hardware supports (bf16 → fp16 → fp32), overridable with --amp {bf16,fp16,off}.

Limits: single-GPU training only for now (no DDP — a v2 item; on multi-GPU instances one GPU is used). See requirements.txt about the torch pin — it is for the original dev box only.

License

Proprietary — licensed, not open source. The software and its architecture remain the exclusive property of the copyright holder. Licensed users may install the binary distribution, train models on their own data, and use the resulting weights and predictions for their own business; copying, redistribution, modification, and reverse engineering are prohibited. See LICENSE.

Hardware notes (8 GB VRAM)

bf16 autocast (no GradScaler), micro-batch 8 × accumulation 4 (effective 32) at 640 by default — measured 5.35 GB peak on an RTX 5060. On a smaller GPU run tools/measure_vram.py and drop to micro-batch 4 × accum 8. Pressure valves: gradient checkpointing around stage 3, then 512 input. Do NOT run this model in fp32 at batch 8/640 on an 8 GB card: it exceeds VRAM and Windows silently pages GPU memory over PCIe (~12× slowdown at 100% reported GPU utilization). (PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True is unsupported on Windows and intentionally not set.)