Metadata-Version: 2.5
Name: telekinesis-trackers
Version: 0.1.3
Summary: Deployment visual tracking runtimes
Author-email: Telekinesis <support@telekinesis.ai>
Requires-Python: >=3.11
Requires-Dist: numpy<3,>=1.26
Requires-Dist: onnxruntime<2,>=1.20
Requires-Dist: pillow<13,>=10
Requires-Dist: requests<3,>=2.31
Provides-Extra: cpu
Provides-Extra: dev
Requires-Dist: build<2,>=1.2; extra == 'dev'
Requires-Dist: onnx<2,>=1.17; extra == 'dev'
Requires-Dist: pyright<2,>=1.1.400; extra == 'dev'
Requires-Dist: pytest<10,>=8; extra == 'dev'
Requires-Dist: ruff<1,>=0.11; extra == 'dev'
Provides-Extra: examples
Requires-Dist: loguru<1,>=0.7; extra == 'examples'
Requires-Dist: opencv-python<5,>=4.8; extra == 'examples'
Provides-Extra: gpu
Requires-Dist: cupy-cuda12x[ctk]<15,>=14; extra == 'gpu'
Requires-Dist: onnxruntime-gpu[cuda,cudnn]<1.27,>=1.21; extra == 'gpu'
Provides-Extra: mcbyte
Requires-Dist: supervision<1,>=0.26.1; extra == 'mcbyte'
Requires-Dist: trackers[mask]==2.6.0; extra == 'mcbyte'
Provides-Extra: tapnextpp
Requires-Dist: einops<1,>=0.8; extra == 'tapnextpp'
Description-Content-Type: text/markdown

<div align="center">
  <p>
    <a href="https://telekinesis.ai">
      <img width="100%" src="https://assets.telekinesis.ai/Telekinesis+Banner.png" alt="Telekinesis">
    </a>
  </p>

![Python](https://img.shields.io/badge/python-3.11%2B-blue)
![Runtime](https://img.shields.io/badge/runtime-ONNX%20Runtime%20%7C%20PyTorch-blue)

</div>

<p align="center">
  <a href="https://github.com/telekinesis-ai">GitHub</a>
  &nbsp;&bull;&nbsp;
  <a href="https://www.linkedin.com/company/telekinesis-ai/">LinkedIn</a>
  &nbsp;&bull;&nbsp;
  <a href="https://x.com/telekinesis_ai">X</a>
  &nbsp;&bull;&nbsp;
  <a href="https://discord.gg/7NnQ3bQHqm">Discord</a>
</p>

# Telekinesis Trackers

Telekinesis Trackers provides visual tracking runtimes for robotics and computer
vision applications within the Telekinesis ecosystem. It runs compatible ONNX
model bundles on CPU, with NVIDIA GPU execution available for CUTIE and TAPIR.
TAPNext++ additionally supports its official PyTorch checkpoint runtime.

It includes:

- Multi-object mask propagation with CUTIE
- Point tracking with TAPIR
- Long-term causal point tracking with TAPNext++
- Click-assisted mask initialization with RITM

## Supported Trackers

| Runtime | Task | Execution |
| --- | --- | --- |
| `CutieTracker` | Multi-object mask propagation | CPU / CUDA |
| `MaskInitializer` | RITM click-assisted masks | CPU only |
| `TapirTracker` | Offline or causal point tracking | CPU / CUDA |
| `TapNextPPTracker` | Causal checkpoint point tracking | PyTorch CPU / CUDA |

Installing the GPU extra does not enable GPU execution for the CPU-only adapters.
The table describes implemented runtime paths, not tracking-quality or performance
benchmarks. CUDA hardware tests are opt-in; see [Tests](DEVELOPMENT.md#tests).

Contributor setup, model exporting, S3 publishing, and tests are documented in
[DEVELOPMENT.md](DEVELOPMENT.md).

## Release Model

Telekinesis Trackers is in active development (pre-1.0). APIs and model bundle
contracts may evolve between releases. See [CHANGELOG.md](CHANGELOG.md) for
release notes and [DEVELOPMENT.md](DEVELOPMENT.md#workflow) for the development
and release workflow.

## Installation

Python 3.11 or newer is required. Create an isolated environment before
installing the package. For example, with Conda:

```sh
conda create -n telekinesis-trackers python=3.11
conda activate telekinesis-trackers
```

Install the package from PyPI:

```sh
pip install telekinesis-trackers
```

The base package includes CPU ONNX Runtime, NumPy, Pillow, and requests, so it
can run inference without an extra. `[cpu]` remains a compatibility alias.

### ONNX trackers on NVIDIA GPU

After installing the package, replace CPU ONNX Runtime with the GPU runtime and
install its CUDA libraries and CuPy:

```sh
pip uninstall -y onnxruntime onnxruntime-gpu
pip install "onnxruntime-gpu[cuda,cudnn]>=1.21,<1.27" "cupy-cuda12x[ctk]>=14,<15"
```

Uninstalling both runtime distributions first also repairs environments where
both were installed, since they share the same Python module directory.
`pip install "telekinesis-trackers[gpu]"` alone installs both runtime packages;
if you use that extra, run the replacement commands above afterward.

Because the package declares CPU ONNX Runtime as a dependency, `pip check` will
report `onnxruntime` as missing after this switch. Pip does not recognize the GPU
distribution as its replacement. Upgrading or reinstalling `telekinesis-trackers`
may restore the CPU dependency; repeat the replacement commands afterward.

ONNX GPU execution needs a CUDA 12-compatible NVIDIA driver. The commands above
install the user-space CUDA libraries, and the ONNX Runtime upper bound keeps it
on CUDA 12. See the [ORT requirements](https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html)
and [CuPy installation guide](https://docs.cupy.dev/en/stable/install.html).

The distribution is named `telekinesis-trackers`; import it in Python as
`telekinesis.trackers`. For installation from the GitLab package registry, see
the [package installation guide](DEVELOPMENT.md#usage).

PyTorch, Transformers, CUTIE, TAPIR, and TAPNext++ are not base dependencies.

### TAPNext++ installation

TAPNext++ ONNX bundles need only the base CPU installation or the optional
`telekinesis-trackers[gpu]` installation for CUDA. The same float32 bundle works
on either device; PyTorch and tapnet are needed only to export or use checkpoints.

For the checkpoint backend, TAPNext++ needs PyTorch, torchvision, einops, and a separately installed tapnet
checkout. Tapnet is deliberately not declared as a package dependency because
deployments may use a locally modified version. PyTorch and torchvision are also
deliberately excluded so pip cannot silently select a CPU build. Install the
matching CUDA wheels explicitly before installing the tracker extra. For CUDA
12.8:

```sh
pip install torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cu128
pip install "telekinesis-trackers[tapnextpp]"
```

The `tapnextpp` extra installs only the small backend-independent dependency
einops. It does not install PyTorch, torchvision, tapnet, CUDA, or a checkpoint.

Install the required tapnet checkout separately. For this repository's sibling
layout, keep the locally modified checkout editable:

```sh
pip install -e "../tapnet"
```

Installing `telekinesis-trackers[tapnextpp]` supplies einops because some tapnet
versions import it without declaring it. The tracker uses the PyTorch,
torchvision, and `tapnet` packages currently installed in the active Python
environment.
Use the [PyTorch installation selector](https://pytorch.org/get-started/locally/)
for a different CUDA or CPU build; do not install ONNX Runtime GPU packages
solely for the PyTorch checkpoint backend.

Verify the selected PyTorch runtime before loading the 2.5 GB checkpoint:

```python
import torch

print(torch.__version__)
print(torch.cuda.is_available())
```

## Verify Installation

Check that the package imports successfully (this does not load models or test
ONNX Runtime/CUDA initialization):

```python
from telekinesis import trackers

print("Available CUTIE tracker:", trackers.CutieTracker)
```

## Usage

### Automatic Model Loading

CUTIE, RITM, and TAPIR download their pretrained bundles on first use:

```python
from telekinesis.trackers import CutieTracker, MaskInitializer

cutie = CutieTracker()
ritm = MaskInitializer()
```

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

```text
~/.cache/telekinesis/trackers/cutie/cutie-dynamic-spatial
~/.cache/telekinesis/trackers/ritm/480x864-20
~/.cache/telekinesis/trackers/tapir-online/causal-2
~/.cache/telekinesis/trackers/tapir-offline/offline-2-2frames
```

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.

The CUTIE default downloads `cutie/cutie-dynamic-spatial.zip`. Its separate
cache directory avoids reusing older fixed-spatial CUTIE downloads.

### 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.

```python
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.

CUTIE requests full FP32 CUDA math (`use_tf32=0`) to limit CPU/GPU prediction
drift across memory updates. This can reduce throughput compared with TF32.

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 bundle supports dynamic frame resolutions and object counts.
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.

Variable frame resolutions work with the default bundle. To preserve the input
resolution, disable the optional internal reduction:

```python
tracker = CutieTracker(max_internal_size=-1)
tracker.seed(seed_labels, first_rgb)
labels, alive_fraction = tracker.execute(next_rgb)
```

This bundle accepts square, landscape, and portrait frames across sessions.
`max_internal_size=-1` preserves the input resolution: a 500×700 frame is padded
to 512×704 for inference and returns a 500×700 mask. The default value, 480,
reduces the short side to at most 480 while preserving aspect ratio, then pads
both dimensions to multiples of 16. It does not stretch frames to the export
example. All frames in a session must retain their original resolution; call
`reset_session()` followed by `seed()`, or directly `seed()` with new labels,
to change resolution. `--dynamic-objects` independently enables different seed
object counts. Existing fixed-size bundles retain their shape requirements.
Custom exports also enable both dynamic modes by default; see the
[export instructions](DEVELOPMENT.md#exporting-bundles) for fixed-axis opt-outs.

### RITM Mask Initialization

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

```python
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.

### TAPIR Point Tracking

TAPIR runs on CPU or CUDA. Online (causal) mode processes frames sequentially and
downloads the default bundle on first use:

```python
from telekinesis.trackers import TapirTracker

tracker = TapirTracker(mode="online", device="auto")  # CUDA when available
tracks = tracker.seed(first_rgb, query_points_xy)
tracks = tracker.step(next_rgb)
```

Offline mode processes a complete clip together:

```python
tracker = TapirTracker(mode="offline")
tracks = tracker.track(video_rgb, query_points_xy)
```

Both default bundles require exactly **two query points** and use 256x256
processing. The default offline bundle requires exactly **two frames**; the
online bundle accepts successive frames through `step()`. Switching mode does
not make these dimensions dynamic.

The point-tracking example can collect those query points interactively from an
OpenCV window. Choose the query frame with `--query-frame`, left-click the
required points, and press Enter:

```console
python examples/track_points.py --video clip.mp4
```

Pass `--points 100,80 240,160` to run the same example without the picker.
The example prints frame progress and selects CUDA automatically when ONNX Runtime
exposes it. Use `--device cuda` to require CUDA (and fail instead of falling back),
or `--device cpu --threads N` to tune CPU inference.

The bundles are cached under `~/.cache/telekinesis/trackers/tapir-online/causal-2`
and `~/.cache/telekinesis/trackers/tapir-offline/offline-2-2frames`. Cached bundles
are reused. Pass `model_dir` to use a custom bundle; its mode is inferred unless
you explicitly specify one. An explicit mode must match the bundle. `"causal"`
is accepted as an alias for `"online"`; `tracker.mode` retains the manifest value
`"causal"` or `"offline"`.

### TAPNext++ Point Tracking

Use an exported ONNX bundle on CPU or CUDA:

```python
from telekinesis.trackers import TapNextPPTracker

tracker = TapNextPPTracker("bundles/tapnextpp/512-dynamic", device="cpu")  # or "cuda"
tracks = tracker.seed(first_rgb, query_points_xy)  # any positive number of points
tracks = tracker.step(next_rgb)
tracker.unload()
```

Build the bundle in an export environment with `scripts/requirements-export.txt`
installed. Export runs on CPU, checks every ONNX graph against PyTorch, and
downloads the official checkpoint if needed:

```console
python scripts/prepare_bundles.py --models tapnextpp --points 2 --smoke-test
python examples/track_points.py --tracker tapnextpp --model-dir bundles/tapnextpp/512-dynamic --device cpu --video clip.mp4
```

Point counts are dynamic: reuse one bundle for different counts by calling
`seed()` again. The count stays constant within a tracking session; adding or
removing points requires reseeding. `--points N` sets only the export trace
example. Existing fixed-count bundles still require their original count.
The model input resolution remains fixed. Use `--tapnextpp-resolution 256` for the 256 checkpoint, or
`--tapnextpp-weights path/to/model.ckpt` for local weights. ONNX execution stores
recurrent state as NumPy arrays between calls; CUDA executes the graphs and
transfers state through host memory. `half_precision` and `compile_model` apply
only to checkpoint loading.

`TapNextPPTracker` loads the official PyTorch checkpoint lazily and keeps its
recurrent state on the selected device. The 512 checkpoint is causal and accepts
any nonzero number of query points:

```python
from telekinesis.trackers import TapNextPPTracker

tracker = TapNextPPTracker("tapnextpp_512.ckpt", device="cuda")
tracks = tracker.seed(first_rgb, query_points_xy)
tracks = tracker.step(next_rgb)
tracker.unload()  # Release weights, recurrent state, and cached CUDA memory.
```

Run the interactive video example with:

```console
python examples/track_points.py --tracker tapnextpp --checkpoint tapnextpp_512.ckpt --video clip.mp4
```

The checkpoint is loaded on the first `load()` or `seed()` call. `device="auto"`
uses CUDA when PyTorch exposes it. TAPNext++ returns thresholded visibility from
its official wrapper, so `PointTracks.visibility` and `occlusion` are binary.

## Resources

- [Development guide](DEVELOPMENT.md): contributor setup, model exports, tests, and CI/CD
- [Changelog](CHANGELOG.md): release notes
- [Telekinesis Examples](https://github.com/telekinesis-ai/telekinesis-examples): examples across the Telekinesis ecosystem
- [Telekinesis Documentation](https://docs.telekinesis.ai): SDK documentation

## Support

For issues and questions:

- Open an issue in this repository with the tracker name, runtime extra, and steps to reproduce.
- Contact the Telekinesis development team at [support@telekinesis.ai](mailto:support@telekinesis.ai) or on [Discord](https://discord.com/invite/7NnQ3bQHqm).
