create_data_lists.py

# standard imports
import os
import json
from typing import List

# third-party imports
from PIL import Image


def create_data_lists(
    train_folders: List[str], test_folders: List[str], min_size: int, output_folder: str
):
    """
    Create lists for images in the training set and each of the test sets.

    :param train_folders: folders containing the training images; these will be merged
    :param test_folders: folders containing the test images; each test folder will form its own test set
    :param min_size: minimum width and height of images to be considered
    :param output_folder: save data lists here
    """
    print("\nCreating data lists... this may take some time.\n")
    train_images = list()
    for d in train_folders:
        for i in os.listdir(d):
            img_path = os.path.join(d, i)
            img = Image.open(img_path, mode="r")
            if img.width >= min_size and img.height >= min_size:
                train_images.append(img_path)
    print("There are %d images in the training data.\n" % len(train_images))
    with open(os.path.join(output_folder, "train_images.json"), "w") as j:
        json.dump(train_images, j)

    for d in test_folders:
        test_images = list()
        test_name = d.split("/")[-1]
        for i in os.listdir(d):
            img_path = os.path.join(d, i)
            img = Image.open(img_path, mode="r")
            if img.width >= min_size and img.height >= min_size:
                test_images.append(img_path)
        print(
            "There are %d images in the %s test data.\n" % (len(test_images), test_name)
        )
        with open(
            os.path.join(output_folder, test_name + "_test_images.json"), "w"
        ) as j:
            json.dump(test_images, j)

    print(
        "JSONS containing lists of Train and Test images have been saved to %s\n"
        % output_folder
    )


if __name__ == "__main__":
    create_data_lists(
        train_folders=[
            "/mnt/d/WatchAndTellCuda/coco/images/train2017",
            "/mnt/d/WatchAndTellCuda/coco/images/val2017",
        ],
        test_folders=["SR/BSDS100", "SR/Set5", "SR/Set14"],
        min_size=100,
        output_folder="./",
    )