prepare_test_images.py

from __future__ import annotations

import argparse
import shutil
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List

from evaluation_utils import (
    KNOWN_SUPPORT_IMAGENET_MAP,
    UNKNOWN_EVAL_IMAGENET_MAP,
    build_stream_manifest,
    write_csv,
    write_json,
)


def _image_files(directory: Path) -> List[Path]:
    return sorted([path for path in directory.iterdir() if path.is_file() and path.suffix.lower() in {".jpg", ".jpeg", ".png"}])


def _support_counts(support_dir: Path) -> Dict[str, int]:
    counts: Dict[str, int] = {}
    for class_dir in sorted([path for path in support_dir.iterdir() if path.is_dir()]):
        counts[class_dir.name] = len(_image_files(class_dir))
    return counts


def _copy_subset(
    source_dir: Path,
    target_dir: Path,
    count: int,
    offset: int = 0,
) -> List[Path]:
    files = _image_files(source_dir)
    selected = files[offset:offset + count]
    if len(selected) < count:
        raise ValueError(
            f"Need {count} images from {source_dir}, but only found {len(selected)} after offset {offset}"
        )

    target_dir.mkdir(parents=True, exist_ok=True)
    copied: List[Path] = []
    for index, src in enumerate(selected, start=1):
        dst = target_dir / f"{index:03d}{src.suffix.lower()}"
        shutil.copy2(src, dst)
        copied.append(dst)
    return copied


def prepare_test_images(
    imagenet_root: str,
    support_dir: str = "support",
    output_dir: str = "test_images",
    known_per_class: int = 4,
    unknown_per_class: int = 4,
    known_offset: int | None = None,
    unknown_offset: int = 0,
) -> Dict[str, Any]:
    imagenet_path = Path(imagenet_root).expanduser().resolve()
    support_path = Path(support_dir).resolve()
    output_path = Path(output_dir).resolve()

    if not imagenet_path.is_dir():
        raise FileNotFoundError(f"ImageNet root not found: {imagenet_path}")
    if not support_path.is_dir():
        raise FileNotFoundError(f"Support folder not found: {support_path}")

    backup_path = None
    if output_path.exists():
        backup_path = output_path.parent / f"{output_path.name}_backup_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        shutil.move(str(output_path), str(backup_path))

    known_root = output_path / "known"
    unknown_root = output_path / "unknown"
    known_root.mkdir(parents=True, exist_ok=True)
    unknown_root.mkdir(parents=True, exist_ok=True)

    support_counts = _support_counts(support_path)
    manifest_items: List[Dict[str, Any]] = []
    copied_summary: Dict[str, Dict[str, int]] = {"known": {}, "unknown": {}}

    live_support_labels = sorted(support_counts.keys())
    missing_mappings = [label for label in live_support_labels if label not in KNOWN_SUPPORT_IMAGENET_MAP]
    if missing_mappings:
        raise ValueError(f"Missing ImageNet source mappings for support classes: {missing_mappings}")

    for label in live_support_labels:
        source_name = KNOWN_SUPPORT_IMAGENET_MAP[label]
        source_dir = imagenet_path / source_name
        if not source_dir.is_dir():
            raise FileNotFoundError(f"Known source directory not found for {label}: {source_dir}")
        class_support_count = int(support_counts.get(label, 0))
        if class_support_count <= 0:
            raise ValueError(f"Support class '{label}' is missing or empty in {support_path}")

        copied = _copy_subset(
            source_dir=source_dir,
            target_dir=known_root / label,
            count=known_per_class,
            offset=class_support_count if known_offset is None else int(known_offset),
        )
        copied_summary["known"][label] = len(copied)
        for idx, copied_path in enumerate(copied):
            manifest_items.append({
                "image_path": str(copied_path),
                "eval_split": "known",
                "expected_label": label,
                "source_label": label,
                "source_imagenet_dir": source_name,
                "round_index": idx,
                "class_index": idx,
            })

    for label, source_name in sorted(UNKNOWN_EVAL_IMAGENET_MAP.items()):
        source_dir = imagenet_path / source_name
        if not source_dir.is_dir():
            raise FileNotFoundError(f"Unknown source directory not found for {label}: {source_dir}")
        copied = _copy_subset(
            source_dir=source_dir,
            target_dir=unknown_root / label,
            count=unknown_per_class,
            offset=int(unknown_offset),
        )
        copied_summary["unknown"][label] = len(copied)
        for idx, copied_path in enumerate(copied):
            manifest_items.append({
                "image_path": str(copied_path),
                "eval_split": "unknown",
                "expected_label": "unknown",
                "source_label": label,
                "source_imagenet_dir": source_name,
                "round_index": idx,
                "class_index": idx,
            })

    ordered_manifest = build_stream_manifest(manifest_items)
    manifest_path = output_path / "manifest.json"
    summary_path = output_path / "summary.json"
    csv_path = output_path / "manifest.csv"

    write_json(manifest_path, ordered_manifest)
    write_csv(csv_path, ordered_manifest)
    write_json(summary_path, {
        "imagenet_root": str(imagenet_path),
        "support_dir": str(support_path),
        "output_dir": str(output_path),
        "backup_dir": str(backup_path) if backup_path else None,
        "known_per_class": known_per_class,
        "unknown_per_class": unknown_per_class,
        "known_offset": "support_count" if known_offset is None else int(known_offset),
        "unknown_offset": int(unknown_offset),
        "known_labels": live_support_labels,
        "support_counts": support_counts,
        "copied": copied_summary,
        "num_manifest_items": len(ordered_manifest),
    })

    return {
        "output_dir": str(output_path),
        "backup_dir": str(backup_path) if backup_path else None,
        "manifest_path": str(manifest_path),
        "summary_path": str(summary_path),
        "num_items": len(ordered_manifest),
        "copied": copied_summary,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="Create a DFOD evaluation test_images folder from ImageNet.")
    parser.add_argument("--imagenet-root", default="/Users/Abhinn/Downloads/val")
    parser.add_argument("--support-dir", default="support")
    parser.add_argument("--output-dir", default="test_images")
    parser.add_argument("--known-per-class", type=int, default=4)
    parser.add_argument("--unknown-per-class", type=int, default=4)
    parser.add_argument("--known-offset", type=int, default=None)
    parser.add_argument("--unknown-offset", type=int, default=0)
    args = parser.parse_args()

    result = prepare_test_images(
        imagenet_root=args.imagenet_root,
        support_dir=args.support_dir,
        output_dir=args.output_dir,
        known_per_class=args.known_per_class,
        unknown_per_class=args.unknown_per_class,
        known_offset=args.known_offset,
        unknown_offset=args.unknown_offset,
    )
    print(f"Prepared test images in {result['output_dir']}")
    print(f"Manifest: {result['manifest_path']}")
    print(f"Total evaluation images: {result['num_items']}")
    if result["backup_dir"]:
        print(f"Previous test_images backed up to {result['backup_dir']}")


if __name__ == "__main__":
    main()