test_dofa.py

#!/usr/bin/env python3
"""Quick CLI smoke test for a self-contained DOFA checkpoint folder."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import torch
from transformers import pipeline


def parse_args() -> argparse.Namespace:
    repo_root = Path(__file__).resolve().parent
    parser = argparse.ArgumentParser(description="Run DOFA feature extraction on dummy or real input.")
    parser.add_argument(
        "--model",
        type=Path,
        default=repo_root / "dofa-base-patch16-224",
        help="Path to a checkpoint folder (default: dofa-base-patch16-224)",
    )
    parser.add_argument(
        "--image",
        type=Path,
        default=None,
        help="Optional image path (.npy HWC array or image readable by rasterio/PIL)",
    )
    parser.add_argument(
        "--pool",
        action="store_true",
        default=True,
        help="Return pooled features (default: True)",
    )
    parser.add_argument(
        "--no-pool",
        action="store_true",
        help="Return sequence features instead of pooled output",
    )
    parser.add_argument(
        "--device",
        default=None,
        help="Torch device, e.g. cuda or cpu (default: auto)",
    )
    return parser.parse_args()


def load_image(path: Path, num_channels: int) -> np.ndarray:
    if path.suffix == ".npy":
        array = np.load(path)
        if array.ndim != 3:
            raise ValueError(f"Expected HWC numpy array, got shape {array.shape}")
        return array

    try:
        import rasterio
    except ImportError as exc:
        raise ImportError("Install rasterio to load geospatial images, or pass a .npy file.") from exc

    with rasterio.open(path) as src:
        array = src.read()
        array = np.transpose(array, (1, 2, 0))
    return array


def main() -> None:
    args = parse_args()
    model_dir = args.model.resolve()
    if not model_dir.is_dir():
        raise SystemExit(f"Model folder not found: {model_dir}")

    config_path = model_dir / "config.json"
    with open(config_path, encoding="utf-8") as handle:
        config = json.load(handle)

    num_channels = config.get("num_channels") or len(config["default_wavelengths"])

    if args.image is None:
        image = np.random.randint(0, 255, (224, 224, num_channels), dtype=np.uint8)
        source = f"random dummy array ({num_channels} channels)"
    else:
        image = load_image(args.image.resolve(), num_channels)
        source = str(args.image)

    device = args.device or ("cuda" if torch.cuda.is_available() else "cpu")
    pool = not args.no_pool

    print(f"model:  {model_dir}")
    print(f"input:  {source}")
    print(f"device: {device}")
    print(f"pool:   {pool}")

    pipe = pipeline(
        task="dofa-feature-extraction",
        model=str(model_dir),
        trust_remote_code=True,
        device=device,
    )

    features = pipe(image, pool=pool, return_tensors=True)
    print(f"output: {tuple(features.shape)} dtype={features.dtype}")
    print("OK")


if __name__ == "__main__":
    main()