dataload.py

from os import listdir
from os.path import join
import torch.utils.data as data
from torchvision.transforms import Compose, ToTensor, Resize,RandomHorizontalFlip,RandomVerticalFlip, Normalize
from PIL import Image


def is_image_file(filename):
    return any(filename.endswith(extension) for extension in [".png", ".jpg", ".jpeg"])


def load_img(filepath):
    img = Image.open(filepath).convert('YCbCr')
    y, _, _ = img.split()
    return y
# def load_img(filepath):
#     return Image.open(filepath).convert('RGB')


class DatasetFromFolder(data.Dataset):
    def __init__(self, image_dir):
        super(DatasetFromFolder, self).__init__()

        # 遍历数据集，获取所有图片的路径
        in_dir = join(image_dir, 'LR')
        self.dataset = []
        for filename in listdir(in_dir):
            input_file = join(in_dir, filename)
            #target_file = join(image_dir, 'H', filename)
            hr_filename = filename.replace('x8', '') if 'x8' in filename else filename
            target_file = join(image_dir, 'HR', hr_filename)

            if all([is_image_file(input_file), is_image_file(target_file)]):
                self.dataset.append((input_file, target_file))

        self.input_transform = Compose([
            Resize((32, 32)),
            # RandomHorizontalFlip(),
            # RandomVerticalFlip(),
            ToTensor(),
            # Normalize(mean=[0.449, 0.438, 0.404], std=[1.0, 1.0, 1.0]),
        ])
        self.target_transform = Compose([
            Resize((256, 256)),
            # RandomHorizontalFlip(),
            # RandomVerticalFlip(),
            ToTensor(),
            # Normalize(mean=[0.449, 0.438, 0.404], std=[1.0, 1.0, 1.0]),
        ])

    def __getitem__(self, index):
        input_image_path, target_image_path = self.dataset[index]
        input_image = load_img(input_image_path)
        target = load_img(target_image_path)

        if self.input_transform:
            input_image = self.input_transform(input_image)
        if self.target_transform:
            target = self.target_transform(target)

        return input_image, target


    def __len__(self):
        return len(self.dataset)