X_SSL_Net/lib/data/utils.py

from __future__ import annotations

from pathlib import Path
import re


IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff"}


def is_image_file(path: Path) -> bool:
    return path.is_file() and path.suffix.lower() in IMAGE_EXTENSIONS


def list_image_files(folder: Path) -> list[Path]:
    if not folder.exists():
        raise FileNotFoundError(f"Folder not found: {folder}")
    return sorted([path for path in folder.iterdir() if is_image_file(path)])


def stem_without_mask_suffix(name: str) -> str:
    stem = Path(name).stem
    stem = re.sub(r"_mask(_\d+)?$", "", stem)
    return stem


def relative_stem(path: Path) -> str:
    return path.stem


__all__ = [
    "IMAGE_EXTENSIONS",
    "is_image_file",
    "list_image_files",
    "stem_without_mask_suffix",
    "relative_stem",
]