전공 공부/파이썬 기초

make_dataset 함수

상솜공방 2023. 10. 2. 20:59
import os
from PIL import Image


IMG_EXTENSIONS = ['.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP', '.tif', '.TIF', '.tiff', '.TIFF',]

def is_image_file(filename):
    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
# 주어진 파일 이름이 이미지 파일인지 여부를 확인하는 함수입니다.
# 파일 이름이 IMG_EXTENSIONS에 정의된 확장자 중 하나로 끝나면 True를 반환하고, 그렇지 않으면 False를 반환합니다.

def make_dataset(dir, max_dataset_size=float("inf")):
# 이미지 파일들의 경로를 수집하여 리스트로 반환하는 함수입니다.
# dir: 이미지 파일을 찾을 디렉터리의 경로입니다.
# max_dataset_size: 반환할 최대 이미지 수입니다. 기본값은 무한대를 의미하며, 이를 통해 데이터셋의 크기를 제한할 수 있습니다.

    images = []
    assert os.path.isdir(dir), '%s is not a valid directory' % dir
    # 주어진 디렉터리 (dir)가 유효한 디렉터리인지 확인합니다.
    
    for root, _, fnames in sorted(os.walk(dir)):
        for fname in fnames:
            if is_image_file(fname):
                path = os.path.join(root, fname)
                images.append(path)
    #디렉터리를 순회하면서 이미지 파일인지 확인하고, 이미지 파일이 맞으면 리스트에 해당 파일 경로를 추가합니다.

    images = sorted(images)
    # 이미지 파일들의 경로를 사전 순서로 정렬합니다.
    return images[:min(max_dataset_size, len(images))]
    # 최대 데이터셋 크기(max_dataset_size)가 정해져 있으면 해당 크기까지만 이미지 파일 리스트를 자릅니다.
    # 최종적으로 이미지 파일들의 경로 리스트를 반환합니다.

디렉토리를 순회하며 해당하는 확장자의 파일을 모두 모아 정렬하여 리턴하는 함수.

출처: https://github.com/xiefan-guo/ctsdg