Skip to content

rslearn.train.tasks.segmentation

segmentation

Segmentation task.

SegmentationTask

Bases: BasicTask

A segmentation (per-pixel classification) task.

Source code in rslearn/train/tasks/segmentation.py 
 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
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
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
class SegmentationTask(BasicTask):
    """A segmentation (per-pixel classification) task."""

    def __init__(
        self,
        num_classes: int,
        class_id_mapping: dict[int, int] | None = None,
        colors: list[tuple[int, int, int]] = DEFAULT_COLORS,
        zero_is_invalid: bool = False,
        nodata_value: int | None = None,
        enable_accuracy_metric: bool = True,
        enable_miou_metric: bool = False,
        enable_f1_metric: bool = False,
        report_metric_per_class: bool = False,
        f1_metric_thresholds: list[list[float]] = [[0.5]],
        metric_kwargs: dict[str, Any] = {},
        miou_metric_kwargs: dict[str, Any] = {},
        prob_scales: list[float] | None = None,
        other_metrics: dict[str, Metric] = {},
        output_probs: bool = False,
        output_class_idx: int | None = None,
        enable_confusion_matrix: bool = False,
        class_names: list[str] | None = None,
        **kwargs: Any,
    ) -> None:
        """Initialize a new SegmentationTask.

        Args:
            num_classes: the number of classes to predict
            colors: optional colors for each class
            zero_is_invalid: whether pixels labeled class 0 should be marked invalid
                Mutually exclusive with nodata_value.
            nodata_value: the value to use for nodata pixels. If None, all pixels are
                considered valid. Mutually exclusive with zero_is_invalid.
            class_id_mapping: optional mapping from original class IDs to new class IDs.
                If provided, class labels will be remapped according to this dictionary.
            enable_accuracy_metric: whether to enable the accuracy metric (default
                true).
            enable_f1_metric: whether to enable the F1 metric (default false).
            report_metric_per_class: whether to report chosen metrics for each class, in
                addition to the average score across classes.
            enable_miou_metric: whether to enable the mean IoU metric (default false).
            f1_metric_thresholds: list of list of thresholds to apply for F1 metric.
                Each inner list is used to initialize a separate F1 metric where the
                best F1 across the thresholds within the inner list is computed. If
                there are multiple inner lists, then multiple F1 scores will be
                reported.
            metric_kwargs: additional arguments to pass to underlying metric, see
                torchmetrics.classification.MulticlassAccuracy.
            miou_metric_kwargs: additional arguments to pass to MeanIoUMetric, if
                enable_miou_metric is passed.
            prob_scales: during inference, scale the output probabilities by this much
                before computing the argmax. There is one scale per class. Note that
                this is only applied during prediction, not when computing val or test
                metrics.
            other_metrics: additional metrics to configure on this task.
            output_probs: if True, output raw softmax probabilities instead of class IDs
                during prediction.
            output_class_idx: if set along with output_probs, only output the probability
                for this specific class index (single-channel output).
            enable_confusion_matrix: whether to compute confusion matrix (default false).
                If true, it requires wandb to be initialized for logging.
            class_names: optional list of class names for labeling confusion matrix axes.
                If not provided, classes will be labeled as "class_0", "class_1", etc.
            kwargs: additional arguments to pass to BasicTask
        """
        super().__init__(**kwargs)
        self.num_classes = num_classes
        self.class_id_mapping = class_id_mapping
        self.colors = colors
        self.nodata_value: int | None

        if zero_is_invalid and nodata_value is not None:
            raise ValueError("zero_is_invalid and nodata_value cannot both be set")
        if zero_is_invalid:
            self.nodata_value = 0
        else:
            self.nodata_value = nodata_value

        self.enable_accuracy_metric = enable_accuracy_metric
        self.enable_f1_metric = enable_f1_metric
        self.enable_miou_metric = enable_miou_metric
        self.report_metric_per_class = report_metric_per_class
        self.f1_metric_thresholds = f1_metric_thresholds
        self.metric_kwargs = metric_kwargs
        self.miou_metric_kwargs = miou_metric_kwargs
        self.prob_scales = prob_scales
        self.other_metrics = other_metrics
        self.output_probs = output_probs
        self.output_class_idx = output_class_idx
        self.enable_confusion_matrix = enable_confusion_matrix
        self.class_names = class_names

    def process_inputs(
        self,
        raw_inputs: Mapping[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
        """
        if not load_targets:
            return {}, {}

        assert isinstance(raw_inputs["targets"], RasterImage)
        labels = raw_inputs["targets"].get_hw_tensor().long()

        if self.class_id_mapping is not None:
            new_labels = labels.clone()
            for old_id, new_id in self.class_id_mapping.items():
                new_labels[labels == old_id] = new_id
            labels = new_labels

        window_valid = self._get_window_valid_mask(labels, metadata)
        if self.nodata_value is not None:
            valid = (labels != self.nodata_value).float() * window_valid
            # Labels, even masked ones, must be in the range 0 to num_classes-1
            if self.nodata_value >= self.num_classes:
                labels[labels == self.nodata_value] = 0
        else:
            valid = window_valid

        # Wrap in RasterImage with CTHW format (C=1, T=1) so classes and valid can be
        # used in image transforms.
        return {}, {
            "classes": RasterImage(labels[None, None, :, :], timestamps=None),
            "valid": RasterImage(valid[None, None, :, :], timestamps=None),
        }

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

        Args:
            raw_output: the output from prediction head, which must be a CHW tensor.
            metadata: metadata about the patch being read

        Returns:
            CHW numpy array. If output_probs is False, returns one channel with
            predicted class IDs. If output_probs is True, returns softmax probabilities
            (num_classes channels, or 1 channel if output_class_idx is set).
        """
        if not isinstance(raw_output, torch.Tensor) or len(raw_output.shape) != 3:
            raise ValueError("the output for SegmentationTask must be a CHW tensor")

        if self.prob_scales is not None:
            raw_output = (
                raw_output
                * torch.tensor(
                    self.prob_scales, device=raw_output.device, dtype=raw_output.dtype
                )[:, None, None]
            )

        if self.output_probs:
            # Return raw softmax probabilities
            probs = raw_output.cpu().numpy()
            if self.output_class_idx is not None:
                # Return only the specified class probability
                return probs[self.output_class_idx : self.output_class_idx + 1, :, :]
            return probs

        classes = raw_output.argmax(dim=0).cpu().numpy()
        return classes[None, :, :]

    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
        """
        image = super().visualize(input_dict, target_dict, output)["image"]
        if target_dict is None:
            raise ValueError("target_dict is required for visualization")
        gt_classes = target_dict["classes"].get_hw_tensor().cpu().numpy()
        pred_classes = output.cpu().numpy().argmax(axis=0)
        gt_vis = np.zeros((gt_classes.shape[0], gt_classes.shape[1], 3), dtype=np.uint8)
        pred_vis = np.zeros(
            (pred_classes.shape[0], pred_classes.shape[1], 3), dtype=np.uint8
        )
        for class_id in range(self.num_classes):
            color = self.colors[class_id % len(self.colors)]
            gt_vis[gt_classes == class_id] = color
            pred_vis[pred_classes == class_id] = color

        return {
            "image": np.array(image),
            "gt": gt_vis,
            "pred": pred_vis,
        }

    def get_metrics(self) -> MetricCollection:
        """Get the metrics for this task."""
        metrics = {}

        if self.enable_accuracy_metric:
            accuracy_metric_kwargs = dict(num_classes=self.num_classes)
            accuracy_metric_kwargs.update(self.metric_kwargs)
            metrics["accuracy"] = SegmentationMetric(
                torchmetrics.classification.MulticlassAccuracy(**accuracy_metric_kwargs)
            )

        if self.enable_f1_metric:
            for thresholds in self.f1_metric_thresholds:
                if len(self.f1_metric_thresholds) == 1:
                    suffix = ""
                else:
                    # Metric name can't contain "." so change to ",".
                    suffix = "_" + str(thresholds[0]).replace(".", ",")

                # Create one metric per type -- it returns either a scalar average or
                # a dict with per-class keys.
                metrics["F1" + suffix] = SegmentationMetric(
                    F1Metric(
                        num_classes=self.num_classes,
                        score_thresholds=thresholds,
                        report_per_class=self.report_metric_per_class,
                    ),
                )
                metrics["precision" + suffix] = SegmentationMetric(
                    F1Metric(
                        num_classes=self.num_classes,
                        score_thresholds=thresholds,
                        metric_mode="precision",
                        report_per_class=self.report_metric_per_class,
                    ),
                )
                metrics["recall" + suffix] = SegmentationMetric(
                    F1Metric(
                        num_classes=self.num_classes,
                        score_thresholds=thresholds,
                        metric_mode="recall",
                        report_per_class=self.report_metric_per_class,
                    ),
                )

        if self.enable_miou_metric:
            miou_metric_kwargs: dict[str, Any] = dict(
                num_classes=self.num_classes,
                report_per_class=self.report_metric_per_class,
            )
            if self.nodata_value is not None:
                miou_metric_kwargs["nodata_value"] = self.nodata_value
            miou_metric_kwargs.update(self.miou_metric_kwargs)
            metrics["mean_iou"] = SegmentationMetric(
                MeanIoUMetric(**miou_metric_kwargs),
                pass_probabilities=False,
            )

        if self.other_metrics:
            metrics.update(self.other_metrics)

        if self.enable_confusion_matrix:
            metrics["confusion_matrix"] = SegmentationMetric(
                ConfusionMatrixMetric(
                    num_classes=self.num_classes,
                    class_names=self.class_names,
                ),
            )

        return MetricCollection(metrics)

process_inputs 

process_inputs(raw_inputs: Mapping[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 Mapping[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/segmentation.py 
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
160
161
162
def process_inputs(
    self,
    raw_inputs: Mapping[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
    """
    if not load_targets:
        return {}, {}

    assert isinstance(raw_inputs["targets"], RasterImage)
    labels = raw_inputs["targets"].get_hw_tensor().long()

    if self.class_id_mapping is not None:
        new_labels = labels.clone()
        for old_id, new_id in self.class_id_mapping.items():
            new_labels[labels == old_id] = new_id
        labels = new_labels

    window_valid = self._get_window_valid_mask(labels, metadata)
    if self.nodata_value is not None:
        valid = (labels != self.nodata_value).float() * window_valid
        # Labels, even masked ones, must be in the range 0 to num_classes-1
        if self.nodata_value >= self.num_classes:
            labels[labels == self.nodata_value] = 0
    else:
        valid = window_valid

    # Wrap in RasterImage with CTHW format (C=1, T=1) so classes and valid can be
    # used in image transforms.
    return {}, {
        "classes": RasterImage(labels[None, None, :, :], timestamps=None),
        "valid": RasterImage(valid[None, None, :, :], timestamps=None),
    }

process_output 

process_output(raw_output: Any, metadata: SampleMetadata) -> NDArray[Any]

Processes an output into raster or vector data.

Parameters:

Name Type Description Default
raw_output Any

the output from prediction head, which must be a CHW tensor.

 required
metadata SampleMetadata

metadata about the patch being read

 required

Returns:

Type Description
NDArray[Any]

CHW numpy array. If output_probs is False, returns one channel with
NDArray[Any]

predicted class IDs. If output_probs is True, returns softmax probabilities
NDArray[Any]

(num_classes channels, or 1 channel if output_class_idx is set).
Source code in rslearn/train/tasks/segmentation.py 
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
def process_output(
    self, raw_output: Any, metadata: SampleMetadata
) -> npt.NDArray[Any]:
    """Processes an output into raster or vector data.

    Args:
        raw_output: the output from prediction head, which must be a CHW tensor.
        metadata: metadata about the patch being read

    Returns:
        CHW numpy array. If output_probs is False, returns one channel with
        predicted class IDs. If output_probs is True, returns softmax probabilities
        (num_classes channels, or 1 channel if output_class_idx is set).
    """
    if not isinstance(raw_output, torch.Tensor) or len(raw_output.shape) != 3:
        raise ValueError("the output for SegmentationTask must be a CHW tensor")

    if self.prob_scales is not None:
        raw_output = (
            raw_output
            * torch.tensor(
                self.prob_scales, device=raw_output.device, dtype=raw_output.dtype
            )[:, None, None]
        )

    if self.output_probs:
        # Return raw softmax probabilities
        probs = raw_output.cpu().numpy()
        if self.output_class_idx is not None:
            # Return only the specified class probability
            return probs[self.output_class_idx : self.output_class_idx + 1, :, :]
        return probs

    classes = raw_output.argmax(dim=0).cpu().numpy()
    return classes[None, :, :]

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/segmentation.py 
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
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
    """
    image = super().visualize(input_dict, target_dict, output)["image"]
    if target_dict is None:
        raise ValueError("target_dict is required for visualization")
    gt_classes = target_dict["classes"].get_hw_tensor().cpu().numpy()
    pred_classes = output.cpu().numpy().argmax(axis=0)
    gt_vis = np.zeros((gt_classes.shape[0], gt_classes.shape[1], 3), dtype=np.uint8)
    pred_vis = np.zeros(
        (pred_classes.shape[0], pred_classes.shape[1], 3), dtype=np.uint8
    )
    for class_id in range(self.num_classes):
        color = self.colors[class_id % len(self.colors)]
        gt_vis[gt_classes == class_id] = color
        pred_vis[pred_classes == class_id] = color

    return {
        "image": np.array(image),
        "gt": gt_vis,
        "pred": pred_vis,
    }

get_metrics 

get_metrics() -> MetricCollection

Get the metrics for this task.

Source code in rslearn/train/tasks/segmentation.py 
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
def get_metrics(self) -> MetricCollection:
    """Get the metrics for this task."""
    metrics = {}

    if self.enable_accuracy_metric:
        accuracy_metric_kwargs = dict(num_classes=self.num_classes)
        accuracy_metric_kwargs.update(self.metric_kwargs)
        metrics["accuracy"] = SegmentationMetric(
            torchmetrics.classification.MulticlassAccuracy(**accuracy_metric_kwargs)
        )

    if self.enable_f1_metric:
        for thresholds in self.f1_metric_thresholds:
            if len(self.f1_metric_thresholds) == 1:
                suffix = ""
            else:
                # Metric name can't contain "." so change to ",".
                suffix = "_" + str(thresholds[0]).replace(".", ",")

            # Create one metric per type -- it returns either a scalar average or
            # a dict with per-class keys.
            metrics["F1" + suffix] = SegmentationMetric(
                F1Metric(
                    num_classes=self.num_classes,
                    score_thresholds=thresholds,
                    report_per_class=self.report_metric_per_class,
                ),
            )
            metrics["precision" + suffix] = SegmentationMetric(
                F1Metric(
                    num_classes=self.num_classes,
                    score_thresholds=thresholds,
                    metric_mode="precision",
                    report_per_class=self.report_metric_per_class,
                ),
            )
            metrics["recall" + suffix] = SegmentationMetric(
                F1Metric(
                    num_classes=self.num_classes,
                    score_thresholds=thresholds,
                    metric_mode="recall",
                    report_per_class=self.report_metric_per_class,
                ),
            )

    if self.enable_miou_metric:
        miou_metric_kwargs: dict[str, Any] = dict(
            num_classes=self.num_classes,
            report_per_class=self.report_metric_per_class,
        )
        if self.nodata_value is not None:
            miou_metric_kwargs["nodata_value"] = self.nodata_value
        miou_metric_kwargs.update(self.miou_metric_kwargs)
        metrics["mean_iou"] = SegmentationMetric(
            MeanIoUMetric(**miou_metric_kwargs),
            pass_probabilities=False,
        )

    if self.other_metrics:
        metrics.update(self.other_metrics)

    if self.enable_confusion_matrix:
        metrics["confusion_matrix"] = SegmentationMetric(
            ConfusionMatrixMetric(
                num_classes=self.num_classes,
                class_names=self.class_names,
            ),
        )

    return MetricCollection(metrics)

SegmentationHead

Bases: Predictor

Head for segmentation task.

Source code in rslearn/train/tasks/segmentation.py 
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
class SegmentationHead(Predictor):
    """Head for segmentation task."""

    def __init__(
        self,
        weights: list[float] | None = None,
        dice_loss: bool = False,
        temperature: float = 1.0,
    ):
        """Initialize a new SegmentationTask.

        Args:
            weights: weights for cross entropy loss (Tensor of size C)
            dice_loss: whether to add dice loss to cross entropy
            temperature: temperature scaling for softmax, does not affect the loss,
                only the predictor outputs
        """
        super().__init__()
        if weights is not None:
            self.register_buffer("weights", torch.Tensor(weights))
        else:
            self.weights = None
        self.dice_loss = dice_loss
        self.temperature = temperature

    def forward(
        self,
        intermediates: Any,
        context: ModelContext,
        targets: list[dict[str, Any]] | None = None,
    ) -> ModelOutput:
        """Compute the segmentation outputs from logits and targets.

        Args:
            intermediates: a FeatureMaps with a single feature map containing the
                segmentation logits.
            context: the model context
            targets: list of target dicts, where each target dict must contain a key
                "classes" containing the per-pixel class labels, along with "valid"
                containing a mask indicating where the example is valid.

        Returns:
            tuple of outputs and loss dict
        """
        if not isinstance(intermediates, FeatureMaps):
            raise ValueError("input to SegmentationHead must be a FeatureMaps")
        if len(intermediates.feature_maps) != 1:
            raise ValueError(
                f"input to SegmentationHead must have one feature map, but got {len(intermediates.feature_maps)}"
            )

        logits = intermediates.feature_maps[0]
        outputs = torch.nn.functional.softmax(logits / self.temperature, dim=1)

        losses = {}
        if targets:
            labels = torch.stack(
                [target["classes"].get_hw_tensor() for target in targets], dim=0
            )
            mask = torch.stack(
                [target["valid"].get_hw_tensor() for target in targets], dim=0
            )
            per_pixel_loss = torch.nn.functional.cross_entropy(
                logits, labels, weight=self.weights, reduction="none"
            )
            mask_sum = torch.sum(mask)
            if mask_sum > 0:
                # Compute average loss over valid pixels.
                losses["cls"] = torch.sum(per_pixel_loss * mask) / torch.sum(mask)
            else:
                # If there are no valid pixels, we avoid dividing by zero and just let
                # the summed mask loss be zero.
                losses["cls"] = torch.sum(per_pixel_loss * mask)
            if self.dice_loss:
                softmax_woT = torch.nn.functional.softmax(logits, dim=1)
                dice_loss = DiceLoss()(softmax_woT, labels, mask)
                losses["dice"] = dice_loss

        return ModelOutput(
            outputs=outputs,
            loss_dict=losses,
        )

forward 

forward(intermediates: Any, context: ModelContext, targets: list[dict[str, Any]] | None = None) -> ModelOutput

Compute the segmentation outputs from logits and targets.

Parameters:

Name Type Description Default
intermediates Any

a FeatureMaps with a single feature map containing the segmentation logits.

 required
context ModelContext

the model context

 required
targets list[dict[str, Any]] | None

list of target dicts, where each target dict must contain a key "classes" containing the per-pixel class labels, along with "valid" containing a mask indicating where the example is valid.

 None

Returns:

Type Description
ModelOutput

tuple of outputs and loss dict
Source code in rslearn/train/tasks/segmentation.py 
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
def forward(
    self,
    intermediates: Any,
    context: ModelContext,
    targets: list[dict[str, Any]] | None = None,
) -> ModelOutput:
    """Compute the segmentation outputs from logits and targets.

    Args:
        intermediates: a FeatureMaps with a single feature map containing the
            segmentation logits.
        context: the model context
        targets: list of target dicts, where each target dict must contain a key
            "classes" containing the per-pixel class labels, along with "valid"
            containing a mask indicating where the example is valid.

    Returns:
        tuple of outputs and loss dict
    """
    if not isinstance(intermediates, FeatureMaps):
        raise ValueError("input to SegmentationHead must be a FeatureMaps")
    if len(intermediates.feature_maps) != 1:
        raise ValueError(
            f"input to SegmentationHead must have one feature map, but got {len(intermediates.feature_maps)}"
        )

    logits = intermediates.feature_maps[0]
    outputs = torch.nn.functional.softmax(logits / self.temperature, dim=1)

    losses = {}
    if targets:
        labels = torch.stack(
            [target["classes"].get_hw_tensor() for target in targets], dim=0
        )
        mask = torch.stack(
            [target["valid"].get_hw_tensor() for target in targets], dim=0
        )
        per_pixel_loss = torch.nn.functional.cross_entropy(
            logits, labels, weight=self.weights, reduction="none"
        )
        mask_sum = torch.sum(mask)
        if mask_sum > 0:
            # Compute average loss over valid pixels.
            losses["cls"] = torch.sum(per_pixel_loss * mask) / torch.sum(mask)
        else:
            # If there are no valid pixels, we avoid dividing by zero and just let
            # the summed mask loss be zero.
            losses["cls"] = torch.sum(per_pixel_loss * mask)
        if self.dice_loss:
            softmax_woT = torch.nn.functional.softmax(logits, dim=1)
            dice_loss = DiceLoss()(softmax_woT, labels, mask)
            losses["dice"] = dice_loss

    return ModelOutput(
        outputs=outputs,
        loss_dict=losses,
    )

SegmentationMetric

Bases: Metric

Metric for segmentation task.

Source code in rslearn/train/tasks/segmentation.py 
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
class SegmentationMetric(Metric):
    """Metric for segmentation task."""

    def __init__(
        self,
        metric: Metric,
        pass_probabilities: bool = True,
        class_idx: int | None = None,
        output_key: str | None = None,
    ):
        """Initialize a new SegmentationMetric.

        Args:
            metric: the metric to wrap. This wrapping class will handle selecting the
                classes from the targets and masking out invalid pixels.
            pass_probabilities: whether to pass predicted probabilities to the metric.
                If False, argmax is applied to pass the predicted classes instead.
            class_idx: if set, return only this class index's value. For backward
                compatibility with configs using standard torchmetrics. Internally
                converted to output_key="cls_{class_idx}".
            output_key: if the wrapped metric returns a dict (or a tensor that gets
                converted to a dict), return only this key's value. For standard
                torchmetrics with average=None, tensors are converted to dicts with
                keys "cls_0", "cls_1", etc. If None, the full dict is returned.
        """
        super().__init__()
        self.metric = metric
        self.pass_probabilities = pass_probabilities
        self.class_idx = class_idx
        self.output_key = output_key

    def update(
        self, preds: list[Any] | torch.Tensor, targets: list[dict[str, Any]]
    ) -> None:
        """Update metric.

        Args:
            preds: the predictions
            targets: the targets
        """
        if not isinstance(preds, torch.Tensor):
            preds = torch.stack(preds)
        labels = torch.stack([target["classes"].get_hw_tensor() for target in targets])

        # Sub-select the valid labels.
        # We flatten the prediction and label images at valid pixels.
        # Prediction is changed from BCHW to BHWC so we can select the valid BHW mask.
        mask = torch.stack([target["valid"].get_hw_tensor() > 0 for target in targets])
        preds = preds.permute(0, 2, 3, 1)[mask]
        labels = labels[mask]
        if len(preds) == 0:
            return

        if not self.pass_probabilities:
            preds = preds.argmax(dim=1)

        self.metric.update(preds, labels)

    def compute(self) -> Any:
        """Returns the computed metric.

        If the wrapped metric returns a multi-element tensor (e.g., standard torchmetrics
        with average=None), it is converted to a dict with keys like "cls_0", "cls_1", etc.
        This allows uniform handling via output_key for both standard torchmetrics and
        custom dict-returning metrics.
        """
        result = self.metric.compute()

        # Convert multi-element tensors to dict for uniform handling.
        # This supports standard torchmetrics with average=None which return per-class tensors.
        if isinstance(result, torch.Tensor) and result.ndim >= 1:
            result = {f"cls_{i}": result[i] for i in range(len(result))}

        if self.output_key is not None:
            if not isinstance(result, dict):
                raise TypeError(
                    f"output_key is set to '{self.output_key}' but metric returned "
                    f"{type(result).__name__} instead of dict"
                )
            return result[self.output_key]
        if self.class_idx is not None:
            # For backward compatibility: class_idx can index into the converted dict
            if isinstance(result, dict):
                return result[f"cls_{self.class_idx}"]
            return result[self.class_idx]

        return result

    def reset(self) -> None:
        """Reset metric."""
        super().reset()
        self.metric.reset()

    def plot(self, *args: list[Any], **kwargs: dict[str, Any]) -> Any:
        """Returns a plot of the metric."""
        return self.metric.plot(*args, **kwargs)

update 

update(preds: list[Any] | Tensor, targets: list[dict[str, Any]]) -> None

Update metric.

Parameters:

Name Type Description Default
preds list[Any] | Tensor

the predictions

 required
targets list[dict[str, Any]]

the targets

 required
Source code in rslearn/train/tasks/segmentation.py 
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
def update(
    self, preds: list[Any] | torch.Tensor, targets: list[dict[str, Any]]
) -> None:
    """Update metric.

    Args:
        preds: the predictions
        targets: the targets
    """
    if not isinstance(preds, torch.Tensor):
        preds = torch.stack(preds)
    labels = torch.stack([target["classes"].get_hw_tensor() for target in targets])

    # Sub-select the valid labels.
    # We flatten the prediction and label images at valid pixels.
    # Prediction is changed from BCHW to BHWC so we can select the valid BHW mask.
    mask = torch.stack([target["valid"].get_hw_tensor() > 0 for target in targets])
    preds = preds.permute(0, 2, 3, 1)[mask]
    labels = labels[mask]
    if len(preds) == 0:
        return

    if not self.pass_probabilities:
        preds = preds.argmax(dim=1)

    self.metric.update(preds, labels)

compute 

compute() -> Any

Returns the computed metric.

If the wrapped metric returns a multi-element tensor (e.g., standard torchmetrics with average=None), it is converted to a dict with keys like "cls_0", "cls_1", etc. This allows uniform handling via output_key for both standard torchmetrics and custom dict-returning metrics.

Source code in rslearn/train/tasks/segmentation.py 
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
def compute(self) -> Any:
    """Returns the computed metric.

    If the wrapped metric returns a multi-element tensor (e.g., standard torchmetrics
    with average=None), it is converted to a dict with keys like "cls_0", "cls_1", etc.
    This allows uniform handling via output_key for both standard torchmetrics and
    custom dict-returning metrics.
    """
    result = self.metric.compute()

    # Convert multi-element tensors to dict for uniform handling.
    # This supports standard torchmetrics with average=None which return per-class tensors.
    if isinstance(result, torch.Tensor) and result.ndim >= 1:
        result = {f"cls_{i}": result[i] for i in range(len(result))}

    if self.output_key is not None:
        if not isinstance(result, dict):
            raise TypeError(
                f"output_key is set to '{self.output_key}' but metric returned "
                f"{type(result).__name__} instead of dict"
            )
        return result[self.output_key]
    if self.class_idx is not None:
        # For backward compatibility: class_idx can index into the converted dict
        if isinstance(result, dict):
            return result[f"cls_{self.class_idx}"]
        return result[self.class_idx]

    return result

reset 

reset() -> None

Reset metric.

Source code in rslearn/train/tasks/segmentation.py 
480
481
482
483
def reset(self) -> None:
    """Reset metric."""
    super().reset()
    self.metric.reset()

plot 

plot(*args: list[Any], **kwargs: dict[str, Any]) -> Any

Returns a plot of the metric.

Source code in rslearn/train/tasks/segmentation.py 
485
486
487
def plot(self, *args: list[Any], **kwargs: dict[str, Any]) -> Any:
    """Returns a plot of the metric."""
    return self.metric.plot(*args, **kwargs)

F1Metric

Bases: Metric

F1 score for segmentation.

It treats each class as a separate prediction task, and computes the maximum F1 score under the different configured thresholds per-class.

Source code in rslearn/train/tasks/segmentation.py 
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
class F1Metric(Metric):
    """F1 score for segmentation.

    It treats each class as a separate prediction task, and computes the maximum F1
    score under the different configured thresholds per-class.
    """

    def __init__(
        self,
        num_classes: int,
        score_thresholds: list[float],
        metric_mode: str = "f1",
        report_per_class: bool = False,
    ):
        """Create a new F1Metric.

        Args:
            num_classes: number of classes.
            score_thresholds: list of score thresholds to check F1 score for. The final
                metric is the best F1 across score thresholds.
            metric_mode: set to "precision" or "recall" to return that instead of F1
                (default "f1")
            report_per_class: whether to return a dict with per-class scores and "avg"
                score instead of just returning the scalar average score.
        """
        super().__init__()
        self.num_classes = num_classes
        self.score_thresholds = score_thresholds
        self.metric_mode = metric_mode
        self.report_per_class = report_per_class

        assert self.metric_mode in ["f1", "precision", "recall"]

        for cls_idx in range(self.num_classes):
            for thr_idx in range(len(self.score_thresholds)):
                cur_prefix = self._get_state_prefix(cls_idx, thr_idx)
                self.add_state(
                    cur_prefix + "tp", default=torch.tensor(0), dist_reduce_fx="sum"
                )
                self.add_state(
                    cur_prefix + "fp", default=torch.tensor(0), dist_reduce_fx="sum"
                )
                self.add_state(
                    cur_prefix + "fn", default=torch.tensor(0), dist_reduce_fx="sum"
                )

    def _get_state_prefix(self, cls_idx: int, thr_idx: int) -> str:
        return f"{cls_idx}_{thr_idx}_"

    def update(self, preds: torch.Tensor, labels: torch.Tensor) -> None:
        """Update metric.

        Args:
            preds: the predictions, NxC.
            labels: the targets, N, with values from 0 to C-1.
        """
        for cls_idx in range(self.num_classes):
            for thr_idx, score_threshold in enumerate(self.score_thresholds):
                pred_bin = preds[:, cls_idx] > score_threshold
                gt_bin = labels == cls_idx

                tp = torch.count_nonzero(pred_bin & gt_bin).item()
                fp = torch.count_nonzero(pred_bin & torch.logical_not(gt_bin)).item()
                fn = torch.count_nonzero(torch.logical_not(pred_bin) & gt_bin).item()

                cur_prefix = self._get_state_prefix(cls_idx, thr_idx)
                setattr(self, cur_prefix + "tp", getattr(self, cur_prefix + "tp") + tp)
                setattr(self, cur_prefix + "fp", getattr(self, cur_prefix + "fp") + fp)
                setattr(self, cur_prefix + "fn", getattr(self, cur_prefix + "fn") + fn)

    def compute(self) -> Any:
        """Compute metric.

        Returns:
            the average F1 score, or, if report_per_class is True, a dict with "f1/avg"
                key containing the average score and "f1/cls_N" keys for each class N.
        """
        cls_best_scores = {}

        for cls_idx in range(self.num_classes):
            best_score = None

            for thr_idx in range(len(self.score_thresholds)):
                cur_prefix = self._get_state_prefix(cls_idx, thr_idx)
                tp = getattr(self, cur_prefix + "tp")
                fp = getattr(self, cur_prefix + "fp")
                fn = getattr(self, cur_prefix + "fn")
                device = tp.device

                if tp + fp == 0:
                    precision = torch.tensor(0, dtype=torch.float32, device=device)
                else:
                    precision = tp / (tp + fp)

                if tp + fn == 0:
                    recall = torch.tensor(0, dtype=torch.float32, device=device)
                else:
                    recall = tp / (tp + fn)

                if precision + recall < 0.001:
                    f1 = torch.tensor(0, dtype=torch.float32, device=device)
                else:
                    f1 = 2 * precision * recall / (precision + recall)

                if self.metric_mode == "f1":
                    score = f1
                elif self.metric_mode == "precision":
                    score = precision
                elif self.metric_mode == "recall":
                    score = recall

                if best_score is None or score > best_score:
                    best_score = score

            cls_best_scores[f"f1/cls_{cls_idx}"] = best_score

        average_score = torch.mean(torch.stack(list(cls_best_scores.values())))

        if self.report_per_class:
            report_scores = {"f1/avg": average_score}
            report_scores.update(cls_best_scores)
            return report_scores
        else:
            return average_score

update 

update(preds: Tensor, labels: Tensor) -> None

Update metric.

Parameters:

Name Type Description Default
preds Tensor

the predictions, NxC.

 required
labels Tensor

the targets, N, with values from 0 to C-1.

 required
Source code in rslearn/train/tasks/segmentation.py 
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
def update(self, preds: torch.Tensor, labels: torch.Tensor) -> None:
    """Update metric.

    Args:
        preds: the predictions, NxC.
        labels: the targets, N, with values from 0 to C-1.
    """
    for cls_idx in range(self.num_classes):
        for thr_idx, score_threshold in enumerate(self.score_thresholds):
            pred_bin = preds[:, cls_idx] > score_threshold
            gt_bin = labels == cls_idx

            tp = torch.count_nonzero(pred_bin & gt_bin).item()
            fp = torch.count_nonzero(pred_bin & torch.logical_not(gt_bin)).item()
            fn = torch.count_nonzero(torch.logical_not(pred_bin) & gt_bin).item()

            cur_prefix = self._get_state_prefix(cls_idx, thr_idx)
            setattr(self, cur_prefix + "tp", getattr(self, cur_prefix + "tp") + tp)
            setattr(self, cur_prefix + "fp", getattr(self, cur_prefix + "fp") + fp)
            setattr(self, cur_prefix + "fn", getattr(self, cur_prefix + "fn") + fn)

compute 

compute() -> Any

Compute metric.

Returns:

Type Description
Any

the average F1 score, or, if report_per_class is True, a dict with "f1/avg" key containing the average score and "f1/cls_N" keys for each class N.
Source code in rslearn/train/tasks/segmentation.py 
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
def compute(self) -> Any:
    """Compute metric.

    Returns:
        the average F1 score, or, if report_per_class is True, a dict with "f1/avg"
            key containing the average score and "f1/cls_N" keys for each class N.
    """
    cls_best_scores = {}

    for cls_idx in range(self.num_classes):
        best_score = None

        for thr_idx in range(len(self.score_thresholds)):
            cur_prefix = self._get_state_prefix(cls_idx, thr_idx)
            tp = getattr(self, cur_prefix + "tp")
            fp = getattr(self, cur_prefix + "fp")
            fn = getattr(self, cur_prefix + "fn")
            device = tp.device

            if tp + fp == 0:
                precision = torch.tensor(0, dtype=torch.float32, device=device)
            else:
                precision = tp / (tp + fp)

            if tp + fn == 0:
                recall = torch.tensor(0, dtype=torch.float32, device=device)
            else:
                recall = tp / (tp + fn)

            if precision + recall < 0.001:
                f1 = torch.tensor(0, dtype=torch.float32, device=device)
            else:
                f1 = 2 * precision * recall / (precision + recall)

            if self.metric_mode == "f1":
                score = f1
            elif self.metric_mode == "precision":
                score = precision
            elif self.metric_mode == "recall":
                score = recall

            if best_score is None or score > best_score:
                best_score = score

        cls_best_scores[f"f1/cls_{cls_idx}"] = best_score

    average_score = torch.mean(torch.stack(list(cls_best_scores.values())))

    if self.report_per_class:
        report_scores = {"f1/avg": average_score}
        report_scores.update(cls_best_scores)
        return report_scores
    else:
        return average_score

MeanIoUMetric

Bases: Metric

Mean IoU for segmentation.

This is the mean of the per-class intersection-over-union scores. The per-class intersection is the number of pixels across all examples where the predicted label and ground truth label are both that class, and the per-class union is defined similarly.

This differs from torchmetrics.segmentation.MeanIoU, where the mean IoU is computed per-image, and averaged across images.

Source code in rslearn/train/tasks/segmentation.py 
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
class MeanIoUMetric(Metric):
    """Mean IoU for segmentation.

    This is the mean of the per-class intersection-over-union scores. The per-class
    intersection is the number of pixels across all examples where the predicted label
    and ground truth label are both that class, and the per-class union is defined
    similarly.

    This differs from torchmetrics.segmentation.MeanIoU, where the mean IoU is computed
    per-image, and averaged across images.
    """

    def __init__(
        self,
        num_classes: int,
        nodata_value: int | None = None,
        ignore_missing_classes: bool = False,
        report_per_class: bool = False,
    ):
        """Create a new MeanIoUMetric.

        Args:
            num_classes: the number of classes for the task.
            nodata_value: the value to treat as nodata/invalid. If set and is one of the
                classes, IoU will not be calculated for it. If None, or not one of the
                classes, IoU is calculated for all classes.
            ignore_missing_classes: whether to ignore classes that don't appear in
                either the predictions or the ground truth. If false, the IoU for a
                missing class will be 0.
            report_per_class: whether to return a dict with per-class mean IoU scores
                and "avg" average across classes instead of returning the scalar
                average only.
        """
        super().__init__()
        self.num_classes = num_classes
        self.nodata_value = nodata_value
        self.ignore_missing_classes = ignore_missing_classes
        self.report_per_class = report_per_class

        self.add_state(
            "intersections", default=torch.zeros(self.num_classes), dist_reduce_fx="sum"
        )
        self.add_state(
            "unions", default=torch.zeros(self.num_classes), dist_reduce_fx="sum"
        )

    def update(self, preds: torch.Tensor, labels: torch.Tensor) -> None:
        """Update metric.

        Like torchmetrics.segmentation.MeanIoU with input_format="index", we expect
        predictions and labels to both be class integers. This is achieved by passing
        pass_probabilities=False to the SegmentationMetric wrapper.

        Args:
            preds: the predictions, (N,), with values from 0 to C-1.
            labels: the targets, (N,), with values from 0 to C-1.
        """
        if preds.min() < 0 or preds.max() >= self.num_classes:
            raise ValueError("predicted class outside of expected range")
        if labels.min() < 0 or labels.max() >= self.num_classes:
            raise ValueError("label class outside of expected range")

        new_intersections = torch.zeros(
            self.num_classes, device=self.intersections.device
        )
        new_unions = torch.zeros(self.num_classes, device=self.unions.device)
        for cls_idx in range(self.num_classes):
            new_intersections[cls_idx] = (
                (preds == cls_idx) & (labels == cls_idx)
            ).sum()
            new_unions[cls_idx] = ((preds == cls_idx) | (labels == cls_idx)).sum()
        self.intersections += new_intersections
        self.unions += new_unions

    def compute(self) -> Any:
        """Compute metric.

        Returns:
            the mean IoU score, or, if report_per_class is true, a dict with
            "mean_iou/avg" key containing the mean IoU across classes and
            "mean_iou/cls_N" keys containing the mean IoU for each class N.
        """
        cls_scores = {}
        valid_scores = []

        for cls_idx in range(self.num_classes):
            # Check if nodata_value is set and is one of the classes
            if self.nodata_value is not None and cls_idx == self.nodata_value:
                continue

            intersection = self.intersections[cls_idx]
            union = self.unions[cls_idx]

            if union == 0 and self.ignore_missing_classes:
                continue

            score = intersection / union
            cls_scores[f"mean_iou/cls_{cls_idx}"] = score
            valid_scores.append(score)

        mean_iou = torch.mean(torch.stack(valid_scores))

        if self.report_per_class:
            report_scores = {"mean_iou/avg": mean_iou}
            report_scores.update(cls_scores)
            return report_scores
        else:
            return mean_iou

update 

update(preds: Tensor, labels: Tensor) -> None

Update metric.

Like torchmetrics.segmentation.MeanIoU with input_format="index", we expect predictions and labels to both be class integers. This is achieved by passing pass_probabilities=False to the SegmentationMetric wrapper.

Parameters:

Name Type Description Default
preds Tensor

the predictions, (N,), with values from 0 to C-1.

 required
labels Tensor

the targets, (N,), with values from 0 to C-1.

 required
Source code in rslearn/train/tasks/segmentation.py 
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
def update(self, preds: torch.Tensor, labels: torch.Tensor) -> None:
    """Update metric.

    Like torchmetrics.segmentation.MeanIoU with input_format="index", we expect
    predictions and labels to both be class integers. This is achieved by passing
    pass_probabilities=False to the SegmentationMetric wrapper.

    Args:
        preds: the predictions, (N,), with values from 0 to C-1.
        labels: the targets, (N,), with values from 0 to C-1.
    """
    if preds.min() < 0 or preds.max() >= self.num_classes:
        raise ValueError("predicted class outside of expected range")
    if labels.min() < 0 or labels.max() >= self.num_classes:
        raise ValueError("label class outside of expected range")

    new_intersections = torch.zeros(
        self.num_classes, device=self.intersections.device
    )
    new_unions = torch.zeros(self.num_classes, device=self.unions.device)
    for cls_idx in range(self.num_classes):
        new_intersections[cls_idx] = (
            (preds == cls_idx) & (labels == cls_idx)
        ).sum()
        new_unions[cls_idx] = ((preds == cls_idx) | (labels == cls_idx)).sum()
    self.intersections += new_intersections
    self.unions += new_unions

compute 

compute() -> Any

Compute metric.

Returns:

Type Description
Any

the mean IoU score, or, if report_per_class is true, a dict with
Any

"mean_iou/avg" key containing the mean IoU across classes and
Any

"mean_iou/cls_N" keys containing the mean IoU for each class N.
Source code in rslearn/train/tasks/segmentation.py 
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
def compute(self) -> Any:
    """Compute metric.

    Returns:
        the mean IoU score, or, if report_per_class is true, a dict with
        "mean_iou/avg" key containing the mean IoU across classes and
        "mean_iou/cls_N" keys containing the mean IoU for each class N.
    """
    cls_scores = {}
    valid_scores = []

    for cls_idx in range(self.num_classes):
        # Check if nodata_value is set and is one of the classes
        if self.nodata_value is not None and cls_idx == self.nodata_value:
            continue

        intersection = self.intersections[cls_idx]
        union = self.unions[cls_idx]

        if union == 0 and self.ignore_missing_classes:
            continue

        score = intersection / union
        cls_scores[f"mean_iou/cls_{cls_idx}"] = score
        valid_scores.append(score)

    mean_iou = torch.mean(torch.stack(valid_scores))

    if self.report_per_class:
        report_scores = {"mean_iou/avg": mean_iou}
        report_scores.update(cls_scores)
        return report_scores
    else:
        return mean_iou

DiceLoss

Bases: Module

Mean Dice Loss for segmentation.

This is the mean of the per-class dice loss (1 - 2*intersection / union scores). The per-class intersection is the number of pixels across all examples where the predicted label and ground truth label are both that class, and the per-class union is defined similarly.

Source code in rslearn/train/tasks/segmentation.py 
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
class DiceLoss(torch.nn.Module):
    """Mean Dice Loss for segmentation.

    This is the mean of the per-class dice loss (1 - 2*intersection / union scores).
    The per-class intersection is the number of pixels across all examples where
    the predicted label and ground truth label are both that class, and the per-class
    union is defined similarly.
    """

    def __init__(self, smooth: float = 1e-7):
        """Initialize a new DiceLoss."""
        super().__init__()
        self.smooth = smooth

    def forward(
        self, inputs: torch.Tensor, targets: torch.Tensor, mask: torch.Tensor
    ) -> torch.Tensor:
        """Compute Dice Loss.

        Returns:
            the mean Dicen Loss across classes
        """
        num_classes = inputs.shape[1]
        targets_one_hot = (
            torch.nn.functional.one_hot(targets, num_classes)
            .permute(0, 3, 1, 2)
            .float()
        )

        # Expand mask to [B, C, H, W]
        mask = mask.unsqueeze(1).expand_as(inputs)

        dice_per_class = []
        for c in range(num_classes):
            pred_c = inputs[:, c] * mask[:, c]
            target_c = targets_one_hot[:, c] * mask[:, c]

            intersection = (pred_c * target_c).sum()
            union = pred_c.sum() + target_c.sum()
            dice_c = (2.0 * intersection + self.smooth) / (union + self.smooth)
            dice_per_class.append(dice_c)

        return 1 - torch.stack(dice_per_class).mean()

forward 

forward(inputs: Tensor, targets: Tensor, mask: Tensor) -> Tensor

Compute Dice Loss.

Returns:

Type Description
Tensor

the mean Dicen Loss across classes
Source code in rslearn/train/tasks/segmentation.py 
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
def forward(
    self, inputs: torch.Tensor, targets: torch.Tensor, mask: torch.Tensor
) -> torch.Tensor:
    """Compute Dice Loss.

    Returns:
        the mean Dicen Loss across classes
    """
    num_classes = inputs.shape[1]
    targets_one_hot = (
        torch.nn.functional.one_hot(targets, num_classes)
        .permute(0, 3, 1, 2)
        .float()
    )

    # Expand mask to [B, C, H, W]
    mask = mask.unsqueeze(1).expand_as(inputs)

    dice_per_class = []
    for c in range(num_classes):
        pred_c = inputs[:, c] * mask[:, c]
        target_c = targets_one_hot[:, c] * mask[:, c]

        intersection = (pred_c * target_c).sum()
        union = pred_c.sum() + target_c.sum()
        dice_c = (2.0 * intersection + self.smooth) / (union + self.smooth)
        dice_per_class.append(dice_c)

    return 1 - torch.stack(dice_per_class).mean()