| 174 | class GaussianBlurBatch(object): |
| 175 | """blur a batch of images on CPU or GPU""" |
| 176 | def __init__(self, kernel_size, use_cuda=False, p=0.5): |
| 177 | radias = kernel_size // 2 |
| 178 | kernel_size = radias * 2 + 1 |
| 179 | self.blur_h = nn.Conv2d(3, 3, kernel_size=(kernel_size, 1), |
| 180 | stride=1, padding=0, bias=False, groups=3) |
| 181 | self.blur_v = nn.Conv2d(3, 3, kernel_size=(1, kernel_size), |
| 182 | stride=1, padding=0, bias=False, groups=3) |
| 183 | self.k = kernel_size |
| 184 | self.r = radias |
| 185 | |
| 186 | self.blur = nn.Sequential( |
| 187 | nn.ReflectionPad2d(radias), |
| 188 | self.blur_h, |
| 189 | self.blur_v |
| 190 | ) |
| 191 | |
| 192 | assert 0 <= p <= 1.0, 'p is out of range [0, 1]' |
| 193 | self.p = p |
| 194 | self.use_cuda = use_cuda |
| 195 | if use_cuda: |
| 196 | self.blur = self.blur.cuda() |
| 197 | |
| 198 | def __call__(self, imgs): |
| 199 | |