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))
