Spaces:
Runtime error
Runtime error
| import pytest | |
| import yaml | |
| from realesrgan.data.realesrgan_dataset import RealESRGANDataset | |
| from realesrgan.data.realesrgan_paired_dataset import RealESRGANPairedDataset | |
| def test_realesrgan_dataset(): | |
| with open('tests/data/test_realesrgan_dataset.yml', mode='r') as f: | |
| opt = yaml.load(f, Loader=yaml.FullLoader) | |
| dataset = RealESRGANDataset(opt) | |
| assert dataset.io_backend_opt['type'] == 'disk' # io backend | |
| assert len(dataset) == 2 # whether to read correct meta info | |
| assert dataset.kernel_list == [ | |
| 'iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso' | |
| ] # correct initialization the degradation configurations | |
| assert dataset.betag_range2 == [0.5, 4] | |
| # test __getitem__ | |
| result = dataset.__getitem__(0) | |
| # check returned keys | |
| expected_keys = ['gt', 'kernel1', 'kernel2', 'sinc_kernel', 'gt_path'] | |
| assert set(expected_keys).issubset(set(result.keys())) | |
| # check shape and contents | |
| assert result['gt'].shape == (3, 400, 400) | |
| assert result['kernel1'].shape == (21, 21) | |
| assert result['kernel2'].shape == (21, 21) | |
| assert result['sinc_kernel'].shape == (21, 21) | |
| assert result['gt_path'] == 'tests/data/gt/baboon.png' | |
| # ------------------ test lmdb backend -------------------- # | |
| opt['dataroot_gt'] = 'tests/data/gt.lmdb' | |
| opt['io_backend']['type'] = 'lmdb' | |
| dataset = RealESRGANDataset(opt) | |
| assert dataset.io_backend_opt['type'] == 'lmdb' # io backend | |
| assert len(dataset.paths) == 2 # whether to read correct meta info | |
| assert dataset.kernel_list == [ | |
| 'iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso' | |
| ] # correct initialization the degradation configurations | |
| assert dataset.betag_range2 == [0.5, 4] | |
| # test __getitem__ | |
| result = dataset.__getitem__(1) | |
| # check returned keys | |
| expected_keys = ['gt', 'kernel1', 'kernel2', 'sinc_kernel', 'gt_path'] | |
| assert set(expected_keys).issubset(set(result.keys())) | |
| # check shape and contents | |
| assert result['gt'].shape == (3, 400, 400) | |
| assert result['kernel1'].shape == (21, 21) | |
| assert result['kernel2'].shape == (21, 21) | |
| assert result['sinc_kernel'].shape == (21, 21) | |
| assert result['gt_path'] == 'comic' | |
| # ------------------ test with sinc_prob = 0 -------------------- # | |
| opt['dataroot_gt'] = 'tests/data/gt.lmdb' | |
| opt['io_backend']['type'] = 'lmdb' | |
| opt['sinc_prob'] = 0 | |
| opt['sinc_prob2'] = 0 | |
| opt['final_sinc_prob'] = 0 | |
| dataset = RealESRGANDataset(opt) | |
| result = dataset.__getitem__(0) | |
| # check returned keys | |
| expected_keys = ['gt', 'kernel1', 'kernel2', 'sinc_kernel', 'gt_path'] | |
| assert set(expected_keys).issubset(set(result.keys())) | |
| # check shape and contents | |
| assert result['gt'].shape == (3, 400, 400) | |
| assert result['kernel1'].shape == (21, 21) | |
| assert result['kernel2'].shape == (21, 21) | |
| assert result['sinc_kernel'].shape == (21, 21) | |
| assert result['gt_path'] == 'baboon' | |
| # ------------------ lmdb backend should have paths ends with lmdb -------------------- # | |
| with pytest.raises(ValueError): | |
| opt['dataroot_gt'] = 'tests/data/gt' | |
| opt['io_backend']['type'] = 'lmdb' | |
| dataset = RealESRGANDataset(opt) | |
| def test_realesrgan_paired_dataset(): | |
| with open('tests/data/test_realesrgan_paired_dataset.yml', mode='r') as f: | |
| opt = yaml.load(f, Loader=yaml.FullLoader) | |
| dataset = RealESRGANPairedDataset(opt) | |
| assert dataset.io_backend_opt['type'] == 'disk' # io backend | |
| assert len(dataset) == 2 # whether to read correct meta info | |
| # test __getitem__ | |
| result = dataset.__getitem__(0) | |
| # check returned keys | |
| expected_keys = ['gt', 'lq', 'gt_path', 'lq_path'] | |
| assert set(expected_keys).issubset(set(result.keys())) | |
| # check shape and contents | |
| assert result['gt'].shape == (3, 128, 128) | |
| assert result['lq'].shape == (3, 32, 32) | |
| assert result['gt_path'] == 'tests/data/gt/baboon.png' | |
| assert result['lq_path'] == 'tests/data/lq/baboon.png' | |
| # ------------------ test lmdb backend -------------------- # | |
| opt['dataroot_gt'] = 'tests/data/gt.lmdb' | |
| opt['dataroot_lq'] = 'tests/data/lq.lmdb' | |
| opt['io_backend']['type'] = 'lmdb' | |
| dataset = RealESRGANPairedDataset(opt) | |
| assert dataset.io_backend_opt['type'] == 'lmdb' # io backend | |
| assert len(dataset) == 2 # whether to read correct meta info | |
| # test __getitem__ | |
| result = dataset.__getitem__(1) | |
| # check returned keys | |
| expected_keys = ['gt', 'lq', 'gt_path', 'lq_path'] | |
| assert set(expected_keys).issubset(set(result.keys())) | |
| # check shape and contents | |
| assert result['gt'].shape == (3, 128, 128) | |
| assert result['lq'].shape == (3, 32, 32) | |
| assert result['gt_path'] == 'comic' | |
| assert result['lq_path'] == 'comic' | |
| # ------------------ test paired_paths_from_folder -------------------- # | |
| opt['dataroot_gt'] = 'tests/data/gt' | |
| opt['dataroot_lq'] = 'tests/data/lq' | |
| opt['io_backend'] = dict(type='disk') | |
| opt['meta_info'] = None | |
| dataset = RealESRGANPairedDataset(opt) | |
| assert dataset.io_backend_opt['type'] == 'disk' # io backend | |
| assert len(dataset) == 2 # whether to read correct meta info | |
| # test __getitem__ | |
| result = dataset.__getitem__(0) | |
| # check returned keys | |
| expected_keys = ['gt', 'lq', 'gt_path', 'lq_path'] | |
| assert set(expected_keys).issubset(set(result.keys())) | |
| # check shape and contents | |
| assert result['gt'].shape == (3, 128, 128) | |
| assert result['lq'].shape == (3, 32, 32) | |
| # ------------------ test normalization -------------------- # | |
| dataset.mean = [0.5, 0.5, 0.5] | |
| dataset.std = [0.5, 0.5, 0.5] | |
| # test __getitem__ | |
| result = dataset.__getitem__(0) | |
| # check returned keys | |
| expected_keys = ['gt', 'lq', 'gt_path', 'lq_path'] | |
| assert set(expected_keys).issubset(set(result.keys())) | |
| # check shape and contents | |
| assert result['gt'].shape == (3, 128, 128) | |
| assert result['lq'].shape == (3, 32, 32) | |