rslearn.train.model_context

model_context

Data classes to provide various context to models.

RasterImage dataclass

A raster image is a torch.tensor containing the images and their associated timestamps.

Source code in rslearn/train/model_context.py

@dataclass
class RasterImage:
    """A raster image is a torch.tensor containing the images and their associated timestamps."""

    # image is a 4D CTHW tensor
    image: torch.Tensor
    # if timestamps is not None, len(timestamps) must match the T dimension of the tensor
    timestamps: list[tuple[datetime, datetime]] | None = None

    def __post_init__(self) -> None:
        """Validate that the image tensor is 4D (CTHW)."""
        if self.image.dim() != 4:
            raise ValueError(
                f"RasterImage expects a 4D CTHW tensor, got {self.image.dim()}D "
                f"with shape {tuple(self.image.shape)}"
            )

    @property
    def shape(self) -> torch.Size:
        """The shape of the image."""
        return self.image.shape

    def dim(self) -> int:
        """The dim of the image."""
        return self.image.dim()

    @property
    def dtype(self) -> torch.dtype:
        """The image dtype."""
        return self.image.dtype

    def single_ts_to_chw_tensor(self) -> torch.Tensor:
        """Single timestep models expect single timestep inputs.

        This function (1) checks this raster image only has 1 timestep and
        (2) returns the tensor for that (single) timestep (going from CTHW to CHW).
        """
        if self.image.shape[1] != 1:
            raise ValueError(f"Expected a single timestep, got {self.image.shape[1]}")
        return self.image[:, 0]

    def get_hw_tensor(self) -> torch.Tensor:
        """Get a 2D HW tensor from a single-channel, single-timestep RasterImage.

        This function checks that C=1 and T=1, then returns the HW tensor.
        Useful for per-pixel labels like segmentation masks.
        """
        if self.image.shape[0] != 1:
            raise ValueError(
                f"Expected single channel (C=1), got {self.image.shape[0]}"
            )
        if self.image.shape[1] != 1:
            raise ValueError(
                f"Expected single timestep (T=1), got {self.image.shape[1]}"
            )
        return self.image[0, 0]

    @classmethod
    def from_raster_array(cls, ra: RasterArray) -> RasterImage:
        """Create a RasterImage from a numpy-based RasterArray.

        Converts the numpy CTHW array to a torch tensor and copies timestamps.

        Args:
            ra: the source RasterArray.

        Returns:
            a new RasterImage with the same data.
        """
        return cls(
            image=torch.as_tensor(ra.array.copy()),
            timestamps=ra.timestamps,
        )

shape property

shape: Size

The shape of the image.

dtype property

dtype: dtype

The image dtype.

dim

dim() -> int

The dim of the image.

Source code in rslearn/train/model_context.py

def dim(self) -> int:
    """The dim of the image."""
    return self.image.dim()

single_ts_to_chw_tensor

single_ts_to_chw_tensor() -> Tensor

Single timestep models expect single timestep inputs.

This function (1) checks this raster image only has 1 timestep and (2) returns the tensor for that (single) timestep (going from CTHW to CHW).

Source code in rslearn/train/model_context.py

def single_ts_to_chw_tensor(self) -> torch.Tensor:
    """Single timestep models expect single timestep inputs.

    This function (1) checks this raster image only has 1 timestep and
    (2) returns the tensor for that (single) timestep (going from CTHW to CHW).
    """
    if self.image.shape[1] != 1:
        raise ValueError(f"Expected a single timestep, got {self.image.shape[1]}")
    return self.image[:, 0]

get_hw_tensor

get_hw_tensor() -> Tensor

Get a 2D HW tensor from a single-channel, single-timestep RasterImage.

This function checks that C=1 and T=1, then returns the HW tensor. Useful for per-pixel labels like segmentation masks.

Source code in rslearn/train/model_context.py

def get_hw_tensor(self) -> torch.Tensor:
    """Get a 2D HW tensor from a single-channel, single-timestep RasterImage.

    This function checks that C=1 and T=1, then returns the HW tensor.
    Useful for per-pixel labels like segmentation masks.
    """
    if self.image.shape[0] != 1:
        raise ValueError(
            f"Expected single channel (C=1), got {self.image.shape[0]}"
        )
    if self.image.shape[1] != 1:
        raise ValueError(
            f"Expected single timestep (T=1), got {self.image.shape[1]}"
        )
    return self.image[0, 0]

from_raster_array classmethod

from_raster_array(ra: RasterArray) -> RasterImage

Create a RasterImage from a numpy-based RasterArray.

Converts the numpy CTHW array to a torch tensor and copies timestamps.

Parameters:

Name	Type	Description	Default
ra	RasterArray	the source RasterArray.	required

Returns:

Type	Description
RasterImage	a new RasterImage with the same data.

Source code in rslearn/train/model_context.py

@classmethod
def from_raster_array(cls, ra: RasterArray) -> RasterImage:
    """Create a RasterImage from a numpy-based RasterArray.

    Converts the numpy CTHW array to a torch tensor and copies timestamps.

    Args:
        ra: the source RasterArray.

    Returns:
        a new RasterImage with the same data.
    """
    return cls(
        image=torch.as_tensor(ra.array.copy()),
        timestamps=ra.timestamps,
    )

SampleMetadata dataclass

Metadata pertaining to an example.

Source code in rslearn/train/model_context.py

@dataclass
class SampleMetadata:
    """Metadata pertaining to an example."""

    window_group: str
    window_name: str
    window_bounds: PixelBounds
    crop_bounds: PixelBounds
    crop_idx: int
    num_crops_in_window: int
    time_range: tuple[datetime, datetime] | None
    projection: Projection

    # Task name to differentiate different tasks.
    dataset_source: str | None

ModelContext dataclass

Context to pass to all model components.

Source code in rslearn/train/model_context.py

@dataclass
class ModelContext:
    """Context to pass to all model components."""

    # One input dict per example in the batch.
    inputs: list[dict[str, torch.Tensor | RasterImage]]
    # One SampleMetadata per example in the batch.
    metadatas: list[SampleMetadata]
    # Arbitrary dict that components can add to.
    context_dict: dict[str, Any] = field(default_factory=lambda: {})

ModelOutput dataclass

The output from the Predictor.

Parameters:

Name	Type	Description	Default
outputs	Iterable[Any]	output compatible with the configured Task.	required
loss_dict	dict[str, Tensor]	map from loss names to scalar tensors.	required
metadata	dict[str, Any]	arbitrary dict that can be used to store other outputs.	(lambda: {})()

Source code in rslearn/train/model_context.py

@dataclass
class ModelOutput:
    """The output from the Predictor.

    Args:
        outputs: output compatible with the configured Task.
        loss_dict: map from loss names to scalar tensors.
        metadata: arbitrary dict that can be used to store other outputs.
    """

    outputs: Iterable[Any]
    loss_dict: dict[str, torch.Tensor]
    metadata: dict[str, Any] = field(default_factory=lambda: {})