def test_sr_folder_gt_dataset(self):
# setup
sr_pipeline = [
dict(type='LoadImageFromFile', io_backend='disk', key='gt'),
dict(type='ImageToTensor', keys=['gt'])
]
target_keys = ['gt_path', 'gt']
gt_folder = self.data_prefix / 'gt'
filename_tmpl = '{}_x4'
# input path is Path object
sr_folder_dataset = SRFolderGTDataset(
gt_folder=gt_folder,
pipeline=sr_pipeline,
scale=4,
filename_tmpl=filename_tmpl)
data_infos = sr_folder_dataset.data_infos
assert data_infos == [dict(gt_path=str(gt_folder / 'baboon.png'))]
result = sr_folder_dataset[0]
assert (len(sr_folder_dataset) == 1)
assert assert_dict_has_keys(result, target_keys)
# input path is str
sr_folder_dataset = SRFolderGTDataset(
gt_folder=str(gt_folder),
pipeline=sr_pipeline,
scale=4,
filename_tmpl=filename_tmpl)
data_infos = sr_folder_dataset.data_infos
assert data_infos == [dict(gt_path=str(gt_folder / 'baboon.png'))]
result = sr_folder_dataset[0]
assert (len(sr_folder_dataset) == 1)
assert assert_dict_has_keys(result, target_keys)
def test_base_vfi_dataset(self):
dataset = BaseVFIDataset(self.pipeline, self.folder, self.ann_file)
dataset.__init__(self.pipeline, self.folder, self.ann_file)
dataset.load_annotations()
assert dataset.folder == self.folder
assert dataset.ann_file == self.ann_file
setattr(dataset, 'data_infos', [
dict(inputs_path=[
'tests/data/vimeo90k/00001/0266/im1.png',
'tests/data/vimeo90k/00001/0266/im3.png'
],
target_path='tests/data/vimeo90k/00001/0266/im2.png',
key='00001/0266')
])
data = dataset.__getitem__(0)
assert_dict_has_keys(data, ['folder', 'ann_file'])
results = [dict(eval_result=dict(psnr=1.1, ssim=0.3))]
eval_result = dataset.evaluate(results)
assert_dict_has_keys(eval_result, ['psnr', 'ssim'])
with pytest.raises(TypeError):
dataset.evaluate(results[0])
with pytest.raises(AssertionError):
dataset.evaluate(results + results)
def test_sr_landmark_dataset(self):
# setup
sr_pipeline = [
dict(type='LoadImageFromFile',
io_backend='disk',
key='gt',
flag='color',
channel_order='rgb',
backend='cv2')
]
target_keys = ['gt_path', 'bbox', 'shape', 'landmark']
gt_folder = self.data_prefix / 'face'
ann_file = self.data_prefix / 'facemark_ann.npy'
# input path is Path object
sr_landmark_dataset = SRFacialLandmarkDataset(gt_folder=gt_folder,
ann_file=ann_file,
pipeline=sr_pipeline,
scale=4)
data_infos = sr_landmark_dataset.data_infos
assert len(data_infos) == 1
result = sr_landmark_dataset[0]
assert len(sr_landmark_dataset) == 1
assert assert_dict_has_keys(result, target_keys)
# input path is str
sr_landmark_dataset = SRFacialLandmarkDataset(gt_folder=str(gt_folder),
ann_file=str(ann_file),
pipeline=sr_pipeline,
scale=4)
data_infos = sr_landmark_dataset.data_infos
assert len(data_infos) == 1
result = sr_landmark_dataset[0]
assert len(sr_landmark_dataset) == 1
assert assert_dict_has_keys(result, target_keys)
def test_vfi_vimeo90k_dataset(self):
dataset_cfg = dict(type='VFIVimeo90KDataset',
folder=self.folder,
ann_file=self.ann_file,
pipeline=self.pipeline)
dataset = build_dataset(dataset_cfg)
data_infos = dataset.data_infos[0]
assert_dict_has_keys(data_infos, ['inputs_path', 'target_path', 'key'])
def test_sr_annotation_dataset(self):
# setup
anno_file_path = self.data_prefix / 'train.txt'
sr_pipeline = [
dict(type='LoadImageFromFile', io_backend='disk', key='lq'),
dict(type='LoadImageFromFile', io_backend='disk', key='gt'),
dict(type='PairedRandomCrop', gt_patch_size=128),
dict(type='ImageToTensor', keys=['lq', 'gt'])
]
target_keys = [
'lq_path', 'gt_path', 'scale', 'lq', 'lq_ori_shape', 'gt',
'gt_ori_shape'
]
# input path is Path object
sr_annotation_dataset = SRAnnotationDataset(
lq_folder=self.data_prefix / 'lq',
gt_folder=self.data_prefix / 'gt',
ann_file=anno_file_path,
pipeline=sr_pipeline,
scale=4,
filename_tmpl='{}_x4')
data_infos = sr_annotation_dataset.data_infos
assert data_infos == [
dict(
lq_path=str(self.data_prefix / 'lq' / 'baboon_x4.png'),
gt_path=str(self.data_prefix / 'gt' / 'baboon.png'))
]
result = sr_annotation_dataset[0]
assert (len(sr_annotation_dataset) == 1)
assert assert_dict_has_keys(result, target_keys)
# input path is str
sr_annotation_dataset = SRAnnotationDataset(
lq_folder=str(self.data_prefix / 'lq'),
gt_folder=str(self.data_prefix / 'gt'),
ann_file=str(anno_file_path),
pipeline=sr_pipeline,
scale=4,
filename_tmpl='{}_x4')
data_infos = sr_annotation_dataset.data_infos
assert data_infos == [
dict(
lq_path=str(self.data_prefix / 'lq' / 'baboon_x4.png'),
gt_path=str(self.data_prefix / 'gt' / 'baboon.png'))
]
result = sr_annotation_dataset[0]
assert (len(sr_annotation_dataset) == 1)
assert assert_dict_has_keys(result, target_keys)
def test_sr_lmdb_dataset(self):
# setup
lq_lmdb_folder = self.data_prefix / 'lq.lmdb'
sr_pipeline = [
dict(
type='LoadImageFromFile',
io_backend='lmdb',
key='lq',
db_path=lq_lmdb_folder),
dict(
type='LoadImageFromFile',
io_backend='lmdb',
key='gt',
db_path=lq_lmdb_folder),
dict(type='ImageToTensor', keys=['lq', 'gt'])
]
target_keys = [
'lq_path', 'gt_path', 'scale', 'lq', 'lq_ori_shape', 'gt',
'gt_ori_shape'
]
# input path is Path object
sr_lmdb_dataset = SRLmdbDataset(
lq_folder=lq_lmdb_folder,
gt_folder=lq_lmdb_folder, # fake gt_folder
pipeline=sr_pipeline,
scale=1)
data_infos = sr_lmdb_dataset.data_infos
assert data_infos == [dict(lq_path='baboon', gt_path='baboon')]
result = sr_lmdb_dataset[0]
assert (len(sr_lmdb_dataset) == 1)
assert assert_dict_has_keys(result, target_keys)
# input path is str
sr_lmdb_dataset = SRLmdbDataset(
lq_folder=str(lq_lmdb_folder),
gt_folder=(lq_lmdb_folder), # fake gt_folder
pipeline=sr_pipeline,
scale=1)
data_infos = sr_lmdb_dataset.data_infos
assert data_infos == [dict(lq_path='baboon', gt_path='baboon')]
result = sr_lmdb_dataset[0]
assert (len(sr_lmdb_dataset) == 1)
assert assert_dict_has_keys(result, target_keys)
with pytest.raises(ValueError):
sr_lmdb_dataset = SRLmdbDataset(
lq_folder=self.data_prefix, # normal folder
gt_folder=lq_lmdb_folder, # fake gt_folder
pipeline=sr_pipeline,
scale=1)
with pytest.raises(ValueError):
sr_lmdb_dataset = SRLmdbDataset(
lq_folder=str(self.data_prefix), # normal folder
gt_folder=lq_lmdb_folder, # fake gt_folder
pipeline=sr_pipeline,
scale=1)
with pytest.raises(ValueError):
sr_lmdb_dataset = SRLmdbDataset(
lq_folder=lq_lmdb_folder,
gt_folder=self.data_prefix, # normal folder
pipeline=sr_pipeline,
scale=1)
with pytest.raises(ValueError):
sr_lmdb_dataset = SRLmdbDataset(
lq_folder=lq_lmdb_folder,
gt_folder=str(self.data_prefix), # normal folder
pipeline=sr_pipeline,
scale=1)
def test_sr_folder_ref_dataset(self):
# setup
sr_pipeline = [
dict(type='LoadImageFromFile', io_backend='disk', key='lq'),
dict(type='LoadImageFromFile', io_backend='disk', key='gt'),
dict(type='LoadImageFromFile', io_backend='disk', key='ref'),
dict(type='PairedRandomCrop', gt_patch_size=128),
dict(type='ImageToTensor', keys=['lq', 'gt', 'ref'])
]
target_keys = [
'lq_path', 'gt_path', 'ref_path', 'scale', 'lq', 'gt', 'ref'
]
lq_folder = self.data_prefix / 'lq'
gt_folder = self.data_prefix / 'gt'
ref_folder = self.data_prefix / 'gt'
filename_tmpl = '{}_x4'
# input path is Path object
sr_folder_ref_dataset = SRFolderRefDataset(
lq_folder=lq_folder,
gt_folder=gt_folder,
ref_folder=str(ref_folder),
pipeline=sr_pipeline,
scale=4,
filename_tmpl_lq=filename_tmpl)
data_infos = sr_folder_ref_dataset.data_infos
assert data_infos == [
dict(
lq_path=str(lq_folder / 'baboon_x4.png'),
gt_path=str(gt_folder / 'baboon.png'),
ref_path=str(ref_folder / 'baboon.png'))
]
result = sr_folder_ref_dataset[0]
assert len(sr_folder_ref_dataset) == 1
assert assert_dict_has_keys(result, target_keys)
# input path is str
sr_folder_ref_dataset = SRFolderRefDataset(
lq_folder=str(lq_folder),
gt_folder=str(gt_folder),
ref_folder=str(ref_folder),
pipeline=sr_pipeline,
scale=4,
filename_tmpl_lq=filename_tmpl)
data_infos = sr_folder_ref_dataset.data_infos
assert data_infos == [
dict(
lq_path=str(lq_folder / 'baboon_x4.png'),
gt_path=str(gt_folder / 'baboon.png'),
ref_path=str(ref_folder / 'baboon.png'))
]
result = sr_folder_ref_dataset[0]
assert len(sr_folder_ref_dataset) == 1
assert assert_dict_has_keys(result, target_keys)
with pytest.raises(AssertionError):
sr_folder_ref_dataset = SRFolderRefDataset(
lq_folder=str(lq_folder),
gt_folder=str(self.data_prefix / 'image'), # fake gt_folder
ref_folder=str(ref_folder),
pipeline=sr_pipeline,
scale=4,
filename_tmpl_lq=filename_tmpl)
with pytest.raises(AssertionError):
sr_folder_ref_dataset = SRFolderRefDataset(
lq_folder=str(self.data_prefix / 'image'), # fake lq_folder
gt_folder=str(gt_folder),
ref_folder=str(ref_folder),
pipeline=sr_pipeline,
scale=4,
filename_tmpl_lq=filename_tmpl)
with pytest.raises(AssertionError):
sr_folder_ref_dataset = SRFolderRefDataset(
lq_folder=str(lq_folder),
gt_folder=str(self.data_prefix / 'bg'), # fake gt_folder
ref_folder=str(ref_folder),
pipeline=sr_pipeline,
scale=4,
filename_tmpl_lq=filename_tmpl)
with pytest.raises(AssertionError):
sr_folder_ref_dataset = SRFolderRefDataset(
lq_folder=str(self.data_prefix / 'bg'), # fake lq_folder
gt_folder=str(gt_folder),
ref_folder=str(ref_folder),
pipeline=sr_pipeline,
scale=4,
filename_tmpl_lq=filename_tmpl)
with pytest.raises(AssertionError):
sr_folder_ref_dataset = SRFolderRefDataset(
lq_folder=None,
gt_folder=None,
ref_folder=str(ref_folder),
pipeline=sr_pipeline,
scale=4,
filename_tmpl_lq=filename_tmpl)
def test_generation_paired_dataset(self):
# setup
img_norm_cfg = dict(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
pipeline = [
dict(type='LoadPairedImageFromFile',
io_backend='disk',
key='pair',
flag='color'),
dict(type='Resize',
keys=['img_a', 'img_b'],
scale=(286, 286),
interpolation='bicubic'),
dict(type='FixedCrop',
keys=['img_a', 'img_b'],
crop_size=(256, 256)),
dict(type='Flip', keys=['img_a', 'img_b'], direction='horizontal'),
dict(type='RescaleToZeroOne', keys=['img_a', 'img_b']),
dict(type='Normalize',
keys=['img_a', 'img_b'],
to_rgb=True,
**img_norm_cfg),
dict(type='ImageToTensor', keys=['img_a', 'img_b']),
dict(type='Collect',
keys=['img_a', 'img_b'],
meta_keys=['img_a_path', 'img_b_path'])
]
target_keys = ['img_a', 'img_b', 'meta']
target_meta_keys = ['img_a_path', 'img_b_path']
pair_folder = self.data_prefix / 'paired'
# input path is Path object
generation_paried_dataset = GenerationPairedDataset(
dataroot=pair_folder, pipeline=pipeline, test_mode=True)
data_infos = generation_paried_dataset.data_infos
assert data_infos == [
dict(pair_path=str(pair_folder / 'test' / '3.jpg'))
]
result = generation_paried_dataset[0]
assert (len(generation_paried_dataset) == 1)
assert assert_dict_has_keys(result, target_keys)
assert assert_dict_has_keys(result['meta'].data, target_meta_keys)
assert (result['meta'].data['img_a_path'] == str(pair_folder / 'test' /
'3.jpg'))
assert (result['meta'].data['img_b_path'] == str(pair_folder / 'test' /
'3.jpg'))
# input path is str
generation_paried_dataset = GenerationPairedDataset(
dataroot=str(pair_folder), pipeline=pipeline, test_mode=True)
data_infos = generation_paried_dataset.data_infos
assert data_infos == [
dict(pair_path=str(pair_folder / 'test' / '3.jpg'))
]
result = generation_paried_dataset[0]
assert (len(generation_paried_dataset) == 1)
assert assert_dict_has_keys(result, target_keys)
assert assert_dict_has_keys(result['meta'].data, target_meta_keys)
assert (result['meta'].data['img_a_path'] == str(pair_folder / 'test' /
'3.jpg'))
assert (result['meta'].data['img_b_path'] == str(pair_folder / 'test' /
'3.jpg'))
# test_mode = False
generation_paried_dataset = GenerationPairedDataset(
dataroot=str(pair_folder), pipeline=pipeline, test_mode=False)
data_infos = generation_paried_dataset.data_infos
assert data_infos == [
dict(pair_path=str(pair_folder / 'train' / '1.jpg')),
dict(pair_path=str(pair_folder / 'train' / '2.jpg'))
]
assert (len(generation_paried_dataset) == 2)
result = generation_paried_dataset[0]
assert assert_dict_has_keys(result, target_keys)
assert assert_dict_has_keys(result['meta'].data, target_meta_keys)
assert (result['meta'].data['img_a_path'] == str(pair_folder /
'train' / '1.jpg'))
assert (result['meta'].data['img_b_path'] == str(pair_folder /
'train' / '1.jpg'))
result = generation_paried_dataset[1]
assert assert_dict_has_keys(result, target_keys)
assert assert_dict_has_keys(result['meta'].data, target_meta_keys)
assert (result['meta'].data['img_a_path'] == str(pair_folder /
'train' / '2.jpg'))
assert (result['meta'].data['img_b_path'] == str(pair_folder /
'train' / '2.jpg'))
