utils/io.py

from __future__ import annotations

from pathlib import Path

import numpy as np
from numpy.typing import DTypeLike
from PIL import Image


def read_image(path: str | Path, dtype: DTypeLike = np.float32) -> np.ndarray:
    """Read an image file and return an HWC NumPy array."""
    path = str(path)

    try:
        import rasterio

        with rasterio.open(path) as src:
            data = src.read()  # (C, H, W)
            return np.moveaxis(data, 0, -1).astype(dtype)
    except ImportError:
        return np.asarray(Image.open(path), dtype=dtype)