mirror of https://github.com/kritiksoman/GIMP-ML
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
77 lines
2.6 KiB
Python
77 lines
2.6 KiB
Python
import os
|
|
import unittest
|
|
from shutil import rmtree
|
|
from tempfile import mkdtemp
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from torch.utils.data import DataLoader
|
|
|
|
from dataset import PairedDataset
|
|
|
|
|
|
def make_img():
|
|
return (np.random.rand(100, 100, 3) * 255).astype('uint8')
|
|
|
|
|
|
class AugTest(unittest.TestCase):
|
|
tmp_dir = mkdtemp()
|
|
raw = os.path.join(tmp_dir, 'raw')
|
|
gt = os.path.join(tmp_dir, 'gt')
|
|
|
|
def setUp(self):
|
|
for d in (self.raw, self.gt):
|
|
os.makedirs(d)
|
|
|
|
for i in range(5):
|
|
for d in (self.raw, self.gt):
|
|
img = make_img()
|
|
cv2.imwrite(os.path.join(d, f'{i}.png'), img)
|
|
|
|
def tearDown(self):
|
|
rmtree(self.tmp_dir)
|
|
|
|
def dataset_gen(self, equal=True):
|
|
base_config = {'files_a': os.path.join(self.raw, '*.png'),
|
|
'files_b': os.path.join(self.raw if equal else self.gt, '*.png'),
|
|
'size': 32,
|
|
}
|
|
for b in ([0, 1], [0, 0.9]):
|
|
for scope in ('strong', 'weak'):
|
|
for crop in ('random', 'center'):
|
|
for preload in (0, 1):
|
|
for preload_size in (0, 64):
|
|
config = base_config.copy()
|
|
config['bounds'] = b
|
|
config['scope'] = scope
|
|
config['crop'] = crop
|
|
config['preload'] = preload
|
|
config['preload_size'] = preload_size
|
|
config['verbose'] = False
|
|
dataset = PairedDataset.from_config(config)
|
|
yield dataset
|
|
|
|
def test_equal_datasets(self):
|
|
for dataset in self.dataset_gen(equal=True):
|
|
dataloader = DataLoader(dataset=dataset,
|
|
batch_size=2,
|
|
shuffle=True,
|
|
drop_last=True)
|
|
dataloader = iter(dataloader)
|
|
batch = next(dataloader)
|
|
a, b = map(lambda x: x.numpy(), map(batch.get, ('a', 'b')))
|
|
|
|
np.testing.assert_allclose(a, b)
|
|
|
|
def test_datasets(self):
|
|
for dataset in self.dataset_gen(equal=False):
|
|
dataloader = DataLoader(dataset=dataset,
|
|
batch_size=2,
|
|
shuffle=True,
|
|
drop_last=True)
|
|
dataloader = iter(dataloader)
|
|
batch = next(dataloader)
|
|
a, b = map(lambda x: x.numpy(), map(batch.get, ('a', 'b')))
|
|
|
|
assert not np.all(a == b), 'images should not be the same'
|