Python API¶

spintrack.track runs what spintrack run does and returns the records instead of writing files:

import spintrack

track = spintrack.track("examples/sample/config.toml")  # the run, without the files
df = track.to_polars()  # columns by name: forward_total, heading, ...
print(track.quality.n_dropped, df.select("forward_total", "heading").tail())

For a live camera, feed frames from your own capture code to a Tracker. The config must already describe the ball, the field of view and the camera position, so track a recorded clip of the rig first and use the config.toml the run writes.

from spintrack import Config, Tracker

cfg = Config.load("config.toml")
tracker = Tracker(cfg, width, height)
for gray, ts_ms in frames():  # your camera SDK: 2-D uint8 images, timestamps in ms
    res = tracker.process_frame(gray, ts_ms)
    if res is not None:  # None: the frame could not be tracked
        print(res.forward, res.side, res.turn)  # rad this frame; side, turn: left +

Reference¶

Track the recording a config describes and return its records.

config is a config file, a Config or None (the defaults); src overrides its video. The run is that of spintrack run, with the same refusals and the same geometry found in the recording, minus the files.

frames opens the frames to track instead of the video, for a caller that decodes them itself: it is called once, twice with two_pass, and the records count its frames from 0. Nothing is found in the recording then, so the config must give the ball, the field of view and the camera position.

Source code in python/spintrack/pipeline.py

def track(
    config: Config | str | Path | None,
    *,
    src: str | int | None = None,
    frames: Callable[[], FrameSource] | None = None,
    two_pass: bool = False,
    max_frames: int | None = None,
) -> Track:
    """Track the recording a config describes and return its records.

    `config` is a config file, a `Config` or None (the defaults); `src` overrides its
    `video`. The run is that of `spintrack run`, with the same refusals and the same
    geometry found in the recording, minus the files.

    `frames` opens the frames to track instead of the video, for a caller that decodes
    them itself: it is called once, twice with `two_pass`, and the records count its
    frames from 0. Nothing is found in the recording then, so the config must give the
    ball, the field of view and the camera position.
    """
    from spintrack.autofit import complete_config

    if frames is None:
        cfg, spec = open_config(config, src, two_pass=two_pass)
        complete_config(cfg, spec)
        frames = functools.partial(open_source, spec)
    else:
        cfg = _load(config)
        camera = cfg.camera
        if not (cfg.ball.rim and camera.vfov_deg and camera.to_animal() is not None):
            raise ValueError(
                "tracking frames rather than a video needs the ball (ball.rim), the "
                "field of view (camera.vfov_deg) and the camera position in the config"
            )
    rows = _Rows()
    source = frames()
    try:
        stats = run(
            cfg,
            source,
            [rows],
            two_pass_source=frames if two_pass else None,
            max_frames=max_frames,
        )
    finally:
        source.close()
    return Track(with_ball(np.array(rows), rows.ball), stats.quality)

What track returns: the records, as in tracks.parquet, and quality.

Source code in python/spintrack/pipeline.py

@dataclass
class Track:
    """What `track` returns: the records, as in `tracks.parquet`, and quality."""

    records: np.ndarray  # (n, 29), one row per tracked frame; `ball_seen` as 0 or 1
    quality: RunQuality | None
    columns: tuple[str, ...] = TABLE_COLUMNS

    def to_polars(self) -> pl.DataFrame:
        """The records as a DataFrame with named, typed columns."""
        from spintrack.io.parquet import to_frame

        return to_frame(self.records)

to_polars ¶

to_polars() -> pl.DataFrame

The records as a DataFrame with named, typed columns.

Source code in python/spintrack/pipeline.py

def to_polars(self) -> pl.DataFrame:
    """The records as a DataFrame with named, typed columns."""
    from spintrack.io.parquet import to_frame

    return to_frame(self.records)

Bases: _Table

A rig and, through video, a recording of it.

Source code in python/spintrack/config.py

class Config(_Table):
    """A rig and, through `video`, a recording of it."""

    video: str | int | None = None  # a path, or a camera index
    camera: CameraConfig = Field(default_factory=CameraConfig)
    ball: BallConfig = Field(default_factory=BallConfig)
    mask: MaskConfig = Field(default_factory=MaskConfig)
    tracking: TrackingConfig = Field(default_factory=TrackingConfig)
    output: OutputConfig = Field(default_factory=OutputConfig)
    stream: StreamConfig = Field(default_factory=StreamConfig)

    @classmethod
    def load(cls, path: str | Path) -> Config:
        """Read a `.toml` config, resolving its relative paths against its folder."""
        path = Path(path)
        if path.suffix.lower() != ".toml":
            raise ValueError(
                f"{path}: configs are TOML files; docs/fictrac.md translates a FicTrac "
                f"config.txt"
            )
        try:
            cfg = cls.model_validate(tomllib.loads(path.read_text(encoding="utf-8")))
        except tomllib.TOMLDecodeError as exc:
            raise ValueError(f"{path}: {exc}") from None
        except ValidationError as exc:
            raise ValueError(_explain(path, exc)) from None
        for table, key in PATHS:
            owner = getattr(cfg, table) if table else cfg
            value = getattr(owner, key)
            if isinstance(value, str):
                setattr(owner, key, os.path.normpath(path.parent / value))
        return cfg

    def save(
        self, path: str | Path, comments: Sequence[str] = (), *, full: bool = False
    ) -> Path:
        """Write the config to `path` as TOML, under `comments` (lines from `#`).

        Only keys that differ from the defaults are written unless `full`. Paths are
        written relative to the file, unless the two share no folder but the root.
        """
        path = Path(path)
        data = self.model_dump(exclude_defaults=not full, exclude_none=True)
        for table, key in PATHS:
            value = getattr(getattr(self, table) if table else self, key)
            if isinstance(value, str):
                (data[table] if table else data)[key] = _relative(value, path.parent)
        header = "".join(f"{c}\n" for c in comments) + ("\n" if comments else "")
        path.write_text(header + _toml(data), "utf-8")
        return path

load `classmethod` ¶

load(path: str | Path) -> Config

Read a .toml config, resolving its relative paths against its folder.

Source code in python/spintrack/config.py

@classmethod
def load(cls, path: str | Path) -> Config:
    """Read a `.toml` config, resolving its relative paths against its folder."""
    path = Path(path)
    if path.suffix.lower() != ".toml":
        raise ValueError(
            f"{path}: configs are TOML files; docs/fictrac.md translates a FicTrac "
            f"config.txt"
        )
    try:
        cfg = cls.model_validate(tomllib.loads(path.read_text(encoding="utf-8")))
    except tomllib.TOMLDecodeError as exc:
        raise ValueError(f"{path}: {exc}") from None
    except ValidationError as exc:
        raise ValueError(_explain(path, exc)) from None
    for table, key in PATHS:
        owner = getattr(cfg, table) if table else cfg
        value = getattr(owner, key)
        if isinstance(value, str):
            setattr(owner, key, os.path.normpath(path.parent / value))
    return cfg

save ¶

save(path: str | Path, comments: Sequence[str] = (), *, full: bool = False) -> Path

Write the config to path as TOML, under comments (lines from #).

Only keys that differ from the defaults are written unless full. Paths are written relative to the file, unless the two share no folder but the root.

Source code in python/spintrack/config.py

def save(
    self, path: str | Path, comments: Sequence[str] = (), *, full: bool = False
) -> Path:
    """Write the config to `path` as TOML, under `comments` (lines from `#`).

    Only keys that differ from the defaults are written unless `full`. Paths are
    written relative to the file, unless the two share no folder but the root.
    """
    path = Path(path)
    data = self.model_dump(exclude_defaults=not full, exclude_none=True)
    for table, key in PATHS:
        value = getattr(getattr(self, table) if table else self, key)
        if isinstance(value, str):
            (data[table] if table else data)[key] = _relative(value, path.parent)
    header = "".join(f"{c}\n" for c in comments) + ("\n" if comments else "")
    path.write_text(header + _toml(data), "utf-8")
    return path