from __future__ import division
import os, glob, shutil, math, random, json
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import basic
from utils import util
eps = 0.0000001
class SPixelLoss:
def __init__(self, psize=8, mpdist=False, gpu_no=0):
self.mpdist = mpdist
self.gpu_no = gpu_no
self.sp_size = psize
def __call__(self, data, epoch_no):
kernel_size = self.sp_size
#pos_weight = 0.003
prob = data['pred_prob']
labxy_feat = data['target_feat']
N,C,H,W = labxy_feat.shape
pooled_labxy = basic.poolfeat(labxy_feat, prob, kernel_size, kernel_size)
reconstr_feat = basic.upfeat(pooled_labxy, prob, kernel_size, kernel_size)
loss_map = reconstr_feat[:,:,:,:] - labxy_feat[:,:,:,:]
featLoss_idx = torch.norm(loss_map[:,:-2,:,:], p=2, dim=1).mean()
posLoss_idx = torch.norm(loss_map[:,-2:,:,:], p=2, dim=1).mean() / kernel_size
totalLoss_idx = 10*featLoss_idx + 0.003*posLoss_idx
return {'totalLoss':totalLoss_idx, 'featLoss':featLoss_idx, 'posLoss':posLoss_idx}
class AnchorColorProbLoss:
def __init__(self, hint2regress=False, enhanced=False, with_grad=False, mpdist=False, gpu_no=0):
self.mpdist = mpdist
self.gpu_no = gpu_no
self.hint2regress = hint2regress
self.enhanced = enhanced
self.with_grad = with_grad
self.rebalance_gradient = basic.RebalanceLoss.apply
self.entropy_loss = nn.CrossEntropyLoss(ignore_index=-1)
if self.enhanced:
self.VGGLoss = VGG19Loss(gpu_no=gpu_no, is_ddp=mpdist)
def _perceptual_loss(self, input_grays, input_colors, pred_colors):
input_RGBs = basic.lab2rgb(torch.cat([input_grays,input_colors], dim=1))
pred_RGBs = basic.lab2rgb(torch.cat([input_grays,pred_colors], dim=1))
## the output of "lab2rgb" just matches the input of "VGGLoss": [0,1]
return self.VGGLoss(input_RGBs, pred_RGBs)
def _laplace_gradient(self, pred_AB, target_AB):
N,C,H,W = pred_AB.shape
kernel = torch.tensor([[1, 1, 1], [1, -8, 1], [1, 1, 1]], device=pred_AB.get_device()).float()
kernel = kernel.view(1, 1, *kernel.size()).repeat(C,1,1,1)
grad_pred = F.conv2d(pred_AB, kernel, groups=C)
grad_trg = F.conv2d(target_AB, kernel, groups=C)
return l1_loss(grad_trg, grad_pred)
def __call__(self, data, epoch_no):
N,C,H,W = data['target_label'].shape
pal_probs = self.rebalance_gradient(data['pal_prob'], data['class_weight'])
#ref_probs = data['ref_prob']
pal_probs = pal_probs.permute(0,2,3,1).contiguous().view(N*H*W, -1)
gt_labels = data['target_label'].permute(0,2,3,1).contiguous().view(N*H*W, -1)
'''
igored_mask = data['empty_entries'].permute(0,2,3,1).contiguous().view(N*H*W, -1)
gt_labels[igored_mask] = -1
gt_labels = gt_probs.squeeze()
'''
palLoss_idx = self.entropy_loss(pal_probs, gt_labels.squeeze(dim=1))
if self.hint2regress:
ref_probs = data['ref_prob']
refLoss_idx = 50 * l2_loss(data['spix_color'], ref_probs)
else:
ref_probs = self.rebalance_gradient(data['ref_prob'], data['class_weight'])
ref_probs = ref_probs.permute(0,2,3,1).contiguous().view(N*H*W, -1)
refLoss_idx = self.entropy_loss(ref_probs, gt_labels.squeeze(dim=1))
reconLoss_idx = torch.zeros_like(palLoss_idx)
if self.enhanced:
scalar = 1.0 if self.hint2regress else 5.0
reconLoss_idx = scalar * self._perceptual_loss(data['input_gray'], data['pred_color'], data['input_color'])
if self.with_grad:
gradient_loss = self._laplace_gradient(data['pred_color'], data['input_color'])
reconLoss_idx += gradient_loss
totalLoss_idx = palLoss_idx + refLoss_idx + reconLoss_idx
#print("loss terms:", palLoss_idx.item(), refLoss_idx.item(), reconLoss_idx.item())
return {'totalLoss':totalLoss_idx, 'palLoss':palLoss_idx, 'refLoss':refLoss_idx, 'recLoss':reconLoss_idx}
def compute_affinity_pos_loss(prob_in, labxy_feat, pos_weight=0.003, kernel_size=16):
S = kernel_size
m = pos_weight
prob = prob_in.clone()
N,C,H,W = labxy_feat.shape
pooled_labxy = basic.poolfeat(labxy_feat, prob, kernel_size, kernel_size)
reconstr_feat = basic.upfeat(pooled_labxy, prob, kernel_size, kernel_size)
loss_map = reconstr_feat[:,:,:,:] - labxy_feat[:,:,:,:]
loss_feat = torch.norm(loss_map[:,:-2,:,:], p=2, dim=1).mean()
loss_pos = torch.norm(loss_map[:,-2:,:,:], p=2, dim=1).mean() * m / S
loss_affinity = loss_feat + loss_pos
return loss_affinity
def l2_loss(y_input, y_target, weight_map=None):
if weight_map is None:
return F.mse_loss(y_input, y_target)
else:
diff_map = torch.mean(torch.abs(y_input-y_target), dim=1, keepdim=True)
batch_dev = torch.sum(diff_map*diff_map*weight_map, dim=(1,2,3)) / (eps+torch.sum(weight_map, dim=(1,2,3)))
return batch_dev.mean()
def l1_loss(y_input, y_target, weight_map=None):
if weight_map is None:
return F.l1_loss(y_input, y_target)
else:
diff_map = torch.mean(torch.abs(y_input-y_target), dim=1, keepdim=True)
batch_dev = torch.sum(diff_map*weight_map, dim=(1,2,3)) / (eps+torch.sum(weight_map, dim=(1,2,3)))
return batch_dev.mean()
def masked_l1_loss(y_input, y_target, outlier_mask):
one = torch.tensor([1.0]).cuda(y_input.get_device())
weight_map = torch.where(outlier_mask, one * 0.0, one * 1.0)
return l1_loss(y_input, y_target, weight_map)
def huber_loss(y_input, y_target, delta=0.01):
mask = torch.zeros_like(y_input)
mann = torch.abs(y_input - y_target)
eucl = 0.5 * (mann**2)
mask[...] = mann < delta
loss = eucl * mask / delta + (mann - 0.5 * delta) * (1 - mask)
return torch.mean(loss)
## Perceptual loss that uses a pretrained VGG network
class VGG19Loss(nn.Module):
def __init__(self, feat_type='liu', gpu_no=0, is_ddp=False, requires_grad=False):
super(VGG19Loss, self).__init__()
os.environ['TORCH_HOME'] = '/apdcephfs/share_1290939/richardxia/Saved/Checkpoints/VGG19'
## data requirement: (N,C,H,W) in RGB format, [0,1] range, and resolution >= 224x224
self.mean = [0.485, 0.456, 0.406]
self.std = [0.229, 0.224, 0.225]
self.feat_type = feat_type
vgg_model = torchvision.models.vgg19(pretrained=True)
## AssertionError: DistributedDataParallel is not needed when a module doesn't have any parameter that requires a gradient
'''
if is_ddp:
vgg_model = vgg_model.cuda(gpu_no)
vgg_model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(vgg_model)
vgg_model = torch.nn.parallel.DistributedDataParallel(vgg_model, device_ids=[gpu_no], find_unused_parameters=True)
else:
vgg_model = vgg_model.cuda(gpu_no)
'''
vgg_model = vgg_model.cuda(gpu_no)
if self.feat_type == 'liu':
## conv1_1, conv2_1, conv3_1, conv4_1, conv5_1
self.slice1 = nn.Sequential(*list(vgg_model.features)[:2]).eval()
self.slice2 = nn.Sequential(*list(vgg_model.features)[2:7]).eval()
self.slice3 = nn.Sequential(*list(vgg_model.features)[7:12]).eval()
self.slice4 = nn.Sequential(*list(vgg_model.features)[12:21]).eval()
self.slice5 = nn.Sequential(*list(vgg_model.features)[21:30]).eval()
self.weights = [1.0/32, 1.0/16, 1.0/8, 1.0/4, 1.0]
elif self.feat_type == 'lei':
## conv1_2, conv2_2, conv3_2, conv4_2, conv5_2
self.slice1 = nn.Sequential(*list(vgg_model.features)[:4]).eval()
self.slice2 = nn.Sequential(*list(vgg_model.features)[4:9]).eval()
self.slice3 = nn.Sequential(*list(vgg_model.features)[9:14]).eval()
self.slice4 = nn.Sequential(*list(vgg_model.features)[14:23]).eval()
self.slice5 = nn.Sequential(*list(vgg_model.features)[23:32]).eval()
self.weights = [1.0/2.6, 1.0/4.8, 1.0/3.7, 1.0/5.6, 10.0/1.5]
else:
## maxpool after conv4_4
self.featureExactor = nn.Sequential(*list(vgg_model.features)[:28]).eval()
'''
for x in range(2):
self.slice1.add_module(str(x), pretrained_features[x])
for x in range(2, 7):
self.slice2.add_module(str(x), pretrained_features[x])
for x in range(7, 12):
self.slice3.add_module(str(x), pretrained_features[x])
for x in range(12, 21):
self.slice4.add_module(str(x), pretrained_features[x])
for x in range(21, 30):
self.slice5.add_module(str(x), pretrained_features[x])
'''
self.criterion = nn.L1Loss()
## fixed parameters
if not requires_grad:
for param in self.parameters():
param.requires_grad = False
self.eval()
print('[*] VGG19Loss init!')
def normalize(self, tensor):
tensor = tensor.clone()
mean = torch.as_tensor(self.mean, dtype=torch.float32, device=tensor.device)
std = torch.as_tensor(self.std, dtype=torch.float32, device=tensor.device)
tensor.sub_(mean[None, :, None, None]).div_(std[None, :, None, None])
return tensor
def forward(self, x, y):
norm_x, norm_y = self.normalize(x), self.normalize(y)
## feature extract
if self.feat_type == 'liu' or self.feat_type == 'lei':
x_relu1, y_relu1 = self.slice1(norm_x), self.slice1(norm_y)
x_relu2, y_relu2 = self.slice2(x_relu1), self.slice2(y_relu1)
x_relu3, y_relu3 = self.slice3(x_relu2), self.slice3(y_relu2)
x_relu4, y_relu4 = self.slice4(x_relu3), self.slice4(y_relu3)
x_relu5, y_relu5 = self.slice5(x_relu4), self.slice5(y_relu4)
x_vgg = [x_relu1, x_relu2, x_relu3, x_relu4, x_relu5]
y_vgg = [y_relu1, y_relu2, y_relu3, y_relu4, y_relu5]
loss = 0
for i in range(len(x_vgg)):
loss += self.weights[i] * self.criterion(x_vgg[i], y_vgg[i].detach())
else:
x_vgg, y_vgg = self.featureExactor(norm_x), self.featureExactor(norm_y)
loss = self.criterion(x_vgg, y_vgg.detach())
return loss