def stitch_pair( self, images_left: torch.Tensor, images_right: torch.Tensor, mask_left: Optional[torch.Tensor] = None, mask_right: Optional[torch.Tensor] = None ) -> Tuple[torch.Tensor, torch.Tensor]: # Compute the transformed images input_dict: Dict[str, torch.Tensor] = self.preprocess( images_left, images_right) out_shape: Tuple[int, int] = (images_left.shape[-2], images_left.shape[-1] + images_right.shape[-1]) correspondences: dict = self.on_matcher(input_dict) h**o: torch.Tensor = self.estimate_transform(**correspondences) src_img = warp_perspective(images_right, h**o, out_shape) dst_img = torch.cat( [images_left, torch.zeros_like(images_right)], dim=-1) # Compute the transformed masks if mask_left is None: mask_left = torch.ones_like(images_left) if mask_right is None: mask_right = torch.ones_like(images_right) # 'nearest' to ensure no floating points in the mask src_mask = warp_perspective(mask_right, h**o, out_shape, mode='nearest') dst_mask = torch.cat([mask_left, torch.zeros_like(mask_right)], dim=-1) return self.blend_image(src_img, dst_img, src_mask), (dst_mask + src_mask).bool().to( src_mask.dtype)
def training_step(self, batch, batch_nb): if self.supervised or len(batch) == 2: x, y = batch y_hat = self.forward(x) l = nn.BCEWithLogitsLoss() loss = 10 * l(y_hat.squeeze(1), y) # fig, ax= plt.subplots(1,x.shape[0]) # for i in range(x.shape[0]): # img1 = (x[i,0,...].cpu().detach().numpy()) # gt1 = (y[i,...].cpu().detach().numpy()) # ax[i].imshow(img1) # ax[i].imshow(gt1,alpha=0.3) # self.logger.experiment.log_image( # 'train images', # fig, # description='x1,x2') else: #angle=d, translate=t, scale=sc, shear=sh x1, x2, affine_matrix = batch y1 = self.forward(x1) y1_transformed = K.warp_perspective(y1, affine_matrix.squeeze(0), y1[0, 0, ...].shape).squeeze(1) # y1_transformed=kornia.compute_affine_transformation(y1,affine_matrix.squeeze(0)) y_teach1 = 1. * (self.teacher.forward(x1)) # y_teach1=torch.from_numpy(np.array(Fvision.affine(Fvision.to_pil_image(y_teach1.squeeze(1).cpu()),angle=d, translate=t, scale=sc, shear=sh))).cuda() y_teach2 = K.warp_perspective(y_teach1, affine_matrix.squeeze(0), y_teach1[0, 0, ...].shape) # y_teach2=kornia.compute_affine_transformation(y_teach1,affine_matrix.squeeze(0)) y_teach1 = 1. * (y_teach1) y2 = self.forward(x2) loss = self.criterion( y2.squeeze(1), y_teach2.squeeze(1) ) #+self.criterion(y1.squeeze(1),y_teach1.squeeze(1))+self.criterion(y1_transformed.squeeze(1),1.*(y2>0.5).squeeze(1)) # self.teacher.load_state_dict(teacher_params) # print(len(teacher_params.keys())) # epo=self.current_epoch # with open(f"models_weights/epoch_{epo}.json",'w+') as f: # json.dump((str(self.teacher.state_dict())),f) # img1 = (x1[0,0,...].cpu().detach().numpy()) # img2= (x1[0,0,...].cpu().detach().numpy()) # gt1 = (y1[0,0,...].cpu().detach().numpy()>0.5) # gt2= (y_teach1[0,0,...].cpu().detach().numpy()>0.5) # fig, (ax1,ax2)= plt.subplots(1,2) # ax1.imshow(img1) # ax1.imshow(gt1,alpha=0.3) # ax2.imshow(img2) # ax2.imshow(gt2,alpha=0.3) # self.logger.experiment.log_image( # 'train images', # fig, # description='x1,x2') # x,y=batch # y_hat = self.forward(x) # loss = self.criterion(y_hat.squeeze(1),y) self.logger.experiment.log_metric('train_loss', loss) return loss
def track_next_frame(self, x: torch.Tensor) -> Tuple[torch.Tensor, bool]: """The frame `x` is prewarped according to the previous frame homography, matched with fast_matcher verified with ransac.""" if self.previous_homography is not None: # mypy, shut up Hwarp = self.previous_homography.clone()[None] # make a bit of border for safety Hwarp[:, 0:2, 0:2] = Hwarp[:, 0:2, 0:2] / 0.8 Hwarp[:, 0:2, 2] -= 10.0 Hinv = torch.inverse(Hwarp) h, w = self.target.shape[2:] frame_warped = warp_perspective(x, Hinv, (h, w)) input_dict: Dict[str, torch.Tensor] = { "image0": self.target, "image1": frame_warped } for k, v in self.target_fast_representation.items(): input_dict[f'{k}0'] = v match_dict = self.fast_matcher(input_dict) keypoints0 = match_dict['keypoints0'][match_dict['batch_indexes'] == 0] keypoints1 = match_dict['keypoints1'][match_dict['batch_indexes'] == 0] keypoints1 = transform_points(Hwarp, keypoints1) if len(keypoints0) < self.minimum_inliers_num: self.reset_tracking() return self.no_match() H, inliers = self.ransac(keypoints0, keypoints1) if inliers.sum().item() < self.minimum_inliers_num: self.reset_tracking() return self.no_match() self.previous_homography = H.clone() return H, True
def apply_transform(self, input: Tensor, params: Dict[str, Tensor], transform: Optional[Tensor] = None) -> Tensor: _, _, height, width = input.shape transform = cast(Tensor, transform) return warp_perspective( input, transform, (height, width), mode=self.flags["resample"].name.lower(), align_corners=self.flags["align_corners"], )
def apply_perspective(input: torch.Tensor, params: Dict[str, torch.Tensor], return_transform: bool = False) -> UnionType: r"""Perform perspective transform of the given torch.Tensor or batch of tensors. Args: input (torch.Tensor): Tensor to be transformed with shape BxCxHxW. start_points (torch.Tensor): Tensor containing [top-left, top-right, bottom-right, bottom-left] of the orignal image with shape Bx4x2. end_points (torch.Tensor): Tensor containing [top-left, top-right, bottom-right, bottom-left] of the transformed image with shape Bx4x2. return_transform (bool): if ``True`` return the matrix describing the transformation applied to each. Default: False. Returns: torch.Tensor: Perspectively transformed tensor. """ if not torch.is_tensor(input): raise TypeError(f"Input type is not a torch.Tensor. Got {type(input)}") device: torch.device = input.device dtype: torch.dtype = input.dtype # arrange input data x_data: torch.Tensor = input.view(-1, *input.shape[-3:]) batch_size, _, height, width = x_data.shape # compute the homography between the input points transform: torch.Tensor = get_perspective_transform( params['start_points'], params['end_points']).to(device, dtype) out_data: torch.Tensor = x_data.clone() # process valid samples mask = params['batch_prob'].to(device) # TODO: look for a workaround for this hack. In CUDA it fails when no elements found. if bool(mask.sum() > 0): # apply the computed transform height, width = x_data.shape[-2:] out_data[mask] = warp_perspective(x_data[mask], transform[mask], (height, width)) if return_transform: return out_data.view_as(input), transform return out_data.view_as(input)
def _apply_affine(input: torch.Tensor, params: Dict[str, torch.Tensor], return_transform: bool = False) -> UnionType: if not torch.is_tensor(input): raise TypeError(f"Input type is not a torch.Tensor. Got {type(input)}") r"""Random affine transformation of the image keeping center invariant Args: input (torch.Tensor): Tensor to be transformed with shape (*, C, H, W). degrees (float or tuple): Range of degrees to select from. If degrees is a number instead of sequence like (min, max), the range of degrees will be (-degrees, +degrees). Set to 0 to deactivate rotations. translate (tuple, optional): tuple of maximum absolute fraction for horizontal and vertical translations. For example translate=(a, b), then horizontal shift is randomly sampled in the range -img_width * a < dx < img_width * a and vertical shift is randomly sampled in the range -img_height * b < dy < img_height * b. Will not translate by default. scale (tuple, optional): scaling factor interval, e.g (a, b), then scale is randomly sampled from the range a <= scale <= b. Will keep original scale by default. shear (sequence or float, optional): Range of degrees to select from. If shear is a number, a shear parallel to the x axis in the range (-shear, +shear) will be applied. Else if shear is a tuple or list of 2 values a shear parallel to the x axis in the range (shear[0], shear[1]) will be applied. Else if shear is a tuple or list of 4 values, a x-axis shear in (shear[0], shear[1]) and y-axis shear in (shear[2], shear[3]) will be applied. Will not apply shear by default return_transform (bool): if ``True`` return the matrix describing the transformation applied to each. Default: False. mode (str): interpolation mode to calculate output values 'bilinear' | 'nearest'. Default: 'bilinear'. padding_mode (str): padding mode for outside grid values 'zeros' | 'border' | 'reflection'. Default: 'zeros'. """ if not torch.is_tensor(input): raise TypeError(f"Input type is not a torch.Tensor. Got {type(input)}") device: torch.device = input.device dtype: torch.dtype = input.dtype # arrange input data x_data: torch.Tensor = input.view(-1, *input.shape[-3:]) height, width = x_data.shape[-2:] transform: torch.Tensor = params['transform'].to(device, dtype) out_data: torch.Tensor = warp_perspective(x_data, transform, (height, width)) if return_transform: return out_data.view_as(input), transform return out_data.view_as(input)