dataset_sampler.py

import glob
import os
import random
import itertools
import shutil
from pathlib import Path


def scan_folders(sort_images=True):
    # Scan folder
    all_images = []
    for directory in os.listdir():
        if os.path.isdir(directory):
            if os.listdir(directory):
                image_paths = sorted(glob.glob(os.path.join(directory, "*.png")))
                selected_images = random.choices(image_paths, k=20)
                print(selected_images)
                all_images.append(selected_images)
    # Sort if you want to
    if sort_images:
        return list(itertools.chain(*all_images))
    else:
        return sorted(list(itertools.chain(*all_images)))


def copy_images(selected_images):
    # Create folder if not exists
    if not Path("sample_dataset").is_dir():
        Path("sample_dataset").mkdir(exist_ok=True)

    # Copy images
    for path in selected_images:
        root, relative = os.path.split(path)
        shutil.copyfile(path, os.path.join("sample_dataset", f"{root}_{relative}"))


if __name__ == "__main__":
    selected_images = scan_folders()
    copy_images(selected_images)