mirror of https://github.com/n00mkrad/flowframes
129 lines
4.5 KiB
Python
129 lines
4.5 KiB
Python
import torch
|
|
import numpy as np
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import torchvision.models as models
|
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
|
|
class EPE(nn.Module):
|
|
def __init__(self):
|
|
super(EPE, self).__init__()
|
|
|
|
def forward(self, flow, gt, loss_mask):
|
|
loss_map = (flow - gt.detach()) ** 2
|
|
loss_map = (loss_map.sum(1, True) + 1e-6) ** 0.5
|
|
return (loss_map * loss_mask)
|
|
|
|
|
|
class Ternary(nn.Module):
|
|
def __init__(self):
|
|
super(Ternary, self).__init__()
|
|
patch_size = 7
|
|
out_channels = patch_size * patch_size
|
|
self.w = np.eye(out_channels).reshape(
|
|
(patch_size, patch_size, 1, out_channels))
|
|
self.w = np.transpose(self.w, (3, 2, 0, 1))
|
|
self.w = torch.tensor(self.w).float().to(device)
|
|
|
|
def transform(self, img):
|
|
patches = F.conv2d(img, self.w, padding=3, bias=None)
|
|
transf = patches - img
|
|
transf_norm = transf / torch.sqrt(0.81 + transf**2)
|
|
return transf_norm
|
|
|
|
def rgb2gray(self, rgb):
|
|
r, g, b = rgb[:, 0:1, :, :], rgb[:, 1:2, :, :], rgb[:, 2:3, :, :]
|
|
gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
|
|
return gray
|
|
|
|
def hamming(self, t1, t2):
|
|
dist = (t1 - t2) ** 2
|
|
dist_norm = torch.mean(dist / (0.1 + dist), 1, True)
|
|
return dist_norm
|
|
|
|
def valid_mask(self, t, padding):
|
|
n, _, h, w = t.size()
|
|
inner = torch.ones(n, 1, h - 2 * padding, w - 2 * padding).type_as(t)
|
|
mask = F.pad(inner, [padding] * 4)
|
|
return mask
|
|
|
|
def forward(self, img0, img1):
|
|
img0 = self.transform(self.rgb2gray(img0))
|
|
img1 = self.transform(self.rgb2gray(img1))
|
|
return self.hamming(img0, img1) * self.valid_mask(img0, 1)
|
|
|
|
|
|
class SOBEL(nn.Module):
|
|
def __init__(self):
|
|
super(SOBEL, self).__init__()
|
|
self.kernelX = torch.tensor([
|
|
[1, 0, -1],
|
|
[2, 0, -2],
|
|
[1, 0, -1],
|
|
]).float()
|
|
self.kernelY = self.kernelX.clone().T
|
|
self.kernelX = self.kernelX.unsqueeze(0).unsqueeze(0).to(device)
|
|
self.kernelY = self.kernelY.unsqueeze(0).unsqueeze(0).to(device)
|
|
|
|
def forward(self, pred, gt):
|
|
N, C, H, W = pred.shape[0], pred.shape[1], pred.shape[2], pred.shape[3]
|
|
img_stack = torch.cat(
|
|
[pred.reshape(N*C, 1, H, W), gt.reshape(N*C, 1, H, W)], 0)
|
|
sobel_stack_x = F.conv2d(img_stack, self.kernelX, padding=1)
|
|
sobel_stack_y = F.conv2d(img_stack, self.kernelY, padding=1)
|
|
pred_X, gt_X = sobel_stack_x[:N*C], sobel_stack_x[N*C:]
|
|
pred_Y, gt_Y = sobel_stack_y[:N*C], sobel_stack_y[N*C:]
|
|
|
|
L1X, L1Y = torch.abs(pred_X-gt_X), torch.abs(pred_Y-gt_Y)
|
|
loss = (L1X+L1Y)
|
|
return loss
|
|
|
|
class MeanShift(nn.Conv2d):
|
|
def __init__(self, data_mean, data_std, data_range=1, norm=True):
|
|
c = len(data_mean)
|
|
super(MeanShift, self).__init__(c, c, kernel_size=1)
|
|
std = torch.Tensor(data_std)
|
|
self.weight.data = torch.eye(c).view(c, c, 1, 1)
|
|
if norm:
|
|
self.weight.data.div_(std.view(c, 1, 1, 1))
|
|
self.bias.data = -1 * data_range * torch.Tensor(data_mean)
|
|
self.bias.data.div_(std)
|
|
else:
|
|
self.weight.data.mul_(std.view(c, 1, 1, 1))
|
|
self.bias.data = data_range * torch.Tensor(data_mean)
|
|
self.requires_grad = False
|
|
|
|
class VGGPerceptualLoss(torch.nn.Module):
|
|
def __init__(self, rank=0):
|
|
super(VGGPerceptualLoss, self).__init__()
|
|
blocks = []
|
|
pretrained = True
|
|
self.vgg_pretrained_features = models.vgg19(pretrained=pretrained).features
|
|
self.normalize = MeanShift([0.485, 0.456, 0.406], [0.229, 0.224, 0.225], norm=True).cuda()
|
|
for param in self.parameters():
|
|
param.requires_grad = False
|
|
|
|
def forward(self, X, Y, indices=None):
|
|
X = self.normalize(X)
|
|
Y = self.normalize(Y)
|
|
indices = [2, 7, 12, 21, 30]
|
|
weights = [1.0/2.6, 1.0/4.8, 1.0/3.7, 1.0/5.6, 10/1.5]
|
|
k = 0
|
|
loss = 0
|
|
for i in range(indices[-1]):
|
|
X = self.vgg_pretrained_features[i](X)
|
|
Y = self.vgg_pretrained_features[i](Y)
|
|
if (i+1) in indices:
|
|
loss += weights[k] * (X - Y.detach()).abs().mean() * 0.1
|
|
k += 1
|
|
return loss
|
|
|
|
if __name__ == '__main__':
|
|
img0 = torch.zeros(3, 3, 256, 256).float().to(device)
|
|
img1 = torch.tensor(np.random.normal(
|
|
0, 1, (3, 3, 256, 256))).float().to(device)
|
|
ternary_loss = Ternary()
|
|
print(ternary_loss(img0, img1).shape)
|