[pytorch] 😎 Custom ImageFolder Class!

😁 안녕하세요, 오늘은 vision 관련 모델 작성시 요긴하게 사용되는 ImageFolder Class 사용법을 간단히 알아보고,
😊 이를 활용하여 Custom Class도 만들어보도록 하겠습니다 :)

ImageFolder

😉 Dataset class의 일종으로서, Data의 경로만 주어지면 Dataset 객체를 간단히 만들 수 있는 클래스입니다.

https://pytorch.org/vision/stable/generated/torchvision.datasets.ImageFolder.html?highlight=imagefolder#torchvision.datasets.ImageFolder

🐶 그리고 저는, 지난번에 사용했던 Stanford Dog Dataset을 활용하도록 하겠습니다. :)

# ubuntu linux
wget http://vision.stanford.edu/aditya86/ImageNetDogs/images.tar
tar -xvf images.tar

위의 명령어를 수행하면, images라는 폴더 이하에 귀여운 강아지 사진 데이터를 얻을 수 있습니다.

😋 제가 좋아하는 골댕이 사진이 잘 다운로드 된 것을 확인하였습니다!
😃 이제는 간단히 ImageFolder class를 활용해보겠습니다. :)

from torchvision import transforms

size = 224
mean = (0.485, 0.456, 0.406)
std = (0.229, 0.224, 0.225)

dog_transform = transforms.Compose([
                transforms.RandomResizedCrop(
                    (size, size), scale=(0.5, 1.0)),  
                transforms.RandomHorizontalFlip(), 
                transforms.ToTensor(),  # 텐서로 변환
                transforms.Normalize(mean, std)  # 표준화
            ])

dog_dataset = torchvision.datasets.ImageFolder('Images/', transform=dog_transform)

dog_dataset.class_to_idx
>>>{'n02085620-Chihuahua': 0,
 'n02085782-Japanese_spaniel': 1,
 'n02085936-Maltese_dog': 2,
 ...

📏 dataset.class_to_index 값을 확인하여, Dataset 객체의 class와 index간 mapping 관계를 확인할 수 있습니다.

data_loader = torch.utils.data.DataLoader(dog_dataset,
                                          batch_size=4,
                                          shuffle=True,
                                          num_workers=2)
                                          
next(iter(data_loader))[0].shape, next(iter(data_loader))[1].shape
>>> (torch.Size([16, 3, 224, 224]), torch.Size([16]))

✨ 그리고, DataLoader class를 활용하여 정상적으로 데이터가 뽑혀나오는 것까지 확인 완료입니다 :)
😜 어때요, 참 쉽죠?

Why Custom Class?

😒 아니 근데 왜 Custom Class를 만드려고 하는거야? 지금도 잘 되는데? - 라고 생각하신다면!
😀 저도 그렇게 생각했습니다만, 위 Class 사용을 위해 하나의 맹점이 있었습니다!

"class_to_idx의 생성 기준" 이 바로 그것이었죠!

🤔 간단하고 강력한 ImageFolder Class는 참 좋지만, class_to_idx는 "alphabet 순서"에 따라서 index가 결정되고 있었습니다.
✔ 만약 apple / banana / cider 3개 label이라면, {"apple" : 0, "banana" : 1, "cider" : 2} 인 것입니다.
😃 물론, 알파벳 순서는 꽤나 보편적인 rule이지만, 실 업무에서는 알파벳 라벨 순서가 아닌 class_to_idx 기준으로도 모델 학습이 필요한 경우가 있었습니다.

새로운 class_to_idx 제공

😋 테스트를 위해 새로운 class_to_idx 객체를 만들었습니다 :)

import os

label_list = os.listdir('Images/')

custom_class_to_idx = {label : idx for idx, label in enumerate(label_list)}

# 알파벳 순서로 idx 지정이 되지 않은 dict
custom_class_to_idx
>>> {'n02111500-Great_Pyrenees': 0,
 'n02111889-Samoyed': 32,
 'n02112018-Pomeranian': 1,
 'n02112137-chow': 97,
 ...
 

😎 이 mapping 관계를 활용하여 Dataset Class를 만들 예정입니다 :)

ImageFolder & DatasetFolder Class의 Source 분석

https://pytorch.org/vision/stable/_modules/torchvision/datasets/folder.html#ImageFolder

✨ 위의 링크는 ImageFolder의 source code입니다.
😎 따로 ImageFolder Class로 작성된 내용보다는 DatasetFolder Class의 로직을 그대로 상속받아 사용하였네요!

class ImageFolder(DatasetFolder):
    def __init__(
        self,
        root: str,
        transform: Optional[Callable] = None,
        target_transform: Optional[Callable] = None,
        loader: Callable[[str], Any] = default_loader,
        is_valid_file: Optional[Callable[[str], bool]] = None,
    ):
        super().__init__(
            root,
            loader,
            IMG_EXTENSIONS if is_valid_file is None else None,
            transform=transform,
            target_transform=target_transform,
            is_valid_file=is_valid_file,
        )
        self.imgs = self.samples

😁 그렇다면 DatasetFolder Class의 source code도 한번 보도록 하죠.
🤣 더헛! VisionDataset Class를 또 상속받았네요! 이렇다면 저 Class까지 분석해야 하나? 싶지만
😎 그럴 필요까지는 없습니다. 저희에게 필요한건 어디까지나 'class_to_idx' attribute의 수정이니까요.

class DatasetFolder(VisionDataset):
    """
    주석 생략
    """

    def __init__(
        self,
        root: str,
        loader: Callable[[str], Any],
        extensions: Optional[Tuple[str, ...]] = None,
        transform: Optional[Callable] = None,
        target_transform: Optional[Callable] = None,
        is_valid_file: Optional[Callable[[str], bool]] = None,
    ) -> None:
        super().__init__(root, transform=transform, target_transform=target_transform)
        classes, class_to_idx = self.find_classes(self.root)
        samples = self.make_dataset(self.root, class_to_idx, extensions, is_valid_file)

        self.loader = loader
        self.extensions = extensions

        self.classes = classes
        self.class_to_idx = class_to_idx
        self.samples = samples
        self.targets = [s[1] for s in samples]

    def find_classes(self, directory: str) -> Tuple[List[str], Dict[str, int]]:
        """Find the class folders in a dataset structured as follows::

            directory/
            ├── class_x
            │   ├── xxx.ext
            │   ├── xxy.ext
            │   └── ...
            │       └── xxz.ext
            └── class_y
                ├── 123.ext
                ├── nsdf3.ext
                └── ...
                └── asd932_.ext

        This method can be overridden to only consider
        a subset of classes, or to adapt to a different dataset directory structure.

        Args:
            directory(str): Root directory path, corresponding to ``self.root``

        Raises:
            FileNotFoundError: If ``dir`` has no class folders.

        Returns:
            (Tuple[List[str], Dict[str, int]]): List of all classes and dictionary mapping each class to an index.
        """
        return find_classes(directory)       

✨ 아하, self.class_to_idx는 self.find_classes 메서드를 통해 결정됩니다.
😋 결국 find_classes 메서드만 수정을 하면 원하는 바를 이룰 수 있을 것 같습니다!

😊 그리고, find_classes 메서드에 사용되는 find_classes 함수의 내용은 아래와 같습니다 :)

def find_classes(directory: str) -> Tuple[List[str], Dict[str, int]]:
    """Finds the class folders in a dataset.

    See :class:`DatasetFolder` for details.
    """
    classes = sorted(entry.name for entry in os.scandir(directory) if entry.is_dir())
    if not classes:
        raise FileNotFoundError(f"Couldn't find any class folder in {directory}.")

    class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}
    return classes, class_to_idx

👌 OK! Dataset Class에 사용될 classes와 class_to_idx를 return합니다.
(classes는 idx 순서대로 class가 나열된 list입니다.)

Custom ImageFolder & DatasetFolder Class 작성!

# from https://pytorch.org/vision/0.11/_modules/torchvision/datasets/folder.html
################################################################################
################################################################################
# copied from folder.py

from typing import Any, Callable, cast, Dict, List, Optional, Tuple
from typing import Union
from PIL import Image
from torchvision.datasets import VisionDataset, DatasetFolder

def has_file_allowed_extension(filename: str, extensions: Tuple[str, ...]) -> bool:
    """Checks if a file is an allowed extension.

    Args:
        filename (string): path to a file
        extensions (tuple of strings): extensions to consider (lowercase)

    Returns:
        bool: True if the filename ends with one of given extensions
    """
    return filename.lower().endswith(extensions)

IMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp')


def pil_loader(path: str) -> Image.Image:
    # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
    with open(path, 'rb') as f:
        img = Image.open(f)
        return img.convert('RGB')


# TODO: specify the return type
def accimage_loader(path: str) -> Any:
    import accimage
    try:
        return accimage.Image(path)
    except IOError:
        # Potentially a decoding problem, fall back to PIL.Image
        return pil_loader(path)


def default_loader(path: str) -> Any:
    from torchvision import get_image_backend
    if get_image_backend() == 'accimage':
        return accimage_loader(path)
    else:
        return pil_loader(path)

def make_dataset(
    directory: str,
    class_to_idx: Optional[Dict[str, int]] = None,
    extensions: Optional[Tuple[str, ...]] = None,
    is_valid_file: Optional[Callable[[str], bool]] = None,
) -> List[Tuple[str, int]]:
    """Generates a list of samples of a form (path_to_sample, class).

    See :class:`DatasetFolder` for details.

    Note: The class_to_idx parameter is here optional and will use the logic of the ``find_classes`` function
    by default.
    """
    directory = os.path.expanduser(directory)

    if class_to_idx is None:
        _, class_to_idx = find_classes(directory)
    elif not class_to_idx:
        raise ValueError("'class_to_index' must have at least one entry to collect any samples.")

    both_none = extensions is None and is_valid_file is None
    both_something = extensions is not None and is_valid_file is not None
    if both_none or both_something:
        raise ValueError("Both extensions and is_valid_file cannot be None or not None at the same time")

    if extensions is not None:

        def is_valid_file(x: str) -> bool:
            return has_file_allowed_extension(x, cast(Tuple[str, ...], extensions))

    is_valid_file = cast(Callable[[str], bool], is_valid_file)

    instances = []
    available_classes = set()
    for target_class in sorted(class_to_idx.keys()):
        class_index = class_to_idx[target_class]
        target_dir = os.path.join(directory, target_class)
        if not os.path.isdir(target_dir):
            continue
        for root, _, fnames in sorted(os.walk(target_dir, followlinks=True)):
            for fname in sorted(fnames):
                if is_valid_file(fname):
                    path = os.path.join(root, fname)
                    item = path, class_index
                    instances.append(item)

                    if target_class not in available_classes:
                        available_classes.add(target_class)

    empty_classes = set(class_to_idx.keys()) - available_classes
    if empty_classes:
        msg = f"Found no valid file for the classes {', '.join(sorted(empty_classes))}. "
        if extensions is not None:
            msg += f"Supported extensions are: {', '.join(extensions)}"
        raise FileNotFoundError(msg)

    return instances

# copied from folder.py END

################################################################################
################################################################################


class CustomDatasetFolder(VisionDataset):
    def __init__(
        self,
        root: str,
        loader: Callable[[str], Any],
        class_list: List[str],
        extensions: Optional[Tuple[str, ...]] = None,
        transform: Optional[Callable] = None,
        target_transform: Optional[Callable] = None,
        is_valid_file: Optional[Callable[[str], bool]] = None,
    ) -> None:
        super().__init__(root, transform=transform, target_transform=target_transform)
        classes, class_to_idx = self.find_classes(class_list)
        samples = self.make_dataset(self.root, class_to_idx, extensions, is_valid_file)

        self.loader = loader
        self.extensions = extensions

        self.classes = classes
        self.class_to_idx = class_to_idx
        self.samples = samples
        self.targets = [s[1] for s in samples]

    @staticmethod
    def make_dataset(
        directory: str,
        class_to_idx: Dict[str, int],
        extensions: Optional[Tuple[str, ...]] = None,
        is_valid_file: Optional[Callable[[str], bool]] = None,
    ) -> List[Tuple[str, int]]:

        if class_to_idx is None:
            raise ValueError("The class_to_idx parameter cannot be None.")
        return make_dataset(directory, class_to_idx, extensions=extensions, is_valid_file=is_valid_file)

    def find_classes(self, class_list: List[str]) -> Tuple[List[str], Dict[str, int]]:
        return class_list, {label : idx for idx, label in enumerate(class_list)}

    def __getitem__(self, index: int) -> Tuple[Any, Any]:

        path, target = self.samples[index]
        sample = self.loader(path)
        if self.transform is not None:
            sample = self.transform(sample)
        if self.target_transform is not None:
            target = self.target_transform(target)

        return sample, target

    def __len__(self) -> int:
        return len(self.samples)


class CustomImageFolder(CustomDatasetFolder):
    def __init__(
        self,
        root: str,
        class_list: List[str],
        transform: Optional[Callable] = None,
        target_transform: Optional[Callable] = None,
        loader: Callable[[str], Any] = default_loader,
        is_valid_file: Optional[Callable[[str], bool]] = None,
    ):
        super().__init__(
        root,
        loader,
        class_list,
        IMG_EXTENSIONS if is_valid_file is None else None,
        transform=transform,
        target_transform=target_transform,
        is_valid_file=is_valid_file,
    )
        self.imgs = self.samples

🤣 헥헥... 좀 작성하다보니 길어졌네요 ;<

custom_dog_dataset = CustomImageFolder('Images/', label_list, transform=dog_transform)

custom_dog_dataset.class_to_idx
>>> {'n02111500-Great_Pyrenees': 0,
 'n02111889-Samoyed': 32,
 'n02112018-Pomeranian': 1,
 'n02112137-chow': 97,
 ...

✌ 이로써 제가 원하는 class_to_idx dict를 기준으로 만들어진 Dataset class가 완성되었습니다!

마치며

😂 오늘 작성한 글은, 생각보다 조금 아쉬운 면이 있습니다.
🤦‍♂️ 단순히 상속과 Overriding을 활용하면 간단히 Custom class를 만들수 있지 않을까? 하였지만..! 컨셉은 간단했으나, 코드는 상당히 길어져버렸네요.
🤔 부족했던 부분도 있을 수 있을것 같아, 코드를 좀더 분석해보고, 간단 명료하게 활용할 수 있는 방안이 있는지 2차 검토가 필요해 보입니다.
😋 뭐... 그래도 원하는 결과는 나왔으니, 어쨌든 된거 아닐까요?! (하.하.하.)
😘 읽어주셔서 감사드리며, 또 뵙도록 하겠습니다!