rslearn 0.0.27__py3-none-any.whl → 0.0.29__py3-none-any.whl

rslearn/train/tasks/segmentation.py

@@ -10,6 +10,7 @@ import torchmetrics.classification
 from torchmetrics import Metric, MetricCollection
 
 from rslearn.models.component import FeatureMaps, Predictor
+from rslearn.train.metrics import ConfusionMatrixMetric
 from rslearn.train.model_context import (
     ModelContext,
     ModelOutput,
@@ -43,6 +44,8 @@ class SegmentationTask(BasicTask):
         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.
@@ -80,6 +83,10 @@ class SegmentationTask(BasicTask):
                 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)
@@ -106,6 +113,8 @@ class SegmentationTask(BasicTask):
         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,
@@ -285,6 +294,14 @@ class SegmentationTask(BasicTask):
         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)
 
 

rslearn/utils/fsspec.py

@@ -2,7 +2,7 @@
 
 import os
 import tempfile
-from collections.abc import Generator
+from collections.abc import Generator, Iterator
 from contextlib import contextmanager
 from typing import Any
 
@@ -16,6 +16,56 @@ from rslearn.log_utils import get_logger
 logger = get_logger(__name__)
 
 
+def iter_nonhidden(path: UPath) -> Iterator[UPath]:
+    """Iterate over non-hidden entries in a directory.
+
+    Hidden entries are those whose basename starts with "." (e.g. ".DS_Store").
+
+    Args:
+        path: the directory to iterate.
+
+    Yields:
+        non-hidden UPath entries in the directory.
+    """
+    try:
+        it = path.iterdir()
+    except (FileNotFoundError, NotADirectoryError):
+        return
+
+    for p in it:
+        if p.name.startswith("."):
+            continue
+        yield p
+
+
+def iter_nonhidden_subdirs(path: UPath) -> Iterator[UPath]:
+    """Iterate over non-hidden subdirectories in a directory.
+
+    Args:
+        path: the directory to iterate.
+
+    Yields:
+        non-hidden subdirectories in the directory.
+    """
+    for p in iter_nonhidden(path):
+        if p.is_dir():
+            yield p
+
+
+def iter_nonhidden_files(path: UPath) -> Iterator[UPath]:
+    """Iterate over non-hidden files in a directory.
+
+    Args:
+        path: the directory to iterate.
+
+    Yields:
+        non-hidden files in the directory.
+    """
+    for p in iter_nonhidden(path):
+        if p.is_file():
+            yield p
+
+
 @contextmanager
 def get_upath_local(
     path: UPath, extra_paths: list[UPath] = []

{rslearn-0.0.27.dist-info → rslearn-0.0.29.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rslearn
-Version: 0.0.27
+Version: 0.0.29
 Summary: A library for developing remote sensing datasets and models
 Author: OlmoEarth Team
 License:                                  Apache License

{rslearn-0.0.27.dist-info → rslearn-0.0.29.dist-info}/RECORD

@@ -9,6 +9,7 @@ rslearn/template_params.py,sha256=Vop0Ha-S44ctCa9lvSZRjrMETznJZlR5y_gJrVIwrPg,79
 rslearn/config/__init__.py,sha256=n1qpZ0ImshTtLYl5mC73BORYyUcjPJyHiyZkqUY1hiY,474
 rslearn/config/dataset.py,sha256=abxIUFDAYmCd4pzGnkPnW_pYyws1yhXFWJ5HnVU4WHo,23942
 rslearn/data_sources/__init__.py,sha256=FZQckYwsnnLokMeYmi0ktUyQd9bAHyLN1_-Xc3qYLag,767
+rslearn/data_sources/aws_google_satellite_embedding_v1.py,sha256=ga5G8uDXdMj2pw8qGC2PD11PBg8Spuf3b4QEwsVJaBY,12805
 rslearn/data_sources/aws_landsat.py,sha256=bJmwBbUV4vjKBNp1MHt4sHhnIjhMis_jOI3FpksQc6w,16435
 rslearn/data_sources/aws_open_data.py,sha256=fum34DqqDiuiiYBfZtGFrNNOLylE9o3o7Cyb2e0Eo0g,29101
 rslearn/data_sources/aws_sentinel1.py,sha256=LfJLDhsd_6h_JinD8PbiiAyxajkIdvAc--5BJryUKlo,4674
@@ -16,7 +17,7 @@ rslearn/data_sources/aws_sentinel2_element84.py,sha256=qeCuiSlvhWChSY3AwYsKT6nZU
 rslearn/data_sources/climate_data_store.py,sha256=mqLJfYubD6m9VwxpLunoIv_MNFN6Ue1hBBVj552e8uQ,18289
 rslearn/data_sources/copernicus.py,sha256=ushAgYGxU2MzPcUNnEvEfPgO0RCC9Rbjzi189xq0jgc,35001
 rslearn/data_sources/data_source.py,sha256=xojlCoAnGTCHKbEx98JkW0oYzAKBbgGMNc0kicEjHWk,4863
-rslearn/data_sources/direct_materialize_data_source.py,sha256=PBmxsLBNakJeX1s92pc4FCuHANfhkYHS2vt60RGdkj0,11276
+rslearn/data_sources/direct_materialize_data_source.py,sha256=UnFuCSJED9-YSFp12-MosV8bMFj6AqCb75a9ADu_Cxw,11030
 rslearn/data_sources/earthdaily.py,sha256=qUtHUG1oV5IlCWXVovUcYxQhqdNDKWaEe-BKnooWX88,14623
 rslearn/data_sources/earthdatahub.py,sha256=KRf1VnxPI9jsT0utEkeYvsCwu7LXo9t-RvMi8gXehag,15889
 rslearn/data_sources/eurocrops.py,sha256=dJ4d0xvt-rID_HuAchyucFJBuAQL-Kk1h_qm6GOH-mE,8641
@@ -32,7 +33,7 @@ rslearn/data_sources/soilgrids.py,sha256=qbnnCIOa6tlN8wxmNCzAj60-pghKEbRxa7lVIgM
 rslearn/data_sources/stac.py,sha256=Gj8TZ5pifVzWPCuzgphrle2ekQ02OET54rj-02sR2nw,10705
 rslearn/data_sources/usda_cdl.py,sha256=3GhcgTB50T7GA44nB9WwItqDJliELquw_YbiAVxh6kc,6808
 rslearn/data_sources/usgs_landsat.py,sha256=IsQOhWY8nwmgixJu1uMSR4CqsC3igcP3TArdBXkETd8,10178
-rslearn/data_sources/utils.py,sha256=NOC0qOyYVS6f8EUrSeP4mH0XZbSrtTLV-gKGbCC6ccg,16586
+rslearn/data_sources/utils.py,sha256=EAVFCYzjFvuHWd7E2ghTub9f-bbDhq83p3x9IJDjgvk,16843
 rslearn/data_sources/vector_source.py,sha256=NCa7CxIrGKe9yRT0NyyFKFQboDGDZ1h7663PV9OfMOM,44
 rslearn/data_sources/worldcereal.py,sha256=OWZA0pvQQiKvuA5AVAc0lw8JStMEeF4DYOh0n2vdg6I,21521
 rslearn/data_sources/worldcover.py,sha256=ahyrGoXMAGWsIUDHSrqPywiK7ycwUD3E3BruNMxpo90,6057
@@ -47,7 +48,7 @@ rslearn/dataset/materialize.py,sha256=VoL5Qf5pGcQV4QMlO5vrcu7w0Sl1NdIRLUVk0kSCMO
 rslearn/dataset/remap.py,sha256=6MaImsY02GNACpvRM81RvWmjZWRfAHxo_R3Ox6XLF6A,2723
 rslearn/dataset/window.py,sha256=X4q8YzcSOTtwKxCPf71QLMoyKUtYMSnZu0kPnmVSUx4,10644
 rslearn/dataset/storage/__init__.py,sha256=R50AVV5LH2g7ol0-jyvGcB390VsclXGbJXz4fmkn9as,52
-rslearn/dataset/storage/file.py,sha256=g9HZ3CD4QcgyVNsBaXhjIKQgDOAeZ4R08sJ7ntx4wo8,6815
+rslearn/dataset/storage/file.py,sha256=GJJgH_eHknLbMQwoC3mOoXJKl6Ha3oNXzz62FIEMWlg,7130
 rslearn/dataset/storage/storage.py,sha256=DxZ7iwV938PiLwdQzb5EXSb4Mj8bRGmOTmA9fzq_Ge8,4840
 rslearn/models/__init__.py,sha256=_vWoF9d2Slah8-6XhYhdU4SRsy_CNxXjCGQTD2yvu3Q,22
 rslearn/models/anysat.py,sha256=nzk6hB83ltNFNXYRNA1rTvq2AQcAhwyvgBaZui1M37o,8107
@@ -61,6 +62,7 @@ rslearn/models/dinov3.py,sha256=Q9X7VTwzjllLSEvc235C9BY_jMnIoSybsiOkeA58uHo,6472
 rslearn/models/faster_rcnn.py,sha256=yOipLPmVHbadvYCR9xfCYgmkU9Mot6fgDK-kKicVTlo,8685
 rslearn/models/feature_center_crop.py,sha256=_Mu3E4iJLBug9I4ZIBIpB_VJo-xGterHmhtIFGaHR34,1808
 rslearn/models/fpn.py,sha256=qm7nKMgsZrCoAdz8ASmNKU2nvZ6USm5CedMfy_w_gwE,2079
+rslearn/models/global_pool.py,sha256=Bl48AVJ7g70hPmVLJbK1y_JN9_FTyANc_7tr6YOHANY,2782
 rslearn/models/module_wrapper.py,sha256=73JspaglnNabUGZB2EiCYF_dZ3-Kicg_OpoTfUWHONk,2271
 rslearn/models/molmo.py,sha256=lXnevwTCNyc1XcnJUB5_pK1G2AJGYMvQYU21mZFf5u0,2246
 rslearn/models/multitask.py,sha256=bpFxvtFowRyT-tvRSdY7AKbEx_i1y7sToEzZgTMcF4s,16264
@@ -112,14 +114,15 @@ rslearn/models/presto/__init__.py,sha256=eZrB-XKi_vYqZhpyAOwppJi4dRuMtYVAdbq7KRy
 rslearn/models/presto/presto.py,sha256=fkyHB85Hfx5L-4yejSFAFv83gk9VFqAR1GTgggtq0EA,11049
 rslearn/models/presto/single_file_presto.py,sha256=-P00xjhj9dx3O6HqWpQmG9dPk_i6bT_t8vhX4uQm5tA,30242
 rslearn/tile_stores/__init__.py,sha256=-cW1J7So60SEP5ZLHCPdaFBV5CxvV3QlOhaFnUkhTJ0,1675
-rslearn/tile_stores/default.py,sha256=PYaDNvBxhJTDKJGw0EjDTSE1OKajR7_iJpMbOjj-mE8,15054
+rslearn/tile_stores/default.py,sha256=AG2j0FCNi_4cnXqLjRIef5wMqMJ5_YtSkTIhk7qJQVQ,15134
 rslearn/tile_stores/tile_store.py,sha256=9AeYduDYPp_Ia2NMlq6osptpz_AFGIOQcLJrqZ_m-z0,10469
 rslearn/train/__init__.py,sha256=fnJyY4aHs5zQqbDKSfXsJZXY_M9fbTsf7dRYaPwZr2M,30
 rslearn/train/all_crops_dataset.py,sha256=CWnqbSjRXJZQsudljvpA07oldiP4fZTmjwrT0sjVnq4,21399
-rslearn/train/data_module.py,sha256=yPShftkHJ2bhJ4carwYYb9c3PkkP7ArzXQyu37EuAxk,23718
-rslearn/train/dataset.py,sha256=AAkdHX3q8VD1Geq9yIatiPkM5blA2luVPPwDmXDB6Z8,43284
+rslearn/train/data_module.py,sha256=G1TRhXg8SPewYy0BTZN5KpeLPK72qIaH15ePfUwrxgM,23865
+rslearn/train/dataset.py,sha256=vCmm6yrW2bAc5A94aBwQe-SOIGdVcZYMM2oBYRq2_sw,45253
 rslearn/train/dataset_index.py,sha256=S5iXhQga5gnnkDqThXXlyjIwkJBPVWiUfDPx3iVs-pw,5306
-rslearn/train/lightning_module.py,sha256=m0aIGk5xO5y12DEiwSl6eAko6X-gQ78_Wsbvz4Hb_NE,15364
+rslearn/train/lightning_module.py,sha256=n4hasJBVlAmMhvf2yaFo0gy1vGz5haQkJpZdCSKlJ8A,17482
+rslearn/train/metrics.py,sha256=RknMf2n09D5XBCf0YM4Zmm0XI-pFbRtsmbY51ipVMPk,4799
 rslearn/train/model_context.py,sha256=8DMWGj5xCRmRDo_38lkhkUMHfK_yg3XZrUJQIz5a1vA,3200
 rslearn/train/optimizer.py,sha256=EKSqkmERalDA0bF32Gey7n6z69KLyaUWKlRsGJfKBmE,927
 rslearn/train/prediction_writer.py,sha256=cRFehEtr0iBuVqzE69a0B4Lvb8ywxLeyon34KWI86H0,16961
@@ -130,13 +133,13 @@ rslearn/train/callbacks/freeze_unfreeze.py,sha256=8fIzBMhCKKjpTffIeAdhdSjsBd8NjT
 rslearn/train/callbacks/gradients.py,sha256=4YqCf0tBb6E5FnyFYbveXfQFlgNPyxIXb2FCWX4-6qs,5075
 rslearn/train/callbacks/peft.py,sha256=wEOKsS3RhsRaZTXn_Kz2wdsZdIiIaZPdCJWtdJBurT8,4156
 rslearn/train/tasks/__init__.py,sha256=dag1u72x1-me6y0YcOubUo5MYZ0Tjf6-dOir9UeFNMs,75
-rslearn/train/tasks/classification.py,sha256=H-Ayqm59IxwrczC8lUV5J5vg-JILhQhTiVlyaTpBs2k,14259
+rslearn/train/tasks/classification.py,sha256=_3cRa8ojd9sG2ELRW_BvByZh2YFdCBaklR8Kv9LAgOY,14864
 rslearn/train/tasks/detection.py,sha256=uDMGtsCMSk9OGXn-vpFKBAyHyVN0ji2NCfqBgg1BQyw,21725
 rslearn/train/tasks/embedding.py,sha256=NdJEAaDWlWYzvOBVf7eIHfFOzqTgavfFH1J1gMbAMVo,3891
 rslearn/train/tasks/multi_task.py,sha256=32hvwyVsHqt7N_M3zXsTErK1K7-0-BPHzt7iGNehyaI,6314
-rslearn/train/tasks/per_pixel_regression.py,sha256=njShN-U9fx3SPcCxGgbDlZAp3DT_GlTt0BRZS416gnw,10387
-rslearn/train/tasks/regression.py,sha256=_TGlj3PA14Iye0duf25TcGZQpFXU9fGfNilzpNuPS78,12693
-rslearn/train/tasks/segmentation.py,sha256=dn1yo1dIArKvW9Giw8-LZyIZ87q76eslL0mk58GyApo,29663
+rslearn/train/tasks/per_pixel_regression.py,sha256=3m_BTP2akadYe3IuAlCG2bd_alfNyom55-pFrI2q4PE,10928
+rslearn/train/tasks/regression.py,sha256=TzHL42gm3aIdev0R7_uz_TSYbAwSvQPjCD42y1p9_7Y,13269
+rslearn/train/tasks/segmentation.py,sha256=b9XS09EQvum89eoW3vWqFMKuCRtznODteKIr1hFnIz4,30531
 rslearn/train/tasks/task.py,sha256=nMPunl9OlnOimr48saeTnwKMQ7Du4syGrwNKVQq4FL4,4110
 rslearn/train/transforms/__init__.py,sha256=BkCAzm4f-8TEhPIuyvCj7eJGh36aMkZFYlq-H_jkSvY,778
 rslearn/train/transforms/concatenate.py,sha256=S8f1svzwb5UmeAgzXe4Af_hFvt5o0tQctIE6t3QYuPI,2625
@@ -153,7 +156,7 @@ rslearn/utils/__init__.py,sha256=GZc1erpEfXTc32yjEDbt5rnMrnXEBY7WVm3v4NlwwWY,620
 rslearn/utils/array.py,sha256=RC7ygtPnQwU6Lb9kwORvNxatJcaJ76JPsykQvndAfes,2444
 rslearn/utils/colors.py,sha256=ELY9_buH06TOVPLrDAyf2S0G--ZiOxnnP8Ujim6_3ig,369
 rslearn/utils/feature.py,sha256=lsg0WThZDJzo1mrbaL04dXYI5G3x-n5FG9aEjj7uUaI,1649
-rslearn/utils/fsspec.py,sha256=h3fER_bkewzR9liEAULXguTIvXLUXA17pC_yZoWN5Tk,5902
+rslearn/utils/fsspec.py,sha256=TcEUgXKvsmtKHv5JVOI2Vp4WNfVNeTok0x4JgZaD1iw,7052
 rslearn/utils/geometry.py,sha256=VzLoxtwdV3uC3szowT-bGuCFF6ge8eK0m01lq8q-01Q,22423
 rslearn/utils/get_utm_ups_crs.py,sha256=kUrcyjCK7KWvuP1XR-nURPeRqYeRO-3L8QUJ1QTF9Ps,3599
 rslearn/utils/grid_index.py,sha256=hRmrtgpqN1pLa-djnZtgSXqKJlbgGyttGnCEmPLD0zo,2347
@@ -175,10 +178,10 @@ rslearn/vis/render_sensor_image.py,sha256=D0ynK6ABPV046970lIKwF98klpSCtrsUvZTwtZ
 rslearn/vis/render_vector_label.py,sha256=ncwgRKCYCJCK1-wTpjgksOiDDebku37LpAyq6wsg4jg,14939
 rslearn/vis/utils.py,sha256=Zop3dEmyaXUYhPiGdYzrTO8BRXWscP2dEZy2myQUnNk,2765
 rslearn/vis/vis_server.py,sha256=kIGnhTy-yfu5lBOVCoo8VVG259i974JPszudCePbzfI,20157
-rslearn-0.0.27.dist-info/licenses/LICENSE,sha256=_99ZWPoLdlUbqZoSC5DF4ihiNwl5rTEmBaq2fACecdg,11352
-rslearn-0.0.27.dist-info/licenses/NOTICE,sha256=wLPr6rwV_jCg-xEknNGwhnkfRfuoOE9MZ-lru2yZyLI,5070
-rslearn-0.0.27.dist-info/METADATA,sha256=_6aHsGpgH_T2MQSb9qKeYshP6zdzwu3CmI2sBhD3dqk,38714
-rslearn-0.0.27.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-rslearn-0.0.27.dist-info/entry_points.txt,sha256=doTBQ57NT7nq-dgYGgTTw6mafcGWb_4PWYtYR4rGm50,46
-rslearn-0.0.27.dist-info/top_level.txt,sha256=XDKo90WBH8P9RQumHxo0giLJsoufT4r9odv-WE6Ahk4,8
-rslearn-0.0.27.dist-info/RECORD,,
+rslearn-0.0.29.dist-info/licenses/LICENSE,sha256=_99ZWPoLdlUbqZoSC5DF4ihiNwl5rTEmBaq2fACecdg,11352
+rslearn-0.0.29.dist-info/licenses/NOTICE,sha256=wLPr6rwV_jCg-xEknNGwhnkfRfuoOE9MZ-lru2yZyLI,5070
+rslearn-0.0.29.dist-info/METADATA,sha256=PMjB15sAZg5VA7qHkMrhG_2hSULc0Wopxr7G3op20Hg,38714
+rslearn-0.0.29.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+rslearn-0.0.29.dist-info/entry_points.txt,sha256=doTBQ57NT7nq-dgYGgTTw6mafcGWb_4PWYtYR4rGm50,46
+rslearn-0.0.29.dist-info/top_level.txt,sha256=XDKo90WBH8P9RQumHxo0giLJsoufT4r9odv-WE6Ahk4,8
+rslearn-0.0.29.dist-info/RECORD,,