Select Git revision
patch_transformer.py 9.71 KiB
import math
from torch.nn.modules.utils import _pair, _quadruple
import torch.nn.functional as F
import numpy as np
import torch
from torch import nn
class PatchApplier(nn.Module):
"""PatchApplier: applies adversarial patches to images.
Module providing the functionality necessary to apply a patch to all detections in all images in the batch.
"""
def __init__(self):
super(PatchApplier, self).__init__()
def forward(self, img, patch):
img = torch.where(torch.tensor(patch < 1e-05), torch.tensor(img)/256, torch.tensor(patch))*256
return img.detach().numpy()
class MedianPool2d(nn.Module):
""" Median pool (usable as median filter when stride=1) module.
Args:
kernel_size: size of pooling kernel, int or 2-tuple
stride: pool stride, int or 2-tuple
padding: pool padding, int or 4-tuple (l, r, t, b) as in pytorch F.pad
same: override padding and enforce same padding, boolean
"""
def __init__(self, kernel_size=3, stride=1, padding=0, same=False):
super(MedianPool2d, self).__init__()
self.k = _pair(kernel_size)
self.stride = _pair(stride)
self.padding = _quadruple(padding) # convert to l, r, t, b
self.same = same
def _padding(self, x):
if self.same:
ih, iw = x.size()[2:]
if ih % self.stride[0] == 0:
ph = max(self.k[0] - self.stride[0], 0)
else:
ph = max(self.k[0] - (ih % self.stride[0]), 0)
if iw % self.stride[1] == 0:
pw = max(self.k[1] - self.stride[1], 0)
else:
pw = max(self.k[1] - (iw % self.stride[1]), 0)
pl = pw // 2
pr = pw - pl
pt = ph // 2
pb = ph - pt
padding = (pl, pr, pt, pb)
else:
padding = self.padding
return padding
def forward(self, x):
# using existing pytorch functions and tensor ops so that we get autograd,
# would likely be more efficient to implement from scratch at C/Cuda level
x = F.pad(x, self._padding(x), mode='reflect')
x = x.unfold(2, self.k[0], self.stride[0]).unfold(3, self.k[1], self.stride[1])
x = x.contiguous().view(x.size()[:4] + (-1,)).median(dim=-1)[0]
return x
class PatchTransformer(nn.Module):
"""PatchTransformer: transforms batch of patches
Module providing the functionality necessary to transform a batch of patches, randomly adjusting brightness and
contrast, adding random amount of noise, and rotating randomly. Resizes patches according to as size based on the
batch of labels, and pads them to the dimension of an image.
"""
def __init__(self):
super(PatchTransformer, self).__init__()
self.min_contrast = 0.8
self.max_contrast = 1.2
self.min_brightness = -0.1
self.max_brightness = 0.1
self.noise_factor = 0.10
self.minangle = -20 / 180 * math.pi
self.maxangle = 20 / 180 * math.pi
self.medianpooler = MedianPool2d(7, same=True)
'''
kernel = torch.cuda.FloatTensor([[0.003765, 0.015019, 0.023792, 0.015019, 0.003765],
[0.015019, 0.059912, 0.094907, 0.059912, 0.015019],
[0.023792, 0.094907, 0.150342, 0.094907, 0.023792],
[0.015019, 0.059912, 0.094907, 0.059912, 0.015019],
[0.003765, 0.015019, 0.023792, 0.015019, 0.003765]])
self.kernel = kernel.unsqueeze(0).unsqueeze(0).expand(3,3,-1,-1)
'''
def forward(self, adv_patch, lab_batch, img_size_x, img_size_y, do_rotate=True, rand_loc=True):
# adv_patch = F.conv2d(adv_patch.unsqueeze(0),self.kernel,padding=(2,2))
adv_patch = self.medianpooler(adv_patch.unsqueeze(0))
# Determine size of padding
pad_x = (img_size_x - adv_patch.size(-1)) / 2
pad_y = (img_size_y - adv_patch.size(-1)) / 2
# Make a batch of patches
adv_patch = adv_patch.unsqueeze(0) # .unsqueeze(0)
adv_batch = adv_patch.expand(lab_batch.size(0), lab_batch.size(1), -1, -1, -1)
batch_size = torch.Size((lab_batch.size(0), lab_batch.size(1)))
# Contrast, brightness and noise transforms
# Create random contrast tensor
contrast = torch.cuda.FloatTensor(batch_size).uniform_(self.min_contrast, self.max_contrast)
contrast = contrast.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
contrast = contrast.expand(-1, -1, adv_batch.size(-3), adv_batch.size(-2), adv_batch.size(-1))
contrast = contrast.cuda()
# Create random brightness tensor
brightness = torch.cuda.FloatTensor(batch_size).uniform_(self.min_brightness, self.max_brightness)
brightness = brightness.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
brightness = brightness.expand(-1, -1, adv_batch.size(-3), adv_batch.size(-2), adv_batch.size(-1))
brightness = brightness.cuda()
# Create random noise tensor
noise = torch.cuda.FloatTensor(adv_batch.size()).uniform_(-1, 1) * self.noise_factor
# Apply contrast/brightness/noise, clamp
adv_batch = adv_batch * contrast + brightness + noise
adv_batch = torch.clamp(adv_batch, 0.000001, 0.99999)
# TODO should that not be class 0 if we do not want to cover person? cls=1 would be bicycle ...
# Where the label class_id is 1 we don't want a patch (padding) --> fill mask with zero's
cls_ids = torch.narrow(lab_batch, 2, 0, 1)
cls_mask = cls_ids.expand(-1, -1, 3)
cls_mask = cls_mask.unsqueeze(-1)
cls_mask = cls_mask.expand(-1, -1, -1, adv_batch.size(3))
cls_mask = cls_mask.unsqueeze(-1)
cls_mask = cls_mask.expand(-1, -1, -1, -1, adv_batch.size(4))
msk_batch = torch.cuda.FloatTensor(cls_mask.size()).fill_(1) - cls_mask
# Pad patch and mask to image dimensions
mypad = nn.ConstantPad2d((int(pad_x + 0.5), int(pad_x), int(pad_y + 0.5), int(pad_y)), 0)
adv_batch = mypad(adv_batch)
msk_batch = mypad(msk_batch)
# Rotation and rescaling transforms
anglesize = (lab_batch.size(0) * lab_batch.size(1))
if do_rotate:
angle = torch.cuda.FloatTensor(anglesize).uniform_(self.minangle, self.maxangle)
else:
angle = torch.cuda.FloatTensor(anglesize).fill_(0)
# Resizes and rotates
current_patch_size = adv_patch.size(-1)
lab_batch_scaled = torch.cuda.FloatTensor(lab_batch.size()).fill_(0)
lab_batch_scaled[:, :, 1] = lab_batch[:, :, 1] * img_size_x
lab_batch_scaled[:, :, 2] = lab_batch[:, :, 2] * img_size_x
lab_batch_scaled[:, :, 3] = lab_batch[:, :, 3] * img_size_y
lab_batch_scaled[:, :, 4] = lab_batch[:, :, 4] * img_size_y
target_size = torch.sqrt(
((lab_batch_scaled[:, :, 3].mul(0.2)) ** 2) + ((lab_batch_scaled[:, :, 4].mul(0.2)) ** 2))
target_x = lab_batch[:, :, 1].view(np.prod(batch_size)) / 2
target_y = lab_batch[:, :, 2].view(np.prod(batch_size)) / 2
targetoff_x = lab_batch[:, :, 3].view(np.prod(batch_size))
targetoff_y = lab_batch[:, :, 4].view(np.prod(batch_size))
if (rand_loc):
off_x = targetoff_x * (torch.cuda.FloatTensor(targetoff_x.size()).uniform_(-0.4, 0.4))
target_x = target_x + off_x
off_y = targetoff_y * (torch.cuda.FloatTensor(targetoff_y.size()).uniform_(-0.4, 0.4))
target_y = target_y + off_y
target_y = target_y - 0.05
scale = target_size / current_patch_size
scale = scale.view(anglesize)
s = adv_batch.size()
adv_batch = adv_batch.view(s[0] * s[1], s[2], s[3], s[4])
msk_batch = msk_batch.view(s[0] * s[1], s[2], s[3], s[4])
tx = (-target_x + 0.5) * 2
ty = (-target_y + 0.5) * 2
sin = torch.sin(angle)
cos = torch.cos(angle)
# Theta = rotation,rescale matrix
theta = torch.cuda.FloatTensor(anglesize, 2, 3).fill_(0)
theta[:, 0, 0] = cos / scale
theta[:, 0, 1] = sin / scale
theta[:, 0, 2] = tx * cos / scale + ty * sin / scale
theta[:, 1, 0] = -sin / scale
theta[:, 1, 1] = cos / scale
theta[:, 1, 2] = -tx * sin / scale + ty * cos / scale
b_sh = adv_batch.shape
grid = F.affine_grid(theta, adv_batch.shape, align_corners=False)
adv_batch_t = F.grid_sample(adv_batch, grid, align_corners=False)
msk_batch_t = F.grid_sample(msk_batch, grid, align_corners=False)
'''
# Theta2 = translation matrix
theta2 = torch.cuda.FloatTensor(anglesize, 2, 3).fill_(0)
theta2[:, 0, 0] = 1
theta2[:, 0, 1] = 0
theta2[:, 0, 2] = (-target_x + 0.5) * 2
theta2[:, 1, 0] = 0
theta2[:, 1, 1] = 1
theta2[:, 1, 2] = (-target_y + 0.5) * 2
grid2 = F.affine_grid(theta2, adv_batch.shape)
adv_batch_t = F.grid_sample(adv_batch_t, grid2)
msk_batch_t = F.grid_sample(msk_batch_t, grid2)
'''
adv_batch_t = adv_batch_t.view(s[0], s[1], s[2], s[3], s[4])
msk_batch_t = msk_batch_t.view(s[0], s[1], s[2], s[3], s[4])
adv_batch_t = torch.clamp(adv_batch_t, 0.000001, 0.999999)
# img = msk_batch_t[0, 0, :, :, :].detach().cpu()
# img = transforms.ToPILImage()(img)
# img.show()
# exit()
return adv_batch_t #* msk_batch_t