utils/preprocessing.py

"""Preprocessing utilities for satellite imagery."""

import numpy as np
import cv2
from typing import Optional


def preprocess_image(
    image: np.ndarray,
    target_size: Optional[tuple] = None,
    normalize: bool = True
) -> np.ndarray:
    """
    Preprocess a satellite image for model input.

    Args:
        image: Input image (H, W, C), uint8 or float
        target_size: Optional (width, height) to resize to
        normalize: If True, output is float32 in [0, 1]

    Returns:
        Preprocessed image as float32 [0,1] or uint8 [0,255]
    """
    if image is None:
        raise ValueError("Input image is None")

    img = image.copy()

    # Ensure 3-channel
    if img.ndim == 2:
        img = np.stack([img, img, img], axis=-1)
    elif img.shape[2] == 1:
        img = np.concatenate([img, img, img], axis=-1)
    elif img.shape[2] > 3:
        img = img[:, :, :3]

    # Resize if requested
    if target_size is not None:
        img = cv2.resize(img, target_size, interpolation=cv2.INTER_LINEAR)

    # Normalise
    if normalize:
        if img.dtype == np.uint8:
            img = img.astype(np.float32) / 255.0
        else:
            img = np.clip(img, 0, 1).astype(np.float32)
    else:
        if img.dtype != np.uint8:
            img = (np.clip(img, 0, 1) * 255).astype(np.uint8)

    return img


def mask_clouds(
    image: np.ndarray,
    cloud_mask: np.ndarray,
    fill_value: float = 0.0
) -> np.ndarray:
    """
    Apply cloud mask to image, replacing cloud pixels with fill_value.

    Args:
        image: Input image (H, W, C)
        cloud_mask: Binary mask (H, W), 1 = cloud
        fill_value: Value to fill masked pixels with

    Returns:
        Masked image same dtype as input
    """
    masked = image.copy().astype(np.float32)
    mask_bool = cloud_mask.astype(bool)

    for c in range(masked.shape[2]):
        masked[:, :, c][mask_bool] = fill_value

    if image.dtype == np.uint8:
        masked = np.clip(masked, 0, 255).astype(np.uint8)

    return masked