-
Notifications
You must be signed in to change notification settings - Fork 187
/
test_image_blur_mapper.py
92 lines (80 loc) · 3.21 KB
/
test_image_blur_mapper.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
import os
import unittest
import numpy as np
from PIL import ImageFilter
from data_juicer.core.data import NestedDataset as Dataset
from data_juicer.ops.mapper.image_blur_mapper import ImageBlurMapper
from data_juicer.utils.mm_utils import load_image
from data_juicer.utils.unittest_utils import DataJuicerTestCaseBase
class ImageBlurMapperTest(DataJuicerTestCaseBase):
data_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), '..',
'data')
img1_path = os.path.join(data_path, 'img1.png')
img2_path = os.path.join(data_path, 'img2.jpg')
img3_path = os.path.join(data_path, 'img3.jpg')
def _get_blur_kernel(self, blur_type='gaussian', radius=2):
if blur_type == 'mean':
return ImageFilter.BLUR
elif blur_type == 'box':
return ImageFilter.BoxBlur(radius)
else:
return ImageFilter.GaussianBlur(radius)
def _run_image_blur_mapper(self, op, source_list, blur_kernel):
dataset = Dataset.from_list(source_list)
dataset = dataset.map(op.process)
res_list = dataset.to_list()
for source, res in zip(source_list, res_list):
for s_path, r_path in zip(source[op.image_key], res[op.image_key]):
s_img = load_image(s_path).convert('RGB').filter(blur_kernel)
t_path = 'temp4test' + os.path.splitext(s_path)[-1]
s_img.save(t_path)
t_img = np.array(load_image(t_path))
r_img = np.array(load_image(r_path))
os.remove(t_path)
np.testing.assert_array_equal(t_img, r_img)
def test(self):
ds_list = [{
'images': [self.img1_path]
}, {
'images': [self.img2_path]
}, {
'images': [self.img3_path]
}]
op = ImageBlurMapper(p=1, blur_type='gaussian', radius=2)
blur_kernel = self._get_blur_kernel('gaussian', 2)
self._run_image_blur_mapper(op, ds_list, blur_kernel)
def test_blur_type(self):
ds_list = [{
'images': [self.img2_path]
}, {
'images': [self.img3_path]
}, {
'images': [self.img1_path]
}]
op = ImageBlurMapper(p=1, blur_type='box', radius=2)
blur_kernel = self._get_blur_kernel('box', 2)
self._run_image_blur_mapper(op, ds_list, blur_kernel)
def test_radius(self):
ds_list = [{
'images': [self.img3_path]
}, {
'images': [self.img2_path]
}, {
'images': [self.img1_path]
}]
op = ImageBlurMapper(p=1, blur_type='gaussian', radius=5)
blur_kernel = self._get_blur_kernel('gaussian', 5)
self._run_image_blur_mapper(op, ds_list, blur_kernel)
def test_multi_img(self):
ds_list = [{
'images': [self.img1_path, self.img2_path, self.img3_path]
}, {
'images': [self.img2_path]
}, {
'images': [self.img3_path, self.img1_path]
}]
op = ImageBlurMapper(p=1, blur_type='gaussian', radius=2)
blur_kernel = self._get_blur_kernel('gaussian', 2)
self._run_image_blur_mapper(op, ds_list, blur_kernel)
if __name__ == '__main__':
unittest.main()