diff --git a/modules/rife/loss.py b/modules/rife/loss.py deleted file mode 100644 index f525ff443..000000000 --- a/modules/rife/loss.py +++ /dev/null @@ -1,121 +0,0 @@ -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F -import torchvision.models as models -from modules import devices - - -class EPE(nn.Module): - def __init__(self): - super().__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().__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(devices.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().__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(devices.device) - self.kernelY = self.kernelY.unsqueeze(0).unsqueeze(0).to(devices.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().__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): # pylint: disable=unused-argument - super().__init__() - 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).to(device=devices.device) - 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 diff --git a/modules/rife/refine.py b/modules/rife/refine.py deleted file mode 100644 index 9ea076d1f..000000000 --- a/modules/rife/refine.py +++ /dev/null @@ -1,90 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F -from modules.rife.warplayer import warp - - -c = 16 - - -def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): - return nn.Sequential( - nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, bias=True), - nn.LeakyReLU(0.2, True) - ) - - -def conv_woact(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): - return nn.Sequential( - nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, bias=True), - ) - - -def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1): # pylint: disable=unused-argument - return nn.Sequential( - torch.nn.ConvTranspose2d(in_channels=in_planes, out_channels=out_planes, kernel_size=4, stride=2, padding=1, bias=True), - nn.LeakyReLU(0.2, True) - ) - - -class Conv2(nn.Module): - def __init__(self, in_planes, out_planes, stride=2): - super().__init__() - self.conv1 = conv(in_planes, out_planes, 3, stride, 1) - self.conv2 = conv(out_planes, out_planes, 3, 1, 1) - - def forward(self, x): - x = self.conv1(x) - x = self.conv2(x) - return x - - -class Contextnet(nn.Module): - def __init__(self): - super().__init__() - self.conv1 = Conv2(3, c) - self.conv2 = Conv2(c, 2*c) - self.conv3 = Conv2(2*c, 4*c) - self.conv4 = Conv2(4*c, 8*c) - - def forward(self, x, flow): - x = self.conv1(x) - flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False) * 0.5 - f1 = warp(x, flow) - x = self.conv2(x) - flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False) * 0.5 - f2 = warp(x, flow) - x = self.conv3(x) - flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False) * 0.5 - f3 = warp(x, flow) - x = self.conv4(x) - flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False) * 0.5 - f4 = warp(x, flow) - return [f1, f2, f3, f4] - - -class Unet(nn.Module): - def __init__(self): - super().__init__() - self.down0 = Conv2(17, 2*c) - self.down1 = Conv2(4*c, 4*c) - self.down2 = Conv2(8*c, 8*c) - self.down3 = Conv2(16*c, 16*c) - self.up0 = deconv(32*c, 8*c) - self.up1 = deconv(16*c, 4*c) - self.up2 = deconv(8*c, 2*c) - self.up3 = deconv(4*c, c) - self.conv = nn.Conv2d(c, 3, 3, 1, 1) - - def forward(self, img0, img1, warped_img0, warped_img1, mask, flow, c0, c1): - s0 = self.down0( - torch.cat((img0, img1, warped_img0, warped_img1, mask, flow), 1)) - s1 = self.down1(torch.cat((s0, c0[0], c1[0]), 1)) - s2 = self.down2(torch.cat((s1, c0[1], c1[1]), 1)) - s3 = self.down3(torch.cat((s2, c0[2], c1[2]), 1)) - x = self.up0(torch.cat((s3, c0[3], c1[3]), 1)) - x = self.up1(torch.cat((x, s2), 1)) - x = self.up2(torch.cat((x, s1), 1)) - x = self.up3(torch.cat((x, s0), 1)) - x = self.conv(x) - return torch.sigmoid(x)