Skip to content

rslearn.train.tasks.task

task

Training tasks.

Task

Represents an ML task like object detection or segmentation.

A task specifies how raster or vector data should be processed into inputs and targets that can be passed to models. It also specifies evaluation functions for computing metrics comparing targets/outputs.

Source code in rslearn/train/tasks/task.py 
14
15
16
17
18
19
20
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
class Task:
    """Represents an ML task like object detection or segmentation.

    A task specifies how raster or vector data should be processed into inputs and
    targets that can be passed to models. It also specifies evaluation functions for
    computing metrics comparing targets/outputs.
    """

    def process_inputs(
        self,
        raw_inputs: dict[str, RasterImage | list[Feature]],
        metadata: SampleMetadata,
        load_targets: bool = True,
    ) -> tuple[dict[str, Any], dict[str, Any]]:
        """Processes the data into targets.

        Args:
            raw_inputs: raster or vector data to process
            metadata: metadata about the patch being read
            load_targets: whether to load the targets or only inputs

        Returns:
            tuple (input_dict, target_dict) containing the processed inputs and targets
                that are compatible with both metrics and loss functions
        """
        raise NotImplementedError

    def process_output(
        self, raw_output: Any, metadata: SampleMetadata
    ) -> npt.NDArray[Any] | list[Feature] | dict[str, Any]:
        """Processes an output into raster or vector data.

        Args:
            raw_output: the output from prediction head.
            metadata: metadata about the patch being read

        Returns:
            raster data, vector data, or multi-task dictionary output.
        """
        raise NotImplementedError

    def visualize(
        self,
        input_dict: dict[str, Any],
        target_dict: dict[str, Any] | None,
        output: Any,
    ) -> dict[str, npt.NDArray[Any]]:
        """Visualize the outputs and targets.

        Args:
            input_dict: the input dict from process_inputs
            target_dict: the target dict from process_inputs
            output: the prediction

        Returns:
            a dictionary mapping image name to visualization image
        """
        raise NotImplementedError

    def get_metrics(self) -> MetricCollection:
        """Get metrics for this task."""
        raise NotImplementedError

process_inputs 

process_inputs(raw_inputs: dict[str, RasterImage | list[Feature]], metadata: SampleMetadata, load_targets: bool = True) -> tuple[dict[str, Any], dict[str, Any]]

Processes the data into targets.

Parameters:

Name Type Description Default
raw_inputs dict[str, RasterImage | list[Feature]]

raster or vector data to process

 required
metadata SampleMetadata

metadata about the patch being read

 required
load_targets bool

whether to load the targets or only inputs

 True

Returns:

Type Description
tuple[dict[str, Any], dict[str, Any]]

tuple (input_dict, target_dict) containing the processed inputs and targets that are compatible with both metrics and loss functions
Source code in rslearn/train/tasks/task.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
def process_inputs(
    self,
    raw_inputs: dict[str, RasterImage | list[Feature]],
    metadata: SampleMetadata,
    load_targets: bool = True,
) -> tuple[dict[str, Any], dict[str, Any]]:
    """Processes the data into targets.

    Args:
        raw_inputs: raster or vector data to process
        metadata: metadata about the patch being read
        load_targets: whether to load the targets or only inputs

    Returns:
        tuple (input_dict, target_dict) containing the processed inputs and targets
            that are compatible with both metrics and loss functions
    """
    raise NotImplementedError

process_output 

process_output(raw_output: Any, metadata: SampleMetadata) -> NDArray[Any] | list[Feature] | dict[str, Any]

Processes an output into raster or vector data.

Parameters:

Name Type Description Default
raw_output Any

the output from prediction head.

 required
metadata SampleMetadata

metadata about the patch being read

 required

Returns:

Type Description
NDArray[Any] | list[Feature] | dict[str, Any]

raster data, vector data, or multi-task dictionary output.
Source code in rslearn/train/tasks/task.py 
41
42
43
44
45
46
47
48
49
50
51
52
53
def process_output(
    self, raw_output: Any, metadata: SampleMetadata
) -> npt.NDArray[Any] | list[Feature] | dict[str, Any]:
    """Processes an output into raster or vector data.

    Args:
        raw_output: the output from prediction head.
        metadata: metadata about the patch being read

    Returns:
        raster data, vector data, or multi-task dictionary output.
    """
    raise NotImplementedError

visualize 

visualize(input_dict: dict[str, Any], target_dict: dict[str, Any] | None, output: Any) -> dict[str, NDArray[Any]]

Visualize the outputs and targets.

Parameters:

Name Type Description Default
input_dict dict[str, Any]

the input dict from process_inputs

 required
target_dict dict[str, Any] | None

the target dict from process_inputs

 required
output Any

the prediction

 required

Returns:

Type Description
dict[str, NDArray[Any]]

a dictionary mapping image name to visualization image
Source code in rslearn/train/tasks/task.py 
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
def visualize(
    self,
    input_dict: dict[str, Any],
    target_dict: dict[str, Any] | None,
    output: Any,
) -> dict[str, npt.NDArray[Any]]:
    """Visualize the outputs and targets.

    Args:
        input_dict: the input dict from process_inputs
        target_dict: the target dict from process_inputs
        output: the prediction

    Returns:
        a dictionary mapping image name to visualization image
    """
    raise NotImplementedError

get_metrics 

get_metrics() -> MetricCollection

Get metrics for this task.

Source code in rslearn/train/tasks/task.py 
73
74
75
def get_metrics(self) -> MetricCollection:
    """Get metrics for this task."""
    raise NotImplementedError

BasicTask

Bases: Task

A task that provides some support for creating visualizations.

Source code in rslearn/train/tasks/task.py 
 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
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
class BasicTask(Task):
    """A task that provides some support for creating visualizations."""

    def __init__(
        self,
        image_bands: tuple[int, ...] = (0, 1, 2),
        remap_values: tuple[tuple[float, float], tuple[int, int]] | None = None,
    ):
        """Initialize a new BasicTask.

        Args:
            image_bands: which bands from the input image to use for the visualization.
            remap_values: if set, remap the values from the first range to the second range
        """
        self.image_bands = image_bands
        self.remap_values = remap_values

    @staticmethod
    def _get_window_valid_mask(
        reference_hw: torch.Tensor, metadata: SampleMetadata
    ) -> torch.Tensor:
        """Return an HW float mask of pixels that fall within window_bounds.

        Raster readers may pad regions outside window_bounds (but inside crop_bounds)
        with zeros when decoding. This mask lets tasks treat those padded pixels as
        invalid, independent of any nodata_value semantics.
        """
        if reference_hw.ndim != 2:
            raise ValueError(
                f"expected an HW tensor for reference_hw, got shape {reference_hw.shape}"
            )

        window_x0, window_y0, window_x1, window_y1 = metadata.window_bounds
        crop_x0, crop_y0, crop_x1, crop_y1 = metadata.crop_bounds
        inter_x0 = max(window_x0, crop_x0)
        inter_y0 = max(window_y0, crop_y0)
        inter_x1 = min(window_x1, crop_x1)
        inter_y1 = min(window_y1, crop_y1)

        window_valid = torch.zeros(
            reference_hw.shape, dtype=torch.float32, device=reference_hw.device
        )
        if inter_x0 < inter_x1 and inter_y0 < inter_y1:
            x0 = int(inter_x0 - crop_x0)
            y0 = int(inter_y0 - crop_y0)
            x1 = int(inter_x1 - crop_x0)
            y1 = int(inter_y1 - crop_y0)
            window_valid[y0:y1, x0:x1] = 1.0

        return window_valid

    def visualize(
        self,
        input_dict: dict[str, Any],
        target_dict: dict[str, Any] | None,
        output: Any,
    ) -> dict[str, npt.NDArray[Any]]:
        """Visualize the outputs and targets.

        Args:
            input_dict: the input dict from process_inputs
            target_dict: the target dict from process_inputs
            output: the prediction

        Returns:
            a dictionary mapping image name to visualization image
        """
        raster_image = input_dict["image"]
        assert isinstance(raster_image, RasterImage)
        # We don't really handle time series here, just use the first timestep.
        image = raster_image.image.cpu()[self.image_bands, 0, :, :]
        if self.remap_values:
            factor = (self.remap_values[1][1] - self.remap_values[1][0]) / (
                self.remap_values[0][1] - self.remap_values[0][0]
            )
            image = (image - self.remap_values[0][0]) * factor + self.remap_values[1][0]
        return {
            "image": torch.clip(image, 0, 255)
            .numpy()
            .transpose(1, 2, 0)
            .astype(np.uint8),
        }

visualize 

visualize(input_dict: dict[str, Any], target_dict: dict[str, Any] | None, output: Any) -> dict[str, NDArray[Any]]

Visualize the outputs and targets.

Parameters:

Name Type Description Default
input_dict dict[str, Any]

the input dict from process_inputs

 required
target_dict dict[str, Any] | None

the target dict from process_inputs

 required
output Any

the prediction

 required

Returns:

Type Description
dict[str, NDArray[Any]]

a dictionary mapping image name to visualization image
Source code in rslearn/train/tasks/task.py 
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
def visualize(
    self,
    input_dict: dict[str, Any],
    target_dict: dict[str, Any] | None,
    output: Any,
) -> dict[str, npt.NDArray[Any]]:
    """Visualize the outputs and targets.

    Args:
        input_dict: the input dict from process_inputs
        target_dict: the target dict from process_inputs
        output: the prediction

    Returns:
        a dictionary mapping image name to visualization image
    """
    raster_image = input_dict["image"]
    assert isinstance(raster_image, RasterImage)
    # We don't really handle time series here, just use the first timestep.
    image = raster_image.image.cpu()[self.image_bands, 0, :, :]
    if self.remap_values:
        factor = (self.remap_values[1][1] - self.remap_values[1][0]) / (
            self.remap_values[0][1] - self.remap_values[0][0]
        )
        image = (image - self.remap_values[0][0]) * factor + self.remap_values[1][0]
    return {
        "image": torch.clip(image, 0, 255)
        .numpy()
        .transpose(1, 2, 0)
        .astype(np.uint8),
    }