auspex

A universal auto-annotator.

Feed it one annotated dataset — CVAT XML, COCO JSON, YOLO, Pascal VOC, or LabelMe — plus the images it describes. It trains up to five per-task models from that single source, then labels unseen images and emits a re-importable CVAT XML and a COCO JSON. No assumptions about domain, labels, or hosting.

Install

pip install auspex-engine

Install as auspex-engine, import as auspex_engine (Python module names can't contain a hyphen, so the underscore form is the import name):

from auspex_engine import Auspex        # train + detect
from auspex_infer import AuspexModel    # detect only (also included)

Requires Python 3.10 / 3.11 / 3.12 on Windows x64, Linux x86_64 / arm64 (manylinux, glibc 2.17+), or macOS Apple Silicon. These are compiled wheels — native binaries with no readable Python source. There is no source distribution, so pip needs a wheel matching your platform; Intel Macs are not supported.

torch >= 2.6.0 is pulled in automatically. If you need a specific CUDA build, install torch first from the PyTorch index — pip then sees it as satisfied and won't replace it:

pip install torch --index-url https://download.pytorch.org/whl/cu121
pip install auspex-engine

Licensing note. auspex-engine itself is proprietary (see LICENSE in the wheel). It depends on Ultralytics, which is AGPL-3.0; using auspex-engine subjects your use of that component to AGPL-3.0 terms. Third-party notices ship in the wheel's .dist-info (THIRD_PARTY_LICENSES).

git clone https://github.com/sohanurislamshuvo/auspex-engine && cd auspex-engine
pip install -e .        # editable — source stays in your working tree; for contributors only

Quick start

Python API

from auspex_engine import Auspex

# TRAIN on your own machine with your dataset
Auspex().train(dataset="ann.xml", images="imgs/", output="runs/exp1", epochs=50)
#   dataset: CVAT XML / COCO JSON / YOLO data.yaml / VOC folder / LabelMe folder
#   -> writes one file: runs/exp1/auspex_model.pt

# DETECT with the trained model
results = Auspex("runs/exp1/auspex_model.pt").predict("test/img.jpg")   # path | folder | ndarray
results.save("out/")            # annotated images + CVAT XML + COCO JSON
results.plot()                  # annotated image (numpy BGR)
results.detections              # list of {type,label,score, box|points|x,y}

CLI

auspex train   --dataset ann.xml --images imgs/ --output runs/exp1 --epochs 50
auspex predict --model runs/exp1/auspex_model.pt --source test/ --output out/
# --tasks bbox,tag to train a subset;  --set KEY=VALUE to override any config knob

Training auto-detects the dataset format, trains only the annotation types present, and bundles every sub-model into one auspex_model.pt (see Single-file model).

Naming the output — by default the bundle is auspex_model.pt. To give each run its own name (handy for an auto-label flywheel that stores one model per id), pass --model-name auspex_lettuce_v3.pt, or the --id shortcut for auspex_<id>.pt:

auspex train --dataset ann.xml --images imgs/ --output runs/exp1 --id lettuce_v3
#   -> runs/exp1/auspex_lettuce_v3.pt

Auspex().train(dataset="ann.xml", images="imgs/", output="runs/exp1", model_id="lettuce_v3")

Orchestrated jobs can set it by env instead: AUSPEX_MODEL_NAME=auspex_${ID}.pt (or AUSPEX_MODEL_ID=${ID}). The name is reduced to a bare filename, so the model always lands inside your output folder.

export DATASET_PATH=/path/to/annotations.xml   # CVAT_XML still works as an alias
export IMG_DIR=/path/to/images
export SAVE_DIR=/path/to/output
python -m auspex_engine.master_train      # trains every task present in the dataset
python -m auspex_engine.master_test       # inference; --images DIR / --output DIR or TEST_IMG_DIR

Train on Google Colab / Kaggle

auspex is cross-platform and auto-uses the GPU (cuda if available, else CPU), so a Colab/Kaggle notebook is a great place to train — Linux + a free GPU, no Windows quirks.

# 1. GPU runtime: Colab -> Runtime -> Change runtime type -> GPU;  Kaggle -> Settings -> Accelerator -> GPU.
#    Kaggle only: Settings -> Internet -> On (needed for pip + pretrained-weight downloads).

# 2. Install (private source repo — use a GitHub token; source visibility is moot in your own notebook)
!git clone https://<TOKEN>@github.com/sohanurislamshuvo/auspex-engine.git
!pip install -q -e ./auspex-engine # or: pip install ./auspex-<ver>-<cpXX>-linux_x86_64.whl

# 3. Train — point at your uploaded / Drive-mounted / Kaggle-attached dataset
from auspex_engine import Auspex
Auspex().train(dataset="/content/data/ann.xml", images="/content/data/images",
               output="/content/runs/exp1", epochs=50)
# -> /content/runs/exp1/auspex_model.pt  (download it or save to Drive)

Gotchas: torch ≥ 2.6 is required (recent Colab/Kaggle images already ship it; older ones get upgraded on install). If a cv2 conflict appears, pip uninstall -y opencv-python opencv-python-headless then reinstall — auspex uses opencv-contrib-python. A compiled wheel's cpXX tag must match the notebook's Python (3.11 today).

Run a trained model without the training code

To let someone run your model without shipping the trainers, give them the one auspex_model.pt plus the lightweight auspex_infer package (a copyable folder, or pip install the auspex-infer wheel built from auspex_infer/):

from auspex_infer import AuspexModel
AuspexModel("auspex_model.pt").predict("img.jpg").save("out/")

auspex_infer has no dependency on the training code — it contains only the model architectures, postprocessing, and export needed to run a model.

Input formats

The dataset format is auto-detected (override with DATASET_FORMAT). Formats that cannot express an annotation type simply contribute none of it — the corresponding model is skipped.

¹ via the non-standard shape_type field (auspex's own COCO output round-trips as input). ² annotations with no geometry keys at all, or a non-standard image-level tags list. ³ LabelMe flags set to true. COCO keypoints use the category's per-keypoint names as classes when defined. Skipped with a counted warning: COCO RLE/iscrowd, malformed/degenerate geometry, YOLO pose lines, LabelMe circle/rotation shapes. VOC segmentation masks are ignored (bbox only).

Environment variables

Every knob above is env-overridable, so a CI or orchestration system can drive per-job runs; standalone runs just use the defaults. Everything else in auspex/config.py is a plain module constant.

Outputs

SAVE_DIR/
  auspex_model.pt                     ← ALL trained models bundled into one file
  data/yolo_bbox/ | yolo_polygon/     YOLO datasets (images are copied, not linked)
  auspex_bb/train/weights/best.pt     YOLO bbox
  auspex_polygon/                     YOLO-seg best.pt, or best_polygon_seg.pt + polygon_classes.json
  auspex_keypoint/best_seg.pt         + kp_classes.json + train_log_kp_seg.csv
  auspex_polyline/best_seg.pt         + train_log_seg.csv
  auspex_tag/best.pt                  + tag_names.json + train_log.csv
  master_train_log.txt
  training_summary.json
  test_output/
    annotated/                        predictions drawn on each image
    predictions_cvat.xml
    predictions_coco.json
    per_image_stats.csv
    polygon_diagnostic.csv
    summary_report.txt

Single-file model

auspex trains up to five independent models, one per annotation type. Every one lands in its own folder above and is bundled into a single portable SAVE_DIR/auspex_model.pt at the end of training — copy or version that one file to move the whole model.

# run inference straight from the bundle (no unpacking needed)
AUSPEX_BUNDLE=auspex_model.pt python -m auspex_engine.master_test --images ./unseen

# or inspect / unpack it back into the per-task layout
python -m auspex_infer.bundle --list   auspex_model.pt
python -m auspex_infer.bundle --unpack auspex_model.pt --into ./models

The bundle embeds each checkpoint's raw bytes, so it is lossless. The per-task files are still written and remain the default for inference; the bundle is used only when AUSPEX_BUNDLE is set.

🔐 Signing bundles (recommended)

A model bundle is executable content. auspex's own loads use weights_only=True, but the bbox/polygon (YOLO) sub-models go through Ultralytics' unrestricted pickle loader — a malicious bundle can execute code when the model is constructed, before predict() runs. weights_only=True does not close that.

Signing does. Sign bundles you produce, and pin the public key when loading: the envelope is read with weights_only=True and its signature is verified before anything is unpacked or handed to a loader, so a tampered or third-party bundle is rejected without ever touching disk. This requires torch ≥ 2.6.0 — on older torch weights_only=True is bypassable (CVE-2025-32434), so auspex refuses to load bundles there rather than give a false guarantee (override with AUSPEX_ALLOW_UNSAFE_TORCH=1 only for bundles you produced yourself).

# once — keep the private key secret, hand out the .pub
auspex keys generate --out mykey            # -> mykey, mykey.pub

# producer: sign at train time (or sign an existing bundle)
auspex train --dataset ann.xml --images imgs/ --output runs/exp1 --sign-key mykey
auspex sign runs/exp1/auspex_model.pt --key mykey

# consumer: pin the public key — refuses anything not signed by it
auspex predict --model auspex_model.pt --source imgs/ --verify-key mykey.pub
auspex verify auspex_model.pt --key mykey.pub

from auspex_engine import Auspex
Auspex().train(dataset="ann.xml", images="imgs/", output="runs/exp1", sign_key="mykey")
Auspex("auspex_model.pt", verify_key="mykey.pub").predict("img.jpg")

The signature covers every embedded checkpoint's bytes plus the metadata, so any tampering — swapping the YOLO weights, injecting a traversal key — breaks it. The public key is pinned by the consumer: an attacker re-signing with their own key is rejected.

auspex predict fails closed: it refuses to load a bundle unless you pass one of --verify-key <key>.pub (verify a trusted signer — recommended), --require-signature (require a valid signature; integrity only, so it rejects a forged/garbage signature block and any post-signing tampering, but does not establish who signed it), or --allow-unsigned (load with no verification — only for bundles you fully trust). The Python AuspexModel/Auspex API stays permissive by default for programmatic use, so pass verify_key= there whenever the bundle is not your own.

Bundle unpacking also rejects path-traversal / absolute / UNC keys (0.1.2), so a tampered bundle cannot write outside its target directory.

Architecture

flowchart TD
  XML[CVAT XML 1.1] --> P["parse_cvat_xml()<br/>parsed ONCE"]
  P --> B[train_bbox<br/>YOLOv8s]
  P --> PG{POLYGON_USE_SEG}
  PG -->|false, default| PY[train_polygon<br/>YOLOv8n-seg]
  PG -->|true| PS[train_polygon_seg<br/>HRNet-W18]
  P --> K[train_keypoint_seg<br/>HRNet + disks]
  P --> PL[train_polyline_seg<br/>HRNet + thick lines]
  P --> T[train_tag<br/>EfficientNet-B0]
  B --> CK[(checkpoints +<br/>sidecar JSON)]
  PY --> CK
  PS --> CK
  K --> CK
  PL --> CK
  T --> CK
  CK --> MT[master_test.py]
  IMGS[unseen images] --> MT
  MT --> OUT[predictions_cvat.xml<br/>predictions_coco.json<br/>annotated/ + CSVs]

Training — master_train.py validates paths, parses the XML once, detects which types are present, then runs five guarded stages. Each stage is wrapped in try/except that records the failure and continues, with _free_gpu() between stages so sequential runs don't accumulate CUDA allocations on small GPUs.

Per-type class lists are derived at the call site with get_labels_for_type(images, TYPE) and passed as an explicit classes= kwarg. The trainers only fall back to cfg.*_CLASSES when that kwarg is absent — nothing mutates the config module at runtime.

Inference — master_test.py loads whichever checkpoints exist on disk, runs each image through every loaded model, then writes the XML, COCO, CSVs and annotated images. All three seg models share one inference shape:

forward → sigmoid → F.interpolate to original resolution → per-class threshold → shape extraction

where shape extraction is findContours + Douglas-Peucker for polygons, connected components + distance NMS for keypoints, and skeletonize + 8-connected trace + Douglas-Peucker for polylines.

Module map

Checkpoint contracts

The three seg checkpoints are self-describing — backbone, class list and input size are embedded in the payload, so inference never depends on the config matching what training used:

{"state_dict": ..., "config": {"backbone", "num_classes", "classes", "input_size", ...},
 "best_miou": ..., "epoch": ...}

The tag path is the exception — it hardcodes TagClassifier and reads its class count from the sidecar JSON rather than the payload.

Why per-class segmentation

Keypoints and polylines were originally single-channel Gaussian-heatmap models trained with MSE, and polylines used a two-stage vertex-heatmap → edge-MLP pipeline. Both trained acceptably and both failed to generalize to unseen images — keypoints missed real points and emitted spurious peaks, polylines never reconnected correctly.

Everything hard was moved to the same recipe: per-class dense masks on an ImageNet-pretrained multi-scale backbone, learned with BCE + soft Dice. On a small dataset that is far more sample-efficient than regressing peak locations — dense supervision gives every pixel a gradient, Dice handles the extreme foreground/background imbalance of thin structures, and the pretrained HRNet carries most of the generalization. Shape recovery is pushed out of the network into deterministic, debuggable postprocessing.

Shared recipe: AdamW, warmup→cosine LR stepped per iteration, AMP autocast + GradScaler, deterministic seeded val split, best checkpoint by val mIoU, per-class IoU logged to CSV every epoch.

The legacy heatmap keypoint trainer, the two-stage polyline trainer and an experimental attention_fusion.py were all removed before this repository's first commit and are not recoverable from this history.

Known issues

Findings from a full read of the codebase. A fix pass has since cleared the tag GPU-crash, the output-format defects (COCO/CVAT), the keypoint-NMS and polygon-score inference bugs, the device/packaging/config-hygiene items, and the stale-checkpoint shadowing. What remains below is open — chiefly the seg-trainer duplication (below). The recent improvement pass fixed the tag-training augmentation and negatives bugs, added COCO co-present-bbox handling and tag input-size persistence, and removed the dead shape-refiner. Tag-model quality still has room (per-class thresholds, pos_weight, best-checkpoint-by-metric — planned).

Medium

• Resolution-dependent seg decode. The seg decoders in auspex_infer.postprocess upsample the model's SxS probability map to the full original resolution and then apply some absolute-pixel thresholds (decode_polygons min_area, plus the polyline MIN_CHAIN_LEN / DP_EPSILON); on very large images those become near-no-ops and the polyline skeleton walk is slow. (eps_ratio and the keypoint NMS are already size-relative.) Decoding on the model grid and scaling points to original coords would fix both — deferred because it changes inference geometry and wants real-data validation.
• tag_names.json is written before training, best.pt only on improvement. A rerun with a different label set that fails early leaves a new names file beside an old checkpoint; load_state_dict then raises a shape mismatch. The checkpoint's own tag_names entry would resolve this but is never read.

Low

• ~730 lines across five files are commented-out mirrors of the pre-XML, COCO-JSON-era implementation: train_tag.py (289/598), split_annotations.py (157/376), coco_to_yolo.py (141/295), train_bbox.py (74/173), train_polygon.py (70/173). _fix_polygon_winding in split_annotations.py is live code with no callers.
• xml_to_yolo calls the global random.seed(), mutating process-wide RNG state, and shutil.copy2s the full image set into both the bbox and polygon dataset directories — three copies on disk.
• master_train.py records per-task failures in training_summary.json but always exits 0, so an calling system cannot distinguish a 5/5 run from a 1/5 run by exit code.
• YOLO and seg trainers use different val-split mechanisms (random.shuffle vs np.random.RandomState), so "the val set" means something different per task.

Designed but never built

• A shared auspex/training_utils.py holding HRNetSegModel, bce_dice_loss, compute_iou, the warmup-cosine lambda, EMA, a mask-aware augmentation pipeline and a deterministic split.
• A tag redesign: stronger backbone read from config, pos_weight from label frequencies, optional asymmetric loss, AMP + warmup + EMA, best-by-macro-F1, and a per-class threshold sweep on val written to tag_thresholds.json. The F1 that is computed today is micro-F1 at a fixed threshold and feeds nothing.
• A fine-tune/resume path (POLY_SEG_RESUME_FROM) for continuing training at a lower LR.

Docker & cloud (AWS / GCP)

auspex ships one env-driven image that both trains and detects, so it drops into local Docker or a cloud orchestrator (AWS ECS / Batch / SageMaker, GCP Vertex AI / GCE / Cloud Run) — the platform just sets env vars and mounts data.

docker build -t auspex:gpu .                                  # GPU (CUDA base)
docker build --build-arg BASE_IMAGE=python:3.11-slim -t auspex:cpu .   # CPU

docker run --rm --gpus all -v "$PWD/data:/data" -v "$PWD/out:/out" auspex:gpu \
  train --dataset /data/ann.xml --images /data/images --output /out --epochs 50

Full build/run recipes, the env-var reference, and per-platform snippets (ECS/Batch/SageMaker, Vertex/GCE/Cloud Run, S3/GCS staging) are in docker/README.md.

Building & distributing

auspex has two build outputs — a compiled per-platform wheel for customers, and a readable pure-Python wheel for internal use. Only the compiled one goes to customers.

Compiled, source-protected (what customers get). Nuitka compiles both packages to native binaries — no readable .py. Platform-specific, so build once per OS × Python (no cross-compile):

pip install nuitka
python setup.py bdist_nuitka                     # -> dist/auspex-<ver>-cp3X-<platform>.whl        (train + detect)
cd auspex_infer && python setup.py bdist_nuitka  # -> dist/auspex_infer-<ver>-cp3X-<platform>.whl  (detect only)

In practice you don't build these by hand: pushing a vX.Y.Z tag runs .github/workflows/build-compiled.yml, which builds both wheels across Windows / Linux / macOS-arm64 × Python 3.10–3.12 and publishes them to the private release repo sohanur-shuvo/auspex-release. Customers then install with gh release download (see Install).

Readable pure-Python wheel (internal / dev only). python -m build produces auspex-<ver>-py3-none-any.whl whose .py is readable in site-packages (like any pip package). Use it for your own environments — never put it in the customer release repo.

Platform caveats. macOS wheels are built on Apple Silicon and are arm64-only (currently tagged universal2, so they will not run on Intel Macs — retag _arm64 or add a macos-13 leg to fix). Linux wheels are plain linux_x86_64 tied to the runner's glibc; build under manylinux + auditwheel repair for broad distro compatibility. Ultralytics (bbox/polygon) is AGPL-3.0 — relevant whenever you distribute or host auspex.

License

Proprietary — Copyright © 2026 Sohanur Islam Shuvo. All rights reserved.

This is not open-source software. No permission is granted to use, copy, modify, or distribute it without prior written consent from the copyright holder. See LICENSE for the full terms. ("Universal" above refers to the range of dataset formats and domains the tool handles — not to any grant of usage rights.)

Third-party components. auspex depends on open-source libraries under their own licenses — see THIRD_PARTY_LICENSES. Note in particular that Ultralytics (the bbox / polygon models) is AGPL-3.0, which carries a source-disclosure obligation when you distribute or network-host auspex; get legal sign-off (or replace/remove Ultralytics) before treating the whole as closed-source.