Python API reference¶

class TokEye:
    """A loaded TokEye segmentation model plus its preprocessing settings.

    Parameters
    ----------
    model
        Registry name (downloaded and cached on first use) or a path to a
        local ``.pt``/``.pt2`` checkpoint. Instance models such as
        ``ae_tf_maskrcnn`` are rejected; use ``tokeye alfvenspec``.
    device
        ``"auto"`` (CUDA, then MPS, then CPU), ``"cpu"``, ``"cuda"``,
        ``"cuda:N"`` or ``"mps"``. An unknown or unavailable device raises
        ``ValueError``.
    n_fft, hop, clip_dc, clip_low, clip_high, log
        Override single fields of ``config`` (see
        :class:`~tokeye.config.SpectrogramConfig`). ``None`` keeps the
        config's value.
    config
        Preprocessing settings (a config or a mapping); defaults when
        omitted.
    fs
        Default sampling rate in Hz for :meth:`segment` (labels axes only).
    tile
        ``"auto"`` (default), ``None`` or an int >= 512; passed to
        :func:`tokeye.inference.infer`.

    Attributes
    ----------
    model_name
        The model's label in results: the registry name, or a local
        checkpoint's file name without its directories
        (:func:`tokeye.hub.model_label`).
    weights
        Which weights were loaded (:func:`tokeye.hub.weights_info`): repo,
        file name, revision and sha256 for a registry model; file name and
        sha256 for a local checkpoint. Every :meth:`segment` result
        records it.
    """

    def __init__(
        self,
        model: str | Path = hub.DEFAULT_MODEL,
        device: str = "auto",
        n_fft: int | None = None,
        hop: int | None = None,
        clip_dc: bool | None = None,
        clip_low: float | None = None,
        clip_high: float | None = None,
        log: bool | None = None,
        *,
        config: SpectrogramConfig | Mapping[str, Any] | None = None,
        fs: float | None = None,
        tile: int | str | None = "auto",
    ) -> None:
        overrides = {
            "n_fft": n_fft,
            "hop": hop,
            "clip_dc": clip_dc,
            "clip_low": clip_low,
            "clip_high": clip_high,
            "log": log,
        }
        cfg = SpectrogramConfig.coerce(config)
        explicit = {key: value for key, value in overrides.items() if value is not None}
        self.config = cfg.replace(**explicit) if explicit else cfg
        self.fs = _check_fs(fs)
        check_tile(tile)
        self.tile = tile

        hub.require_task(model, "segmentation")
        self.model_name = hub.model_label(model)
        self.channels = hub.channels_for(model)
        self.model = hub.load_model(model, device)
        self.weights = hub.weights_info(model)

    def __repr__(self) -> str:
        device = str(_device_of(self.model))
        return (
            f"TokEye(model={self.model_name!r}, device={device!r}, tile={self.tile!r})"
        )

    @property
    def log(self) -> bool:
        """Whether 2D inputs are ``log1p``-scaled (``config.log``)."""
        return self.config.log

    @log.setter
    def log(self, value: bool) -> None:
        self.config = self.config.replace(log=value)

    def _prepare(
        self,
        data: Any,
        *,
        fs: float | None = None,
        reference: Any = None,
        log: bool | None = None,
    ) -> Spectrogram:
        cfg = self.config if log is None else self.config.replace(log=log)
        return prepare(data, cfg, fs=self.fs if fs is None else fs, reference=reference)

    def spectrogram(self, data: Any, log: bool | None = None) -> np.ndarray:
        """The ``(H, W)`` float64 spectrogram the model will see (pre-1.0).

        1D input goes through the STFT (which log-scales internally); 2D
        input is used as-is, ``log1p``-scaled first when ``log`` is on.
        ``log=None`` defers to the instance setting. The result is always a
        new array. The model sees these values cast to float32; use
        :meth:`segment` for that float32 input together with its axes.
        """
        cfg = self.config if log is None else self.config.replace(log=log)
        values = _values64(data, cfg)[0]
        if np.shares_memory(values, data):  # 2D float64 input, log off
            values = values.copy()
        return values

    def segment(
        self,
        data: Any,
        *,
        fs: float | None = None,
        reference: Any = None,
        log: bool | None = None,
    ) -> Segmentation:
        """Segment a signal or spectrogram.

        Parameters
        ----------
        data
            1D signal or 2D spectrogram.
        fs
            Sampling rate in Hz (overrides the instance ``fs``).
        reference
            Second 1D signal for a cross-power spectrum.
        log
            Per-call override of ``config.log``.

        Returns
        -------
        Segmentation
            Mask, the spectrogram it was computed from, and its axes.
        """
        spec = self._prepare(data, fs=fs, reference=reference, log=log)
        mask = infer(self.model, spec.values, tile=self.tile)
        channels = resolve_channels(self.channels, mask.shape[0])
        return Segmentation(mask, spec, channels, self.model_name, self.weights)

    def predict(self, data: Any, log: bool | None = None) -> np.ndarray:
        """Run inference; returns a float32 mask of shape ``(C, H, W)``.

        For ``big_tf_unet``, channel 0 = coherent and channel 1 =
        transient activity, both sigmoid scores in ``[0, 1]``.
        Standardization is applied internally. A one-channel model gives
        ``(H, W)``, as before 1.0; :meth:`segment` always gives
        ``(C, H, W)``.
        """
        mask = self.segment(data, log=log).mask
        return mask[0] if mask.shape[0] == 1 else mask

    def __call__(self, data: Any, log: bool | None = None) -> np.ndarray:
        return self.predict(data, log=log)

@dataclass(frozen=True)
class SpectrogramConfig:
    """How a raw signal (or a stored spectrogram) becomes model input.

    Parameters
    ----------
    n_fft
        STFT window length in samples.
    hop
        STFT hop in samples. 128 matches the released model's training.
    window
        Window name understood by :func:`scipy.signal.get_window`.
    clip_dc
        Drop the DC (0 Hz) row of the STFT.
    clip_low, clip_high
        Percentiles the log-magnitude spectrogram is clipped to.
    log
        Apply ``log1p`` to 2D inputs stored in linear scale. 1D inputs
        are always log-scaled by the STFT, so this only affects 2D input.

    Notes
    -----
    ``clip_dc`` and ``log`` take ``True`` or ``False`` (NumPy booleans too);
    the integers ``0`` and ``1`` still work, with a ``DeprecationWarning``.
    """

    n_fft: int = field(default=1024, metadata={"help": "STFT window length [samples]"})
    hop: int = field(default=128, metadata={"help": "STFT hop [samples]"})
    window: str = field(default="hann", metadata={"help": "STFT window name"})
    clip_dc: bool = field(default=True, metadata={"help": "drop the DC row"})
    clip_low: float = field(default=1.0, metadata={"help": "lower clip percentile"})
    clip_high: float = field(default=99.0, metadata={"help": "upper clip percentile"})
    log: bool = field(
        default=False,
        metadata={"help": "log1p 2D inputs stored in linear scale"},
    )

    def __post_init__(self) -> None:
        for name in _INT_FIELDS:
            value = getattr(self, name)
            if isinstance(value, bool) or not isinstance(value, numbers.Integral):
                raise TypeError(f"{name} must be an integer, got {value!r}")
            object.__setattr__(self, name, int(value))
        for name in _FLOAT_FIELDS:
            value = getattr(self, name)
            if isinstance(value, bool) or not isinstance(value, numbers.Real):
                raise TypeError(f"{name} must be a number, got {value!r}")
            object.__setattr__(self, name, float(value))
        for name in _BOOL_FIELDS:
            value = getattr(self, name)
            if isinstance(value, bool):
                continue
            # Without numpy imported, the value cannot be a NumPy bool (and
            # importing it here would load numpy with this module).
            np_mod = sys.modules.get("numpy")
            if np_mod is not None and isinstance(value, np_mod.bool_):
                object.__setattr__(self, name, bool(value))
            elif isinstance(value, numbers.Integral) and value in (0, 1):
                warnings.warn(
                    f"{name}={value!r}: pass True or False (integers are "
                    "deprecated and rejected in 2.0)",
                    DeprecationWarning,
                    stacklevel=_user_stacklevel(),
                )
                object.__setattr__(self, name, bool(value))
            else:
                raise TypeError(f"{name} must be True or False, got {value!r}")
        if not isinstance(self.window, str):
            raise TypeError(f"window must be a string, got {self.window!r}")

        if not self.window:
            raise ValueError("window must be a non-empty string, got ''")
        if self.n_fft < 2:
            raise ValueError(f"n_fft must be >= 2, got {self.n_fft}")
        if self.hop < 1:
            raise ValueError(f"hop must be >= 1, got {self.hop}")
        if not 0.0 <= self.clip_low < self.clip_high <= 100.0:
            raise ValueError(
                "need 0 <= clip_low < clip_high <= 100, got "
                f"clip_low={self.clip_low}, clip_high={self.clip_high}"
            )

    def to_dict(self) -> dict[str, Any]:
        """Plain-``dict`` form (JSON-serializable)."""
        return dataclasses.asdict(self)

    @classmethod
    def from_dict(cls, data: Mapping[str, Any]) -> SpectrogramConfig:
        """Build from a mapping; unknown keys raise ``ValueError``."""
        valid = [f.name for f in dataclasses.fields(cls)]
        unknown = sorted(set(data) - set(valid), key=repr)
        if unknown:
            hint = ""
            if "fs" in unknown:
                hint = (
                    " (the sampling rate is not a preprocessing setting; "
                    "pass fs= separately)"
                )
            raise ValueError(
                f"unknown SpectrogramConfig field(s) {unknown}; valid: {valid}{hint}"
            )
        return cls(**dict(data))

    def replace(self, **changes: Any) -> SpectrogramConfig:
        """Copy with ``changes`` applied (validated like the constructor)."""
        return self.from_dict({**self.to_dict(), **changes})

    @classmethod
    def coerce(
        cls, value: SpectrogramConfig | Mapping[str, Any] | None
    ) -> SpectrogramConfig:
        """Accept ``None`` (defaults), a config, or a mapping of fields."""
        if value is None:
            return cls()
        if isinstance(value, cls):
            return value
        if isinstance(value, Mapping):
            return cls.from_dict(value)
        raise TypeError(
            f"expected SpectrogramConfig, a mapping or None, got {type(value).__name__}"
        )

    def stft_kwargs(self) -> dict[str, Any]:
        """Keyword arguments for :func:`tokeye.transforms.compute_stft`."""
        return {
            "n_fft": self.n_fft,
            "hop": self.hop,
            "window": self.window,
            "clip_dc": self.clip_dc,
            "clip_low": self.clip_low,
            "clip_high": self.clip_high,
        }

@dataclass(frozen=True, eq=False)
class Segmentation:
    """The result of :meth:`tokeye.TokEye.segment`.

    Attributes
    ----------
    mask
        ``(C, H, W)`` float32 scores in ``[0, 1]``; for ``big_tf_unet``,
        channel 0 is coherent and channel 1 transient activity.
    spectrogram
        The :class:`~tokeye.preprocess.Spectrogram` the model saw.
    channels
        One name per mask channel.
    model
        Registry name or file name of the model that produced the mask.
    weights
        Which weights produced it (see :func:`tokeye.hub.weights_info`):
        ``repo``, ``filename``, ``revision`` and ``sha256`` for a registry
        model, ``name`` and ``sha256`` for a local checkpoint; ``None`` when
        unknown. A mapping is stored as a plain ``dict``.
    """

    mask: np.ndarray
    spectrogram: Spectrogram
    channels: tuple[str, ...] = DEFAULT_CHANNELS
    model: str = "big_tf_unet"
    weights: Mapping[str, str | None] | None = None

    def __post_init__(self) -> None:
        if not isinstance(self.spectrogram, Spectrogram):
            raise TypeError(
                "spectrogram must be a tokeye Spectrogram, got "
                f"{type(self.spectrogram).__name__}"
            )
        if self.weights is not None:
            if not isinstance(self.weights, Mapping):
                raise TypeError(
                    "weights must be a mapping or None, got "
                    f"{type(self.weights).__name__}"
                )
            object.__setattr__(self, "weights", dict(self.weights))
        mask = np.asarray(self.mask, dtype=np.float32)
        if mask.ndim != 3 or mask.shape[1:] != self.spectrogram.shape:
            raise ValueError(
                f"mask shape {mask.shape} does not match spectrogram shape "
                f"{self.spectrogram.shape}; expected (C, H, W)"
            )
        channels = tuple(self.channels)
        if len(channels) != mask.shape[0]:
            raise ValueError(
                f"{len(channels)} channel name(s) for a {mask.shape[0]}-channel mask"
            )
        object.__setattr__(self, "mask", mask)
        object.__setattr__(self, "channels", channels)

    def __repr__(self) -> str:
        return (
            f"Segmentation(model={self.model!r}, mask={self.mask.shape}, "
            f"channels={self.channels!r}, "
            f"{_axes_text(self.fs, self.freqs, self.times)})"
        )

    def __getitem__(self, name: str) -> np.ndarray:
        """The ``(H, W)`` scores of channel ``name``."""
        try:
            return self.mask[self.channels.index(name)]
        except ValueError:
            raise KeyError(
                f"no channel {name!r}; channels: {list(self.channels)}"
            ) from None

    def _channel_attribute(self, name: str) -> np.ndarray:
        # AttributeError, so hasattr()/getattr(..., default) work when the
        # channel is absent; seg[name] keeps raising KeyError.
        try:
            return self[name]
        except KeyError as exc:
            raise AttributeError(exc.args[0]) from None

    @property
    def coherent(self) -> np.ndarray:
        """``(H, W)`` coherent-activity scores (``AttributeError`` if absent)."""
        return self._channel_attribute("coherent")

    @property
    def transient(self) -> np.ndarray:
        """``(H, W)`` transient-activity scores (``AttributeError`` if absent)."""
        return self._channel_attribute("transient")

    @property
    def freqs(self) -> np.ndarray:
        """Row centres: Hz when ``fs`` is known, else bin index."""
        return self.spectrogram.freqs

    @property
    def times(self) -> np.ndarray:
        """Column centres: seconds when ``fs`` is known, else frame index."""
        return self.spectrogram.times

    @property
    def fs(self) -> float | None:
        """Sampling rate in Hz, or ``None``."""
        return self.spectrogram.fs

    def threshold(self, t: float = 0.5) -> np.ndarray:
        """Boolean ``(C, H, W)`` mask, ``True`` where the score is ``>= t``."""
        return self.mask >= t

    def save(self, path: str | Path) -> Path:
        """Save as a ``tokeye-analysis/v1`` ``.npz`` (the app's schema).

        Returns the path written (``.npz`` is appended when missing). With
        ``fs`` known, the result's own axes are stored (in ms and kHz).
        Without it the axes are bin and frame indices and are not stored;
        other axes then warn (``UserWarning``), since they reload as indices.
        """
        spec = self.spectrogram
        axes = None
        if spec.fs is not None:
            axes = (
                np.asarray(spec.times, dtype=np.float64) * 1e3,
                np.asarray(spec.freqs, dtype=np.float64) / 1e3,
            )
        params = {
            "config": spec.config.to_dict(),
            "model": self.model,
            "weights": self.weights,
            "tokeye_version": __version__,
            "kind": spec.kind,
            "fs": spec.fs,
            "n_samples": None if spec.n_samples is None else int(spec.n_samples),
            "channels": list(self.channels),
        }
        bundle = export.analysis_bundle(
            spectrogram=spec.values,
            mask=self.mask,
            axes=axes,
            params=params,
            source="segment",
        )
        written = export.save_npz(path, bundle)
        if spec.fs is None:
            freqs, times = _axes(spec.shape, spec.config, None)
            if not (
                np.array_equal(spec.freqs, freqs) and np.array_equal(spec.times, times)
            ):
                warnings.warn(
                    f"{written}: axes without fs are not saved; they reload as "
                    "bin/frame indices",
                    UserWarning,
                    stacklevel=_user_stacklevel(),
                )
        return written

    @classmethod
    def load(cls, path: str | Path) -> Segmentation:
        """Load a ``tokeye-analysis/v1`` bundle that contains a mask.

        Bundles written by the app (flat ``n_fft``/``hop``/... params, no
        axes) load too; missing settings fall back to the defaults, silently.
        A recorded value that is invalid (an unknown or out-of-range setting,
        a bad ``fs``, ``kind``, ``n_samples``, channel list, weights record or
        stored axis) is not used: it falls back to its default with a
        ``UserWarning`` naming the file. Bundles without a weights record
        load with ``weights=None``, and without ``n_samples`` with
        ``n_samples=None``.

        Bundles do not record their STFT framing, so their stored axes
        (``time_ms``/``freq_khz``) are the only record of it: with a valid
        ``fs``, each valid stored axis is used as stored, and only a missing
        or invalid one is rebuilt from ``fs`` and the settings.

        Raises
        ------
        ValueError
            The file is not a ``tokeye-analysis/v1`` bundle (a bare ``.npy``
            mask, say: ``tokeye run --format npz`` and :meth:`save` write
            bundles), or has no mask.
        """
        loaded = np.load(Path(path), allow_pickle=False)
        if isinstance(loaded, np.ndarray):
            raise ValueError(
                f"{path} is a bare array (such as the mask `tokeye run` writes by "
                f"default), not a {export.SCHEMA_ANALYSIS} bundle; "
                "`tokeye run --format npz` or Segmentation.save() writes a bundle "
                "that loads"
            )
        with loaded as data:
            files = set(data.files)
            schema = str(data["schema"]) if "schema" in files else None
            if schema != export.SCHEMA_ANALYSIS:
                raise ValueError(
                    f"{path} is not a {export.SCHEMA_ANALYSIS} bundle "
                    f"(schema={schema!r})"
                )
            if "mask" not in files:
                raise ValueError(f"{path} has no mask (saved before inference?)")
            values = np.asarray(data["spectrogram"], dtype=np.float32)
            mask = np.asarray(data["mask"], dtype=np.float32)
            params_json = str(data["params_json"]) if "params_json" in files else "{}"
            time_ms = data["time_ms"] if "time_ms" in files else None
            freq_khz = data["freq_khz"] if "freq_khz" in files else None

        if mask.ndim == 2:
            mask = mask[np.newaxis]
        params = _params_from_json(path, params_json)
        cfg = _config_from_params(params, path)

        kind = params.get("kind")
        if kind is None:
            kind = "spectrogram"
        elif kind not in KINDS:
            _warn_ignored(path, "kind", kind, f"not one of {KINDS}", "'spectrogram'")
            kind = "spectrogram"

        fs = params.get("fs")
        if fs is not None and not _positive_finite(fs):
            _warn_ignored(
                path,
                "fs",
                fs,
                "not a positive finite number",
                "no fs: the axes are rebuilt as bin/frame indices",
            )
            fs = None
        fs = None if fs is None else float(fs)

        freqs, times = _axes(values.shape, cfg, fs)
        if fs is not None:
            stored_times = _stored_axis(path, "time_ms", time_ms, values.shape[1])
            if stored_times is not None:
                times = stored_times / 1e3
            stored_freqs = _stored_axis(path, "freq_khz", freq_khz, values.shape[0])
            if stored_freqs is not None:
                freqs = stored_freqs * 1e3
        n_samples = _n_samples_from(path, params.get("n_samples"))
        spectrogram = Spectrogram(values, freqs, times, kind, cfg, fs, n_samples)

        default_channels = resolve_channels(DEFAULT_CHANNELS, mask.shape[0])
        channels = params.get("channels")
        if channels is None:
            channels = default_channels
        elif not (
            isinstance(channels, list)
            and len(channels) == mask.shape[0]
            and all(isinstance(name, str) for name in channels)
        ):
            _warn_ignored(
                path,
                "channels",
                channels,
                f"not a list of {mask.shape[0]} names",
                f"{list(default_channels)}",
            )
            channels = default_channels
        model = params.get("model") or "unknown"
        weights = params.get("weights")
        if weights is not None and not _is_weights_record(weights):
            _warn_ignored(
                path,
                "weights",
                weights,
                "not a mapping of strings to strings or null",
                "None",
            )
            weights = None
        return cls(mask, spectrogram, tuple(channels), str(model), weights)

    def plot(
        self, ax: Axes | None = None, *, threshold: float = 0.5, alpha: float = 0.4
    ) -> Axes:
        """Draw the spectrogram with the thresholded mask overlaid.

        Uses ``origin="lower"`` (low frequencies at the bottom) and labelled
        axes. Creates a new pyplot figure when ``ax`` is ``None``.
        """
        from ._plotting import draw_segmentation

        if ax is None:
            import matplotlib.pyplot as plt

            _, ax = plt.subplots(figsize=(8, 4), layout="constrained")
        return draw_segmentation(ax, self, threshold=threshold, alpha=alpha)