rslearn 0.0.1__py3-none-any.whl → 0.0.21__py3-none-any.whl

rslearn/train/tasks/classification.py

@@ -15,6 +15,13 @@ from torchmetrics.classification import (
     MulticlassRecall,
 )
 
+from rslearn.models.component import FeatureVector, Predictor
+from rslearn.train.model_context import (
+    ModelContext,
+    ModelOutput,
+    RasterImage,
+    SampleMetadata,
+)
 from rslearn.utils import Feature, STGeometry
 
 from .task import BasicTask
@@ -27,7 +34,7 @@ class ClassificationTask(BasicTask):
         self,
         property_name: str,
         classes: list[str],
-        filters: list[tuple[str, str]] | None = None,
+        filters: list[tuple[str, str]] = [],
         read_class_id: bool = False,
         allow_invalid: bool = False,
         skip_unknown_categories: bool = False,
@@ -37,7 +44,7 @@ class ClassificationTask(BasicTask):
         f1_metric_kwargs: dict[str, Any] = {},
         positive_class: str | None = None,
         positive_class_threshold: float = 0.5,
-        **kwargs,
+        **kwargs: Any,
     ):
         """Initialize a new ClassificationTask.
 
@@ -49,8 +56,8 @@ class ClassificationTask(BasicTask):
                 features with matching properties.
             read_class_id: whether to read an integer class ID instead of the class
                 name.
-            allow_invalid: instead of throwing error when no regression label is found
-                at a window, simply mark the example invalid for this task
+            allow_invalid: instead of throwing error when no classification label is
+                found at a window, simply mark the example invalid for this task
             skip_unknown_categories: whether to skip examples with categories that are
                 not passed via classes, instead of throwing error
             prob_property: when predicting, write probabilities in addition to class ID
@@ -95,13 +102,10 @@ class ClassificationTask(BasicTask):
             else:
                 self.positive_class_id = self.classes.index(self.positive_class)
 
-        if not self.filters:
-            self.filters = []
-
     def process_inputs(
         self,
-        raw_inputs: dict[str, torch.Tensor | list[Feature]],
-        metadata: dict[str, Any],
+        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.
@@ -119,7 +123,10 @@ class ClassificationTask(BasicTask):
             return {}, {}
 
         data = raw_inputs["targets"]
+        assert isinstance(data, list)
         for feat in data:
+            if feat.properties is None:
+                continue
             for property_name, property_value in self.filters:
                 if feat.properties.get(property_name) != property_value:
                     continue
@@ -155,17 +162,25 @@ class ClassificationTask(BasicTask):
         }
 
     def process_output(
-        self, raw_output: Any, metadata: dict[str, Any]
-    ) -> npt.NDArray[Any] | list[Feature]:
+        self, raw_output: Any, metadata: SampleMetadata
+    ) -> list[Feature]:
         """Processes an output into raster or vector data.
 
         Args:
-            raw_output: the output from prediction head.
+            raw_output: the output from prediction head, which must be a tensor
+                containing output probabilities (one dimension).
             metadata: metadata about the patch being read
 
         Returns:
-            either raster or vector data.
+            a list with one Feature corresponding to the input patch extent with a
+                property name containing the predicted class. It will have another
+                property containing the probabilities if prob_property was set.
         """
+        if not isinstance(raw_output, torch.Tensor) or len(raw_output.shape) != 1:
+            raise ValueError(
+                "expected output for ClassificationTask to be a Tensor with one dimension"
+            )
+
         probs = raw_output.cpu().numpy()
         if len(self.classes) == 2 and self.positive_class_threshold != 0.5:
             positive_class_prob = probs[self.positive_class_id]
@@ -175,24 +190,25 @@ class ClassificationTask(BasicTask):
                 class_idx = 1 - self.positive_class_id
         else:
             # For multiclass classification or when using the default threshold
-            class_idx = probs.argmax()
+            class_idx = probs.argmax().item()
 
+        value: str | int
         if not self.read_class_id:
-            value = self.classes[class_idx]
+            value = self.classes[class_idx]  # type: ignore
         else:
             value = class_idx
 
         feature = Feature(
             STGeometry(
-                metadata["projection"],
-                shapely.Point(metadata["bounds"][0], metadata["bounds"][1]),
+                metadata.projection,
+                shapely.Point(metadata.patch_bounds[0], metadata.patch_bounds[1]),
                 None,
             ),
             {
                 self.property_name: value,
             },
         )
-        if self.prob_property:
+        if self.prob_property is not None and feature.properties is not None:
             feature.properties[self.prob_property] = probs.tolist()
         return [feature]
 
@@ -215,6 +231,8 @@ class ClassificationTask(BasicTask):
         image = super().visualize(input_dict, target_dict, output)["image"]
         image = Image.fromarray(image)
         draw = ImageDraw.Draw(image)
+        if target_dict is None:
+            raise ValueError("target_dict is required for visualization")
         target_class = self.classes[target_dict["class"]]
         output_class = self.classes[output.argmax()]
         text = f"Label: {target_class}\nOutput: {output_class}"
@@ -263,28 +281,34 @@ class ClassificationTask(BasicTask):
         return MetricCollection(metrics)
 
 
-class ClassificationHead(torch.nn.Module):
+class ClassificationHead(Predictor):
     """Head for classification task."""
 
     def forward(
         self,
-        logits: torch.Tensor,
-        inputs: list[dict[str, Any]],
+        intermediates: Any,
+        context: ModelContext,
         targets: list[dict[str, Any]] | None = None,
-    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+    ) -> ModelOutput:
         """Compute the classification outputs and loss from logits and targets.
 
         Args:
-            logits: tensor that is (BatchSize, NumClasses) in shape.
-            inputs: original inputs (ignored).
-            targets: should contain class key that stores the class label.
+            intermediates: output from the previous model component, it should be a
+                FeatureVector with a tensor that is (BatchSize, NumClasses) in shape.
+            context: the model context.
+            targets: must contain "class" key that stores the class label, along with
+                "valid" key indicating whether the label is valid for each example.
 
         Returns:
             tuple of outputs and loss dict
         """
+        if not isinstance(intermediates, FeatureVector):
+            raise ValueError("the input to ClassificationHead must be a FeatureVector")
+
+        logits = intermediates.feature_vector
         outputs = torch.nn.functional.softmax(logits, dim=1)
 
-        loss = None
+        losses = {}
         if targets:
             class_labels = torch.stack([target["class"] for target in targets], dim=0)
             mask = torch.stack([target["valid"] for target in targets], dim=0)
@@ -294,9 +318,12 @@ class ClassificationHead(torch.nn.Module):
                 )
                 * mask
             )
-            loss = torch.mean(loss)
+            losses["cls"] = torch.mean(loss)
 
-        return outputs, {"cls": loss}
+        return ModelOutput(
+            outputs=outputs,
+            loss_dict=losses,
+        )
 
 
 class ClassificationMetric(Metric):

rslearn/train/tasks/detection.py

@@ -12,26 +12,27 @@ import torchmetrics.classification
 import torchvision
 from torchmetrics import Metric, MetricCollection
 
+from rslearn.train.model_context import RasterImage, SampleMetadata
 from rslearn.utils import Feature, STGeometry
 
 from .task import BasicTask
 
 DEFAULT_COLORS = [
-    [255, 0, 0],
-    [0, 255, 0],
-    [0, 0, 255],
-    [255, 255, 0],
-    [0, 255, 255],
-    [255, 0, 255],
-    [0, 128, 0],
-    [255, 160, 122],
-    [139, 69, 19],
-    [128, 128, 128],
-    [255, 255, 255],
-    [143, 188, 143],
-    [95, 158, 160],
-    [255, 200, 0],
-    [128, 0, 0],
+    (255, 0, 0),
+    (0, 255, 0),
+    (0, 0, 255),
+    (255, 255, 0),
+    (0, 255, 255),
+    (255, 0, 255),
+    (0, 128, 0),
+    (255, 160, 122),
+    (139, 69, 19),
+    (128, 128, 128),
+    (255, 255, 255),
+    (143, 188, 143),
+    (95, 158, 160),
+    (255, 200, 0),
+    (128, 0, 0),
 ]
 
 
@@ -53,14 +54,30 @@ class DetectionTask(BasicTask):
         score_threshold: float = 0.5,
         enable_map_metric: bool = True,
         enable_f1_metric: bool = False,
+        enable_precision_recall: bool = False,
+        f1_metric_thresholds: list[list[float]] = [
+            [
+                0.05,
+                0.1,
+                0.2,
+                0.3,
+                0.4,
+                0.5,
+                0.6,
+                0.7,
+                0.8,
+                0.9,
+                0.95,
+            ]
+        ],
         f1_metric_kwargs: dict[str, Any] = {},
-        **kwargs,
-    ):
-        """Initialize a new SegmentationTask.
+        **kwargs: Any,
+    ) -> None:
+        """Initialize a new DetectionTask.
 
         Args:
-            property_name: the property from which to extract the class name. The class
-                is read from the first matching feature.
+            property_name: the property from which to extract the class name. Features
+                without this property name are ignored.
             classes: a list of class names.
             filters: optional list of (property_name, property_value) to only consider
                 features with matching properties.
@@ -70,14 +87,20 @@ class DetectionTask(BasicTask):
                 not passed via classes, instead of throwing error
             skip_empty_examples: whether to skip examples with zero labels.
             colors: optional colors for each class
-            box_size: force all boxes to be this size, centered at the centroid of the
-                geometry. Required for Point geometries.
+            box_size: force all boxes to be two times this size, centered at the
+                centroid of the geometry. Required for Point geometries.
             clip_boxes: whether to clip boxes to the image bounds.
             exclude_by_center: before optionally clipping boxes, exclude boxes if the
                 center is outside the image bounds.
             score_threshold: confidence threshold for visualization and prediction.
             enable_map_metric: whether to compute mAP (default true)
             enable_f1_metric: whether to compute F1 (default false)
+            enable_precision_recall: whether to compute precision and recall.
+            f1_metric_thresholds: list of list of thresholds to apply for F1 metric, as
+                well as for precision and recall if enabled. 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.
             f1_metric_kwargs: extra arguments to pass to F1 metric.
             kwargs: additional arguments to pass to BasicTask
         """
@@ -95,6 +118,8 @@ class DetectionTask(BasicTask):
         self.score_threshold = score_threshold
         self.enable_map_metric = enable_map_metric
         self.enable_f1_metric = enable_f1_metric
+        self.enable_precision_recall = enable_precision_recall
+        self.f1_metric_thresholds = f1_metric_thresholds
         self.f1_metric_kwargs = f1_metric_kwargs
 
         if not self.filters:
@@ -102,8 +127,8 @@ class DetectionTask(BasicTask):
 
     def process_inputs(
         self,
-        raw_inputs: dict[str, torch.Tensor | list[Feature]],
-        metadata: dict[str, Any],
+        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.
@@ -120,15 +145,21 @@ class DetectionTask(BasicTask):
         if not load_targets:
             return {}, {}
 
+        bounds = metadata.patch_bounds
+
         boxes = []
         class_labels = []
         valid = 1
 
         data = raw_inputs["targets"]
+        assert isinstance(data, list)
         for feat in data:
-            for property_name, property_value in self.filters:
-                if feat.properties.get(property_name) != property_value:
-                    continue
+            if feat.properties is None:
+                continue
+            if self.filters is not None:
+                for property_name, property_value in self.filters:
+                    if feat.properties.get(property_name) != property_value:
+                        continue
             if self.property_name not in feat.properties:
                 continue
 
@@ -159,39 +190,33 @@ class DetectionTask(BasicTask):
             else:
                 box = [int(val) for val in shp.bounds]
 
-            if box[0] >= metadata["bounds"][2] or box[2] <= metadata["bounds"][0]:
+            if box[0] >= bounds[2] or box[2] <= bounds[0]:
                 continue
-            if box[1] >= metadata["bounds"][3] or box[3] <= metadata["bounds"][1]:
+            if box[1] >= bounds[3] or box[3] <= bounds[1]:
                 continue
 
             if self.exclude_by_center:
                 center_col = (box[0] + box[2]) // 2
                 center_row = (box[1] + box[3]) // 2
-                if (
-                    center_col <= metadata["bounds"][0]
-                    or center_col >= metadata["bounds"][2]
-                ):
+                if center_col <= bounds[0] or center_col >= bounds[2]:
                     continue
-                if (
-                    center_row <= metadata["bounds"][1]
-                    or center_row >= metadata["bounds"][3]
-                ):
+                if center_row <= bounds[1] or center_row >= bounds[3]:
                     continue
 
             if self.clip_boxes:
                 box = [
-                    np.clip(box[0], metadata["bounds"][0], metadata["bounds"][2]),
-                    np.clip(box[1], metadata["bounds"][1], metadata["bounds"][3]),
-                    np.clip(box[2], metadata["bounds"][0], metadata["bounds"][2]),
-                    np.clip(box[3], metadata["bounds"][1], metadata["bounds"][3]),
+                    np.clip(box[0], bounds[0], bounds[2]),
+                    np.clip(box[1], bounds[1], bounds[3]),
+                    np.clip(box[2], bounds[0], bounds[2]),
+                    np.clip(box[3], bounds[1], bounds[3]),
                 ]
 
             # Convert to relative coordinates.
             box = [
-                box[0] - metadata["bounds"][0],
-                box[1] - metadata["bounds"][1],
-                box[2] - metadata["bounds"][0],
-                box[3] - metadata["bounds"][1],
+                box[0] - bounds[0],
+                box[1] - bounds[1],
+                box[2] - bounds[0],
+                box[3] - bounds[1],
             ]
 
             boxes.append(box)
@@ -211,16 +236,12 @@ class DetectionTask(BasicTask):
             "valid": torch.tensor(valid, dtype=torch.int32),
             "boxes": boxes,
             "labels": class_labels,
-            "width": torch.tensor(
-                metadata["bounds"][2] - metadata["bounds"][0], dtype=torch.float32
-            ),
-            "height": torch.tensor(
-                metadata["bounds"][3] - metadata["bounds"][1], dtype=torch.float32
-            ),
+            "width": torch.tensor(bounds[2] - bounds[0], dtype=torch.float32),
+            "height": torch.tensor(bounds[3] - bounds[1], dtype=torch.float32),
         }
 
     def process_output(
-        self, raw_output: Any, metadata: dict[str, Any]
+        self, raw_output: Any, metadata: SampleMetadata
     ) -> npt.NDArray[Any] | list[Feature]:
         """Processes an output into raster or vector data.
 
@@ -240,13 +261,13 @@ class DetectionTask(BasicTask):
         features = []
         for box, class_id, score in zip(boxes, class_ids, scores):
             shp = shapely.box(
-                metadata["bounds"][0] + float(box[0]),
-                metadata["bounds"][1] + float(box[1]),
-                metadata["bounds"][0] + float(box[2]),
-                metadata["bounds"][1] + float(box[3]),
+                metadata.patch_bounds[0] + float(box[0]),
+                metadata.patch_bounds[1] + float(box[1]),
+                metadata.patch_bounds[0] + float(box[2]),
+                metadata.patch_bounds[1] + float(box[3]),
             )
-            geom = STGeometry(metadata["projection"], shp, None)
-            properties = {
+            geom = STGeometry(metadata.projection, shp, None)
+            properties: dict[str, Any] = {
                 "score": float(score),
             }
 
@@ -278,7 +299,9 @@ class DetectionTask(BasicTask):
         """
         image = super().visualize(input_dict, target_dict, output)["image"]
 
-        def draw_boxes(image: npt.NDArray[Any], d: dict[str, torch.Tensor]):
+        def draw_boxes(
+            image: npt.NDArray[Any], d: dict[str, torch.Tensor]
+        ) -> npt.NDArray[Any]:
             boxes = d["boxes"].cpu().numpy()
             class_ids = d["labels"].cpu().numpy()
             if "scores" in d:
@@ -299,6 +322,8 @@ class DetectionTask(BasicTask):
 
             return image
 
+        if target_dict is None:
+            raise ValueError("target_dict is required for visualization")
         return {
             "gt": draw_boxes(image.copy(), target_dict),
             "pred": draw_boxes(image.copy(), output),
@@ -307,17 +332,46 @@ class DetectionTask(BasicTask):
     def get_metrics(self) -> MetricCollection:
         """Get the metrics for this task."""
         metrics = {}
+
         if self.enable_map_metric:
             metrics["mAP"] = DetectionMetric(
                 torchmetrics.detection.mean_ap.MeanAveragePrecision(),
                 output_key="map",
             )
-        if self.enable_f1_metric:
+
+        if self.enable_f1_metric or self.enable_precision_recall:
             kwargs = dict(
                 num_classes=len(self.classes),
             )
             kwargs.update(self.f1_metric_kwargs)
-            metrics["F1"] = DetectionMetric(F1Metric(**kwargs))
+
+            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(".", ",")
+
+                if self.enable_f1_metric:
+                    metrics["F1" + suffix] = DetectionMetric(
+                        F1Metric(score_thresholds=thresholds, **kwargs)  # type: ignore
+                    )
+                if self.enable_precision_recall:
+                    metrics["precision" + suffix] = DetectionMetric(
+                        F1Metric(
+                            score_thresholds=thresholds,
+                            metric_mode="precision",
+                            **kwargs,  # type: ignore
+                        )
+                    )
+                    metrics["recall" + suffix] = DetectionMetric(
+                        F1Metric(
+                            score_thresholds=thresholds,
+                            metric_mode="recall",
+                            **kwargs,  # type: ignore
+                        )
+                    )
+
         return MetricCollection(metrics)
 
 
@@ -377,22 +431,11 @@ class F1Metric(Metric):
     def __init__(
         self,
         num_classes: int,
+        score_thresholds: list[float],
         cmp_mode: str = "iou",
         cmp_threshold: float = 0.5,
-        score_thresholds: list[float] = [
-            0.05,
-            0.1,
-            0.2,
-            0.3,
-            0.4,
-            0.5,
-            0.6,
-            0.7,
-            0.8,
-            0.9,
-            0.95,
-        ],
         flatten_classes: bool = False,
+        metric_mode: str = "f1",
     ):
         """Create a new F1Metric.
 
@@ -406,6 +449,8 @@ class F1Metric(Metric):
             flatten_classes: sum true positives, false positives, and false negatives
                 across classes and report combined F1 instead of computing F1 score for
                 each class and then reporting the average.
+            metric_mode: set to "precision" or "recall" to return that instead of F1
+                (default "f1")
         """
         super().__init__()
         self.num_classes = num_classes
@@ -413,6 +458,10 @@ class F1Metric(Metric):
         self.cmp_threshold = cmp_threshold
         self.score_thresholds = score_thresholds
         self.flatten_classes = flatten_classes
+        self.metric_mode = metric_mode
+
+        assert self.cmp_mode in ["iou", "distance"]
+        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)):
@@ -531,8 +580,15 @@ class F1Metric(Metric):
                 else:
                     f1 = 2 * precision * recall / (precision + recall)
 
-                if best_score is None or f1 > best_score:
-                    best_score = f1
+                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
 
             best_scores.append(best_score)
 

rslearn/train/tasks/embedding.py

@@ -0,0 +1,120 @@
+"""Embedding task."""
+
+from typing import Any
+
+import numpy.typing as npt
+import torch
+from torchmetrics import MetricCollection
+
+from rslearn.models.component import FeatureMaps
+from rslearn.train.model_context import ModelContext, ModelOutput, SampleMetadata
+from rslearn.utils import Feature
+
+from .task import Task
+
+
+class EmbeddingTask(Task):
+    """A dummy task for computing embeddings.
+
+    This task does not compute any targets or loss. Instead, it is just set up for
+    inference, to save embeddings from the configured model.
+    """
+
+    def process_inputs(
+        self,
+        raw_inputs: dict[str, torch.Tensor],
+        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
+        """
+        return {}, {}
+
+    def process_output(
+        self, raw_output: Any, metadata: SampleMetadata
+    ) -> npt.NDArray[Any] | list[Feature]:
+        """Processes an output into raster or vector data.
+
+        Args:
+            raw_output: the output from prediction head, which must be a CxHxW tensor.
+            metadata: metadata about the patch being read
+
+        Returns:
+            either raster or vector data.
+        """
+        if not isinstance(raw_output, torch.Tensor) or len(raw_output.shape) != 3:
+            raise ValueError(
+                "output for EmbeddingTask must be a tensor with three dimensions"
+            )
+
+        # Just convert the raw output to numpy array that can be saved to GeoTIFF.
+        return raw_output.cpu().numpy()
+
+    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
+        """
+        # EmbeddingTask is only set up to support `model predict`.
+        raise NotImplementedError
+
+    def get_metrics(self) -> MetricCollection:
+        """Get the metrics for this task."""
+        return MetricCollection({})
+
+
+class EmbeddingHead:
+    """Head for embedding task.
+
+    It just adds a dummy loss to act as a Predictor.
+    """
+
+    def forward(
+        self,
+        intermediates: Any,
+        context: ModelContext,
+        targets: list[dict[str, Any]] | None = None,
+    ) -> ModelOutput:
+        """Return the feature map along with a dummy loss.
+
+        Args:
+            intermediates: output from the previous model component, which must be a
+                FeatureMaps consisting of a single feature map.
+            context: the model context.
+            targets: the targets (ignored).
+
+        Returns:
+            model output with the feature map that was input to this component along
+                with a dummy loss.
+        """
+        if not isinstance(intermediates, FeatureMaps):
+            raise ValueError("input to EmbeddingHead must be a FeatureMaps")
+        if len(intermediates.feature_maps) != 1:
+            raise ValueError(
+                f"input to EmbeddingHead must have one feature map, but got {len(intermediates.feature_maps)}"
+            )
+
+        return ModelOutput(
+            outputs=intermediates.feature_maps[0],
+            loss_dict={"loss": 0},
+        )