imagenet_parser.py

import shutil
from pathlib import Path

def flatten_imagenet(imagenet_root: str, output_dir: str, move=False):
    """
    Flatten imagenet/train/* and imagenet/val/* into a single directory.
    
    Parameters:
        imagenet_root (str): Path to imagenet directory containing train/ and val/.
        output_dir (str): Destination folder for all images.
        move (bool): If True, move files instead of copying.
    """
    root = Path(imagenet_root)
    train_dir = root / "train"
    val_dir = root / "val"

    out = Path(output_dir)
    out.mkdir(parents=True, exist_ok=True)

    def move_or_copy(src, dst):
        if move:
            shutil.move(src, dst)
        else:
            shutil.copy(src, dst)

    for split_dir in [train_dir, val_dir]:
        if not split_dir.exists():
            continue
        for class_dir in split_dir.iterdir():
            if not class_dir.is_dir():
                continue
            for img_path in class_dir.iterdir():
                if not img_path.is_file():
                    continue

                new_name = f"{class_dir.name}_{img_path.name}"
                dst = out / new_name
                move_or_copy(str(img_path), str(dst))

    print(f"All images moved into: {output_dir}")


if __name__ == "__main__":
    flatten_imagenet(
        "raw_data/imagenet",
        "raw_data/imagenet_parsed"
    )