model checkpoints, sample input, readme

README.md

@@ -1,6 +1,35 @@
----
-license: mit
----
-
-
-Python 3.12
+---
+license: mit
+datasets:
+- csaybar/CloudSEN12-high
+language:
+- en
+base_model:
+- NickWright/OmniCloudMask
+tags:
+- remote-sensing
+- cloud-detection
+---
+
+
+# Cloud Detection Model
+
+This model is based on NickWright/OmniCloudMask for cloud detection in satellite imagery. It provides pixel-level segmentation with the following classes:
+
+0 = Clear
+1 = Thick Cloud
+2 = Thin Cloud
+3 = Cloud Shadow
+
+## Usage
+
+The model requires Python 3.10 or higher. To use this model:
+```bash
+pip install -r requirements.txt
+```
+```bash
+python3 model.py
+```
+
+Below is a visualization of the cloud mask generated by the model:
+![Cloud Mask Visualization](cloud_mask_visualization.png)

omnicloudmask/__init__.py

@@ -0,0 +1,16 @@
+from .__version__ import __version__
+from .cloud_mask import predict_from_array, predict_from_load_func
+from .data_loaders import (
+    load_ls8,
+    load_multiband,
+    load_s2,
+)
+
+__all__ = [
+    "predict_from_load_func",
+    "predict_from_array",
+    "load_ls8",
+    "load_multiband",
+    "load_s2",
+    "__version__",
+]

omnicloudmask/__version__.py

@@ -0,0 +1 @@
+__version__ = "1.0.9"

omnicloudmask/cloud_mask.py

@@ -0,0 +1,493 @@
+import warnings
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
+from threading import Thread
+from typing import Callable, Generator, Optional, Union
+
+import numpy as np
+import torch
+from rasterio.profiles import Profile
+from tqdm.auto import tqdm
+
+from .__version__ import __version__
+from .download_models import get_models
+from .model_utils import (
+    create_gradient_mask,
+    default_device,
+    get_torch_dtype,
+    inference_and_store,
+    load_model_from_weights,
+)
+from .raster_utils import (
+    get_patch,
+    make_patch_indexes,
+    mask_prediction,
+    save_prediction,
+)
+
+
+def compile_batches(
+    batch_size: int,
+    patch_size: int,
+    patch_indexes: list[tuple[int, int, int, int]],
+    input_array: np.ndarray,
+    no_data_value: int,
+    inference_device: torch.device,
+    inference_dtype: torch.dtype,
+) -> Generator[tuple[torch.Tensor, list[tuple[int, int, int, int]]], None, None]:
+    """Used to compile batches of patches from the input array and return them as a generator."""
+
+    with ThreadPoolExecutor(max_workers=batch_size) as executor:
+        futures = [
+            executor.submit(get_patch, input_array, index, no_data_value)
+            for index in patch_indexes
+        ]
+
+        total_futures = len(futures)
+        all_indexes = set()
+        index_batch = []
+        patch_batch_array = np.zeros(
+            (batch_size, input_array.shape[0], patch_size, patch_size), dtype=np.float32
+        )
+
+        for index, future in enumerate(as_completed(futures)):
+            patch, new_index = future.result()
+
+            if patch is not None and new_index not in all_indexes:
+                index_batch.append(new_index)
+                patch_batch_array[len(index_batch) - 1] = patch
+                all_indexes.add(new_index)
+
+            if len(index_batch) == batch_size or index == total_futures - 1:
+                if len(index_batch) == 0:
+                    continue
+                input_tensor = (
+                    torch.tensor(patch_batch_array[: len(index_batch)])
+                    .to(inference_device)
+                    .to(inference_dtype)
+                )
+                yield input_tensor, index_batch
+                index_batch = []
+
+
+def run_models_on_array(
+    models: list[torch.nn.Module],
+    input_array: np.ndarray,
+    pred_tracker: torch.Tensor,
+    grad_tracker: Union[torch.Tensor, None],
+    patch_size: int,
+    patch_overlap: int,
+    inference_device: torch.device,
+    batch_size: int = 2,
+    inference_dtype: torch.dtype = torch.float32,
+    no_data_value: int = 0,
+) -> None:
+    """Used to execute the model on the input array, in patches. Predictions are stored in pred_tracker and grad_tracker, updated in place."""
+    patch_indexes = make_patch_indexes(
+        array_height=input_array.shape[1],
+        array_width=input_array.shape[2],
+        patch_size=patch_size,
+        patch_overlap=patch_overlap,
+    )
+
+    gradient = create_gradient_mask(
+        patch_size, patch_overlap, device=inference_device, dtype=inference_dtype
+    )
+
+    input_tensor_gen = compile_batches(
+        batch_size=batch_size,
+        patch_size=patch_size,
+        patch_indexes=patch_indexes,
+        input_array=input_array,
+        no_data_value=no_data_value,
+        inference_device=inference_device,
+        inference_dtype=inference_dtype,
+    )
+
+    for patch_batch, index_batch in input_tensor_gen:
+        inference_and_store(
+            models=models,
+            patch_batch=patch_batch,
+            index_batch=index_batch,
+            pred_tracker=pred_tracker,
+            gradient=gradient,
+            grad_tracker=grad_tracker,
+        )
+
+
+def check_patch_size(
+    input_array: np.ndarray, no_data_value: int, patch_size: int, patch_overlap: int
+) -> tuple[int, int]:
+    """Used to check the inputs and adjust the patch size and overlap if necessary."""
+    # check the shape of the input array
+    if len(input_array.shape) != 3:
+        raise ValueError(
+            f"Input array must have 3 dimensions, found {len(input_array.shape)}. The input should be in format (bands (red,green,NIR), height, width)."
+        )
+
+    # check the width and height are greater than 10 pixels
+    if min(input_array.shape[1], input_array.shape[2]) < 10:
+        raise ValueError(
+            f"Input array must have a width and height greater than 10 pixels, found shape {input_array.shape}. The input should be in format (bands (red,green,NIR), height, width)."
+        )
+    if min(input_array.shape[1], input_array.shape[2]) < 50:
+        warnings.warn(
+            f"Input width or height is less than 50 pixels, found shape {input_array.shape}. Such a small image may not provide adequate spatial context for the model."
+        )
+
+    # if the input has a lot of no data values and the patch size is larger than half the image size, we reduce the patch size and overlap
+    if np.count_nonzero(input_array == no_data_value) / input_array.size > 0.3:
+        if patch_size > min(input_array.shape[1], input_array.shape[2]) / 2:
+            patch_size = min(input_array.shape[1], input_array.shape[2]) // 2
+            if patch_size // 2 < patch_overlap:
+                patch_overlap = patch_size // 2
+
+            warnings.warn(
+                f"Significant no-data areas detected. Adjusting patch size to {patch_size}px and overlap to {patch_overlap}px to minimize no-data patches."
+            )
+
+    # if the patch size is larger than the image size, we reduce the patch size and overlap
+    if patch_size > min(input_array.shape[1], input_array.shape[2]):
+        patch_size = min(input_array.shape[1], input_array.shape[2])
+        if patch_size // 2 < patch_overlap:
+            patch_overlap = patch_size // 2
+        warnings.warn(
+            f"Patch size too large, reducing to {patch_size} and overlap to {patch_overlap}."
+        )
+
+    # if the patch overlap is larger than the patch size, raise an error
+    if patch_overlap >= patch_size:
+        raise ValueError(
+            f"Patch overlap {patch_overlap}px must be less than patch size {patch_size}px."
+        )
+    return patch_overlap, patch_size
+
+
+def coordinator(
+    input_array: np.ndarray,
+    models: list[torch.nn.Module],
+    inference_dtype: torch.dtype,
+    export_confidence: bool,
+    softmax_output: bool,
+    inference_device: torch.device,
+    mosaic_device: torch.device,
+    patch_size: int,
+    patch_overlap: int,
+    batch_size: int,
+    profile: Profile = Profile(),
+    output_path: Path = Path(""),
+    no_data_value: int = 0,
+    pbar: Optional[tqdm] = None,
+    apply_no_data_mask: bool = False,
+    export_to_disk: bool = True,
+    save_executor: Optional[ThreadPoolExecutor] = None,
+    pred_classes: int = 4,
+) -> np.ndarray:
+    """Used to coordinate the process of predicting from an input array."""
+
+    patch_overlap, patch_size = check_patch_size(
+        input_array, no_data_value, patch_size, patch_overlap
+    )
+
+    pred_tracker = torch.zeros(
+        (pred_classes, *input_array.shape[1:3]),
+        dtype=inference_dtype,
+        device=mosaic_device,
+    )
+
+    grad_tracker = (
+        torch.zeros(input_array.shape[1:3], dtype=inference_dtype, device=mosaic_device)
+        if export_confidence
+        else None
+    )
+
+    run_models_on_array(
+        models=models,
+        input_array=input_array,
+        pred_tracker=pred_tracker,
+        grad_tracker=grad_tracker,
+        inference_device=inference_device,
+        inference_dtype=inference_dtype,
+        no_data_value=no_data_value,
+        patch_size=patch_size,
+        patch_overlap=patch_overlap,
+        batch_size=batch_size,
+    )
+
+    if export_confidence:
+        pred_tracker_norm = pred_tracker / grad_tracker
+        if softmax_output:
+            pred_tracker = torch.clip(
+                (torch.nn.functional.softmax(pred_tracker_norm, 0) + 0.001),
+                0.001,
+                0.999,
+            )
+        else:
+            pred_tracker = pred_tracker_norm
+
+        pred_tracker_np = pred_tracker.float().numpy(force=True)
+
+    else:
+        pred_tracker_np = (
+            torch.argmax(pred_tracker, 0, keepdim=True)
+            .numpy(force=True)
+            .astype(np.uint8)
+        )
+
+    if apply_no_data_mask:
+        pred_tracker_np = mask_prediction(input_array, pred_tracker_np, no_data_value)
+
+    if export_to_disk:
+        export_profile = profile.copy()
+        export_profile.update(
+            dtype=pred_tracker_np.dtype,
+            count=pred_tracker_np.shape[0],
+            compress="lzw",
+            nodata=0,
+            driver="GTiff",
+        )
+        # if executer has been passed, submit the save_prediction function to it, to avoid blocking the main thread
+        if save_executor:
+            save_executor.submit(
+                save_prediction, output_path, export_profile, pred_tracker_np
+            )
+        # otherwise save the prediction directly
+
+        else:
+            save_prediction(output_path, export_profile, pred_tracker_np)
+
+    if pbar:
+        pbar.update(1)
+    return pred_tracker_np
+
+
+def collect_models(
+    custom_models: Union[list[torch.nn.Module], torch.nn.Module],
+    inference_device: torch.device,
+    inference_dtype: torch.dtype,
+    source: str,
+    destination_model_dir: Union[str, Path, None] = None,
+) -> list[torch.nn.Module]:
+    if not custom_models:
+        models = []
+        for model_details in get_models(model_dir=destination_model_dir, source=source):
+            models.append(
+                load_model_from_weights(
+                    model_name=model_details["timm_model_name"],
+                    weights_path=model_details["Path"],
+                    device=inference_device,
+                    dtype=inference_dtype,
+                )
+            )
+    else:
+        # if not a list, make it a list of models
+        if not isinstance(custom_models, list):
+            custom_models = [custom_models]
+
+        models = [
+            model.to(inference_dtype).to(inference_device) for model in custom_models
+        ]
+    return models
+
+
+def predict_from_array(
+    input_array: np.ndarray,
+    patch_size: int = 1000,
+    patch_overlap: int = 300,
+    batch_size: int = 1,
+    inference_device: Union[str, torch.device] = default_device(),
+    mosaic_device: Optional[Union[str, torch.device]] = None,
+    inference_dtype: Union[torch.dtype, str] = torch.float32,
+    export_confidence: bool = False,
+    softmax_output: bool = True,
+    no_data_value: int = 0,
+    apply_no_data_mask: bool = True,
+    custom_models: Union[list[torch.nn.Module], torch.nn.Module] = [],
+    pred_classes: int = 4,
+    destination_model_dir: Union[str, Path, None] = None,
+    model_download_source: str = "google_drive",
+) -> np.ndarray:
+    """Predict a cloud and cloud shadow mask from a Red, Green and NIR numpy array, with a spatial res between 10 m and 50 m.
+
+    Args:
+        input_array (np.ndarray): A numpy array with shape (3, height, width) representing the Red, Green and NIR bands.
+        patch_size (int, optional): Size of the patches for inference. Defaults to 1000.
+        patch_overlap (int, optional): Overlap between patches for inference. Defaults to 300.
+        batch_size (int, optional): Number of patches to process in a batch. Defaults to 1.
+        inference_device (Union[str, torch.device], optional): Device to use for inference (e.g., 'cpu', 'cuda', 'mps'). Defaults to the device returned by default_device().
+        mosaic_device (Union[str, torch.device], optional): Device to use for mosaicking patches. Defaults to inference device.
+        inference_dtype (Union[torch.dtype, str], optional): Data type for inference. Defaults to torch.float32.
+        export_confidence (bool, optional): If True, exports confidence maps instead of predicted classes. Defaults to False.
+        softmax_output (bool, optional): If True, applies a softmax to the output, only used if export_confidence = True. Defaults to True.
+        no_data_value (int, optional): Value within input scenes that specifies no data region. Defaults to 0.
+        apply_no_data_mask (bool, optional): If True, applies a no-data mask to the predictions. Defaults to True.
+        custom_models Union[list[torch.nn.Module], torch.nn.Module], optional): A list or singular custom torch models to use for prediction. Defaults to [].
+        pred_classes (int, optional): Number of classes to predict. Defaults to 4, to be used with custom models.
+        destination_model_dir Union[str, Path, None]: Directory to save the model weights. Defaults to None.
+        model_download_source (str, optional): Source from which to download the model weights. Defaults to "google_drive", can also be "hugging_face".
+    Returns:
+        np.ndarray: A numpy array with shape (1, height, width) or (4, height, width if export_confidence = True) representing the predicted cloud and cloud shadow mask.
+
+    """
+
+    inference_device = torch.device(inference_device)
+    if mosaic_device is None:
+        mosaic_device = inference_device
+    else:
+        mosaic_device = torch.device(mosaic_device)
+
+    inference_dtype = get_torch_dtype(inference_dtype)
+    # if no custom model paths are provided, use the default models
+    models = collect_models(
+        custom_models=custom_models,
+        inference_device=inference_device,
+        inference_dtype=inference_dtype,
+        source=model_download_source,
+        destination_model_dir=destination_model_dir,
+    )
+
+    pred_tracker = coordinator(
+        input_array=input_array,
+        models=models,
+        inference_device=inference_device,
+        mosaic_device=mosaic_device,
+        inference_dtype=inference_dtype,
+        export_confidence=export_confidence,
+        softmax_output=softmax_output,
+        patch_size=patch_size,
+        patch_overlap=patch_overlap,
+        batch_size=batch_size,
+        no_data_value=no_data_value,
+        export_to_disk=False,
+        apply_no_data_mask=apply_no_data_mask,
+        pred_classes=pred_classes,
+    )
+
+    return pred_tracker
+
+
+def predict_from_load_func(
+    scene_paths: Union[list[Path], list[str]],
+    load_func: Callable,
+    patch_size: int = 1000,
+    patch_overlap: int = 300,
+    batch_size: int = 1,
+    inference_device: Union[str, torch.device] = default_device(),
+    mosaic_device: Optional[Union[str, torch.device]] = None,
+    inference_dtype: Union[torch.dtype, str] = torch.float32,
+    export_confidence: bool = False,
+    softmax_output: bool = True,
+    no_data_value: int = 0,
+    overwrite: bool = True,
+    apply_no_data_mask: bool = True,
+    output_dir: Optional[Union[Path, str]] = None,
+    custom_models: Union[list[torch.nn.Module], torch.nn.Module] = [],
+    destination_model_dir: Union[str, Path, None] = None,
+    model_download_source: str = "google_drive",
+) -> list[Path]:
+    """
+    Predicts cloud and cloud shadow masks for a list of scenes using a specified loading function.
+
+    Args:
+        scene_paths (Union[list[Path], list[str]]): A list of paths to the scene files to be processed.
+        load_func (Callable): A function to load the scene data. This function should take an input_path parameter and return a R,G,NIR numpy array and a rasterio for export profile, several load func are provided within data_loaders.py
+        patch_size (int, optional): Size of the patches for inference. Defaults to 1000.
+        patch_overlap (int, optional): Overlap between patches for inference. Defaults to 300.
+        batch_size (int, optional): Number of patches to process in a batch. Defaults to 1.
+        inference_device (Union[str, torch.device], optional): Device to use for inference (e.g., 'cpu', 'cuda', 'mps'). Defaults to the device returned by default_device().
+        mosaic_device (Union[str, torch.device], optional): Device to use for mosaicking patches. Defaults to inference device.
+        inference_dtype (Union[torch.dtype, str], optional): Data type for inference. Defaults to torch.float32.
+        export_confidence (bool, optional): If True, exports confidence maps instead of predicted classes. Defaults to False.
+        softmax_output (bool, optional): If True, applies a softmax to the output, only used if export_confidence = True. Defaults to True.
+        no_data_value (int, optional): Value within input scenes that specifies no data region. Defaults to 0.
+        overwrite (bool, optional): If False, skips scenes that already have a prediction file. Defaults to True.
+        apply_no_data_mask (bool, optional): If True, applies a no-data mask to the predictions. Defaults to True.
+        output_dir (Optional[Union[Path, str]], optional): Directory to save the prediction files. Defaults to None. If None, the predictions will be saved in the same directory as the input scene.
+        custom_models Union[list[torch.nn.Module], torch.nn.Module], optional): A list or singular custom torch models to use for prediction. Defaults to [].
+        destination_model_dir Union[str, Path, None]: Directory to save the model weights. Defaults to None.
+        model_download_source (str, optional): Source from which to download the model weights. Defaults to "google_drive", can also be "hugging_face".
+
+    Returns:
+        list[Path]: A list of paths to the output prediction files.
+
+    """
+    pred_paths = []
+    inf_thread = Thread()
+    save_executor = ThreadPoolExecutor(max_workers=1)
+
+    inference_device = torch.device(inference_device)
+    if mosaic_device is None:
+        mosaic_device = inference_device
+    else:
+        mosaic_device = torch.device(mosaic_device)
+
+    inference_dtype = get_torch_dtype(inference_dtype)
+
+    models = collect_models(
+        custom_models=custom_models,
+        inference_device=inference_device,
+        inference_dtype=inference_dtype,
+        destination_model_dir=destination_model_dir,
+        source=model_download_source,
+    )
+
+    pbar = tqdm(
+        total=len(scene_paths),
+        desc=f"Running inference using {inference_device.type} {str(inference_dtype).split('.')[-1]}",
+    )
+
+    for scene_path in scene_paths:
+        scene_path = Path(scene_path)
+        file_name = f"{scene_path.stem}_OCM_v{__version__.replace('.','_')}.tif"
+
+        if output_dir is None:
+            output_path = scene_path.parent / file_name
+        else:
+            Path(output_dir).mkdir(parents=True, exist_ok=True)
+            output_path = Path(output_dir) / file_name
+
+        pred_paths.append(output_path)
+
+        if output_path.exists() and not overwrite:
+            pbar.update(1)
+            pbar.refresh()
+            continue
+
+        input_array, profile = load_func(input_path=scene_path)
+
+        while inf_thread.is_alive():
+            inf_thread.join()
+
+        inf_thread = Thread(
+            target=coordinator,
+            kwargs={
+                "input_array": input_array,
+                "profile": profile,
+                "output_path": output_path,
+                "models": models,
+                "inference_dtype": inference_dtype,
+                "export_confidence": export_confidence,
+                "softmax_output": softmax_output,
+                "inference_device": inference_device,
+                "mosaic_device": mosaic_device,
+                "patch_size": patch_size,
+                "patch_overlap": patch_overlap,
+                "batch_size": batch_size,
+                "no_data_value": no_data_value,
+                "pbar": pbar,
+                "apply_no_data_mask": apply_no_data_mask,
+                "save_executor": save_executor,
+            },
+        )
+        inf_thread.start()
+
+    while inf_thread.is_alive():
+        inf_thread.join()
+
+    if inference_device.type.startswith("cuda"):
+        torch.cuda.empty_cache()
+
+    save_executor.shutdown(wait=True)
+    pbar.refresh()
+
+    return pred_paths

omnicloudmask/data_loaders.py

@@ -0,0 +1,159 @@
+import warnings
+from pathlib import Path
+from typing import Optional, Union
+
+import numpy as np
+import rasterio as rio
+from rasterio.profiles import Profile
+
+
+def load_s2(
+    input_path: Union[Path, str],
+    resolution: float = 10.0,
+    required_bands: list[str] = ["B04", "B03", "B8A"],
+) -> tuple[np.ndarray, Profile]:
+    """Load a Sentinel-2 (L1C or L2A) image from a SAFE folder containing the bands"""
+    if not 10 <= resolution <= 50:
+        raise ValueError("Resolution must be between 10 and 50")
+    input_path = Path(input_path)
+    processing_level = find_s2_processing_level(input_path)
+    return open_s2_bands(input_path, processing_level, resolution, required_bands)
+
+
+def find_s2_processing_level(
+    input_path: Path,
+) -> str:
+    """Derive the processing level of a Sentinel-2 image from the folder name."""
+
+    folder_name = Path(input_path).name
+    processing_level = folder_name.split("_")[1][3:6]
+
+    if processing_level not in ["L1C", "L2A"]:
+        raise ValueError(
+            f"Processing level {processing_level} not recognized, expected L1C or L2A"
+        )
+    return processing_level
+
+
+def open_s2_bands(
+    input_path: Path,
+    processing_level: str,
+    resolution: float,
+    required_bands: list[str],
+) -> tuple[np.ndarray, Profile]:
+    bands = []
+    for band_name in required_bands:
+        if processing_level == "L1C":
+            try:
+                band = list(input_path.rglob(f"*IMG_DATA/*{band_name}.jp2"))[0]
+
+            except IndexError:
+                raise ValueError(f"Band {band_name} not found in {input_path}")
+        else:
+            band = None
+            for search_resolution in [10, 20, 60]:
+                band_paths = list(
+                    input_path.rglob(f"*{band_name}_{search_resolution}m.jp2")
+                )
+                if band_paths:
+                    band = band_paths[0]
+                    break
+            if not band:
+                raise ValueError(f"Band {band_name} not found in {input_path}")
+
+        with rio.open(band) as src:
+            profile = src.profile
+            native_resolution = int(src.res[0])
+            scale_factor = native_resolution / resolution
+            if native_resolution == resolution:
+                bands.append(src.read(1))
+            else:
+                bands.append(
+                    src.read(
+                        1,
+                        out_shape=(
+                            int(src.height * scale_factor),
+                            int(src.width * scale_factor),
+                        ),
+                    )
+                )
+    profile["transform"] = rio.transform.from_origin(  # type: ignore
+        profile["transform"][2],
+        profile["transform"][5],
+        resolution,
+        resolution,
+    )
+    data = np.array(bands)
+    profile["height"] = data.shape[1]
+    profile["width"] = data.shape[2]
+    return data, profile
+
+
+def load_multiband(
+    input_path: Union[Path, str],
+    resample_res: Optional[float] = None,
+    band_order: Optional[list[int]] = None,
+) -> tuple[np.ndarray, Profile]:
+    """Load a multiband image and resample it to requested resolution."""
+    if band_order is None:
+        warnings.warn(
+            "No band order provided, using default [1, 2, 3] (RGN)", UserWarning
+        )
+        band_order = [1, 2, 3]
+    input_path = Path(input_path)
+
+    with rio.open(input_path) as src:
+        if resample_res:
+            current_res = src.res
+            desired_res = (resample_res, resample_res)
+            scale_factor = (
+                current_res[0] / desired_res[0],
+                current_res[1] / desired_res[1],
+            )
+        else:
+            scale_factor = (1, 1)
+
+        data = src.read(
+            band_order,
+            out_shape=(
+                len(band_order),
+                int(src.height * scale_factor[0]),
+                int(src.width * scale_factor[1]),
+            ),
+            resampling=rio.enums.Resampling.nearest,  # type: ignore
+        )
+        profile = src.profile
+
+        return data, profile
+
+
+def load_ls8(
+    input_path: Union[Path, str],
+    resolution: int = 30,
+    required_bands=["B4", "B3", "B5"],
+) -> tuple[np.ndarray, Profile]:
+    """Load a Landsat 8 image from a folder containing the bands"""
+    if resolution != 30:
+        raise ValueError("Resolution must be 30")
+
+    input_path = Path(input_path)
+
+    band_files = {}
+    for band_name in required_bands:
+        try:
+            band = list(input_path.rglob(f"*{band_name}.TIF"))[0]
+
+        except IndexError:
+            raise ValueError(f"Band {band_name} not found in {input_path}")
+        band_files[band_name] = band
+
+    data = []
+    profile = Profile()
+    for band_name in required_bands:
+        with rio.open(band_files[band_name]) as src:
+            if not profile:
+                profile = src.profile
+            data.append(src.read(1))
+
+    data = np.array(data)
+    return data, profile

omnicloudmask/download_models.py

@@ -0,0 +1,92 @@
+from pathlib import Path
+from typing import Union
+
+import gdown
+import pandas as pd
+import torch
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+
+
+def download_file_from_google_drive(file_id: str, destination: Path) -> None:
+    """
+    Downloads a file from Google Drive and saves it at the given destination using gdown.
+
+    Args:
+        file_id (str): The ID of the file on Google Drive.
+        destination (Path): The local path where the file should be saved.
+    """
+    url = f"https://drive.google.com/uc?id={file_id}"
+    gdown.download(url, str(destination), quiet=False)
+
+
+def download_file_from_hugging_face(destination: Path) -> None:
+    """
+    Downloads a file from Hugging Face and saves it at the given destination using hf_hub_download.
+    Loads the resulting safetensors file and saves it as a PyTorch model state for compatibility with the rest of the codebase.
+
+    Args:
+        file_id (str): The ID of the file on Hugging Face.
+        destination (Path): The local path where the file should be saved.
+    """
+    file_name = destination.stem
+    safetensor_path = hf_hub_download(
+        repo_id="NickWright/OmniCloudMask",
+        filename=f"{file_name}.safetensors",
+        force_download=True,
+        cache_dir=destination.parent,
+    )
+    model_state = load_file(safetensor_path)
+    torch.save(model_state, destination)
+
+
+def download_file(file_id: str, destination: Path, source: str) -> None:
+    if source == "google_drive":
+        download_file_from_google_drive(file_id, destination)
+    elif source == "hugging_face":
+        download_file_from_hugging_face(destination)
+    else:
+        raise ValueError(
+            "Invalid source. Supported sources are 'google_drive' and 'hugging_face'."
+        )
+
+
+def get_models(
+    force_download: bool = False,
+    model_dir: Union[str, Path, None] = None,
+    source: str = "google_drive",
+) -> list[dict]:
+    """
+    Downloads the model weights from Google Drive and saves them locally.
+
+    Args:
+        force_download (bool): Whether to force download the model weights even if they already exist locally.
+        model_dir (Union[str, Path, None]): The directory where the model weights should be saved.
+        source (str): The source from which the model weights should be downloaded. Currently, only "google_drive" or "hugging_face" are supported.
+    """
+
+    df = pd.read_csv(
+        Path(__file__).resolve().parent / "models/model_download_links.csv"
+    )
+    model_paths = []
+
+    for _, row in df.iterrows():
+        file_id = str(row["google_drive_id"])
+
+        if model_dir is not None:
+            model_dir = Path(model_dir)
+        else:
+            model_dir = Path(__file__).resolve().parent / "models"
+
+        model_dir.mkdir(exist_ok=True)
+        destination = model_dir / str(row["file_name"])
+        timm_model_name = row["timm_model_name"]
+
+        if not destination.exists() or force_download:
+            download_file(file_id=file_id, destination=destination, source=source)
+
+        elif destination.stat().st_size <= 1024 * 1024:
+            download_file(file_id=file_id, destination=destination, source=source)
+
+        model_paths.append({"Path": destination, "timm_model_name": timm_model_name})
+    return model_paths

omnicloudmask/model_utils.py

@@ -0,0 +1,208 @@
+from functools import partial
+from pathlib import Path
+from typing import Optional, Union
+
+import numpy as np
+import timm
+import torch
+from fastai.vision.learner import create_unet_model
+
+
+def get_torch_dtype(dtype: Union[torch.dtype, str]) -> torch.dtype:
+    """Return a torch.dtype from a string or torch.dtype."""
+    if isinstance(dtype, str):
+        dtype_mapping = {
+            "float16": torch.float16,
+            "half": torch.float16,
+            "fp16": torch.float16,
+            "float32": torch.float32,
+            "float": torch.float32,
+            "bfloat16": torch.bfloat16,
+            "bf16": torch.bfloat16,
+        }
+        try:
+            return dtype_mapping[dtype.lower()]
+        except KeyError:
+            raise ValueError(
+                f"Invalid dtype: {dtype}. Must be one of {list(dtype_mapping.keys())}"
+            )
+    elif isinstance(dtype, torch.dtype):
+        return dtype
+    else:
+        raise TypeError(
+            f"Expected dtype to be a str or torch.dtype, but got {type(dtype)}"
+        )
+
+
+def create_gradient_mask(
+    patch_size: int, patch_overlap: int, device: torch.device, dtype: torch.dtype
+) -> torch.Tensor:
+    """Create a gradient mask for a given patch size and overlap."""
+    if patch_overlap > 0:
+        if patch_overlap * 2 > patch_size:
+            patch_overlap = patch_size // 2
+
+        gradient_strength = 1
+        gradient = (
+            torch.ones((patch_size, patch_size), dtype=torch.int, device=device)
+            * patch_overlap
+        )
+        gradient[:, :patch_overlap] = torch.tile(
+            torch.arange(1, patch_overlap + 1),
+            (patch_size, 1),
+        )
+        gradient[:, -patch_overlap:] = torch.tile(
+            torch.arange(patch_overlap, 0, -1),
+            (patch_size, 1),
+        )
+        gradient = gradient / patch_overlap
+        rotated_gradient = torch.rot90(gradient)
+        combined_gradient = rotated_gradient * gradient
+
+        combined_gradient = (combined_gradient * gradient_strength) + (
+            1 - gradient_strength
+        )
+    else:
+        combined_gradient = torch.ones(
+            (patch_size, patch_size), dtype=torch.int, device=device
+        )
+    return combined_gradient.to(dtype)
+
+
+def channel_norm(patch: np.ndarray, nodata_value: Optional[int] = 0) -> np.ndarray:
+    """Normalize each band of the input array by subtracting the nonzero mean and dividing
+    by the nonzero standard deviation then fill nodata values with 0."""
+    out_array = np.zeros(patch.shape).astype(np.float32)
+    for id, band in enumerate(patch):
+        # Mask for non-zero values
+        mask = band != nodata_value
+        # Check if there are any non-zero values
+        if np.any(mask):
+            mean = band[mask].mean()
+            std = band[mask].std()
+            if std == 0:
+                std = 1  # Prevent division by zero
+            # Normalize only non-zero values
+            out_array[id][mask] = (band[mask] - mean) / std
+        else:
+            continue
+        # Fill original nodata values with 0
+        out_array[id][~mask] = 0
+    return out_array
+
+
+def store_results(
+    pred_batch: torch.Tensor,
+    index_batch: list[tuple],
+    pred_tracker: torch.Tensor,
+    gradient: torch.Tensor,
+    grad_tracker: Optional[torch.Tensor] = None,
+) -> None:
+    """Store the results of the model inference in the pred_tracker and grad_tracker tensors."""
+    # Store the predictions in the pred_tracker tensor
+    assert pred_batch.ndim == 4, "pred_batch must have 4 dimensions, (B, class, H, W)"
+    assert pred_batch.shape[0] == len(index_batch), "Batch size must match index_batch"
+    assert pred_batch.shape[1] == pred_tracker.shape[0], "Number of classes must match"
+    assert pred_batch.shape[2] == gradient.shape[0], "Height must match gradient"
+    assert pred_batch.shape[3] == gradient.shape[1], "Width must match gradient"
+
+    pred_batch *= gradient[None, None, :, :]
+
+    for pred, index in zip(pred_batch.to(pred_tracker.device), index_batch):
+        pred_tracker[:, index[0] : index[1], index[2] : index[3]] += pred
+        if grad_tracker is not None:
+            grad_tracker[index[0] : index[1], index[2] : index[3]] += gradient.to(
+                grad_tracker.device
+            )
+
+
+def inference_and_store(
+    models: list[torch.nn.Module],
+    patch_batch: torch.Tensor,
+    index_batch: list[tuple],
+    pred_tracker: torch.Tensor,
+    gradient: torch.Tensor,
+    grad_tracker: Optional[torch.Tensor] = None,
+) -> None:
+    """Perform inference on the patch_batch and store the results in the pred_tracker and grad_tracker tensors."""
+    # pre-initialize the all_preds tensor to store the predictions from each model
+    all_preds = torch.zeros(
+        len(models),
+        patch_batch.shape[0],
+        pred_tracker.shape[0],
+        patch_batch.shape[2],
+        patch_batch.shape[3],
+        device=patch_batch.device,
+        dtype=patch_batch.dtype,
+    )
+    for index, model in enumerate(models):
+        with torch.no_grad():
+            all_preds[index] = model(patch_batch)
+
+    mean_preds = all_preds.mean(dim=0)
+
+    store_results(
+        pred_batch=mean_preds,
+        index_batch=index_batch,
+        pred_tracker=pred_tracker,
+        gradient=gradient,
+        grad_tracker=grad_tracker,
+    )
+
+
+def default_device() -> torch.device:
+    """Return the default device for model inference"""
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    elif torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+def load_model(
+    model_path: Union[Path, str],
+    device: torch.device,
+    dtype: torch.dtype = torch.float32,
+) -> torch.nn.Module:
+    """Load a PyTorch model from a file and move it to the specified device and dtype."""
+    model_path = Path(model_path)
+    if not model_path.is_file():
+        raise FileNotFoundError(f"Model file not found at: {model_path}")
+
+    try:
+        model = torch.load(model_path, map_location="cpu")
+    except Exception as e:
+        raise RuntimeError(f"Error loading model: {e}")
+
+    model.eval()
+    return model.to(dtype).to(device)
+
+
+def load_model_from_weights(
+    model_name: str,
+    weights_path: Union[Path, str],
+    device: torch.device,
+    dtype: torch.dtype = torch.float32,
+    in_chans: int = 3,
+    n_out: int = 4,
+) -> torch.nn.Module:
+    """Build Fastai DynamicUnet model from timm model and load weights from file"""
+    timm_model = partial(
+        timm.create_model,
+        model_name=model_name,
+        pretrained=False,
+        in_chans=in_chans,
+    )
+
+    model = create_unet_model(
+        arch=timm_model,
+        n_out=n_out,
+        img_size=(509, 509),
+        act_cls=torch.nn.Mish,
+        pretrained=False,
+    )
+
+    model.load_state_dict(torch.load(weights_path, weights_only=True))
+    model.eval()
+
+    return model.to(dtype).to(device)

omnicloudmask/models/PM_model_2.2.10_RG_NIR_509_convnextv2_nano.fcmae_ft_in1k_PT_state.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7d83ddef55797fd443cb30fdd545edf4f070c76cfb031ab53af2cd01f51d6d0f
+size 130226202

omnicloudmask/models/PM_model_2.2.10_RG_NIR_509_regnety_004.pycls_in1k_PT_state.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:be9e29fa69464a286d40e71b5af10894cabe2a258a6f1de4e869500ae704c7bd
+size 72458313

omnicloudmask/models/model_download_links.csv

@@ -0,0 +1,3 @@
+file_name,timm_model_name,google_drive_id
+PM_model_2.2.10_RG_NIR_509_regnety_004.pycls_in1k_PT_state.pth,regnety_004,1tGJh9nnrH-apjmV70AcK8VtXnbBtRb67
+PM_model_2.2.10_RG_NIR_509_convnextv2_nano.fcmae_ft_in1k_PT_state.pth,convnextv2_nano,1QXQ_oPhLKEowC9fxlZGLOACt8gCNMbWP

omnicloudmask/raster_utils.py

@@ -0,0 +1,118 @@
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+import rasterio as rio
+from rasterio.profiles import Profile
+
+from .model_utils import channel_norm
+
+
+def get_patch(
+    input_array: np.ndarray,
+    index: tuple,
+    no_data_value: Optional[int] = 0,
+) -> tuple[Optional[np.ndarray], Optional[tuple[int, int, int, int]]]:
+    """Extract a patch from a 3D array and normalize it. If the patch is entirely nodata, return None.
+    If the patch contains nodata, try to move patches to reduce nodata regions in patches.
+    """
+    assert input_array.ndim == 3, "Input array must have 3 dimensions"
+
+    top, bottom, left, right = index
+    patch = input_array[:, top:bottom, left:right].astype(np.float32)
+
+    if patch.sum() == 0:
+        return None, None
+
+    if no_data_value is None:
+        if np.all(patch == no_data_value):
+            return None, None
+
+    if np.any(patch == 0):
+        max_bottom, max_right = input_array.shape[1:3]
+
+        if np.any(patch[:, 0, :]) or np.any(patch[:, -1, :]):
+            while not np.any(patch[:, 0, :]) and bottom < max_bottom:  # check top row
+                patch = patch[:, 1:, :]
+                top += 1
+                bottom += 1
+
+            while not np.any(patch[:, -1, :]) and top > 0:
+                patch = patch[:, :-1, :]
+                bottom -= 1
+                top -= 1
+
+        # Both sides are not zero-filled
+        if np.any(patch[:, :, 0]) or np.any(patch[:, :, -1]):
+            while not np.any(patch[:, :, 0]) and right < max_right:  # check left column
+                patch = patch[:, :, 1:]
+                left += 1
+                right += 1
+
+            while not np.any(patch[:, :, -1]) and left > 0:  # check right column
+                patch = patch[:, :, :-1]
+                right -= 1
+                left -= 1
+        patch = input_array[:, top:bottom, left:right].astype(np.float32)
+        index = (top, bottom, left, right)
+
+    # trim index bottom and right to match patch shape
+    index = (top, top + patch.shape[1], left, left + patch.shape[2])
+    return channel_norm(patch, no_data_value), index
+
+
+def mask_prediction(
+    scene: np.ndarray, pred_tracker_np: np.ndarray, no_data_value: int = 0
+) -> np.ndarray:
+    """Create a no data mask from a raster scene."""
+    assert scene.ndim == 3, "Scene must have 3 dimensions"
+    assert pred_tracker_np.ndim == 3, "Prediction tracker must have 3 dimensions"
+    assert (
+        scene.shape[1:] == pred_tracker_np.shape[1:]
+    ), "Scene and prediction tracker must have the same shape"
+    mask = np.all(scene != no_data_value, axis=0).astype(np.uint8)
+    pred_tracker_np *= mask
+    return pred_tracker_np
+
+
+def make_patch_indexes(
+    array_width: int,
+    array_height: int,
+    patch_size: int = 1000,
+    patch_overlap: int = 300,
+) -> list[tuple[int, int, int, int]]:
+    """Create a list of patch indexes for a given shape and patch size."""
+    assert patch_size > patch_overlap, "Patch size must be greater than patch overlap"
+    assert patch_overlap >= 0, "Patch overlap must be greater than or equal to 0"
+    assert patch_size > 0, "Patch size must be greater than 0"
+    assert (
+        patch_size <= array_width
+    ), "Patch size must be less than or equal to array width"
+    assert (
+        patch_size <= array_height
+    ), "Patch size must be less than or equal to array height"
+
+    stride = patch_size - patch_overlap
+
+    max_bottom = array_height - patch_size
+    max_right = array_width - patch_size
+
+    patch_indexes = []
+    for top in range(0, array_height, stride):
+        if top > max_bottom:
+            top = max_bottom
+        bottom = top + patch_size
+        for left in range(0, array_width, stride):
+            if left > max_right:
+                left = max_right
+            right = left + patch_size
+            patch_indexes.append((top, bottom, left, right))
+
+    return patch_indexes
+
+
+def save_prediction(
+    output_path: Path, export_profile: Profile, pred_tracker_np: np.ndarray
+) -> None:
+    with rio.open(output_path, "w", **export_profile) as dst:
+        dst.write(pred_tracker_np)