rslearn.train.tasks

tasks

rslearn 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

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

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

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

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

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