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