Skip to content

rslearn.train.transforms.concatenate

concatenate

Concatenate bands across multiple image inputs.

ConcatenateDim

Bases: Enum

Enum for concatenation dimensions.

Source code in rslearn/train/transforms/concatenate.py 
14
15
16
17
18
class ConcatenateDim(Enum):
    """Enum for concatenation dimensions."""

    CHANNEL = 0
    TIME = 1

Concatenate

Bases: Transform

Concatenate bands across multiple image inputs.

Source code in rslearn/train/transforms/concatenate.py 
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
class Concatenate(Transform):
    """Concatenate bands across multiple image inputs."""

    def __init__(
        self,
        selections: dict[str, list[int]],
        output_selector: str,
        concatenate_dim: ConcatenateDim | int = ConcatenateDim.TIME,
        skip_missing: bool = False,
    ):
        """Initialize a new Concatenate.

        Args:
            selections: map from selector to list of band indices in that input to
                retain, or empty list to use all bands.
            output_selector: the output selector under which to save the concatenate image.
            concatenate_dim: the dimension against which to concatenate the inputs
            skip_missing: if True, selectors that are absent from the input/target dicts
                are silently skipped instead of raising. Useful when concatenating
                optional inputs (e.g. per-month layers where some months are missing).
        """
        super().__init__(skip_missing=skip_missing)
        self.selections = selections
        self.output_selector = output_selector
        self.concatenate_dim = (
            concatenate_dim.value
            if isinstance(concatenate_dim, ConcatenateDim)
            else concatenate_dim
        )

    def forward(
        self, input_dict: dict[str, Any], target_dict: dict[str, Any]
    ) -> tuple[dict[str, Any], dict[str, Any]]:
        """Apply concatenation over the inputs and targets.

        Args:
            input_dict: the input
            target_dict: the target

        Returns:
            (input_dicts, target_dicts) where the entry corresponding to
            output_selector contains the concatenated RasterImage.
        """
        concatenate_time = self.concatenate_dim == ConcatenateDim.TIME.value

        tensors: list[torch.Tensor] = []
        # For CHANNEL concatenation, all inputs share the same timesteps, so we keep the
        # first available timestamps. For TIME concatenation, timestamps are collected
        # across all inputs so that len(timestamps) matches the concatenated time
        # dimension (required by models that consume real per-timestep timestamps).
        channel_timestamps: list[tuple[datetime, datetime]] | None = None
        time_timestamps: list[tuple[datetime, datetime]] = []
        time_has_all_timestamps = True

        for selector, wanted_bands in self.selections.items():
            if self.skip_missing and not selector_exists(
                input_dict, target_dict, selector
            ):
                continue
            image = read_selector(input_dict, target_dict, selector)
            if wanted_bands:
                tensors.append(image.image[wanted_bands, :, :])
            else:
                tensors.append(image.image)

            if concatenate_time:
                if image.timestamps is not None:
                    time_timestamps.extend(image.timestamps)
                else:
                    time_has_all_timestamps = False
            elif channel_timestamps is None and image.timestamps is not None:
                channel_timestamps = image.timestamps

        if not tensors:
            raise ValueError(
                f"Concatenate produced no inputs for output_selector "
                f"'{self.output_selector}' (all selectors missing with skip_missing)."
            )

        if concatenate_time:
            timestamps = time_timestamps if time_has_all_timestamps else None
        else:
            timestamps = channel_timestamps

        result = RasterImage(
            torch.concatenate(tensors, dim=self.concatenate_dim),
            timestamps=timestamps,
        )
        write_selector(input_dict, target_dict, self.output_selector, result)
        return input_dict, target_dict

forward 

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

Apply concatenation over the inputs and targets.

Parameters:

Name Type Description Default
input_dict dict[str, Any]

the input

 required
target_dict dict[str, Any]

the target

 required

Returns:

Type Description
dict[str, Any]

(input_dicts, target_dicts) where the entry corresponding to
dict[str, Any]

output_selector contains the concatenated RasterImage.
Source code in rslearn/train/transforms/concatenate.py 
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def forward(
    self, input_dict: dict[str, Any], target_dict: dict[str, Any]
) -> tuple[dict[str, Any], dict[str, Any]]:
    """Apply concatenation over the inputs and targets.

    Args:
        input_dict: the input
        target_dict: the target

    Returns:
        (input_dicts, target_dicts) where the entry corresponding to
        output_selector contains the concatenated RasterImage.
    """
    concatenate_time = self.concatenate_dim == ConcatenateDim.TIME.value

    tensors: list[torch.Tensor] = []
    # For CHANNEL concatenation, all inputs share the same timesteps, so we keep the
    # first available timestamps. For TIME concatenation, timestamps are collected
    # across all inputs so that len(timestamps) matches the concatenated time
    # dimension (required by models that consume real per-timestep timestamps).
    channel_timestamps: list[tuple[datetime, datetime]] | None = None
    time_timestamps: list[tuple[datetime, datetime]] = []
    time_has_all_timestamps = True

    for selector, wanted_bands in self.selections.items():
        if self.skip_missing and not selector_exists(
            input_dict, target_dict, selector
        ):
            continue
        image = read_selector(input_dict, target_dict, selector)
        if wanted_bands:
            tensors.append(image.image[wanted_bands, :, :])
        else:
            tensors.append(image.image)

        if concatenate_time:
            if image.timestamps is not None:
                time_timestamps.extend(image.timestamps)
            else:
                time_has_all_timestamps = False
        elif channel_timestamps is None and image.timestamps is not None:
            channel_timestamps = image.timestamps

    if not tensors:
        raise ValueError(
            f"Concatenate produced no inputs for output_selector "
            f"'{self.output_selector}' (all selectors missing with skip_missing)."
        )

    if concatenate_time:
        timestamps = time_timestamps if time_has_all_timestamps else None
    else:
        timestamps = channel_timestamps

    result = RasterImage(
        torch.concatenate(tensors, dim=self.concatenate_dim),
        timestamps=timestamps,
    )
    write_selector(input_dict, target_dict, self.output_selector, result)
    return input_dict, target_dict