def make_homography_kornia(shifts, horizon_line):
horizon_line = horizon_line * 2 - 1
left_top_dx, left_top_dy = shifts[0][0]
right_top_dx, right_top_dy = shifts[0][1]
left_hor_dx, left_hor_dy = shifts[1][0]
right_hor_dx, right_hor_dy = shifts[1][1]
if horizon_line > -1 + 1e-4:
points_src = torch.tensor([[
[-1, -1], # left top
[1, -1], # right top
[-1, horizon_line], # left bottom
[1, horizon_line]
]]).float() # right bottom
points_tgt = torch.tensor([[
[-1 + left_top_dx, -1 + left_top_dy], # left top
[1 + right_top_dx, -1 + right_top_dy], # right top
[-1 + left_hor_dx, horizon_line + left_hor_dy], # left bottom
[1 + right_hor_dx, horizon_line + right_hor_dy]
]]).float() # right bottom
sky_transform = kornia.get_perspective_transform(
points_src, points_tgt)
else:
points_src = torch.tensor([[
[-1, -1], # left top
[1, -1], # right top
[-1, 1], # left bottom
[1, 1]
]]).float() # right bottom
sky_transform = kornia.get_perspective_transform(
points_src, points_src)
if horizon_line <= 1 - 1e-4:
points_src = torch.tensor([[
[-1, horizon_line], # left horizon
[1, horizon_line], # right horizon
[-1, 1], # left bottom
[1, 1]
]]).float() # right bottom
points_tgt = torch.tensor([[
[-1 + left_hor_dx, horizon_line - left_hor_dy], # left horizon
[1 + right_hor_dx, horizon_line - right_hor_dy], # right horizon
[-1 + left_top_dx, 1 - left_top_dy], # left bottom
[1 + right_top_dx, 1 - right_top_dy]
]]).float() # right bottom
earth_transform = kornia.get_perspective_transform(
points_src, points_tgt)
else:
points_src = torch.tensor([[
[-1, -1], # left horizon
[1, -1], # right horizon
[-1, 1], # left bottom
[1, 1]
]]).float() # right bottom
earth_transform = kornia.get_perspective_transform(
points_src, points_src)
return sky_transform, earth_transform
def test_crop(self, batch_size, channels, device, dtype):
# generate input data
src_h, src_w = 3, 3
dst_h, dst_w = 3, 3
# [x, y] origin
# top-left, top-right, bottom-right, bottom-left
points_src = torch.tensor(
[[[0, 0], [0, src_w - 1], [src_h - 1, src_w - 1], [src_h - 1, 0]]],
device=device,
dtype=dtype)
# [x, y] destination
# top-left, top-right, bottom-right, bottom-left
points_dst = torch.tensor(
[[[0, 0], [0, dst_w - 1], [dst_h - 1, dst_w - 1], [dst_h - 1, 0]]],
device=device,
dtype=dtype)
# compute transformation between points
dst_trans_src = kornia.get_perspective_transform(
points_src, points_dst).expand(batch_size, -1, -1)
# warp tensor
patch = torch.tensor([[[[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12],
[13, 14, 15, 16]]]],
device=device,
dtype=dtype).expand(batch_size, channels, -1, -1)
expected = patch[..., :3, :3]
# warp and assert
patch_warped = kornia.warp_perspective(patch, dst_trans_src,
(dst_h, dst_w))
assert_close(patch_warped, expected)
def _convert(self):
src_img = kornia.image_to_tensor(self.src_img, keepdim=False)
dst_h, dst_w = 800, 800
# Compute perspective transform
M = kornia.get_perspective_transform(self.points_src, self.points_dst)
# Image to BEV transformation
dst_img = kornia.warp_perspective(src_img.float(),
M,
dsize=(dst_h, dst_w),
flags='bilinear',
border_mode='zeros')
# remove unwanted portion of BEV image. e.g for FRONT view dst point should not be higher than 450.
if 'FRONT' in self.path:
dst_img[:, :, 400:, :] = 0
if 'BACK' in self.path:
dst_img[:, :, :400, :] = 0
if 'LEFT' in self.path:
dst_img[:, :, :, 400:] = 0
if 'RIGHT' in self.path:
dst_img[:, :, :, :400] = 0
dst_img = kornia.tensor_to_image(dst_img.byte())
return dst_img
def _homography_floor_svd(self,
top_corners: torch.Tensor, # in [-1, 1]
bottom_corners: torch.Tensor, # in [-1, 1]
floor_z: float=-1.6,
):
b, N, _ = top_corners.size()
u = bottom_corners[:, :, 0] * np.pi
v = bottom_corners[:, :, 1] * (-0.5 * np.pi)
c = floor_z / torch.tan(v)
x = c * torch.sin(u)
y = -c * torch.cos(u)
floor_xy = torch.stack([x, y], dim=-1)
scale = self._get_scale_all(floor_xy)
scale = scale / 2.0
centroid = floor_xy.mean(dim=1)
c = torch.linalg.norm(floor_xy, ord=2, dim=-1)
v = top_corners[:, :, 1] * (-0.5 * np.pi)
ceil_z = (c * torch.tan(v)).mean(dim=1, keepdim=True)
ceil_z = ceil_z.unsqueeze(1).expand(b, 4, 1).contiguous()
floor_xy = floor_xy - centroid.unsqueeze(1)
inds = torch.sort(torch.atan2(floor_xy[..., 0], floor_xy[..., 1] + 1e-12))[1]
axes = self.cuboid_axes[:, inds.squeeze(), :]
homography = kornia.get_perspective_transform(floor_xy, axes)
homogeneous = torch.cat([floor_xy, torch.ones_like(floor_xy[..., -1:])], dim=2)
xformed = (homography @ homogeneous.transpose(1, 2)).transpose(1, 2)
xformed = xformed[:, :, :2] / xformed[:, :, 2].unsqueeze(-1)
rect_floor_xy = xformed * scale.unsqueeze(1) + centroid.unsqueeze(1)
original_xy = floor_xy + centroid.unsqueeze(1)
R, t, s = self._svd(rect_floor_xy, original_xy[:, inds.squeeze(), :])
rect_floor_xy = self._transform_points(rect_floor_xy, R, t, s)
bottom_points = torch.cat([rect_floor_xy, floor_z * torch.ones_like(c.unsqueeze(-1))], dim=-1)
top_points = torch.cat([rect_floor_xy, ceil_z], dim=-1)
return top_points, bottom_points
def photometric_loss(delta, img_a, patch_b, corners):
corners_hat = corners + delta
# in order to apply transform and center crop,
# subtract points by top-left corner (corners[N, 0])
# print("ori_corners:",corners[0:1])
corners = corners - corners[:, 0].view(-1, 1,
2) #关键！！！！区分大变换还是小变换！！！！！！！！！！！！！
h = kornia.get_perspective_transform(corners, corners_hat)
h_inv = torch.inverse(h) #求逆矩阵
# patch_b_hat = kornia.warp_perspective(img_a, h_inv, (128, 128))
patch_b_hat = kornia.warp_perspective(
img_a, h_inv, (patch_b.shape[-2], patch_b.shape[-1]))
# patch_b_hat2 = kornia.warp_perspective(img_a, h_inv, (img_a.shape[-2],img_a.shape[-1]))
# print("corners:",corners[0:1])
# print("corners_hat:",corners_hat[0:1])
# print("H:",h[0:1])
# print(img_a.size())
# print(patch_b.size())
# print(patch_b_hat.size())
#
# save_pic(img_a[0:1, :], "img_a.png")
# save_pic(patch_b[0:1, :], "patch_b.png")
# save_pic(patch_b_hat[0:1, :], "patch_b_hat.png")
# save_pic(patch_b_hat2[0:1, :], "patch_b_hat2.png")
# raise ValueError("print size")
return F.l1_loss(patch_b_hat, patch_b)
def warp(pert_tensor, bbox_src, bbox_dest):
'''
Input: pert_tensor : Tensor (3, W, H)
bbox_src and bbox_dest: (B, 4)
Output: Tensor (B, 3, W, H)
'''
if type(bbox_src) == torch.Tensor:
bbox_src = bbox_src.cpu()
bbox_dest = bbox_dest.cpu()
bbox_src = np.array(bbox_src).reshape(-1, 4)
bbox_dest = np.array(bbox_dest).reshape(-1, 4)
masks = list()
for i in range(bbox_src.shape[0]):
x, y, w, h = bbox_src[i]
points_src = torch.FloatTensor([[
[x, y],
[x + w, y],
[x, y + h],
[x + w, y + h],
]])
x, y, w, h = bbox_dest[i]
points_dst = torch.FloatTensor([[
[x, y],
[x + w, y],
[x, y + h],
[x + w, y + h],
]])
M = kornia.get_perspective_transform(points_src,
points_dst).to(pert_tensor.device)
size = pert_tensor.shape[-2:]
masks.append(kornia.warp_perspective(pert_tensor, M, size))
return torch.cat(masks)
def test_get_perspective_transform(batch_size, device):
# generate input data
h_max, w_max = 64, 32 # height, width
h = torch.ceil(h_max * torch.rand(batch_size)).to(device)
w = torch.ceil(w_max * torch.rand(batch_size)).to(device)
norm = torch.rand(batch_size, 4, 2).to(device)
points_src = torch.zeros_like(norm)
points_src[:, 1, 0] = h
points_src[:, 2, 1] = w
points_src[:, 3, 0] = h
points_src[:, 3, 1] = w
points_dst = points_src + norm
# compute transform from source to target
dst_homo_src = kornia.get_perspective_transform(points_src, points_dst)
assert_allclose(kornia.transform_points(dst_homo_src, points_src),
points_dst)
# compute gradient check
points_src = utils.tensor_to_gradcheck_var(points_src) # to var
points_dst = utils.tensor_to_gradcheck_var(points_dst) # to var
assert gradcheck(kornia.get_perspective_transform, (
points_src,
points_dst,
),
raise_exception=True)
def main(args):
if args.resume !="":
model = HomographyModel.load_from_checkpoint(args.resume)
else:
model_dir = 'lightning_logs/version*'
model_dir_list = sorted(glob.glob(model_dir))
model_dir = model_dir_list[-1]
model_path = osp.join(model_dir, "checkpoints", "*.ckpt")
model_path_list = sorted(glob.glob(model_path))
if len(model_path_list) > 0:
model_path = model_path_list[-1]
print(model_path)
# model = HomographyModel.load_from_checkpoint(model_path)
model = HomographyModel() #test专用，内部重新定义精简版的类
model_old = torch.load(model_path, map_location=lambda storage, loc: storage)
# print(model_old.keys())
# net.load_state_dict(torch.load('path/params.pkl'))
model.load_state_dict(model_old['state_dict'])
print(model_path)
print("model loaded.")
else:
raise ValueError(f'No load model!') #raise Error
model.eval() #不训练
test_set = SyntheticDataset(args.test_path, rho=args.rho, filetype=args.filetype,pic_size=pic_size,patch_size=patch_size)
#clear last output
last_output = "figures/*"
os.system("rm "+last_output)
print('clear last ok.')
for i in range(args.n):
img_a,img_b, patch_a, patch_b, corners, delta = test_set[i]
# after unsqueeze to save
# tensors_to_gif(patch_a, patch_b, f"figures/input_{i}.gif")
tensors_to_gif(img_a, img_b, f"figures/input_{i}.gif")
#add
img_a = img_a.unsqueeze(0)
img_b = img_b.unsqueeze(0)
##
patch_a = patch_a.unsqueeze(0)
patch_b = patch_b.unsqueeze(0)
corners = corners.unsqueeze(0)
corners = corners - corners[:, 0].view(-1, 1, 2)
delta_hat = model(patch_a, patch_b)
corners_hat = corners + delta_hat
#获取h
h = kornia.get_perspective_transform(corners, corners_hat)
h_inv = torch.inverse(h)
patch_b_hat = kornia.warp_perspective(img_a, h_inv, (patch_a.shape[-2],patch_a.shape[-1])) #128 最初设置 #注意，用img_a / patch_a
img_b_hat = kornia.warp_perspective(img_a, h_inv, (img_a.shape[-2],img_a.shape[-1]))
#输出
tensors_to_gif(patch_b_hat[0], patch_b[0], f"figures/output_patch{i}.gif")
tensors_to_gif(img_b_hat[0], img_b[0], f"figures/output_{i}.gif")
def four_point_to_homography(corners, deltas, crop=False):
"""
Args:
corners ():
deltas ():
crop (bool): If set to true, homography will aready contain cropping part.
"""
assert len(corners.shape) == 3, 'corners should be of size B, 4, 2'
assert len(deltas.shape) == 3, 'deltas should be of size B, 4, 2'
# in order to apply transform and center crop,
# subtract points by top-left corner (corners[N, 0])
if 'torch' in str(type(corners)):
if crop:
corners = corners - corners[:, 0].view(-1, 1, 2)
corners_hat = corners + deltas
return kornia.get_perspective_transform(corners, corners_hat)
elif 'numpy' in str(type(corners)):
if crop:
corners = corners - corners[:, 0].reshape(-1, 1, 2)
corners_hat = corners + deltas
return cv2.getPerspectiveTransform(np.float32(corners), np.float32(corners_hat))
else:
assert False, 'Wrong type?'
def test_crop(self, batch_size, channels, device, dtype):
# generate input data
src_h, src_w = 3, 3
dst_h, dst_w = 3, 3
# [x, y] origin
# top-left, top-right, bottom-right, bottom-left
points_src = torch.tensor([[
[0, 0],
[0, src_w - 1],
[src_h - 1, src_w - 1],
[src_h - 1, 0],
]],
device=device,
dtype=dtype)
# [x, y] destination
# top-left, top-right, bottom-right, bottom-left
points_dst = torch.tensor([[
[0, 0],
[0, dst_w - 1],
[dst_h - 1, dst_w - 1],
[dst_h - 1, 0],
]],
device=device,
dtype=dtype)
# compute transformation between points
dst_trans_src = kornia.get_perspective_transform(
points_src, points_dst).expand(batch_size, -1, -1)
# warp tensor
patch = torch.tensor([[[
[1, 2, 3, 4],
[5, 6, 7, 8],
[9, 10, 11, 12],
[13, 14, 15, 16],
]]],
device=device,
dtype=dtype).expand(batch_size, channels, -1, -1)
expected = torch.tensor(
[[[[0.2500, 0.9167, 1.5833], [2.1667, 5.1667, 6.5000],
[4.8333, 10.5000, 11.8333]]]],
device=device,
dtype=dtype).repeat(batch_size, channels, 1, 1)
# warp and assert
patch_warped = kornia.warp_perspective(patch, dst_trans_src,
(dst_h, dst_w))
assert_allclose(patch_warped, expected, rtol=1e-4, atol=1e-4)
# check jit
patch_warped_jit = kornia.jit.warp_perspective(patch, dst_trans_src,
(dst_h, dst_w))
assert_allclose(patch_warped, patch_warped_jit, rtol=1e-4, atol=1e-4)
def test_crop_center_resize(self, device, dtype):
# generate input data
dst_h, dst_w = 4, 4
# [x, y] origin
# top-left, top-right, bottom-right, bottom-left
points_src = torch.tensor([[
[1, 1],
[1, 2],
[2, 2],
[2, 1],
]],
device=device,
dtype=dtype)
# [x, y] destination
# top-left, top-right, bottom-right, bottom-left
points_dst = torch.tensor([[
[0, 0],
[0, dst_w - 1],
[dst_h - 1, dst_w - 1],
[dst_h - 1, 0],
]],
device=device,
dtype=dtype)
# compute transformation between points
dst_trans_src = kornia.get_perspective_transform(
points_src, points_dst)
# warp tensor
patch = torch.tensor([[[
[1, 2, 3, 4],
[5, 6, 7, 8],
[9, 10, 11, 12],
[13, 14, 15, 16],
]]],
device=device,
dtype=dtype)
expected = torch.tensor([[[[5.1667, 5.6111, 6.0556, 6.5000],
[6.9444, 7.3889, 7.8333, 8.2778],
[8.7222, 9.1667, 9.6111, 10.0556],
[10.5000, 10.9444, 11.3889, 11.8333]]]],
device=device,
dtype=dtype)
# warp and assert
patch_warped = kornia.warp_perspective(patch, dst_trans_src,
(dst_h, dst_w))
assert_allclose(patch_warped, expected, rtol=1e-4, atol=1e-4)
# check jit
patch_warped_jit = kornia.jit.warp_perspective(patch, dst_trans_src,
(dst_h, dst_w))
assert_allclose(patch_warped, patch_warped_jit, rtol=1e-4, atol=1e-4)
def sample(self, idx, points_source):
r = self.dataloader._all_records.iloc[idx]
image_id = r['Image']
_device = points_source.device
img = self.dataloader.read_image(image_id, r['flipped'])
img = torchvision.transforms.functional.to_tensor(img).to(_device)
px = points_source[:,0] * img.shape[1]
py = points_source[:,1] * img.shape[2]
p = torch.stack([px, py], dim=1)
M = kornia.get_perspective_transform(p.unsqueeze(0), self.points_dest.to(_device))
return kornia.warp_perspective(img.unsqueeze(0), M, dsize=(self.h, self.w))[0]
def _homography_joint_svd(self,
top_corners: torch.Tensor, # in [-1, 1]
bottom_corners: torch.Tensor, # in [-1, 1]
floor_z: float=-1.6,
ceil_z: float=1.6,
):
b, N, _ = top_corners.size()
floor_u = bottom_corners[:, :, 0] * np.pi
floor_v = bottom_corners[:, :, 1] * (-0.5 * np.pi)
floor_c = floor_z / torch.tan(floor_v)
floor_x = floor_c * torch.sin(floor_u)
floor_y = -floor_c * torch.cos(floor_u)
floor_xy = torch.stack([floor_x, floor_y], dim=-1)
floor_scale = self._get_scale_all(floor_xy)
floor_scale = floor_scale / 2.0
floor_ceil_c = torch.linalg.norm(floor_xy, ord=2, dim=-1)
floor_ceil_v = top_corners[:, :, 1] * (-0.5 * np.pi)
floor_ceil_z = (floor_ceil_c * torch.tan(floor_ceil_v)).mean(dim=1, keepdim=True)
floor_ceil_z = floor_ceil_z.unsqueeze(1).expand(b, 4, 1).contiguous()
ceil_u_t = top_corners[:, :, 0] * np.pi
ceil_v_t = top_corners[:, :, 1] * (-0.5 * np.pi)
ceil_c = ceil_z / torch.tan(ceil_v_t)
ceil_x = ceil_c * torch.sin(ceil_u_t)
ceil_y = -ceil_c * torch.cos(ceil_u_t)
ceil_xy = torch.stack([ceil_x, ceil_y], dim=-1)
ceil_floor_c = torch.linalg.norm(ceil_xy, ord=2, dim=-1)
ceil_v_b = bottom_corners[:, :, 1] * (-0.5 * np.pi)
ceil_floor_z = (ceil_floor_c * torch.tan(ceil_v_b)).mean(dim=1, keepdim=True)
fix_ceil = -ceil_z / ceil_floor_z
ceil_z_fix = ceil_z * fix_ceil
ceil_z_fix = ceil_z_fix.unsqueeze(1).expand(b, 4, 1).contiguous()
ceil_floor_fixed_c = ceil_z_fix.squeeze(-1) / torch.tan(ceil_v_t)
ceil_x = ceil_floor_fixed_c * torch.sin(ceil_u_t)
ceil_y = -ceil_floor_fixed_c * torch.cos(ceil_u_t)
ceil_xy = torch.stack([ceil_x, ceil_y], dim=-1)
ceil_scale = self._get_scale_all(ceil_xy)
ceil_scale = ceil_scale / 2.0
joint_xy = 0.5 * (floor_xy + ceil_xy)
joint_scale = 0.5 * (floor_scale + ceil_scale)
joint_centroid = joint_xy.mean(dim=1)
joint_xy = joint_xy - joint_centroid.unsqueeze(1)
inds = torch.sort(torch.atan2(joint_xy[..., 0], joint_xy[..., 1] + 1e-12))[1]
axes = self.cuboid_axes[:, inds.squeeze(), :]
homography = kornia.get_perspective_transform(joint_xy, axes)
homogeneous = torch.cat([joint_xy, torch.ones_like(joint_xy[..., -1:])], dim=2)
xformed = (homography @ homogeneous.transpose(1, 2)).transpose(1, 2)
xformed = xformed[:, :, :2] / xformed[:, :, 2].unsqueeze(-1)
rect_joint_xy = xformed * joint_scale.unsqueeze(1) + joint_centroid.unsqueeze(1)
original_xy = joint_xy + joint_centroid.unsqueeze(1)
R, t, s = self._svd(rect_joint_xy, original_xy[:, inds.squeeze(), :])
rect_joint_xy = self._transform_points(rect_joint_xy, R, t, s)
bottom_points = torch.cat([rect_joint_xy, floor_z * torch.ones_like(floor_c.unsqueeze(-1))], dim=-1)
top_points = torch.cat([rect_joint_xy, ceil_z_fix], dim=-1)
return top_points, bottom_points
def get_perspective_transform_torch(src, dst):
"""Get the homography matrix between src and dst
Arguments:
src: Tensor of shape (B,4,2), the four original points per image
dst: Tensor of shape (B,4,2), the four corresponding points per image
Returns:
A tensor of shape (B,3,3), each homography per image
Raises:
"""
return kornia.get_perspective_transform(src, dst)
def get_h(self, a, b, corners):
x = torch.cat((a, b), dim=1) # combine two images in channel dimension
x = self.cnn(x)
x = self.fc(x)
delta = x.view(-1, 4, 2)
corners_hat = corners + delta
# 获取h
h = kornia.get_perspective_transform(corners, corners_hat)
h_inv = torch.inverse(h)
return h_inv
def photometric_loss(delta, img_a, patch_b, corners):
corners_hat = corners + delta
# in order to apply transform and center crop,
# subtract points by top-left corner (corners[N, 0])
corners = corners - corners[:, 0].view(-1, 1, 2)
h = kornia.get_perspective_transform(corners, corners_hat)
h_inv = torch.inverse(h)
patch_b_hat = kornia.warp_perspective(img_a, h_inv, (128, 128))
return F.l1_loss(patch_b_hat, patch_b)
def test_crop_center_resize(self, device, dtype):
# generate input data
dst_h, dst_w = 4, 4
# [x, y] origin
# top-left, top-right, bottom-right, bottom-left
points_src = torch.tensor([[
[1, 1],
[1, 2],
[2, 2],
[2, 1],
]],
device=device,
dtype=dtype)
# [x, y] destination
# top-left, top-right, bottom-right, bottom-left
points_dst = torch.tensor([[
[0, 0],
[0, dst_w - 1],
[dst_h - 1, dst_w - 1],
[dst_h - 1, 0],
]],
device=device,
dtype=dtype)
# compute transformation between points
dst_trans_src = kornia.get_perspective_transform(
points_src, points_dst)
# warp tensor
patch = torch.tensor([[[
[1, 2, 3, 4],
[5, 6, 7, 8],
[9, 10, 11, 12],
[13, 14, 15, 16],
]]],
device=device,
dtype=dtype)
expected = torch.tensor([[[[6.0000, 6.3333, 6.6667, 7.0000],
[7.3333, 7.6667, 8.0000, 8.3333],
[8.6667, 9.0000, 9.3333, 9.6667],
[10.0000, 10.3333, 10.6667, 11.0000]]]],
device=device,
dtype=dtype)
# warp and assert
patch_warped = kornia.warp_perspective(patch, dst_trans_src,
(dst_h, dst_w))
assert_allclose(patch_warped, expected)
def stn(self, x, ret_theta=False):
points_src = self.locnet(x)
points_src = torch.clamp(points_src, -1, 1)
points_src = (points_src + 1.) * 63.5 # rescale to actual coordinate
points_src = points_src.view(-1, 4, 2)
B = points_src.shape[0]
points_dst = self.points_dst.repeat(B, 1, 1)
M = kornia.get_perspective_transform(points_src, points_dst)
x = kornia.warp_perspective(x, M, dsize=(self.h, self.w))
if ret_theta: return x, points_src
return x
def __getitem__(self, index):
img_a = Image.open(self.fnames[index]).convert('RGB')
img_a = self.transforms(img_a)
# grayscale
img_a = torch.mean(img_a, dim=0, keepdim=True)
# pick top left corner
x = random.randint(self.rho,
self.img_size - self.rho - self.patch_size)
y = random.randint(self.rho,
self.img_size - self.rho - self.patch_size)
corners = torch.tensor([
[x, y],
[x + self.patch_size, y],
[x + self.patch_size, y + self.patch_size],
[x, y + self.patch_size],
])
delta = torch.randint_like(corners, -self.rho, self.rho)
perturbed_corners = corners + delta
try:
# compute homography from points
h = kornia.get_perspective_transform(
corners.unsqueeze(0).float(),
perturbed_corners.unsqueeze(0).float())
h_inv = torch.inverse(h)
# apply homography to single img
img_b = kornia.warp_perspective(img_a.unsqueeze(0), h_inv,
(self.img_size, self.img_size))[0]
except:
# either matrix could not be solved or inverted
# this will show up as None, so use safe_collate in train.py
return None
patch_a = img_a[:, y:y + self.patch_size, x:x + self.patch_size]
patch_b = img_b[:, y:y + self.patch_size, x:x + self.patch_size]
return img_a, patch_a, patch_b, corners.float(), delta.float()
def forward(self, data):
corners = data['corners']
delta = data['delta']
img_a = data['img_a']
patch_b = data['patch_b']
corners_hat = corners + delta
# in order to apply transform and center crop,
# subtract points by top-left corner (corners[N, 0])
corners = corners - corners[:, 0].view(-1, 1, 2)
h = kornia.get_perspective_transform(corners, corners_hat)
h_inv = torch.inverse(h)
patch_b_hat = kornia.warp_perspective(
img_a, h_inv, (self.patch_size, self.patch_size))
return F.l1_loss(patch_b_hat, patch_b)
def warp_patch(patch_tensor, img_tensor, bbox_dest):
'''
Apply the patch to images.
Input: patch_tensor : Tensor ([B, ]3, h0, w0)
img_tensor: Tensor(B, 3, H, W)
bbox_dest: Tensor(B, 4)
Output: Tensor (B, 3, H, W)
'''
B = bbox_dest.shape[0]
x, y, w, h = 0, 0, patch_tensor.shape[-1], patch_tensor.shape[-2]
points_src = torch.FloatTensor([[
[x, y],
[x + w - 1, y],
[x, y + h - 1],
[x + w - 1, y + h - 1],
]])
points_src = points_src.expand(bbox_dest.shape[0], 4,
2).to(img_tensor.device)
xy = torch.stack((bbox_dest[:, 0], bbox_dest[:, 1]), dim=1)
x2y = torch.stack((bbox_dest[:, 0] + bbox_dest[:, 2] - 1, bbox_dest[:, 1]),
dim=1)
xy2 = torch.stack((bbox_dest[:, 0], bbox_dest[:, 1] + bbox_dest[:, 3] - 1),
dim=1)
x2y2 = torch.stack((bbox_dest[:, 0] + bbox_dest[:, 2] - 1,
bbox_dest[:, 1] + bbox_dest[:, 3] - 1),
dim=1)
points_dst = torch.stack([xy, x2y, xy2, x2y2], dim=1).to(torch.float32)
M = kornia.get_perspective_transform(points_src,
points_dst).to(img_tensor.device)
if len(patch_tensor.shape) == 3:
patch_tensor = patch_tensor.expand(B, -1, -1, -1)
patch_warped = kornia.warp_perspective(
patch_tensor, M, (img_tensor.shape[2], img_tensor.shape[3]))
# res_img = torch.where((patch_warped==0), img_tensor, patch_warped)
return patch_warped
def main(args):
# model = HomographyModel.load_from_checkpoint(args.checkpoint)
model_dir = 'lightning_logs/version*'
model_dir_list = sorted(glob.glob(model_dir))
model_dir = model_dir_list[-1]
model_path = osp.join(model_dir, "checkpoints", "*.ckpt")
model_path_list = sorted(glob.glob(model_path))
if len(model_path_list) > 0:
model_path = model_path_list[-1]
model = HomographyModel.load_from_checkpoint(model_path)
print(model_path)
print("model loaded.")
else:
raise ValueError(f'No load model!') #raise Error
model.eval()
test_set = SyntheticDataset(args.test_path,
rho=args.rho,
filetype=args.filetype)
#clear last output
last_output = "figures/*"
os.system("rm " + last_output)
print('clear last ok.')
for i in range(args.n):
img_a, patch_a, patch_b, corners, delta = test_set[i]
tensors_to_gif(patch_a, patch_b, f"figures/input_{i}.gif")
patch_a = patch_a.unsqueeze(0)
patch_b = patch_b.unsqueeze(0)
corners = corners.unsqueeze(0)
corners = corners - corners[:, 0].view(-1, 1, 2)
delta_hat = model(patch_a, patch_b)
corners_hat = corners + delta_hat
h = kornia.get_perspective_transform(corners, corners_hat)
h_inv = torch.inverse(h)
patch_b_hat = kornia.warp_perspective(patch_a, h_inv, (128, 128))
tensors_to_gif(patch_b_hat[0], patch_b[0], f"figures/output_{i}.gif")
def main(args):
model = HomographyModel.load_from_checkpoint(args.checkpoint)
model.eval()
test_set = SyntheticDataset(args.test_path,
rho=args.rho,
filetype=args.filetype)
for i in range(args.n):
img_a, patch_a, patch_b, corners, delta = test_set[i]
tensors_to_gif(patch_a, patch_b, f"figures/input_{i}.gif")
patch_a = patch_a.unsqueeze(0)
patch_b = patch_b.unsqueeze(0)
corners = corners.unsqueeze(0)
corners = corners - corners[:, 0].view(-1, 1, 2)
delta_hat = model(patch_a, patch_b)
corners_hat = corners + delta_hat
h = kornia.get_perspective_transform(corners, corners_hat)
h_inv = torch.inverse(h)
patch_b_hat = kornia.warp_perspective(patch_a, h_inv, (128, 128))
tensors_to_gif(patch_b_hat[0], patch_b[0], f"figures/output_{i}.gif")
def get_img_crop(img_batch, target_kps, name, spatial_size, box_factor):
"""
Crops an image around a set of keypoints or a single keypoint
"""
# leave out batch dimension
if name == "head":
# kepyoints are assumed to be (rshoulder, lshoulder, head)
assert target_kps.shape[1] == 3
necks = 0.5 * (target_kps[:, 0] + target_kps[:, 1])
necks_to_noses = target_kps[:, 2] - necks
up_heads = necks + 2 * necks_to_noses
segments = necks - up_heads
normals = torch.stack([-segments[:, 1], segments[:, 0]], dim=-1)
alpha = 0.5
a = up_heads + alpha * normals
b = up_heads - alpha * normals
c = necks - alpha * normals
d = necks + alpha * normals
elif name == "hand":
assert target_kps.shape[1] == 2
# target keypoints are assumed to be (wrist,hand) --> segments point from wrist to hand
segments = target_kps[:, 1] - target_kps[:, 0]
# s_norm = torch.norm(segments, dim=1, p=2).unsqueeze(dim=1)
# # normals are rotated in mathematical positive direction
normals = torch.stack([-segments[:, 1], segments[:, 0]], dim=-1)
# n_norm = torch.norm(normals, dim=1, p=2).unsqueeze(dim=1)
# # bisector of segments and vectors
# bisectors = torch.mul(normals,s_norm) + torch.mul(segments,n_norm)
# # should have same norm as normals
# bisectors = torch.div(bisectors, 2 * s_norm)
# alpha = 0.5
# rot_90 = torch.stack([-bisectors[:, 1], bisectors[:, 0]], dim=-1)
# a = target_kps[:,0] - alpha * bisectors
# b = target_kps[:,0] + alpha * rot_90
# c = target_kps[:,1] + alpha * bisectors
# d = target_kps[:,1] - alpha * rot_90
alpha = 1.0
beta = 0.25
a = target_kps[:, 0] + alpha * normals - beta * segments
b = target_kps[:, 0] - alpha * normals - beta * segments
c = target_kps[:, 1] - alpha * normals + beta * segments
d = target_kps[:, 1] + alpha * normals + beta * segments
else:
raise ValueError("Invalid ids or keypoints.")
src_windows = torch.stack([a, b, c, d], dim=1).to(torch.float)
dev = src_windows.get_device() if src_windows.get_device() > 0 else "cpu"
dst = torch.tensor(
[
[0.0, 0.0],
[0.0, spatial_size // (2**box_factor) - 1.0],
[
spatial_size // (2**box_factor) - 1.0,
spatial_size // (2**box_factor) - 1.0,
],
[spatial_size // (2**box_factor) - 1.0, 0.0],
],
dtype=torch.float,
device=dev,
)
dst_windows = torch.stack([dst] * src_windows.shape[0],
dim=0).to(torch.float)
M = kornia.get_perspective_transform(src_windows, dst_windows)
if dev != "cpu":
with torch.cuda.device(dev):
crop = kornia.warp_perspective(
img_batch,
M,
dsize=(
spatial_size // (2**box_factor),
spatial_size // (2**box_factor),
),
)
else:
crop = kornia.warp_perspective(
img_batch,
M,
dsize=(
spatial_size // (2**box_factor),
spatial_size // (2**box_factor),
),
)
return crop
# convert to torch tensor
data: torch.tensor = kornia.image_to_tensor(img, keepdim=False) # BxCxHxW
# the source points are the region to crop corners
points_src = torch.tensor([[
[125., 150.], [562., 40.], [562., 282.], [54., 328.],
]])
# the destination points are the image vertexes
h, w = 64, 128 # destination size
points_dst = torch.tensor([[
[0., 0.], [w - 1., 0.], [w - 1., h - 1.], [0., h - 1.],
]])
# compute perspective transform
M: torch.tensor = kornia.get_perspective_transform(points_src, points_dst)
# warp the original image by the found transform
data_warp: torch.tensor = kornia.warp_perspective(data.float(), M, dsize=(h, w))
# convert back to numpy
img_warp: np.ndarray = kornia.tensor_to_image(data_warp.byte())
# draw points into original image
for i in range(4):
center = tuple(points_src[0, i].long().numpy())
img = cv2.circle(img.copy(), center, 5, (0, 255, 0), -1)
# create the plot
fig, axs = plt.subplots(1, 2, figsize=(16, 10))
axs = axs.ravel()
def test_epoch(epoch, test_dataloader, modelhomo, model, model2, criterion):
modelhomo.eval() #h**o
model.eval()
model2.eval()
device = next(model.parameters()).device
loss = AverageMeter()
bpp_loss = AverageMeter()
mse_loss = AverageMeter()
aux_loss = AverageMeter()
ssim_loss = AverageMeter()
ssim_loss1 = AverageMeter()
ssim_loss2 = AverageMeter()
psnr1 = AverageMeter()
psnr2 = AverageMeter()
bpp1 = AverageMeter()
bpp2 = AverageMeter()
with torch.no_grad():
for d in test_dataloader:
d1 = d[0].to(device)
d2 = d[1].to(device)
# h_matrix = d[2].to(device)
#通过h**o获取h_matrix 注意要加逆变换
# print(len(d))
homo_img1 = d[-3].to(device)
homo_img2 = d[-2].to(device)
homo_corners = d[-1].to(device)
homo_corners = homo_corners - homo_corners[:, 0].view(-1, 1, 2)
delta_hat = modelhomo(homo_img1, homo_img2)
homo_corners_hat = homo_corners + delta_hat
h = kornia.get_perspective_transform(homo_corners,
homo_corners_hat)
h_matrix = torch.inverse(h)
h_matrix = h_adjust(d1.shape[-2], d1.shape[-1], pic_size, pic_size,
h_matrix)
out_net = model(d1, d2, h_matrix)
out_net2 = model2(out_net['x1_hat'], out_net['x2_hat'], h_matrix)
# out_net['x1_hat'] = out_net2['x1_hat']
# out_net['x2_hat'] = out_net2['x2_hat']
out_criterion1 = criterion(out_net, d1, d2, need_bpp=True)
out_criterion = criterion(out_net2, d1, d2)
# out_criterion = criterion(out_net, d1,d2)
aux_loss.update(model.aux_loss())
# bpp_loss.update(out_criterion['bpp_loss'])
loss.update(out_criterion['loss'])
mse_loss.update(out_criterion['mse_loss'])
ssim_loss.update(out_criterion['ms_ssim']) # 已除2
ssim_loss1.update(out_criterion['ms_ssim1']) # 已除2
ssim_loss2.update(out_criterion['ms_ssim2']) # 已除2
psnr1.update(out_criterion['psnr1'])
psnr2.update(out_criterion['psnr2'])
bpp1.update(out_criterion1['bpp1'])
bpp2.update(out_criterion1['bpp2'])
ssim_val = out_criterion['ms_ssim']
psnr1_val = out_criterion['psnr1']
psnr2_val = out_criterion['psnr2']
bpp1_val = out_criterion1['bpp1']
bpp2_val = out_criterion1['bpp2']
print_context = (
str(d[3]).split("'")[1] +
f'\tPSNR (dB): {(psnr1_val + psnr2_val) / 2:.3f} |' # 平均一张图的PSNR
f'\tBPP: {(bpp1_val+bpp2_val)/2:.4f} |'
f'\tMS-SSIM: {ssim_val:.4f} |' # 已除2，相加时候便除了2
f'\tPSNR1: {psnr1_val:.3f} |'
f'\tPSNR2: {psnr2_val:.3f} |'
f'\tBPP1: {bpp1_val:.4f} |'
f'\tBPP2: {bpp2_val:.4f}\n')
out_root_path_file.write(print_context)
print(print_context)
##save pic
save_pic(out_net2['x1_hat'],
osp.join(left_save_path,
str(d[3]).split("'")[1]))
save_pic(out_net2['x2_hat'],
osp.join(right_save_path,
str(d[3]).split("'")[1]))
# x1_warp = kornia.warp_perspective(out_net2['x1_hat'], h_matrix, (out_net2['x1_hat'].size()[-2], out_net2['x1_hat'].size()[-1]))
# save_pic(x1_warp, "demo_warp.png")
# raise ValueError("stop test warp save")
print(str(d[3]).split("'")[1])
####
out_root_path_file.close()
print(f'Test epoch {epoch}: Average losses:'
f'\tTime: {time.strftime("%Y-%m-%d %H:%M:%S")} |'
f'\tLoss: {loss.val:.3f} |'
f'\tMSE loss: {mse_loss.val:.4f} |'
f'\tPSNR (dB): {(psnr1.val+psnr2.val)/2:.3f} |' #平均一张图的PSNR
f'\tMS-SSIM: {ssim_loss.val:.4f} |' #已除2，相加时候便除了2
f'\tMS-SSIM1: {ssim_loss1.val:.4f} |'
f'\tMS-SSIM2: {ssim_loss2.val:.4f} |'
f'\tPSNR1: {psnr1.val:.3f} |'
f'\tPSNR2: {psnr2.val:.3f} \n')
return loss.val
def _run_encoder(self, batch):
src, src_lengths = batch.src if isinstance(batch.src, tuple) \
else (batch.src, None)
# import pdb
# pdb.set_trace()
import skimage.io
notusestn = False
if self.use_stn and notusestn:
# input images are downsampled before being fed into stn_head.
# stn_input = F.interpolate(src, self.tps_inputsize, mode='bilinear', align_corners=True)
h, w = src.shape[2:4]
padW = w // 6
padH = h // 6
src = torch.nn.functional.pad(src, (padW, padW, padH, padH),
mode='constant',
value=1)
stn_input = kornia.geometry.transform.resize(
src, self.tps_inputsize, interpolation='bicubic')
stn_input = stn_input.clamp(0, 1)
stn_input = (stn_input - 0.5) / 0.5
_, ctrl_points = self.model.model_stn_head(stn_input)
self.cnt += 1
skimage.io.imsave(
str(self.cnt) + '-1.jpg', src[0].squeeze().cpu().numpy())
# src, _ = self.model.model_tps(src, ctrl_points)
############### using kornia ##################
dst_h, dst_w = src.shape[2:4]
points_dst = torch.tensor([[
[padW, padH],
[dst_w - 1 - padW, padH],
[dst_w - 1 - padW, dst_h - 1 - padH],
[padW, dst_h - 1 - padH],
]])
points_dst = points_dst.float()
points_dst = points_dst.repeat(src.shape[0], 1, 1)
points_dst = points_dst.cuda()
ctrl_points2 = ctrl_points.clone()
ctrl_points[:, 1, :] = ctrl_points2[:, 3, :]
ctrl_points[:, 3, :] = ctrl_points2[:, 1, :]
ctrl_points[:, :, 0] = ctrl_points[:, :, 0] / 192 * dst_w
ctrl_points[:, :, 1] = ctrl_points[:, :, 1] / 48 * dst_h
M = kornia.get_perspective_transform(ctrl_points, points_dst)
src = kornia.warp_perspective(src,
M,
dsize=(dst_h, dst_w),
border_mode='border')
skimage.io.imsave(
str(self.cnt) + '-2.jpg', src[0].squeeze().cpu().numpy())
enc_states, memory_bank, src_lengths = self.model.encoder(
src, src_lengths)
if src_lengths is None:
assert not isinstance(memory_bank, tuple), \
'Ensemble decoding only supported for text data'
src_lengths = torch.Tensor(batch.batch_size) \
.type_as(memory_bank) \
.long() \
.fill_(memory_bank.size(0))
return src, enc_states, memory_bank, src_lengths
def test_epoch(epoch, test_dataloader, modelhomo, model, model2, criterion):
modelhomo.eval() #h**o
model.eval()
model2.eval()
device = next(model.parameters()).device
loss = AverageMeter()
bpp_loss = AverageMeter()
mse_loss = AverageMeter()
aux_loss = AverageMeter()
ssim_loss = AverageMeter()
ssim_loss1 = AverageMeter()
ssim_loss2 = AverageMeter()
psnr1 = AverageMeter()
psnr2 = AverageMeter()
bpp1 = AverageMeter()
bpp2 = AverageMeter()
with torch.no_grad():
for d in test_dataloader:
d1 = d[0].to(device)
d2 = d[1].to(device)
# h_matrix = d[2].to(device)
#通过h**o获取h_matrix 注意要加逆变换
# print(len(d))
homo_img1 = d[-3].to(device)
homo_img2 = d[-2].to(device)
homo_corners = d[-1].to(device)
homo_corners = homo_corners - homo_corners[:, 0].view(-1, 1, 2)
delta_hat = modelhomo(homo_img1, homo_img2)
homo_corners_hat = homo_corners + delta_hat
h = kornia.get_perspective_transform(homo_corners,
homo_corners_hat)
h_matrix = torch.inverse(h)
h_matrix = h_adjust(d1.shape[-2], d1.shape[-1], pic_size, pic_size,
h_matrix)
out_net = model(d1, d2, h_matrix)
out_net2 = model2(out_net['x1_hat'], out_net['x2_hat'], h_matrix)
# out_net['x1_hat'] = out_net2['x1_hat']
# out_net['x2_hat'] = out_net2['x2_hat']
out_criterion = criterion(out_net2, d1, d2)
#+ 计算bpp
bpp_out_criterion = criterion(out_net, d1, d2, kind=1)
aux_loss.update(model.aux_loss())
loss.update(out_criterion['loss'])
mse_loss.update(out_criterion['mse_loss'])
ssim_loss.update(out_criterion['ms_ssim']) # 已除2
ssim_loss1.update(out_criterion['ms_ssim1']) # 已除2
ssim_loss2.update(out_criterion['ms_ssim2']) # 已除2
psnr1.update(out_criterion['psnr1'])
psnr2.update(out_criterion['psnr2'])
bpp_loss.update(bpp_out_criterion['bpp_loss'])
bpp1.update(bpp_out_criterion['bpp1'])
bpp2.update(bpp_out_criterion['bpp2'])
print(f'Test epoch {epoch}: Average losses:'
f'\tTime: {time.strftime("%Y-%m-%d %H:%M:%S")} |'
f'\tLoss: {loss.val:.3f} |'
f'\tBPP: {bpp_loss.val:.3f} |'
f'\tBPP1: {bpp1.val:.3f} |'
f'\tBPP2: {bpp2.val:.3f} |'
f'\tMSE loss: {mse_loss.val:.4f} |'
f'\tPSNR (dB): {(psnr1.val+psnr2.val)/2:.3f} |' #平均一张图的PSNR
f'\tMS-SSIM: {ssim_loss.val:.4f} |' #已除2，相加时候便除了2
f'\tMS-SSIM1: {ssim_loss1.val:.4f} |'
f'\tMS-SSIM2: {ssim_loss2.val:.4f} |'
f'\tPSNR1: {psnr1.val:.3f} |'
f'\tPSNR2: {psnr2.val:.3f} \n')
return loss.val
def forward(self, src, tgt, lengths, bptt=False, with_align=False):
"""Forward propagate a `src` and `tgt` pair for training.
Possible initialized with a beginning decoder state.
Args:
src (Tensor): A source sequence passed to encoder.
typically for inputs this will be a padded `LongTensor`
of size ``(len, batch, features)``. However, may be an
image or other generic input depending on encoder.
tgt (LongTensor): A target sequence passed to decoder.
Size ``(tgt_len, batch, features)``.
lengths(LongTensor): The src lengths, pre-padding ``(batch,)``.
bptt (Boolean): A flag indicating if truncated bptt is set.
If reset then init_state
with_align (Boolean): A flag indicating whether output alignment,
Only valid for transformer decoder.
Returns:
(FloatTensor, dict[str, FloatTensor]):
* decoder output ``(tgt_len, batch, hidden)``
* dictionary attention dists of ``(tgt_len, batch, src_len)``
"""
dec_in = tgt[:-1] # exclude last target from inputs
if self.rectifier is not None:
# input images are downsampled before being fed into stn_head.
stn_input = kornia.geometry.transform.resize(
src, self.tps_inputsize, interpolation='bicubic')
stn_input = stn_input.clamp(0, 1)
stn_input = (stn_input - 0.5) / 0.5
_, ctrl_points = self.model_stn_head(stn_input)
# self.cnt += 1
# skimage.io.imsave(str(self.cnt)+'-1.jpg', src[0].squeeze().cpu().numpy())
# src, _ = self.model_tps(src, ctrl_points)
############### using kornia ##################
dst_h, dst_w = src.shape[2:4]
points_dst = torch.tensor([[
[0., 0.],
[dst_w - 1., 0.],
[dst_w - 1., dst_h - 1.],
[0., dst_h - 1.],
]])
points_dst = points_dst.repeat(src.shape[0], 1, 1)
points_dst = points_dst.cuda()
ctrl_points2 = ctrl_points.clone()
ctrl_points[:, 1, :] = ctrl_points2[:, 3, :]
ctrl_points[:, 3, :] = ctrl_points2[:, 1, :]
ctrl_points[:, :, 0] = ctrl_points[:, :, 0] / 192 * dst_w
ctrl_points[:, :, 1] = ctrl_points[:, :, 1] / 48 * dst_h
M = kornia.get_perspective_transform(ctrl_points, points_dst)
src = kornia.warp_perspective(src,
M,
dsize=(dst_h, dst_w),
border_mode='border')
# skimage.io.imsave(str(self.cnt)+'-kornia.jpg', src[0].squeeze().cpu().detach().numpy())
enc_state, memory_bank, lengths = self.encoder(src, lengths)
if bptt is False:
self.decoder.init_state(src, memory_bank, enc_state)
dec_out, attns = self.decoder(dec_in,
memory_bank,
memory_lengths=lengths,
with_align=with_align)
return dec_out, attns
def forward(self,
adv_patch,
lab_batch,
height,
width,
do_rotate=True,
rand_loc=False,
scale_factor=0.25,
cls_label=1):
adv_patch = self.medianpooler(adv_patch.unsqueeze(0))
patch_size = adv_patch.size(-1)
# Determine size of padding
pad_width = (width - adv_patch.size(-1)) / 2
pad_height = (height - adv_patch.size(-2)) / 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)
# perspective transformations
dis_scale = (lab_batch.size(0) * lab_batch.size(1))
distortion = torch.empty(dis_scale).uniform_(0, self.distortion_max)
adv_height = adv_batch.size(-1)
adv_width = adv_batch.size(-2)
# tps transformation
if self.augment:
theta, dst = self.get_tps_thetas(dis_scale)
img = adv_batch.clone()
img = img.view(-1, 3, adv_width, adv_height)
grid = tps.torch.tps_grid(theta.cuda(), dst.cuda(),
(img.size(0), 1, adv_width, adv_height))
adv_batch = F.grid_sample(img, grid.cuda(), padding_mode='border')
adv_batch = adv_batch.view(lab_batch.size(0), lab_batch.size(1), 3,
adv_width, adv_height)
# perpective transformations with random distortion scales
start_end = [torch.tensor(self.get_perspective_params(adv_width, adv_height, x), \
dtype=torch.float).unsqueeze(0) for x in distortion]
start_end = torch.cat(start_end, 0)
start_points = start_end[:, 0, :, :].squeeze()
end_points = start_end[:, 1, :, :].squeeze()
if dis_scale == 1:
start_points = start_points.unsqueeze(0)
end_points = end_points.unsqueeze(0)
try:
M = kornia.get_perspective_transform(start_points, end_points)
img = adv_batch.clone()
img = img.view(-1, 3, adv_width, adv_height)
adv_batch = kornia.warp_perspective(img,
M.cuda(),
dsize=(adv_width, adv_height))
adv_batch = adv_batch.view(lab_batch.size(0), lab_batch.size(1), 3,
adv_width, adv_height)
except:
print('hihi')
cls_ids = torch.narrow(lab_batch, 2, 0, 1)
cls_mask = 1.0 * (cls_ids.expand(-1, -1, 3) != cls_label)
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.float()
mypad = nn.ConstantPad2d((int(pad_width + 0.5), int(pad_width),
int(pad_height + 0.5), int(pad_height)), 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.FloatTensor(anglesize).uniform_(
self.minangle, self.maxangle)
else:
angle = torch.FloatTensor(anglesize).fill_(0)
angle = angle.cuda()
# Resizes and rotates
batch_size = torch.Size((lab_batch.size(0), lab_batch.size(1)))
lab_batch_scaled = torch.zeros(lab_batch.size())
lab_batch_scaled[:, :, 1] = lab_batch[:, :, 1] * width
lab_batch_scaled[:, :, 2] = lab_batch[:, :, 2] * height
lab_batch_scaled[:, :, 3] = lab_batch[:, :, 3] * width
lab_batch_scaled[:, :, 4] = lab_batch[:, :, 4] * height
# roughly estimate the size and compute the scale
target_size = scale_factor * torch.sqrt(
(lab_batch_scaled[:, :, 3])**2 + (lab_batch_scaled[:, :, 4])**2)
target_x = lab_batch[:, :, 1].view(
np.prod(batch_size)) # (batch_size, num_objects)
target_y = lab_batch[:, :, 2].view(np.prod(batch_size))
# target_y = target_y - 0.0
# shift a bit from the center
targetoff_x = lab_batch[:, :, 3].view(np.prod(batch_size))
targetoff_y = lab_batch[:, :, 4].view(np.prod(batch_size))
off_y = (torch.FloatTensor(targetoff_y.size()).uniform_(
self.sliding, 0))
off_y = targetoff_y * off_y.cuda()
target_y = target_y + off_y
scale = target_size / patch_size
scale = scale.view(anglesize)
scale = scale.cuda()
s = adv_batch.size()
# print(adv_batch.size())
adv_batch = adv_batch.view(-1, 3, height, width)
msk_batch = msk_batch.view_as(adv_batch)
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.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
theta = theta.cuda()
grid = F.affine_grid(theta, adv_batch.shape)
adv_batch = F.grid_sample(adv_batch, grid)
msk_batch = F.grid_sample(msk_batch, grid)
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])
adv_batch = adv_batch * msk_batch
return adv_batch
