pixano_inference.models.tracking

Tracking model base class.

The I/O and prompt types are the wire contract and live in :mod:pixano_inference.schemas.tracking; they are re-exported here so from pixano_inference.models.tracking import TrackingInput (and the prompt types) keep working.

TrackingModel(config)

Bases: InferenceModel[TrackingInput, TrackingOutput]

Base class for video tracking models.

predict receives a :class:TrackingInput and returns a :class:TrackingOutput: for each frame, the objects tracked in it. A model is either prompted (the request names the objects with points, boxes or masks, and the model returns their masks, as SAM2 does) or prompt-free (tracking by detection: the request has no object, and the model creates the tracks and returns their boxes and scores, as ByteTrack does). The interface property, when overridden, returns a :class:TrackingInterface: which prompts the model takes, whether it accepts a prompt-free request, its class set and what each tracked object carries.

Example

@register_model("my-tracker")
class MyTracker(TrackingModel):
    def load_model(self):
        self.detector, self.tracker = load_detector(...), load_tracker(...)

    def predict(self, input: TrackingInput) -> TrackingOutput:
        frames = []
        for index, image in enumerate(load_frames(input.video)):
            tracks = self.tracker.update(self.detector(image))
            frames.append(
                TrackedFrame(
                    frame_index=index,
                    objects=[TrackedObject(track_id=t.id, box=t.xyxy, score=t.score) for t in tracks],
                )
            )
        return TrackingOutput(frames=frames)

Source code in pixano_inference/models/base.py

def __init__(self, config: ModelDeploymentConfig) -> None:
    """Initialize the model with deployment config.

    Args:
        config: Model deployment configuration.
    """
    self._config = config