utils/env_checker.py

"""Check importable packages for the CD-Models training environment."""

from __future__ import annotations


REQUIRED_PACKAGES = {
    ("torch", "torch", "all models"),
    ("torchvision", "torchvision", "all models"),
    ("numpy", "numpy", "all models"),
    ("PIL", "Pillow", "all models"),
    ("cv2", "opencv-python", "all models"),
    ("yaml", "pyyaml", "all models"),
    ("tqdm", "tqdm", "weight_downloader"),
    ("timm", "timm", "TinyCD, ELGCNet, ChangeFormer"),
    ("einops", "einops", "BIT_CD, ChangeFormer"),
    ("mmengine", "mmengine==0.10.1", "ChangeMamba, Changer"),
    ("mmcv", "mmcv==2.1.0", "ChangeMamba, Changer"),
    ("mmseg", "mmsegmentation==1.2.2", "ChangeMamba, Changer"),
    ("gdown", "gdown", "weight_downloader GDrive fallback"),
}


def check_env(verbose: bool = True) -> bool:
    missing = []
    for import_name, pip_name, required_by in sorted(REQUIRED_PACKAGES):
        try:
            __import__(import_name)
            if verbose:
                print(f"  [OK] {import_name}")
        except ImportError:
            print(f"  [MISSING] {import_name} (install: pip install {pip_name}) - needed by: {required_by}")
            missing.append((import_name, pip_name))

    if missing:
        print(f"\n{len(missing)} packages missing. Install with:")
        for _, pip_name in missing:
            print(f"  pip install {pip_name}")
    else:
        print("\nAll packages present.")
    return not missing


if __name__ == "__main__":
    import sys

    sys.exit(0 if check_env(verbose=True) else 1)