Python data.util方法代码示例

# 需要导入模块: from torch.utils import data [as 别名]
# 或者: from torch.utils.data import util [as 别名]
def __init__(self, opt):
        super(LRHRSeg_BG_Dataset, self).__init__()
        self.opt = opt
        self.paths_LR = None
        self.paths_HR = None
        self.paths_HR_bg = None  # HR images for background scenes
        self.LR_env = None  # environment for lmdb
        self.HR_env = None
        self.HR_env_bg = None

        # read image list from lmdb or image files
        self.HR_env, self.paths_HR = util.get_image_paths(opt['data_type'], opt['dataroot_GT'])
        self.LR_env, self.paths_LR = util.get_image_paths(opt['data_type'], opt['dataroot_LR'])
        self.HR_env_bg, self.paths_HR_bg = util.get_image_paths(opt['data_type'],
                                                                opt['dataroot_GT_bg'])

        assert self.paths_HR, 'Error: HR path is empty.'
        if self.paths_LR and self.paths_HR:
            assert len(self.paths_LR) == len(self.paths_HR), \
                'HR and LR datasets have different number of images - {}, {}.'.format(
                len(self.paths_LR), len(self.paths_HR))

        self.random_scale_list = [1, 0.9, 0.8, 0.7, 0.6, 0.5]
        self.ratio = 10  # 10 OST data samples and 1 DIV2K general data samples(background) 

# 需要导入模块: from torch.utils import data [as 别名]
# 或者: from torch.utils.data import util [as 别名]
def __getitem__(self, index):
        LR_path = None

        # get LR image
        LR_path = self.paths_LR[index]
        img_LR = util.read_img(self.LR_env, LR_path)
        H, W, C = img_LR.shape

        # change color space if necessary
        if self.opt['color']:
            img_LR = util.channel_convert(C, self.opt['color'], [img_LR])[0]

        # BGR to RGB, HWC to CHW, numpy to tensor
        if img_LR.shape[2] == 3:
            img_LR = img_LR[:, :, [2, 1, 0]]
        img_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()

        return {'LR': img_LR, 'LR_path': LR_path} 

# 需要导入模块: from torch.utils import data [as 别名]
# 或者: from torch.utils.data import util [as 别名]
def __init__(self, opt):
        super(LQGTDataset, self).__init__()
        self.opt = opt
        self.data_type = self.opt['data_type']
        self.paths_LQ, self.paths_GT = None, None
        self.sizes_LQ, self.sizes_GT = None, None
        self.LQ_env, self.GT_env = None, None  # environment for lmdb

        self.paths_GT, self.sizes_GT = util.get_image_paths(self.data_type, opt['dataroot_GT'])
        self.paths_LQ, self.sizes_LQ = util.get_image_paths(self.data_type, opt['dataroot_LQ'])
        assert self.paths_GT, 'Error: GT path is empty.'
        if self.paths_LQ and self.paths_GT:
            assert len(self.paths_LQ) == len(
                self.paths_GT
            ), 'GT and LQ datasets have different number of images - {}, {}.'.format(
                len(self.paths_LQ), len(self.paths_GT))
        self.random_scale_list = [1] 

# 需要导入模块: from torch.utils import data [as 别名]
# 或者: from torch.utils.data import util [as 别名]
def __init__(self, opt):
        super(LQDataset, self).__init__()
        self.opt = opt
        self.opt_P = opt
        self.LR_paths = None
        self.LR_sizes = None  # environment for lmdb
        self.LR_env = None
        self.LR_size = opt['LR_size']

        self.real_ker_path = '/mnt/yjchai/SR_data/Flickr2K/kermap.pt'
        self.real_ker_map_list = util.load_ker_map_list(self.real_ker_path)

        # read image list from lmdb or image files
        if opt['data_type'] == 'lmdb':
            self.LR_paths, self.LR_sizes = util.get_image_paths(opt['data_type'], opt['dataroot_LQ'])
        elif opt['data_type'] == 'img':
            self.LR_paths = util.get_image_paths(opt['data_type'], opt['dataroot_LQ']) #LR_list
        else:
            print('Error: data_type is not matched in Dataset')
        assert self.LR_paths, 'Error: LR paths are empty.' 

# 需要导入模块: from torch.utils import data [as 别名]
# 或者: from torch.utils.data import util [as 别名]
def __init__(self, opt):
        super(LQGTKerDataset, self).__init__()
        self.opt = opt
        self.opt_F = opt
        self.opt_P = opt
        self.opt_C = opt
        self.LR_paths, self.GT_paths = None, None
        self.LR_env, self.GT_env = None, None  # environment for lmdb
        self.LR_size, self.GT_size = opt['LR_size'], opt['GT_size']

        # read image list from lmdb or image files
        if opt['data_type'] == 'lmdb':
            self.LR_paths, self.LR_sizes = util.get_image_paths(opt['data_type'], opt['dataroot_LQ'])
            self.GT_paths, self.GT_sizes = util.get_image_paths(opt['data_type'], opt['dataroot_GT'])
        elif opt['data_type'] == 'img':
            self.LR_paths = util.get_image_paths(opt['data_type'], opt['dataroot_LQ']) # LR list
            self.GT_paths = util.get_image_paths(opt['data_type'], opt['dataroot_GT']) # GT list
        else:
            print('Error: data_type is not matched in Dataset')
        assert self.GT_paths, 'Error: GT paths are empty.'
        if self.LR_paths and self.GT_paths:
            assert len(self.LR_paths) == len(self.GT_paths), 'GT and LR datasets have different number of images - {}, {}.'.format(len(self.LR_paths), len(self.GT_paths))
        self.random_scale_list = [1] 

# 需要导入模块: from torch.utils import data [as 别名]
# 或者: from torch.utils.data import util [as 别名]
def __getitem__(self, index):
        if self.data_type == 'lmdb' and self.LQ_env is None:
            self._init_lmdb()
        LQ_path = None

        # get LQ image
        LQ_path = self.LQ_path[index]
        resolution = [int(s) for s in self.sizes_LQ[index].split('_')
                      ] if self.data_type == 'lmdb' else None
        img_LQ = util.read_img(self.LQ_env, LQ_path, resolution)
        H, W, C = img_LQ.shape

        if self.opt['color']:  # change color space if necessary
            img_LQ = util.channel_convert(C, self.opt['color'], [img_LQ])[0]

        # BGR to RGB, HWC to CHW, numpy to tensor
        if img_LQ.shape[2] == 3:
            img_LQ = img_LQ[:, :, [2, 1, 0]]
        img_LQ = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LQ, (2, 0, 1)))).float()

        return {'LQ': img_LQ, 'LQ_path': LQ_path} 

# 需要导入模块: from torch.utils import data [as 别名]
# 或者: from torch.utils.data import util [as 别名]
def __getitem__(self, index):
        # path_LQ = self.data_info['path_LQ'][index]
        # path_GT = self.data_info['path_GT'][index]
        folder = self.data_info['folder'][index]
        idx, max_idx = self.data_info['idx'][index].split('/')
        idx, max_idx = int(idx), int(max_idx)
        border = self.data_info['border'][index]

        if self.cache_data:
            select_idx = util.index_generation(idx, max_idx, self.opt['N_frames'],
                                               padding=self.opt['padding'])
            imgs_LQ = self.imgs_LQ[folder].index_select(0, torch.LongTensor(select_idx))
            img_GT = self.imgs_GT[folder][idx]
        else:
            pass  # TODO

        return {
            'LQs': imgs_LQ,
            'GT': img_GT,
            'folder': folder,
            'idx': self.data_info['idx'][index],
            'border': border
        } 

# 需要导入模块: from torch.utils import data [as 别名]
# 或者: from torch.utils.data import util [as 别名]
def __getitem__(self, index):
        if self.data_type == 'lmdb' and self.LQ_env is None:
            self._init_lmdb()
        LQ_path = None

        # get LQ image
        LQ_path = self.paths_LQ[index]
        resolution = [int(s) for s in self.sizes_LQ[index].split('_')
                      ] if self.data_type == 'lmdb' else None
        img_LQ = util.read_img(self.LQ_env, LQ_path, resolution)
        H, W, C = img_LQ.shape

        if self.opt['color']:  # change color space if necessary
            img_LQ = util.channel_convert(C, self.opt['color'], [img_LQ])[0]

        # BGR to RGB, HWC to CHW, numpy to tensor
        if img_LQ.shape[2] == 3:
            img_LQ = img_LQ[:, :, [2, 1, 0]]
        img_LQ = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LQ, (2, 0, 1)))).float()

        return {'LQ': img_LQ, 'LQ_path': LQ_path} 

# 需要导入模块: from torch.utils import data [as 别名]
# 或者: from torch.utils.data import util [as 别名]
def __init__(self, opt):
        super(LRDataset, self).__init__()
        self.opt = opt
        self.paths_LR = None
        self.LR_env = None  # environment for lmdb

        # read image list from lmdb or image files
        self.LR_env, self.paths_LR = util.get_image_paths(opt['data_type'], opt['dataroot_LR'])
        assert self.paths_LR, 'Error: LR paths are empty.' 

# 需要导入模块: from torch.utils import data [as 别名]
# 或者: from torch.utils.data import util [as 别名]
def __getitem__(self, index):
        if self.opt['data_type'] == 'lmdb':
            if self.LR_env is None:
                self._init_lmdb()

        LR_size = self.LR_size

        # get real kernel map
        real_ker_map = self.real_ker_map_list[index].float()

        # get LR image
        LR_path = self.LR_paths[index]
        if self.opt['data_type'] == 'lmdb':
            resolution = [int(s) for s in self.LR_sizes[index].split('_')]
        else:
            resolution = None
        img_LR = util.read_img(self.LR_env, LR_path, resolution)
        H, W, C = img_LR.shape

        if self.opt['phase'] == 'train':
            #randomly crop
            rnd_h = random.randint(0, max(0, H - LR_size))
            rnd_w = random.randint(0, max(0, W - LR_size))
            img_LR = img_LR[rnd_h:rnd_h + LR_size, rnd_w:rnd_w + LR_size, :]

            # augmentation - flip, rotate
            img_LR = util.augment(img_LR, self.opt['use_flip'], self.opt['use_rot'], self.opt['mode'])

        # change color space if necessary
        if self.opt['color']:
            img_LR = util.channel_convert(C, self.opt['color'], [img_LR])[0]

        # BGR to RGB, HWC to CHW, numpy to tensor
        if img_LR.shape[2] == 3:
            img_LR = img_LR[:, :, [2, 1, 0]]
        img_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()

        return {'LQ': img_LR, 'LQ_path': LR_path, 'real_ker': real_ker_map} 

# 需要导入模块: from torch.utils import data [as 别名]
# 或者: from torch.utils.data import util [as 别名]
def __init__(self, opt, ker_map_list, SR_img_list):
        super(SRKerDataset, self).__init__()
        self.opt = opt
        self.opt_C = opt

        self.LR_paths = None
        self.LR_sizes = None  # environment for lmdb
        self.LR_env = None
        self.LR_size = opt['LR_size']

        self.SR_env = None
        self.SR_img_list = SR_img_list
        self.SR_size = opt['GT_size']

        self.ker_map_list = ker_map_list
        self.real_ker_path = '/mnt/yjchai/SR_data/Flickr2K/kermap.pt'
        self.real_ker_map_list = util.load_ker_map_list(self.real_ker_path)

        # read image list from lmdb or image files
        #if opt['data_type'] == 'lmdb':
        #    self.LR_paths, self.LR_sizes = util.get_image_paths(opt['data_type'], opt['dataroot_LQ'])
        #elif opt['data_type'] == 'img':
        #    self.LR_paths = util.get_image_paths(opt['data_type'], opt['dataroot_LQ']) #LR_list
        #else:
        #    print('Error: data_type is not matched in Dataset')
        #assert self.LR_paths, 'Error: LR paths are empty.' 

# 需要导入模块: from torch.utils import data [as 别名]
# 或者: from torch.utils.data import util [as 别名]
def __getitem__(self, index):
        if self.opt['data_type'] == 'lmdb':
            if self.LR_env is None:
                self._init_lmdb()

        LR_size = self.LR_size

        # get LR image, kernel map
        LR_path = self.LR_paths[index]
        ker_map = self.ker_maps[index]
        if self.opt['data_type'] == 'lmdb':
            resolution = [int(s) for s in self.LR_sizes[index].split('_')]
        else:
            resolution = None
        img_LR = util.read_img(self.LR_env, LR_path, resolution)
        H, W, C = img_LR.shape

        if self.opt['phase'] == 'train':
            #randomly crop
            rnd_h = random.randint(0, max(0, H - LR_size))
            rnd_w = random.randint(0, max(0, W - LR_size))
            img_LR = img_LR[rnd_h:rnd_h + LR_size, rnd_w:rnd_w + LR_size, :]

            # augmentation - flip, rotate
            img_LR = util.augment(img_LR, self.opt['use_flip'], self.opt['use_rot'], self.opt['mode'])

        # change color space if necessary
        if self.opt['color']:
            img_LR = util.channel_convert(C, self.opt['color'], [img_LR])[0]

        # BGR to RGB, HWC to CHW, numpy to tensor
        if img_LR.shape[2] == 3:
            img_LR = img_LR[:, :, [2, 1, 0]]
        img_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()

        return {'LQ': img_LR, 'ker': ker_map, 'LQ_path': LR_path} 

# 需要导入模块: from torch.utils import data [as 别名]
# 或者: from torch.utils.data import util [as 别名]
def __init__(self, opt):
        super(LQDataset, self).__init__()
        self.opt = opt
        self.paths_LQ, self.paths_GT = None, None
        self.LQ_env = None  # environment for lmdb

        self.paths_LQ, self.sizes_LQ = util.get_image_paths(self.data_type, opt['dataroot_LQ'])
        assert self.paths_LQ, 'Error: LQ paths are empty.' 

# 需要导入模块: from torch.utils import data [as 别名]
# 或者: from torch.utils.data import util [as 别名]
def __init__(self, opt):
        super(VideoTestDataset, self).__init__()
        self.opt = opt
        self.cache_data = opt['cache_data']
        self.half_N_frames = opt['N_frames'] // 2
        self.GT_root, self.LQ_root = opt['dataroot_GT'], opt['dataroot_LQ']
        self.data_type = self.opt['data_type']
        self.data_info = {'path_LQ': [], 'path_GT': [], 'folder': [], 'idx': [], 'border': []}
        if self.data_type == 'lmdb':
            raise ValueError('No need to use LMDB during validation/test.')
        #### Generate data info and cache data
        self.imgs_LQ, self.imgs_GT = {}, {}
        if opt['name'].lower() in ['vid4', 'reds4']:
            subfolders_LQ = util.glob_file_list(self.LQ_root)
            subfolders_GT = util.glob_file_list(self.GT_root)
            for subfolder_LQ, subfolder_GT in zip(subfolders_LQ, subfolders_GT):
                subfolder_name = osp.basename(subfolder_GT)
                img_paths_LQ = util.glob_file_list(subfolder_LQ)
                img_paths_GT = util.glob_file_list(subfolder_GT)
                max_idx = len(img_paths_LQ)
                assert max_idx == len(
                    img_paths_GT), 'Different number of images in LQ and GT folders'
                self.data_info['path_LQ'].extend(img_paths_LQ)
                self.data_info['path_GT'].extend(img_paths_GT)
                self.data_info['folder'].extend([subfolder_name] * max_idx)
                for i in range(max_idx):
                    self.data_info['idx'].append('{}/{}'.format(i, max_idx))
                border_l = [0] * max_idx
                for i in range(self.half_N_frames):
                    border_l[i] = 1
                    border_l[max_idx - i - 1] = 1
                self.data_info['border'].extend(border_l)

                if self.cache_data:
                    self.imgs_LQ[subfolder_name] = util.read_img_seq(img_paths_LQ)
                    self.imgs_GT[subfolder_name] = util.read_img_seq(img_paths_GT)
        elif opt['name'].lower() in ['vimeo90k-test']:
            pass  # TODO
        else:
            raise ValueError(
                'Not support video test dataset. Support Vid4, REDS4 and Vimeo90k-Test.') 

# 需要导入模块: from torch.utils import data [as 别名]
# 或者: from torch.utils.data import util [as 别名]
def __init__(self, opt):
        super(REDSDataset, self).__init__()
        self.opt = opt
        # temporal augmentation
        self.interval_list = opt['interval_list']
        self.random_reverse = opt['random_reverse']
        logger.info('Temporal augmentation interval list: [{}], with random reverse is {}.'.format(
            ','.join(str(x) for x in opt['interval_list']), self.random_reverse))

        self.half_N_frames = opt['N_frames'] // 2
        self.GT_root, self.LQ_root = opt['dataroot_GT'], opt['dataroot_LQ']
        self.data_type = self.opt['data_type']
        self.LR_input = False if opt['GT_size'] == opt['LQ_size'] else True  # low resolution inputs
        #### directly load image keys
        if self.data_type == 'lmdb':
            self.paths_GT, _ = util.get_image_paths(self.data_type, opt['dataroot_GT'])
            logger.info('Using lmdb meta info for cache keys.')
        elif opt['cache_keys']:
            logger.info('Using cache keys: {}'.format(opt['cache_keys']))
            self.paths_GT = pickle.load(open(opt['cache_keys'], 'rb'))['keys']
        else:
            raise ValueError(
                'Need to create cache keys (meta_info.pkl) by running [create_lmdb.py]')

        # remove the REDS4 for testing
        self.paths_GT = [
            v for v in self.paths_GT if v.split('_')[0] not in ['000', '011', '015', '020']
        ]
        assert self.paths_GT, 'Error: GT path is empty.'

        if self.data_type == 'lmdb':
            self.GT_env, self.LQ_env = None, None
        elif self.data_type == 'mc':  # memcached
            self.mclient = None
        elif self.data_type == 'img':
            pass
        else:
            raise ValueError('Wrong data type: {}'.format(self.data_type))