def compute_loss(pred_pha, pred_fgr, pred_err, true_pha, true_fgr): true_err = torch.abs(pred_pha.detach() - true_pha) true_msk = true_pha != 0 return F.l1_loss(pred_pha, true_pha) + \ F.l1_loss(kornia.sobel(pred_pha), kornia.sobel(true_pha)) + \ F.l1_loss(pred_fgr * true_msk, true_fgr * true_msk) + \ F.mse_loss(pred_err, true_err)
def sobel_shift_loss(self, img, img_shifted): """Edge pairwise loss""" edge = K.sobel(utils.rgb2gray(img).unsqueeze(0)) # extract edges edge_shifted = K.sobel( utils.rgb2gray(img_shifted).unsqueeze(0)) # extract edges return torch.mean( (edge - edge_shifted)**2) # encourage edges to be the same
def compute_mattingbase_loss(pred_pha: torch.Tensor, pred_fgr: torch.Tensor, pred_err: torch.Tensor, true_pha: torch.Tensor, true_fgr: torch.Tensor): true_err = torch.abs(pred_pha.detach() - true_pha) true_msk = true_pha != 0 matting_loss = F.l1_loss(pred_pha, true_pha) + \ F.l1_loss(kornia.sobel(pred_pha), kornia.sobel(true_pha)) + \ F.l1_loss(pred_fgr * true_msk, true_fgr * true_msk) + \ F.mse_loss(pred_err, true_err) return matting_loss
def compute_loss(pred_pha_lg, pred_fgr_lg, pred_pha_sm, pred_fgr_sm, pred_err_sm, true_pha_lg, true_fgr_lg): true_pha_sm = kornia.resize(true_pha_lg, pred_pha_sm.shape[2:]) true_fgr_sm = kornia.resize(true_fgr_lg, pred_fgr_sm.shape[2:]) true_msk_lg = true_pha_lg != 0 true_msk_sm = true_pha_sm != 0 return F.l1_loss(pred_pha_lg, true_pha_lg) + \ F.l1_loss(pred_pha_sm, true_pha_sm) + \ F.l1_loss(kornia.sobel(pred_pha_lg), kornia.sobel(true_pha_lg)) + \ F.l1_loss(kornia.sobel(pred_pha_sm), kornia.sobel(true_pha_sm)) + \ F.l1_loss(pred_fgr_lg * true_msk_lg, true_fgr_lg * true_msk_lg) + \ F.l1_loss(pred_fgr_sm * true_msk_sm, true_fgr_sm * true_msk_sm) + \ F.mse_loss(kornia.resize(pred_err_sm, true_pha_lg.shape[2:]), kornia.resize(pred_pha_sm, true_pha_lg.shape[2:]).sub(true_pha_lg).abs())
def compute_loss(pred_pha: torch.Tensor, pred_fgr: torch.Tensor, pred_err: torch.Tensor, pred_depth: torch.Tensor, true_pha: torch.Tensor, true_fgr: torch.Tensor, true_depth: torch.Tensor): true_err = torch.abs(pred_pha.detach() - true_pha) true_msk = true_pha != 0 matting_loss = F.l1_loss(pred_pha, true_pha) + \ F.l1_loss(kornia.sobel(pred_pha), kornia.sobel(true_pha)) + \ F.l1_loss(pred_fgr * true_msk, true_fgr * true_msk) + \ F.mse_loss(pred_err, true_err) true_depth = true_depth.detach() depth_abs = (pred_depth - true_depth).abs() * (true_depth >= 0) c = depth_abs.max() / 5 mask = depth_abs <= c depth_loss = ((mask * depth_abs).sum() + ((~mask * depth_abs)**2 + c**2).sum() / (2 * c)) / (mask.shape[2] * mask.shape[3]) return 0 + depth_loss
def edge_aware_smoothness_reg(*, im, depth, flow, mask, depth_coeff=0., flow_coeff=0., mask_coeff=0.): """ Computes the total optical flow, depth, and mask edge-aware regualizer. The 2nd order derivates of depth are penalized Parameters: ----------- im: the batch of images (BxCxHxW) depth: depth values (BxHxW) flow: flow values (BxHxWx2) mask: mask values (BxKxHxW) Output: ------- Regularizer, scalar """ B, C, H, W = im.shape K = mask.shape[1] assert depth.shape == (B, H, W), f'depth is wrong shape {depth.shape}' assert flow.shape == (B, H, W, 2), f'flow is wrong shape {flow.shape}' assert mask.shape == (B, K, H, W), f'mask is wrong shape {mask.shape}' d_im = kornia.sobel(im) # BxCxHxW norm_d_im = torch.norm(d_im, p=1, dim=1) # BxHxW exp_norm = torch.exp(-norm_d_im) # BxHxW d_flow = kornia.sobel(flow) # BxHxWx2 flow_smooth_reg = torch.mean(exp_norm * torch.norm(flow, p=1, dim=-1)) lap_depth = kornia.laplacian(depth.unsqueeze(1), kernel_size=3) depth_smooth_reg = torch.mean(exp_norm * lap_depth) if K != 0: d_mask = kornia.sobel(mask) mask_smooth_reg = torch.mean(exp_norm * torch.norm(d_mask, p=1, dim=1)) else: mask_smooth_reg = 0. return flow_coeff * flow_smooth_reg + depth_coeff * depth_smooth_reg + mask_coeff * mask_smooth_reg
def forward(self, data: dict) -> list: # data["result"][0] = kn.normalize_min_max(data["result"][0]) edgeData = {"result": [], "gts": []} for inp in data["result"]: edgeData["result"].append(kn.normalize_min_max(kn.sobel(inp))) mean = edgeData["result"][-1].mean() std = edgeData["result"][-1].std() # edgeData["result"][-1][edgeData["result"][-1] <= mean + 1.2 * std] = 0 for inp in data["gts"]: edgeData["gts"].append(kn.normalize_min_max(kn.sobel(inp))) mean = edgeData["gts"][-1].mean() std = edgeData["gts"][-1].std() # edgeData["gts"][-1][edgeData["gts"][-1] <= mean + 1.2 * std] = 0 # make ir edges more stronger # mean = edgeData["gts"][0].mean() # std = edgeData["gts"][0].std() # edgeData["gts"][0][edgeData["gts"][0] > mean + 4 * std] = 1 edgeLoss = self.qLoss(edgeData) for i in range(len(edgeLoss)): loss = edgeLoss[i] loss = torch.pow(loss, 0.7) edgeLoss[i] = loss # there is bug in the following line # edgeLoss_tmp = [torch.pow(edgeLoss[i], torch.tensor(0.7)) for i in range(len(edgeLoss))] # debugging data["edges"] = edgeData # debugging # inplace 1 - overallQscore_abf for i in range(len(edgeLoss)): edgeLoss[i] = torch.add(edgeLoss[i], -1) edgeLoss[i] = torch.mul(edgeLoss[i], -1) return edgeLoss
def train_on_loader(self, dataloader): self.backbone.train() losses = [] timer_data, timer_model = utility.timer(), utility.timer() for batch, (lr, hr, idx_scale) in enumerate(tqdm(dataloader)): lr, hr = self.prepare(lr, hr) self.optimizer.zero_grad() sr = self.backbone(lr, 0) loss = F.l1_loss(sr, hr) * self.exp_dict["loss_l1_weight"] print(loss.item()) if self.exp_dict["loss_sobel_weight"] > 0: # https://kornia.readthedocs.io/en/latest/filters.html#kornia.filters.sobel # normalized is True by default sr_sobel = K.sobel(sr, normalized=True) hr_sobel = K.sobel(hr, normalized=True) loss_sobel = F.l1_loss( sr_sobel, hr_sobel) * self.exp_dict["loss_sobel_weight"] loss += loss_sobel print(loss_sobel.item()) loss.backward() losses.append(float(loss)) # if self.args.gclip > 0: # utils.clip_grad_value_( # self.model.parameters(), # self.args.gclip # ) self.optimizer.step() if (batch + 1) % 20 == 0: print("loss", float(loss)) return dict(train_loss=float(np.mean(losses)))
grads_x = grads[:, :, 0] grads_y = grads[:, :, 1] ################################# # Show first derivatives in x imshow(grads_x) ################################# # Show first derivatives in y imshow(grads_y) ################################# # Sobel Edges # Once with the gradients in the two directions we can computet the Sobel edges. # However, in kornia we already have it implemented. x_sobel: torch.Tensor = kornia.sobel(x_gray) imshow(x_sobel) ################################# # Compute the 2nd order derivates grads: torch.Tensor = kornia.spatial_gradient(x_gray, order=2) # BxCx2xHxW grads_x = grads[:, :, 0] grads_y = grads[:, :, 1] ################################# # Show second derivatives in x imshow(grads_x) ################################# # Show second derivatives in y imshow(grads_y)
def fit_on_edges(self, root_dir, num_epochs, batch_size, val_split=0.1, num_workers=4): image_datasets = datasets.ImageFolder(root_dir, self.data_transforms) image_datasets.classes = [ self.translate[i] for i in image_datasets.classes ] train_size = int((1 - val_split) * len(image_datasets)) train, valid = torch.utils.data.random_split( image_datasets, [train_size, len(image_datasets) - train_size]) dataloaders = { 'train': torch.utils.data.DataLoader(train, batch_size=batch_size, shuffle=True, num_workers=num_workers), 'val': torch.utils.data.DataLoader(valid, batch_size=batch_size, shuffle=True, num_workers=num_workers) } criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(self.parameters(), lr=0.0001) epoch_bar = tqdm.tqdm(range(num_epochs), desc='Epochs---', position=0) self.train_epoch_loss = [] self.val_epoch_loss = [] train_acc = 0. val_acc = 0. # self.to(self.device) for epoch in epoch_bar: self.train() pbar = tqdm.tqdm(dataloaders['train'], desc='Epoch= {}'.format(epoch + 1), position=1) running_loss = 0.0 val_loss = 0.0 total = 0 correct = 0 for inputs, labels in pbar: inputs = inputs.to(self.device) inputs = torch.stack([ kornia.sobel(inputs[i].view(1, 3, 224, 224)).view(3, 224, 224) for i in range(len(inputs)) ]) labels = labels.to(self.device) # zero the parameter gradients optimizer.zero_grad() labels_hat = self(inputs) _, predicted = torch.max(labels_hat.data, 1) loss = criterion(labels_hat, labels) total += labels.size(0) correct += (predicted == labels).sum().item() pbar.set_postfix(train_run_loss='{:.4f}'.format(loss.item())) running_loss += loss.item() loss.backward() optimizer.step() self.train_epoch_loss.append(running_loss / len(dataloaders['train'])) train_acc = 100 * correct / total epoch_bar.set_postfix(train_acc='{:.4f}'.format(train_acc), val_acc='{:.4f}'.format(val_acc)) # pbar = tqdm.tqdm(dataloaders['val'],desc='Epoch={}, train_epoch_loss= {}'.format(epoch+1,epoch_loss)) total = 0 correct = 0 with torch.no_grad(): self.eval() for inputs, labels in dataloaders['val']: inputs = inputs.to(self.device) inputs = torch.stack([ kornia.sobel(inputs[i].view(1, 3, 224, 224)).view(3, 224, 224) for i in range(len(inputs)) ]) labels = labels.to(self.device) labels_hat = self(inputs) _, predicted = torch.max(labels_hat.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() loss = criterion(labels_hat, labels) val_loss += loss.item() self.val_epoch_loss.append(val_loss / len(dataloaders['val'])) val_acc = 100 * correct / total epoch_bar.set_postfix(train_acc='{:.4f}'.format(train_acc), val_acc='{:.4f}'.format(val_acc))