ReflectionEraser/ReflectionEraserApp
0
1from __future__ import division #This ensures true division (i.e., division results in floating-point numbers) in Python 2 otherwise we get //
2
3import math
4import random
5
6import kornia #for computer vision operations in PyTorch.
7import torch
8from PIL import Image
9
10try:
11 import accimage #third-party Python library designed to provide accelerated image loading and processing capabilities, primarily intended for use with the PyTorch deep learning framework.Alternative of PIL/pillow
12except ImportError:
13 accimage = None
14import numpy as np
15import scipy.stats as st
16import cv2
17import collections
18import torchvision.transforms as transforms # for common image transformations.
19import util.util as util
20from scipy.signal import convolve2d
21
22
23# utility
24#
25# PIL and accimage for loading
26def _is_pil_image(img): #if accimage is not supported in environment
27 if accimage is not None:
28 return isinstance(img, (Image.Image, accimage.Image))
29 else:
30 return isinstance(img, Image.Image)
31
32#in transformation step for augmentation and feeding into a model.
33def _is_tensor_image(img): #to handle different types of image representations that might be encountered in various stages of image processing
34 return torch.is_tensor(img) and img.ndimension() == 3 # is 3D tensor
35
36# for certain preprocessing steps.
37def _is_numpy_image(img): #to handle different types of image representations that might be encountered in various stages of image processing
38 return isinstance(img, np.ndarray) and (img.ndim in {2, 3}) #is numpy array of 2 or 3 dimension
39
40
41def arrshow(arr):
42 Image.fromarray(arr.astype(np.uint8)).show()#converts a NumPy array to a PIL image and then displays it.
43 #arr.astype(np.uint8) converts the NumPy array arr to have a data type of np.uint8
44 # PIL images expect pixel values to be in the range [0, 255],
45 #Image.fromarray() function is a method provided by the PIL (Pillow) library to convert a pixel values to a PIL image.
46 #.show() is a method of the PIL Image class to open image in image viewer of system
47
48'''"transformations" refer to a series of operations applied to the images to preprocess or augment them. These transformations can include resizing, cropping, flipping, rotating, and many others. They help in normalizing the data, enhancing certain features, or augmenting the dataset to improve the robustness and performance of the model.'''
49
50def get_transform(opt): #based on option given, a list of transformations are generated
51 transform_list = []
52 osizes = util.parse_args(opt.loadSize)
53 fineSize = util.parse_args(opt.fineSize)
54 if opt.resize_or_crop == 'resize_and_crop':
55 transform_list.append(
56 transforms.RandomChoice([
57 transforms.Resize([osize, osize], Image.BICUBIC) for osize in osizes
58 ]))
59 transform_list.append(transforms.RandomCrop(fineSize))
60 elif opt.resize_or_crop == 'crop':
61 transform_list.append(transforms.RandomCrop(fineSize))
62 elif opt.resize_or_crop == 'scale_width':
63 transform_list.append(transforms.Lambda(
64 lambda img: __scale_width(img, fineSize)))
65 elif opt.resize_or_crop == 'scale_width_and_crop':
66 transform_list.append(transforms.Lambda(
67 lambda img: __scale_width(img, opt.loadSize)))
68 transform_list.append(transforms.RandomCrop(opt.fineSize))
69
70 if opt.isTrain and not opt.no_flip:
71 transform_list.append(transforms.RandomHorizontalFlip())
72
73 return transforms.Compose(transform_list)
74
75
76to_norm_tensor = transforms.Compose([
77 transforms.ToTensor(),
78 transforms.Normalize(
79 (0.5, 0.5, 0.5),
80 (0.5, 0.5, 0.5)
81 )
82])
83
84to_tensor = transforms.ToTensor()
85
86
87def __scale_width(img, target_width):
88 ow, oh = img.size
89 if (ow == target_width):
90 return img
91 w = target_width
92 h = int(target_width * oh / ow)
93 h = math.ceil(h / 2.) * 2 # round up to even
94 return img.resize((w, h), Image.BICUBIC)
95
96
97# functional
98def gaussian_blur(img, kernel_size, sigma):
99 if not _is_pil_image(img):
100 raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
101
102 img = np.asarray(img)
103 # the 3rd dimension (i.e. inter-band) would be filtered which is unwanted for our purpose
104 # new = gaussian_filter(img, sigma=sigma, truncate=truncate)
105 if isinstance(kernel_size, int):
106 kernel_size = (kernel_size, kernel_size)
107 elif isinstance(kernel_size, collections.Sequence):
108 assert len(kernel_size) == 2
109 new = cv2.GaussianBlur(img, kernel_size, sigma) # apply gaussian filter band by band
110 return Image.fromarray(new)
111
112
113# transforms
114class GaussianBlur(object):
115 def __init__(self, kernel_size=11, sigma=3):
116 self.kernel_size = kernel_size
117 self.sigma = sigma
118
119 def __call__(self, img):
120 return gaussian_blur(img, self.kernel_size, self.sigma)
121
122
123class ReflectionSythesis_0(object):
124 """Reflection image data synthesis for weakly-supervised learning
125 of ICCV 2017 paper *"A Generic Deep Architecture for Single Image Reflection Removal and Image Smoothing"*
126 """
127
128 def __init__(self, kernel_sizes=None, low_sigma=2, high_sigma=5, low_gamma=1.3,
129 high_gamma=1.3, low_delta=0.4, high_delta=1.8):
130 self.kernel_sizes = kernel_sizes or [11]
131 self.low_sigma = low_sigma
132 self.high_sigma = high_sigma
133 self.low_gamma = low_gamma
134 self.high_gamma = high_gamma
135 self.low_delta = low_delta
136 self.high_delta = high_delta
137 print('[i] reflection sythesis model: {}'.format({
138 'kernel_sizes': kernel_sizes, 'low_sigma': low_sigma, 'high_sigma': high_sigma,
139 'low_gamma': low_gamma, 'high_gamma': high_gamma}))
140
141 def __call__(self, B, R):
142 if not _is_pil_image(B):
143 raise TypeError('B should be PIL Image. Got {}'.format(type(B)))
144 if not _is_pil_image(R):
145 raise TypeError('R should be PIL Image. Got {}'.format(type(R)))
146 B_ = np.asarray(B, np.float32)
147 if random.random() < 0.4:
148 B_ = np.tile(np.random.uniform(0, 30, (1, 1, 1)), B_.shape) / 255.
149 else:
150 B_ = np.tile(np.random.normal(50, 50, (1, 1, 3)), (B_.shape[0], B_.shape[1], 1)).clip(0, 255) / 255.
151 R_ = np.asarray(R, np.float32) / 255.
152
153 kernel_size = np.random.choice(self.kernel_sizes)
154 sigma = np.random.uniform(self.low_sigma, self.high_sigma)
155 gamma = np.random.uniform(self.low_gamma, self.high_gamma)
156 delta = np.random.uniform(self.low_delta, self.high_delta)
157 R_blur = R_
158 kernel = cv2.getGaussianKernel(11, sigma)
159 kernel2d = np.dot(kernel, kernel.T)
160
161 for i in range(3):
162 R_blur[..., i] = convolve2d(R_blur[..., i], kernel2d, mode='same')
163
164 R_blur = np.clip(R_blur - np.mean(R_blur) * gamma, 0, 1)
165 R_blur = np.clip(R_blur * delta, 0, 1)
166 M_ = np.clip(R_blur + B_, 0, 1)
167
168 return B_, R_blur, M_
169
170
171class ReflectionSythesis_1(object):
172 """Reflection image data synthesis for weakly-supervised learning
173 of ICCV 2017 paper *"A Generic Deep Architecture for Single Image Reflection Removal and Image Smoothing"*
174 """
175
176 def __init__(self, kernel_sizes=None, low_sigma=2, high_sigma=5, low_gamma=1.3, high_gamma=1.3):
177 self.kernel_sizes = kernel_sizes or [11]
178 self.low_sigma = low_sigma
179 self.high_sigma = high_sigma
180 self.low_gamma = low_gamma
181 self.high_gamma = high_gamma
182 print('[i] reflection sythesis model: {}'.format({
183 'kernel_sizes': kernel_sizes, 'low_sigma': low_sigma, 'high_sigma': high_sigma,
184 'low_gamma': low_gamma, 'high_gamma': high_gamma}))
185
186 def __call__(self, B, R):
187 if not _is_pil_image(B):
188 raise TypeError('B should be PIL Image. Got {}'.format(type(B)))
189 if not _is_pil_image(R):
190 raise TypeError('R should be PIL Image. Got {}'.format(type(R)))
191
192 B_ = np.asarray(B, np.float32) / 255.
193 R_ = np.asarray(R, np.float32) / 255.
194
195 kernel_size = np.random.choice(self.kernel_sizes)
196 sigma = np.random.uniform(self.low_sigma, self.high_sigma)
197 gamma = np.random.uniform(self.low_gamma, self.high_gamma)
198 R_blur = R_
199 kernel = cv2.getGaussianKernel(11, sigma)
200 kernel2d = np.dot(kernel, kernel.T)
201
202 for i in range(3):
203 R_blur[..., i] = convolve2d(R_blur[..., i], kernel2d, mode='same')
204
205 M_ = B_ + R_blur
206
207 if np.max(M_) > 1:
208 m = M_[M_ > 1]
209 m = (np.mean(m) - 1) * gamma
210 R_blur = np.clip(R_blur - m, 0, 1)
211 M_ = np.clip(R_blur + B_, 0, 1)
212
213 return B_, R_blur, M_
214
215
216class NoiseReflectionSythesis(object):
217 """Reflection image data synthesis for weakly-supervised learning
218 of ICCV 2017 paper *"A Generic Deep Architecture for Single Image Reflection Removal and Image Smoothing"*
219 """
220
221 def __init__(self, kernel_sizes=None, low_sigma=2, high_sigma=5, low_gamma=1.3, high_gamma=1.3):
222 self.kernel_sizes = kernel_sizes or [11]
223 self.low_sigma = low_sigma
224 self.high_sigma = high_sigma
225 self.low_gamma = low_gamma
226 self.high_gamma = high_gamma
227 print('[i] reflection sythesis model: {}'.format({
228 'kernel_sizes': kernel_sizes, 'low_sigma': low_sigma, 'high_sigma': high_sigma,
229 'low_gamma': low_gamma, 'high_gamma': high_gamma}))
230
231 def __call__(self, B, R, N):
232 if not _is_pil_image(B):
233 raise TypeError('B should be PIL Image. Got {}'.format(type(B)))
234 if not _is_pil_image(R):
235 raise TypeError('R should be PIL Image. Got {}'.format(type(R)))
236
237 B_ = np.asarray(B, np.float32) / 255. + N
238 R_ = np.asarray(R, np.float32) / 255.
239
240 kernel_size = np.random.choice(self.kernel_sizes)
241 sigma = np.random.uniform(self.low_sigma, self.high_sigma)
242 gamma = np.random.uniform(self.low_gamma, self.high_gamma)
243 R_blur = R_
244 kernel = cv2.getGaussianKernel(11, sigma)
245 kernel2d = np.dot(kernel, kernel.T)
246
247 for i in range(3):
248 R_blur[..., i] = convolve2d(R_blur[..., i], kernel2d, mode='same')
249
250 M_ = B_ * 0.5 + R_blur * 0.5
251
252 return B_.astype(np.float32), R_blur.astype(np.float32), M_.astype(np.float32)
253
254
255class NoiseReflectionSythesisTorch(object):
256 """Reflection image data synthesis for weakly-supervised learning
257 of ICCV 2017 paper *"A Generic Deep Architecture for Single Image Reflection Removal and Image Smoothing"*
258 """
259
260 def __init__(self, kernel_sizes=None, low_sigma=2, high_sigma=5, low_gamma=1.3, high_gamma=1.3):
261 self.kernel_sizes = kernel_sizes or [11]
262 self.low_sigma = low_sigma
263 self.high_sigma = high_sigma
264 self.low_gamma = low_gamma
265 self.high_gamma = high_gamma
266 print('[i] reflection sythesis model: {}'.format({
267 'kernel_sizes': kernel_sizes, 'low_sigma': low_sigma, 'high_sigma': high_sigma,
268 'low_gamma': low_gamma, 'high_gamma': high_gamma}))
269
270 def __call__(self, B, R):
271 stdN = np.random.uniform(15, 55)
272 noise = torch.zeros(B.size()).normal_(mean=0, std=stdN / 255.)
273 B = B + noise
274
275 sigma = np.random.uniform(self.low_sigma, self.high_sigma)
276 R = kornia.gaussian_blur2d(R.unsqueeze(0), (11, 11), (sigma, sigma), border_type='replicate').squeeze(0)
277 M = B * 0.5 + R * 0.5
278
279 return B, R, M
280
281
282class Sobel(object):
283 def __call__(self, img):
284 if not _is_pil_image(img):
285 raise TypeError('img should be PIL Image. Got {}'.format(type(img)))
286
287 gray_img = np.array(img.convert('L'))
288 x = cv2.Sobel(gray_img, cv2.CV_16S, 1, 0)
289 y = cv2.Sobel(gray_img, cv2.CV_16S, 0, 1)
290
291 absX = cv2.convertScaleAbs(x)
292 absY = cv2.convertScaleAbs(y)
293
294 dst = cv2.addWeighted(absX, 0.5, absY, 0.5, 0)
295 return Image.fromarray(dst)
296
297
298class ReflectionSythesis_2(object):
299 """Reflection image data synthesis for weakly-supervised learning
300 of CVPR 2018 paper *"Single Image Reflection Separation with Perceptual Losses"*
301 """
302
303 def __init__(self, kernel_sizes=None):
304 self.kernel_sizes = kernel_sizes or np.linspace(1, 5, 80)
305
306 @staticmethod
307 def gkern(kernlen=100, nsig=1):
308 """Returns a 2D Gaussian kernel array."""
309 interval = (2 * nsig + 1.) / (kernlen)
310 x = np.linspace(-nsig - interval / 2., nsig + interval / 2., kernlen + 1)
311 kern1d = np.diff(st.norm.cdf(x))
312 kernel_raw = np.sqrt(np.outer(kern1d, kern1d))
313 kernel = kernel_raw / kernel_raw.sum()
314 kernel = kernel / kernel.max()
315 return kernel
316
317 def __call__(self, t, r):
318 t = np.float32(t) / 255.
319 r = np.float32(r) / 255.
320 ori_t = t
321 # create a vignetting mask
322 g_mask = self.gkern(560, 3)
323 g_mask = np.dstack((g_mask, g_mask, g_mask))
324 sigma = self.kernel_sizes[np.random.randint(0, len(self.kernel_sizes))]
325
326 t = np.power(t, 2.2)
327 r = np.power(r, 2.2)
328
329 sz = int(2 * np.ceil(2 * sigma) + 1)
330
331 r_blur = cv2.GaussianBlur(r, (sz, sz), sigma, sigma, 0)
332 blend = r_blur + t
333
334 att = 1.08 + np.random.random() / 10.0
335
336 for i in range(3):
337 maski = blend[:, :, i] > 1
338 mean_i = max(1., np.sum(blend[:, :, i] * maski) / (maski.sum() + 1e-6))
339 r_blur[:, :, i] = r_blur[:, :, i] - (mean_i - 1) * att
340 r_blur[r_blur >= 1] = 1
341 r_blur[r_blur <= 0] = 0
342
343 h, w = r_blur.shape[0:2]
344 neww = np.random.randint(0, 560 - w - 10)
345 newh = np.random.randint(0, 560 - h - 10)
346 alpha1 = g_mask[newh:newh + h, neww:neww + w, :]
347 alpha2 = 1 - np.random.random() / 5.0
348 r_blur_mask = np.multiply(r_blur, alpha1)
349 blend = r_blur_mask + t * alpha2
350
351 t = np.power(t, 1 / 2.2)
352 r_blur_mask = np.power(r_blur_mask, 1 / 2.2)
353 blend = np.power(blend, 1 / 2.2)
354 blend[blend >= 1] = 1
355 blend[blend <= 0] = 0
356
357 return np.float32(ori_t), np.float32(r_blur_mask), np.float32(blend)
358
359
360# Examples
361if __name__ == '__main__':
362 """cv2 imread"""
363 # img = cv2.imread('testdata_reflection_real/19-input.png')
364 # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
365 # img2 = cv2.GaussianBlur(img, (11,11), 3)
366
367 """Sobel Operator"""
368 # img = np.array(Image.open('datasets/VOC224/train/B/2007_000250.png').convert('L'))
369
370 """Reflection Sythesis"""
371 b = Image.open('')
372 r = Image.open('')
373 G = ReflectionSythesis_0()
374 m, r = G(b, r)
375 r.show()
376