rslearn.train.transforms.random_time_dropping

random_time_dropping

Random temporal dropping augmentation transform.

RandomTimeDropping

Bases: Transform

Randomly drop entire timesteps from a CTHW tensor.

Dropped timesteps are removed from both image.image (shrinking T) and image.timestamps.

Source code in rslearn/train/transforms/random_time_dropping.py

class RandomTimeDropping(Transform):
    """Randomly drop entire timesteps from a CTHW tensor.

    Dropped timesteps are removed from both ``image.image`` (shrinking T)
    and ``image.timestamps``.
    """

    def __init__(
        self,
        selectors: list[str] = ["image"],
        drop_ratio: float = 0.2,
        min_keep: int = 1,
        skip_missing: bool = False,
    ):
        """Initialize RandomTimeDropping.

        Args:
            selectors: inputs to augment.
            drop_ratio: expected fraction of timesteps to drop.
            min_keep: minimum number of timesteps to keep (at least 1).
            skip_missing: skip missing selectors.
        """
        super().__init__(skip_missing=skip_missing)
        if not 0 <= drop_ratio <= 1:
            raise ValueError(f"drop_ratio must be between 0 and 1, got {drop_ratio}")
        self.selectors = selectors
        self.drop_ratio = drop_ratio
        self.min_keep = max(min_keep, 1)

    def apply_image(self, image: RasterImage) -> RasterImage:
        """Drop random timesteps from the image."""
        T = image.shape[1]
        if T <= self.min_keep:
            return image

        # Decide which timesteps to keep.
        keep = torch.rand(T) >= self.drop_ratio

        # Guarantee we keep at least min_keep timesteps.
        if keep.sum() < self.min_keep:
            # Randomly pick min_keep indices to force-keep.
            indices = torch.randperm(T)[: self.min_keep]
            keep[:] = False
            keep[indices] = True

        # If nothing was dropped, return unchanged.
        if keep.all():
            return image

        image.image = image.image[:, keep, :, :]

        if image.timestamps is not None:
            image.timestamps = [
                ts for ts, k in zip(image.timestamps, keep.tolist()) if k
            ]

        return image

    def forward(
        self, input_dict: dict[str, Any], target_dict: dict[str, Any]
    ) -> tuple[dict[str, Any], dict[str, Any]]:
        """Apply transform."""
        self.apply_fn(self.apply_image, input_dict, target_dict, self.selectors)
        return input_dict, target_dict

apply_image

apply_image(image: RasterImage) -> RasterImage

Drop random timesteps from the image.

Source code in rslearn/train/transforms/random_time_dropping.py

def apply_image(self, image: RasterImage) -> RasterImage:
    """Drop random timesteps from the image."""
    T = image.shape[1]
    if T <= self.min_keep:
        return image

    # Decide which timesteps to keep.
    keep = torch.rand(T) >= self.drop_ratio

    # Guarantee we keep at least min_keep timesteps.
    if keep.sum() < self.min_keep:
        # Randomly pick min_keep indices to force-keep.
        indices = torch.randperm(T)[: self.min_keep]
        keep[:] = False
        keep[indices] = True

    # If nothing was dropped, return unchanged.
    if keep.all():
        return image

    image.image = image.image[:, keep, :, :]

    if image.timestamps is not None:
        image.timestamps = [
            ts for ts, k in zip(image.timestamps, keep.tolist()) if k
        ]

    return image

forward

forward(input_dict: dict[str, Any], target_dict: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]

Apply transform.

Source code in rslearn/train/transforms/random_time_dropping.py

def forward(
    self, input_dict: dict[str, Any], target_dict: dict[str, Any]
) -> tuple[dict[str, Any], dict[str, Any]]:
    """Apply transform."""
    self.apply_fn(self.apply_image, input_dict, target_dict, self.selectors)
    return input_dict, target_dict