Telekinesis trackers

Deployment visual trackers backed by ONNX Runtime.

Contributor setup, model exporting, S3 publishing, and tests are documented in DEVELOPMENT.md.

Install

Python 3.11 or newer is required. Choose one runtime extra (the base package includes NumPy, Pillow, and requests):

pip install ".[cpu]"

# NVIDIA GPU: CUDA 12 runtime libraries, cuDNN 9, and CuPy
pip install ".[gpu]"

# Optional OpenCV examples
pip install ".[cpu,examples]"

# McByte association and detection data dependencies
pip install ".[cpu,mcbyte]"

PyTorch, Transformers, CUTIE, and TAPIR are not runtime dependencies.

Use a fresh environment when switching runtimes, or uninstall both onnxruntime and onnxruntime-gpu before installing the chosen extra. They provide the same Python module and must not be installed together. The GPU extra needs a CUDA 12-compatible NVIDIA driver; it installs the user-space CUDA libraries. Its ONNX Runtime upper bound keeps it on CUDA 12. See the ORT requirements and CuPy installation guide.

Automatic model loading

CUTIE and RITM download their pretrained bundles on first use:

from telekinesis.trackers import CutieTracker, MaskInitializer

cutie = CutieTracker()
ritm = MaskInitializer()

Models are fetched from https://assets.telekinesis.ai/trackers/, validated, and extracted under:

~/.cache/telekinesis/trackers/cutie/480x864-dynamic
~/.cache/telekinesis/trackers/ritm/480x864-20

The cached bundle is reused without another network request. The asset origin, cache location, progress display, and network settings are fixed. Pass model_dir only to use a local bundle instead.

CUTIE mask tracking

Frames are uint8 RGB arrays shaped (H,W,3). Seed labels are integer arrays shaped (H,W) where zero is background and values 1–255 are object IDs.

from telekinesis.trackers import CutieTracker

tracker = CutieTracker()
tracker.seed(seed_labels, first_rgb)
labels, alive_fraction = tracker.execute(next_rgb)

Use CutieTracker(device="cuda") (or "cuda:N") to select a GPU explicitly, and device="cpu" to force CPU execution. The default "auto" selects CUDA when ONNX Runtime advertises it, otherwise CPU. A selected CUDA backend that cannot initialize raises an error; it does not silently run the tracker on CPU. tracker.device reports the selected device.

On CUDA, graph features and recurrent state stay on the GPU through I/O binding; CuPy performs memory attention and output resizing there. Input validation and frame preprocessing remain on CPU; returned labels are NumPy arrays. Existing float32 bundles work without re-exporting. Unsupported graph operators may still use ONNX Runtime's CPU provider.

The pretrained graph has 480x864 padded dimensions and a dynamic object axis. Frames can be reduced with max_internal_size; objects are selected from the IDs present in the seed mask. Use correct(image_rgb, labels) to replace the current mask and memory state, and reset_session() to clear the tracker.

RITM mask initialization

MaskInitializer creates one object's mask from positive and negative (x, y) clicks:

from telekinesis.trackers import MaskInitializer

initializer = MaskInitializer()
probability = initializer.predict_proba(
    image_rgb,
    positive_points=[[320, 200]],
)
mask = initializer.predict(
    image_rgb,
    positive_points=[[320, 200]],
    negative_points=[[20, 20]],
    previous_mask=probability,
)

Images and previous-mask probabilities are resized internally, and results are returned at the original image size. Calls are stateless; pass the preceding probability map when refining the same object. The pretrained bundle accepts up to 20 positive and 20 negative clicks.

Live camera

Install the example dependencies and run:

python examples/track_live.py

# Use RITM click-assisted initialization
python examples/track_live.py --ritm

# After installing .[gpu,examples]
python examples/track_live.py --device cuda

Use 1–9 or n/p to select an object ID. Press Enter after defining the objects, c during tracking to correct the current masks, and q to stop. Left-click adds a positive RITM point and right-click adds a negative point.

Camera frames preserve aspect ratio with letterboxing. Use --resize-mode cover to crop or --resize-mode stretch to distort. CUTIE memory settings can be evaluated with --mem-every, --max-mem-frames, and --top-k.

SAM3

SAM3 requires a compatible local recurrent ONNX bundle:

from telekinesis.trackers import Sam3Tracker

tracker = Sam3Tracker("bundles/sam3/1080x1080-4", device="cpu")
tracker.seed(seed_labels, first_rgb)
labels, alive_fraction = tracker.execute(next_rgb)

The adapter accepts device="auto", "cpu", "cuda", or "cuda:N". The cpu and gpu extras select the ONNX Runtime installation. See DEVELOPMENT.md for the required bundle contract.

Run a video with an existing bundle:

python examples/track_video.py --tracker sam3 \
  --model-dir bundles/sam3/1080x1080-4 \
  --video clip.mp4 --masks object-a.png object-b.png \
  --resize 1080 1080 --out tracked.mp4

TAPIR point tracking

from telekinesis.trackers import TapirTracker

tracker = TapirTracker("bundles/tapir/causal-2")
tracks = tracker.seed(first_rgb, query_points_xy)
tracks = tracker.step(next_rgb)

Bundle mode, frame count, resolution, and query count are fixed in its manifest.

McByte association

from telekinesis.trackers import McByteTracker

tracker = McByteTracker(
    "bundles/mcbyte/sam/vit-b",
    "bundles/mcbyte/cutie/480x864",
)
tracked_detections = tracker.update(detections, frame_rgb)

McByte does not include a detector. Supply supervision.Detections on every frame. Its CUTIE bundle must use the dynamic-object graph layout.