pixel_cursor/cursor.py

"""PixelCursor: the spatio-temporal-semantic-identity selector that owns pixel mutation.

Every studio operation that touches pixels routes through a PixelCursor. Artifacts
are immutable from outside this package; the only way to produce a new Artifact is
via cursor.write_* methods.

Op-set (closed):
    seek_*           — reposition (returns new cursor)
    range / at_frame / all_frames / propagate_temporal — temporal traversal
    lock_identity    — pin subject or motion identity
    bind_motion_exemplar — few-shot motion attach (returns new cursor)
    read_*           — non-mutating queries (returns data)
    write_*          — mutation (returns new Artifact)
    tell             — introspect cursor state
"""
from __future__ import annotations
from dataclasses import dataclass, replace, field
from typing import Optional, Sequence, Tuple

import numpy as np

from .artifact import Artifact
from .selectors import Selector, Point, Region, Mask, Token, Layer, Anchor3D
from .scope import (
    TemporalScope,
    FrameIndex,
    FrameRange,
    ALL_FRAMES,
    DepthScope,
    IdentityLock,
    OpacityProfile,
)


@dataclass(frozen=True, slots=True)
class PixelCursor:
    artifact: Artifact
    selector: Selector
    temporal: TemporalScope
    depth: Optional[DepthScope] = None
    identity_lock: Optional[IdentityLock] = None
    opacity: OpacityProfile = field(default_factory=lambda: OpacityProfile(kind="hard"))

    # ── seek_* : reposition without mutating the artifact ──

    def seek_point(self, x: float, y: float, z: Optional[float] = None) -> "PixelCursor":
        return replace(self, selector=Point(x=x, y=y, z=z))

    def seek_region(self, x: int, y: int, w: int, h: int) -> "PixelCursor":
        return replace(self, selector=Region(x=x, y=y, w=w, h=h))

    def seek_mask(self, bitmap: np.ndarray, is_soft: bool = False) -> "PixelCursor":
        return replace(self, selector=Mask(bitmap=bitmap, is_soft=is_soft))

    def seek_token(self, text: str, embedding_id: Optional[str] = None) -> "PixelCursor":
        return replace(self, selector=Token(text=text, embedding_id=embedding_id))

    def seek_layer(self, layer_id: str) -> "PixelCursor":
        return replace(self, selector=Layer(layer_id=layer_id))

    def seek_3d_anchor(
        self, anchor_id: str, bbox: Tuple[float, float, float, float, float, float]
    ) -> "PixelCursor":
        return replace(self, selector=Anchor3D(anchor_id=anchor_id, bbox=bbox))

    # ── temporal ──

    def at_frame(self, index: int) -> "PixelCursor":
        return replace(self, temporal=FrameIndex(index=index))

    def range(self, start: int, end: int) -> "PixelCursor":
        return replace(self, temporal=FrameRange(start=start, end=end))

    def all_frames(self) -> "PixelCursor":
        return replace(self, temporal=ALL_FRAMES)

    def propagate_temporal(self, end_frame: int, mode: str = "forward") -> "PixelCursor":
        """Extend cursor temporal extent. Backend (e.g. CoTracker3, SAM2) does the actual traversal."""
        cur = self.temporal
        if isinstance(cur, FrameIndex):
            if mode == "backward":
                new: TemporalScope = FrameRange(start=end_frame, end=cur.index + 1)
            else:
                new = FrameRange(start=cur.index, end=end_frame)
        elif isinstance(cur, FrameRange):
            new = FrameRange(start=cur.start, end=max(cur.end, end_frame))
        else:
            new = cur
        return replace(self, temporal=new)

    # ── identity ──

    def lock_identity(self, lock: IdentityLock) -> "PixelCursor":
        return replace(self, identity_lock=lock)

    def bind_motion_exemplar(
        self, exemplar_paths: Sequence[str], backend: str = "auto"
    ) -> "PixelCursor":
        from .ops import _dispatch

        return _dispatch("bind_motion_exemplar", backend, self, exemplar_paths)

    # ── read ops (non-mutating) ──

    def read_state(self) -> dict:
        return self.tell()

    def read_pixels(self) -> np.ndarray:
        from .artifact import Image, FrameStack

        a = self.artifact
        if isinstance(a, Image):
            return a.pixels
        if isinstance(a, FrameStack):
            t = self.temporal
            if isinstance(t, FrameIndex):
                return a.frames[t.index]
            if isinstance(t, FrameRange):
                return a.frames[t.start : t.end]
            return a.frames
        raise TypeError(f"read_pixels: unsupported artifact type {type(a).__name__}")

    # ── write ops (return new Artifact) ──

    def write_motion(self, motion_spec, backend: str = "auto") -> Artifact:
        from .ops import _dispatch

        return _dispatch("write_motion", backend, self, motion_spec)

    def write_displacement(self, target_xy: Tuple[float, float], backend: str = "auto") -> Artifact:
        from .ops import _dispatch

        return _dispatch("write_displacement", backend, self, target_xy)

    def write_inpaint(self, prompt: str, backend: str = "auto") -> Artifact:
        from .ops import _dispatch

        return _dispatch("write_inpaint", backend, self, prompt)

    def write_keyframe(self, frame_index: int, pose_state, backend: str = "auto") -> Artifact:
        from .ops import _dispatch

        return _dispatch("write_keyframe", backend, self, frame_index, pose_state)

    def write_condition(self, control_signal, backend: str = "auto") -> Artifact:
        from .ops import _dispatch

        return _dispatch("write_condition", backend, self, control_signal)

    # ── introspection ──

    def tell(self) -> dict:
        return {
            "selector": type(self.selector).__name__,
            "selector_value": self.selector,
            "temporal": self.temporal,
            "depth": self.depth,
            "identity_locked": self.identity_lock is not None,
            "identity_source": (self.identity_lock.source if self.identity_lock else None),
            "opacity_kind": self.opacity.kind,
            "artifact": type(self.artifact).__name__,
        }


def open_cursor(
    artifact: Artifact,
    selector: Optional[Selector] = None,
    temporal: Optional[TemporalScope] = None,
) -> PixelCursor:
    if selector is None:
        selector = Token(text="*all*")
    if temporal is None:
        temporal = ALL_FRAMES
    return PixelCursor(artifact=artifact, selector=selector, temporal=temporal)