initial implementation of raster datasets

Author: PaulDanielML

README.md

@@ -1,20 +1,32 @@
-# eotorch: Template Python repository
-
-This repository serves as a template for creating a Python library
-
-## How do I use this?
-
-1. Create a new repository in GitHub with this repo as a starting point
-    ![](images/new_repo.png)
-
-2. Change all occurences of `eotorch` to match the name of your new library
-
-3. Consider if the [license](LICENSE) should be modified.
-
-
-## Additional resources
-
-If you're interested in learning more about best practices for developing Python packages, check out the following resources:
-
-- [Python Package Development at DHI](https://dhi.github.io/python-package-development/)
-- [Scientific Python Library Development Guide](https://learn.scientific-python.org/development/)
+## Example of how to use SegmentationRasterDataset:
+
+```python
+from eotorch.datasets.geo import SegmentationRasterDataset
+from eotorch.plot.plot import plot_samples
+
+class_mapping = {
+    1: "Baresoil",
+    2: "Buildings",
+    3: "Coniferous Trees",
+    4: "Deciduous Trees",
+    5: "Grass",
+    6: "Impervious",
+    7: "Water",
+}
+
+bla = SegmentationRasterDataset.create(
+    images_dir="dev_data/sr_data",
+    labels_dir="dev_data/labels",
+    image_kwargs=dict(
+        all_bands=("B02", "B03", "B04", "B08", "B11", "B12"),
+        rgb_bands=("B04", "B03", "B02"),
+    ),
+    label_kwargs=dict(
+        class_mapping=class_mapping,
+    ),
+)
+plot_samples(bla, n=2, patch_size=256)
+```
+
+![alt text](media/sample_1.png)
+![alt text](media/sample_2.png)

eotorch/datasets/geo.py

@@ -0,0 +1,107 @@
+from pathlib import Path
+from typing import Any
+
+import matplotlib.patches as mpatches
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib import cm
+from rasterio.plot import show
+from torchgeo.datasets import IntersectionDataset, RasterDataset
+
+
+class PlottableRasterDataset(RasterDataset):
+    colormap = cm.tab20
+    nodata_value = 0
+    class_mapping = None
+
+    def plot(self, sample, ax=None, **kwargs):
+        if ax is None:
+            _, ax = plt.subplots()
+
+        if self.is_image:
+            rgb_indices = []
+            for band in self.rgb_bands:
+                rgb_indices.append(self.all_bands.index(band))
+
+            image = sample["image"][rgb_indices]
+            return show(image.numpy(), ax=ax, adjust=True, **kwargs)
+
+        vals = sample["mask"].numpy()
+        ax = show(vals, ax=ax, cmap=self.colormap, **kwargs)
+        values = np.unique(vals.ravel()).tolist()
+
+        class_mapping = self.class_mapping or {v: str(v) for v in values}
+
+        if (self.nodata_value in values) and (self.nodata_value not in class_mapping):
+            class_mapping[self.nodata_value] = "No Data"
+
+        im = ax.get_images()[0]
+        colors = [im.cmap(im.norm(value)) for value in range(len(self.colormap.colors))]
+
+        patches = [
+            mpatches.Patch(color=colors[i], label=class_mapping[i]) for i in values
+        ]
+        ax.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0)
+
+
+class LabelledRasterDataset(IntersectionDataset):
+    def plot(self, sample: dict[str, Any], **kwargs):
+        fig, axes = plt.subplots(1, 2, figsize=(8, 4))
+
+        for i, dataset in enumerate(self.datasets):
+            if isinstance(dataset, PlottableRasterDataset):
+                dataset.plot(sample, ax=axes[i], **kwargs)
+            else:
+                raise NotImplementedError("Dataset must be plottable")
+
+        return fig
+
+
+class SegmentationRasterDataset:
+    """
+    A dataset for semantic segmentation of raster data.
+
+    Args:
+        images_dir: A directory containing the image files.
+        labels_dir: A directory containing the label files.
+        image_glob: A glob pattern to match the image files.
+        label_glob: A glob pattern to match the label files.
+        image_kwargs: Keyword arguments to pass to the image dataset.
+        label_kwargs: Keyword arguments to pass to the label dataset.
+
+    This dataset is an intersection of two raster datasets: one for the images and one for the labels.
+    Both those datasets are based on TorchGeo's RasterDataset (which itself is a subclass of TorchGeo's GeoDataset).
+    Therefore it comes with strong capabilities for spatial and temporal indexing, however it is also limited
+    in how it handles global datasets spanning across multiple crs.
+
+    """
+
+    @classmethod
+    def create(
+        cls,
+        images_dir: str | Path,
+        labels_dir: str | Path = None,
+        image_glob="*.tif",
+        label_glob="*.tif",
+        image_kwargs: dict[str, Any] = None,
+        label_kwargs: dict[str, Any] = None,
+    ):
+        image_kwargs = image_kwargs or {}
+        label_kwargs = label_kwargs or {}
+
+        image_ds = PlottableRasterDataset(paths=images_dir)
+        image_ds.filename_glob = image_glob
+        for k, v in image_kwargs.items():
+            setattr(image_ds, k, v)
+
+        if labels_dir:
+            label_ds = PlottableRasterDataset(paths=labels_dir)
+            label_ds.is_image = False
+            label_ds.filename_glob = label_glob
+            for k, v in label_kwargs.items():
+                setattr(label_ds, k, v)
+
+            return LabelledRasterDataset(image_ds, label_ds)
+
+        else:
+            return image_ds

eotorch/inference/inference.py

@@ -8,6 +8,7 @@
 from rasterio.io import BufferedDatasetWriter
 
 from eotorch.inference import inference_utils as iu
+from eotorch.plot import plot
 
 
 def predict_on_tif(
@@ -101,5 +102,5 @@ def predict_on_tif(
 
     if show_results:
         print(f"Showing results for {out_file_path}")
-        return iu.plot_predictions_pyplot(out_file_path, classes, ax=ax)
+        return plot.plot_predictions_pyplot(out_file_path, classes, ax=ax)
     return out_file_path

eotorch/inference/inference_utils.py

@@ -4,54 +4,9 @@
 import numpy as np
 import rasterio as rst
 from affine import Affine
-from matplotlib import pyplot as plt
 from rasterio.windows import Window, from_bounds
 
 
-def plot_predictions_pyplot(
-    predictions_path: str | Path, classes: dict[int, str], ax: plt.Axes = None
-):
-    """
-    Plot predictions using matplotlib.pyplot.
-
-    Parameters
-    ----------
-        predictions_path: str or Path, path to the predictions file
-        classes: dict, mapping of class indices to class names
-        ax: matplotlib.axes.Axes, axes to plot on
-    """
-
-    import matplotlib.patches as mpatches
-    import matplotlib.pyplot as plt
-
-    if ax is None:
-        _, ax = plt.subplots(nrows=1, ncols=1)
-
-    with rst.open(predictions_path) as src:
-        predictions = src.read(1)
-        no_data_val = src.nodata
-    im = ax.imshow(predictions, interpolation="none")
-    values = np.unique(predictions.ravel()).tolist()
-
-    if no_data_val not in values:
-        plot_values = values + [no_data_val]
-    else:
-        plot_values = values
-
-    if not classes:
-        classes = {v: str(v) for v in values}
-    if (no_data_val in values) and (no_data_val not in classes):
-        classes[no_data_val] = "No Data"
-
-    colors = [im.cmap(im.norm(value)) for value in plot_values]
-    patches = [
-        mpatches.Patch(color=colors[i], label=classes[i])
-        for i in values
-        # mpatches.Patch(color=colors[i - 1], label=classes[i - 1]) for i in values
-    ]
-    ax.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0)
-
-
 def crop_np_to_window(arr: np.ndarray, w_buffered: Window, w_unbuffered: Window):
     left = int(round(w_unbuffered.col_off - w_buffered.col_off))
     top = int(round(w_unbuffered.row_off - w_buffered.row_off))

eotorch/plot/plot.py

@@ -0,0 +1,66 @@
+from pathlib import Path
+
+import numpy as np
+import rasterio as rst
+from matplotlib import pyplot as plt
+from torch.utils.data import DataLoader
+from torchgeo.datasets import GeoDataset, stack_samples, unbind_samples
+from torchgeo.samplers import RandomGeoSampler
+
+
+def plot_predictions_pyplot(
+    predictions_path: str | Path, classes: dict[int, str], ax: plt.Axes = None
+):
+    """
+    Plot predictions using matplotlib.pyplot.
+
+    Parameters
+    ----------
+        predictions_path: str or Path, path to the predictions file
+        classes: dict, mapping of class indices to class names
+        ax: matplotlib.axes.Axes, axes to plot on
+    """
+
+    import matplotlib.patches as mpatches
+    import matplotlib.pyplot as plt
+
+    if ax is None:
+        _, ax = plt.subplots(nrows=1, ncols=1)
+
+    with rst.open(predictions_path) as src:
+        predictions = src.read(1)
+        no_data_val = src.nodata
+    im = ax.imshow(predictions, interpolation="none")
+    values = np.unique(predictions.ravel()).tolist()
+
+    if no_data_val not in values:
+        plot_values = values + [no_data_val]
+    else:
+        plot_values = values
+
+    if not classes:
+        classes = {v: str(v) for v in values}
+    if (no_data_val in values) and (no_data_val not in classes):
+        classes[no_data_val] = "No Data"
+
+    colors = [im.cmap(im.norm(value)) for value in plot_values]
+    patches = [
+        mpatches.Patch(color=colors[i], label=classes[i])
+        for i in values
+        # mpatches.Patch(color=colors[i - 1], label=classes[i - 1]) for i in values
+    ]
+    ax.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0)
+
+
+def plot_samples(dataset: GeoDataset, n: int = 3, patch_size: int = 256):
+    assert hasattr(dataset, "plot"), "Dataset must have a plot method"
+    # sampler = GridGeoSampler(dataset, size=patch_size, stride=100)
+    sampler = RandomGeoSampler(dataset, size=patch_size)
+    dataloader = DataLoader(
+        dataset, sampler=sampler, collate_fn=stack_samples, batch_size=n
+    )
+    batch = next(iter(dataloader))
+    samples = unbind_samples(batch)
+    for i, sample in enumerate(samples):
+        dataset.plot(sample, title=f"Sample {i + 1}")
+        plt.show()

media/sample_1.png

@@ -16,6 +16,7 @@ exclude = ["notebooks", "tests", "images"]
 name = "eotorch"
 version = "0.0.1"
 dependencies = [
+    "alive-progress>=3.2.0",
     "torchgeo>=0.6.2",
 ]