omnicloudmask/cloud_mask.py

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