AI-KSK-ComfyUI-AnimaForge-Windows/dataset_utils.py

from __future__ import annotations

import os
import re
from pathlib import Path
from typing import Any, Dict, Iterable, List, Sequence, Tuple

try:
    from natsort import natsorted
except Exception:  # pragma: no cover - optional dependency fallback
    natsorted = sorted


DEFAULT_IMAGE_EXTENSIONS = ".png,.jpg,.jpeg,.webp,.bmp"
PLUGIN_DIR = Path(__file__).resolve().parent
WORKSPACE_DIR = PLUGIN_DIR.parent
LEGACY_WORKSPACE_PREFIX = "/workspace"


def resolve_container_path(value: str) -> Path:
    raw = os.path.expandvars(str(value or "").strip())
    if raw == LEGACY_WORKSPACE_PREFIX:
        raw = str(WORKSPACE_DIR)
    elif raw.startswith(f"{LEGACY_WORKSPACE_PREFIX}/"):
        raw = str(WORKSPACE_DIR / raw.removeprefix(f"{LEGACY_WORKSPACE_PREFIX}/"))
    return Path(os.path.expanduser(raw)).resolve()


def parse_extensions(value: str | Sequence[str] | None) -> Tuple[str, ...]:
    if value is None:
        value = DEFAULT_IMAGE_EXTENSIONS
    if isinstance(value, str):
        raw_items = re.split(r"[,;\s]+", value)
    else:
        raw_items = list(value)

    extensions: List[str] = []
    for item in raw_items:
        item = str(item).strip().lower()
        if not item:
            continue
        if not item.startswith("."):
            item = f".{item}"
        extensions.append(item)
    return tuple(dict.fromkeys(extensions)) or tuple(DEFAULT_IMAGE_EXTENSIONS.split(","))


def is_image_file(path: Path, extensions: Iterable[str]) -> bool:
    return path.is_file() and path.suffix.lower() in set(extensions)


def list_images(folder: Path, extensions: Iterable[str], max_images: int = 0) -> List[str]:
    images = [path for path in folder.iterdir() if is_image_file(path, extensions)]
    images = natsorted(images, key=lambda p: p.name.lower())
    if max_images and max_images > 0:
        images = images[:max_images]
    return [str(path) for path in images]


def normalize_trigger_word(name: str) -> str:
    trigger = name.strip()
    trigger = re.sub(r"\s+", "_", trigger)
    trigger = re.sub(r"[^0-9A-Za-z_\-]+", "_", trigger)
    trigger = re.sub(r"_+", "_", trigger).strip("_")
    return trigger or "anima_subject"


def matches_folder_filter(folder_name: str, folder_filter: str) -> bool:
    folder_filter = (folder_filter or "").strip()
    if not folder_filter:
        return True

    items = [item.strip() for item in re.split(r"[,;\n]+", folder_filter) if item.strip()]
    if not items:
        return True

    for item in items:
        if item == folder_name:
            return True
        try:
            if re.search(item, folder_name):
                return True
        except re.error:
            if item.lower() in folder_name.lower():
                return True
    return False


def count_existing_captions(images: Sequence[str]) -> int:
    count = 0
    for image in images:
        caption_path = Path(image).with_suffix(".txt")
        if caption_path.exists() and caption_path.read_text(encoding="utf-8", errors="ignore").strip():
            count += 1
    return count


def scan_task_folders(
    dataset_root: str,
    image_extensions: str = DEFAULT_IMAGE_EXTENSIONS,
    folder_filter: str = "",
    include_root_images: bool = False,
    max_images_per_folder: int = 0,
) -> Tuple[List[Dict[str, Any]], str]:
    root = resolve_container_path(dataset_root)
    logs: List[str] = []
    if not dataset_root.strip():
        return [], "ERROR: dataset_root is empty."
    if not root.exists():
        return [], f"ERROR: dataset_root does not exist: {root}"
    if not root.is_dir():
        return [], f"ERROR: dataset_root is not a directory: {root}"

    extensions = parse_extensions(image_extensions)
    candidate_folders: List[Path] = []
    child_folders = [path for path in root.iterdir() if path.is_dir()]
    child_folders = natsorted(child_folders, key=lambda p: p.name.lower())

    for folder in child_folders:
        if matches_folder_filter(folder.name, folder_filter):
            candidate_folders.append(folder)

    if include_root_images:
        candidate_folders.insert(0, root)

    tasks: List[Dict[str, Any]] = []
    for folder in candidate_folders:
        images = list_images(folder, extensions, max_images_per_folder)
        if not images:
            logs.append(f"SKIP: no images found in {folder}")
            continue

        trigger_word = normalize_trigger_word(folder.name if folder != root else root.name)
        caption_count = count_existing_captions(images)
        task = {
            "dataset_root": str(root),
            "task_dir": str(folder),
            "folder_name": folder.name,
            "trigger_word": trigger_word,
            "images": images,
            "image_count": len(images),
            "caption_count": caption_count,
        }
        tasks.append(task)
        logs.append(
            f"OK: {folder.name} -> trigger={trigger_word}, images={len(images)}, captions={caption_count}"
        )

    if not tasks:
        logs.append("ERROR: no valid LoRA task folders were found.")
    else:
        logs.insert(0, f"Found {len(tasks)} task folder(s) under {root}.")
    return tasks, "\n".join(logs)


def iter_task_images(task_list: Sequence[Dict[str, Any]]):
    for task in task_list or []:
        trigger_word = normalize_trigger_word(str(task.get("trigger_word") or task.get("folder_name") or ""))
        for image_path in task.get("images") or []:
            yield task, Path(image_path), trigger_word


def refresh_task_caption_counts(task_list: Sequence[Dict[str, Any]]) -> List[Dict[str, Any]]:
    refreshed: List[Dict[str, Any]] = []
    for task in task_list or []:
        task_copy = dict(task)
        images = list(task_copy.get("images") or [])
        task_copy["caption_count"] = count_existing_captions(images)
        refreshed.append(task_copy)
    return refreshed