使用ImageFolder加载转换后的原始图像之外的原始图像

root/dog/xxx.png
root/dog/xxy.png
root/dog/xxz.png

root/cat/123.png
root/cat/nsdf3.png
root/cat/asd932_.png

from torchvision.datasets.folder import default_loader
from torchvision.datasets.vision import VisionDataset

class CustomDataset(VisionDataset):
    def __init__(self, root, transform=None, target_transform=None, loader=default_loader):
        super(CustomDataset, self).__init__(root, transform=transform, target_transform=target_transform)
        self.loader = loader
        self.extensions = ['.jpg', '.jpeg', '.png']  # 支持的图像格式

        self.samples = self._load_samples()

    def _load_samples(self):
        samples = []
        classes = []
        class_to_idx = {}

        for root, _, filenames in self._walk(self.root):
            for filename in filenames:
                if self._has_valid_extension(filename):
                    path = self._join_path(root, filename)
                    classes.append(self._get_class(root))
                    samples.append((path, self._get_class(root)))

        classes = sorted(set(classes))
        class_to_idx = {cls: i for i, cls in enumerate(classes)}

        return samples

    def _walk(self, dir):
        for root, _, filenames in os.walk(dir):
            yield root, _, filenames

    def _has_valid_extension(self, filename):
        return any(filename.endswith(ext) for ext in self.extensions)

    def _join_path(self, root, filename):
        return os.path.join(root, filename)

    def _get_class(self, dir):
        return os.path.basename(dir)

    def __getitem__(self, index):
        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):
        return len(self.samples)