nodes/jonstreeter-ComfyUI-Reference-Based-Video-Colorization
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor

Browse Source
This commit is contained in:
zhangmozhe
2020-12-01 14:06:29 +08:00
parent 5118d07b0f
commit 0c1389853a
246 changed files with 186 additions and 3834 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/ColorVidNet.py
+2 -123
View File
@@ -1,127 +1,11 @@
from os import path as osp
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.parallel
def type_name(m):
name = type(m).__name__
if name == "TorchObject":
name = m._typename.replace("cudnn.", "")
return name
def model_tm2pm(model):
layer_map = {}
for m in model.modules:
layer_map[m.name] = tm2pm(m)
return layer_map
"""'
removed layers: [[ 'ab_conv1_1', 'bw_conv1_1' ]]
"""
ref_color_layernames = [
"relu1_1",
"conv1_2",
"relu1_2",
"conv1_2norm",
"conv1_2norm_ss",
"conv2_1",
"relu2_1",
"conv2_2",
"relu2_2",
"conv2_2norm",
"conv2_2norm_ss",
"conv3_1",
"relu3_1",
"conv3_2",
"relu3_2",
"conv3_3",
"relu3_3",
"conv3_3norm",
"conv3_3norm_ss",
"conv4_1",
"relu4_1",
"conv4_2",
"relu4_2",
"conv4_3",
"relu4_3",
"conv4_3norm",
"conv5_1",
"relu5_1",
"conv5_2",
"relu5_2",
"conv5_3",
"relu5_3",
"conv5_3norm",
"conv6_1",
"relu6_1",
"conv6_2",
"relu6_2",
"conv6_3",
"relu6_3",
"conv6_3norm",
"conv7_1",
"relu7_1",
"conv7_2",
"relu7_2",
"conv7_3",
"relu7_3",
"conv7_3norm",
"conv8_1",
"conv3_3_short",
"relu8_1_comb",
"conv8_2",
"relu8_2",
"conv8_3",
"relu8_3",
"conv8_3norm",
"conv9_1",
"conv2_2_short",
"relu9_1_comb",
"conv9_2",
"relu9_2",
"conv9_2norm",
"conv10_1",
"conv1_2_short",
"relu10_1_comb",
"conv10_2",
"relu10_2",
"conv10_ab",
]
class ColorVidNet(nn.Module):
def __init__(
self, ic, ref_color_torch_model="convert_caffe_model/torch_model/reference_model/color_nodist_model.t7"
):
def __init__(self, ic):
super(ColorVidNet, self).__init__()
"""torch_model = load_lua(ref_color_torch_model, unknown_classes=True)
if type(torch_model).__name__ == 'hashable_uniq_dict':
torch_model = torch_model.model
# fix the ugly typo error
for m in torch_model.modules:
if m.name == 'conv8_1_comb' and type_name(m) == 'ReLU':
m.name = 'relu8_1_comb'
pytorch_model = model_tm2pm(torch_model)
self.add_module('conv1_1', nn.Sequential(nn.Conv2d(ic, 32, 3, 1, 1), nn.ReLU(), nn.Conv2d(32, 64, 3, 1, 1), nn.ReLU()))
for l_name in ref_color_layernames:
self.add_module(l_name, pytorch_model[l_name])
### replace all deconvolution with neareast upsample + convolution
assert type(self.conv8_1).__name__=='ConvTranspose2d' and type(self.conv9_1).__name__=='ConvTranspose2d' and \
type(self.conv10_1).__name__=='ConvTranspose2d', 'conv8_1, conv9_1, conv10_1 should be all ConvTranspose2d'
"""
##### Written by Ak. #####
self.conv1_1 = nn.Sequential(nn.Conv2d(ic, 32, 3, 1, 1), nn.ReLU(), nn.Conv2d(32, 64, 3, 1, 1))
self.conv1_2 = nn.Conv2d(64, 64, 3, 1, 1)
self.conv1_2norm = nn.BatchNorm2d(64, affine=False)
@@ -191,11 +75,8 @@ class ColorVidNet(nn.Module):
self.relu9_1_comb = nn.ReLU()
self.relu9_2 = nn.ReLU()
self.relu10_1_comb = nn.ReLU()
# self.relu10_2 = nn.ReLU()
self.relu10_2 = nn.LeakyReLU(0.2, True)
##### Above #####
print("replace all deconv with [nearest + conv]")
self.conv8_1 = nn.Sequential(nn.Upsample(scale_factor=2, mode="nearest"), nn.Conv2d(512, 256, 3, 1, 1))
self.conv9_1 = nn.Sequential(nn.Upsample(scale_factor=2, mode="nearest"), nn.Conv2d(256, 128, 3, 1, 1))
@@ -211,7 +92,6 @@ class ColorVidNet(nn.Module):
self.conv7_3norm = nn.InstanceNorm2d(512)
self.conv8_3norm = nn.InstanceNorm2d(256)
self.conv9_2norm = nn.InstanceNorm2d(128)
# print(self)
def forward(self, x):
""" x: gray image (1 channel), ab(2 channel), ab_err, ba_err"""
@@ -260,6 +140,5 @@ class ColorVidNet(nn.Module):
conv10_1_comb = self.relu10_1_comb(conv10_1 + conv1_2_short)
conv10_2 = self.relu10_2(self.conv10_2(conv10_1_comb))
conv10_ab = self.conv10_ab(conv10_2)
pred_ab = torch.tanh(conv10_ab) * 128
return pred_ab
return torch.tanh(conv10_ab) * 128
models/FrameColor.py
+11 -24
View File
@@ -1,21 +1,9 @@
# from imports import *
import torch
from utils.tensor_lab2rgb import tensor_lab2rgb
from utils.util import *
if torch.cuda.is_available():
device = torch.device("cuda")
else:
device = troch.device("cpu")
def warp_color(IA_l, IB_lab, features_B, vggnet, nonlocal_net, colornet, feature_noise=0, temperature=0.01):
"""
joint_train=True, enable grad; otherwise disable grad
"""
# change to rgb for feature extraction
IA_rgb_from_gray = gray2rgb_batch(IA_l)
with torch.no_grad():
A_relu1_1, A_relu2_1, A_relu3_1, A_relu4_1, A_relu5_1 = vggnet(
IA_rgb_from_gray, ["r12", "r22", "r32", "r42", "r52"], preprocess=True
@@ -34,18 +22,17 @@ def warp_color(IA_l, IB_lab, features_B, vggnet, nonlocal_net, colornet, feature
B_relu4_1 = feature_normalize(B_relu4_1)
B_relu5_1 = feature_normalize(B_relu5_1)
# if feature_noise > 0:
# A_relu2_1 = A_relu2_1 + torch.randn_like(A_relu2_1, requires_grad=False) * feature_noise
# B_relu2_1 = B_relu2_1 + torch.randn_like(B_relu2_1, requires_grad=False) * feature_noise
# A_relu3_1 = A_relu3_1 + torch.randn_like(A_relu3_1, requires_grad=False) * feature_noise
# B_relu3_1 = B_relu3_1 + torch.randn_like(B_relu3_1, requires_grad=False) * feature_noise
# A_relu4_1 = A_relu4_1 + torch.randn_like(A_relu4_1, requires_grad=False) * feature_noise
# B_relu4_1 = B_relu4_1 + torch.randn_like(B_relu4_1, requires_grad=False) * feature_noise
# A_relu5_1 = A_relu5_1 + torch.randn_like(A_relu5_1, requires_grad=False) * feature_noise
# B_relu5_1 = B_relu5_1 + torch.randn_like(B_relu5_1, requires_grad=False) * feature_noise
nonlocal_BA_lab, similarity_map = nonlocal_net(
IB_lab, A_relu2_1, A_relu3_1, A_relu4_1, A_relu5_1, B_relu2_1, B_relu3_1, B_relu4_1, B_relu5_1, temperature=0.01
IB_lab,
A_relu2_1,
A_relu3_1,
A_relu4_1,
A_relu5_1,
B_relu2_1,
B_relu3_1,
B_relu4_1,
B_relu5_1,
temperature=temperature,
)
return nonlocal_BA_lab, similarity_map, features_A
@@ -64,7 +51,7 @@ def frame_colorization(
luminance_noise=0,
temperature=0.01,
):
# change to rgb for feature extraction
IA_l = IA_lab[:, 0:1, :, :]
if luminance_noise:
IA_l = IA_l + torch.randn_like(IA_l, requires_grad=False) * luminance_noise
models/GAN_models.py
+11 -31
View File
@@ -81,11 +81,11 @@ class Self_Attention(nn.Module):
def forward(self, x):
"""
inputs :
x : input feature maps( B X C X W X H)
returns :
out : self attention value + input feature
attention: B X N X N (N is Width*Height)
inputs :
x : input feature maps( B X C X W X H)
returns :
out : self attention value + input feature
attention: B X N X N (N is Width*Height)
"""
m_batchsize, C, width, height = x.size()
proj_query = self.query_conv(x).view(m_batchsize, -1, width * height).permute(0, 2, 1) # B X CX(N)
@@ -98,13 +98,12 @@ class Self_Attention(nn.Module):
out = out.view(m_batchsize, C, width, height)
out = self.gamma * out + x
# return out, attention
return out
class Discriminator_WITHOUT_FC_x64_video(nn.Module):
"""
Discriminative Network
Discriminative Network
"""
def __init__(self, in_size=6, ndf=64):
@@ -137,18 +136,7 @@ class Discriminator_WITHOUT_FC_x64_video(nn.Module):
nn.LeakyReLU(0.2, inplace=True),
)
# sate size: (ndf * 8) x 4 x 4
# Self_Attention(self.ndf * 8, 'relu')
self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, [3, 6], 1, 0))
# nn.Sigmoid()
# state size: 1 x 1 x 1
# for m in self.modules():
# if isinstance(m, nn.Conv2d):
# m.weight.data.normal_(0.0, 0.02)
# if m.bias is not None:
# m.bias.data.zero_()
def forward(self, input):
feature1 = self.layer1(input)
@@ -163,7 +151,7 @@ class Discriminator_WITHOUT_FC_x64_video(nn.Module):
class Discriminator_WITHOUT_FC_x64_video_BIG(nn.Module):
"""
Discriminative Network
Discriminative Network
"""
def __init__(self, in_size=6, ndf=64):
@@ -219,7 +207,7 @@ class Discriminator_WITHOUT_FC_x64_video_BIG(nn.Module):
class Discriminator_WITHOUT_FC_x64(nn.Module):
"""
Discriminative Network
Discriminative Network
"""
def __init__(self, in_size=3, ndf=64):
@@ -248,14 +236,6 @@ class Discriminator_WITHOUT_FC_x64(nn.Module):
)
self.last = SpectralNorm(nn.Conv2d(self.ndf * 8, 1, 4, 1, 0))
# nn.Sigmoid()
# state size: 1 x 1 x 1
# for m in self.modules():
# if isinstance(m, nn.Conv2d):
# m.weight.data.normal_(0.0, 0.02)
# if m.bias is not None:
# m.bias.data.zero_()
def forward(self, input):
@@ -268,7 +248,7 @@ class Discriminator_WITHOUT_FC_x64(nn.Module):
class Discriminator_WITHOUT_FC_x64_BIG(nn.Module):
"""
Discriminative Network
Discriminative Network
"""
def __init__(self, in_size=3, ndf=64):
@@ -316,7 +296,7 @@ class Discriminator_WITHOUT_FC_x64_BIG(nn.Module):
class Discriminator_WITHOUT_FC_x128(nn.Module):
"""
Discriminative Network
Discriminative Network
"""
def __init__(self, in_size=3, ndf=64):
@@ -359,7 +339,7 @@ class Discriminator_WITHOUT_FC_x128(nn.Module):
class Discriminator_WITHOUT_FC_x256(nn.Module):
"""
Discriminative Network
Discriminative Network
"""
def __init__(self, in_size=3, ndf=64):
models/NonlocalNet.py
+55 -124
View File
@@ -14,12 +14,14 @@ def find_local_patch(x, patch_size):
x_unfold = F.unfold(
x, kernel_size=(patch_size, patch_size), padding=(patch_size // 2, patch_size // 2), stride=(1, 1)
)
out = x_unfold.view(N, x_unfold.shape[1], H, W)
return out
return x_unfold.view(N, x_unfold.shape[1], H, W)
class WeightedAverage(nn.Module):
def __init__(self,):
def __init__(
self,
):
super(WeightedAverage, self).__init__()
def forward(self, x_lab, patch_size=3, alpha=1, scale_factor=1):
@@ -34,14 +36,13 @@ class WeightedAverage(nn.Module):
local_difference_l = (local_l - l) ** 2
correlation = nn.functional.softmax(-1 * local_difference_l / alpha, dim=1)
weighted_ab = torch.cat(
return torch.cat(
(
torch.sum(correlation * local_a, dim=1, keepdim=True),
torch.sum(correlation * local_b, dim=1, keepdim=True),
),
1,
)
return weighted_ab
class WeightedAverage_color(nn.Module):
@@ -49,11 +50,13 @@ class WeightedAverage_color(nn.Module):
smooth the image according to the color distance in the LAB space
"""
def __init__(self,):
def __init__(
self,
):
super(WeightedAverage_color, self).__init__()
def forward(self, x_lab, x_lab_predict, patch_size=3, alpha=1, scale_factor=1):
# alpha=0: less smooth; alpha=inf: smoother
""" alpha=0: less smooth; alpha=inf: smoother """
x_lab = F.interpolate(x_lab, scale_factor=scale_factor)
l = uncenter_l(x_lab[:, 0:1, :, :])
a = x_lab[:, 1:2, :, :]
@@ -71,18 +74,19 @@ class WeightedAverage_color(nn.Module):
-1 * local_color_difference / alpha, dim=1
) # so that sum of weights equal to 1
weighted_ab = torch.cat(
return torch.cat(
(
torch.sum(correlation * local_a_predict, dim=1, keepdim=True),
torch.sum(correlation * local_b_predict, dim=1, keepdim=True),
),
1,
)
return weighted_ab
class NonlocalWeightedAverage(nn.Module):
def __init__(self,):
def __init__(
self,
):
super(NonlocalWeightedAverage, self).__init__()
def forward(self, x_lab, feature, patch_size=3, alpha=0.1, scale_factor=1):
@@ -113,56 +117,36 @@ class CorrelationLayer(nn.Module):
self.search_range = search_range
def forward(self, x1, x2, alpha=1, raw_output=False, metric="similarity"):
# args = self.args
shape = list(x1.size())
shape[1] = (self.search_range * 2 + 1) ** 2
cv = torch.zeros(shape).to(torch.device("cuda"))
if metric == "similarity":
for i in range(-self.search_range, self.search_range + 1):
for j in range(-self.search_range, self.search_range + 1):
if i < 0:
slice_h, slice_h_r = slice(None, i), slice(-i, None)
elif i > 0:
slice_h, slice_h_r = slice(i, None), slice(None, -i)
else:
slice_h, slice_h_r = slice(None), slice(None)
for i in range(-self.search_range, self.search_range + 1):
for j in range(-self.search_range, self.search_range + 1):
if i < 0:
slice_h, slice_h_r = slice(None, i), slice(-i, None)
elif i > 0:
slice_h, slice_h_r = slice(i, None), slice(None, -i)
else:
slice_h, slice_h_r = slice(None), slice(None)
if j < 0:
slice_w, slice_w_r = slice(None, j), slice(-j, None)
elif j > 0:
slice_w, slice_w_r = slice(j, None), slice(None, -j)
else:
slice_w, slice_w_r = slice(None), slice(None)
if j < 0:
slice_w, slice_w_r = slice(None, j), slice(-j, None)
elif j > 0:
slice_w, slice_w_r = slice(j, None), slice(None, -j)
else:
slice_w, slice_w_r = slice(None), slice(None)
# storage sequence (eg. search_range=3): -24, -23, .., -1, 0, 1, ..., 23, 24
if metric == "similarity":
cv[:, (self.search_range * 2 + 1) * i + j, slice_h, slice_w] = (
x1[:, :, slice_h, slice_w] * x2[:, :, slice_h_r, slice_w_r]
).sum(1)
else: # patchwise subtraction
for i in range(-self.search_range, self.search_range + 1):
for j in range(-self.search_range, self.search_range + 1):
if i < 0:
slice_h, slice_h_r = slice(None, i), slice(-i, None)
elif i > 0:
slice_h, slice_h_r = slice(i, None), slice(None, -i)
else:
slice_h, slice_h_r = slice(None), slice(None)
if j < 0:
slice_w, slice_w_r = slice(None, j), slice(-j, None)
elif j > 0:
slice_w, slice_w_r = slice(j, None), slice(None, -j)
else:
slice_w, slice_w_r = slice(None), slice(None)
# storage sequence (eg. search_range=3): -24, -23, .., -1, 0, 1, ..., 23, 24
else: # patchwise subtraction
cv[:, (self.search_range * 2 + 1) * i + j, slice_h, slice_w] = -(
(x1[:, :, slice_h, slice_w] - x2[:, :, slice_h_r, slice_w_r]) ** 2
).sum(1)
## TODO sigmoid?
# TODO sigmoid?
if raw_output:
return cv
else:
@@ -185,11 +169,11 @@ class Self_Attn(nn.Module):
def forward(self, x):
"""
inputs :
x : input feature maps(B X C X W X H)
returns :
out : self attention value + input feature
attention: B X N X N (N is Width*Height)
inputs :
x : input feature maps(B X C X W X H)
returns :
out : self attention value + input feature
attention: B X N X N (N is Width*Height)
"""
m_batchsize, C, width, height = x.size()
proj_query = self.query_conv(x).view(m_batchsize, -1, width * height).permute(0, 2, 1) # B X N X C
@@ -206,7 +190,7 @@ class Self_Attn(nn.Module):
class VGG19_feature_color_torchversion(nn.Module):
"""
"""
NOTE: no need to pre-process the input; input tensor should range in [0,1]
"""
@@ -242,7 +226,7 @@ class VGG19_feature_color_torchversion(nn.Module):
self.pool5 = nn.AvgPool2d(kernel_size=2, stride=2)
def forward(self, x, out_keys, preprocess=True):
"""
"""
NOTE: input tensor should range in [0,1]
"""
out = {}
@@ -283,8 +267,8 @@ class VGG19_feature_color(nn.Module):
"""Extract multiple convolutional feature maps."""
features = []
for name, layer in self.vgg._modules.items():
x = layer(x)
if name in self.select:
x = layer(x)
features.append(x)
return features
@@ -301,70 +285,20 @@ class VGG19_feature(nn.Module):
return A_relu3_1, A_relu4_1, A_relu5_1, B_relu3_1, B_relu4_1, B_relu5_1
class VGG19_feature_new(nn.Module):
# input: [LLL] channels, range=[0,255]
def __init__(self, gpu_ids):
super(VGG19_feature_new, self).__init__()
self.vgg19_gray_new = vgg19_gray_new().cuda()
def forward(self, A_l, B_l):
A_relu2_1, A_relu3_1, A_relu4_1, A_relu5_1 = self.vgg19_gray_new(A_l)
B_relu2_1, B_relu3_1, B_relu4_1, B_relu5_1 = self.vgg19_gray_new(B_l)
return A_relu2_1, A_relu3_1, A_relu4_1, A_relu5_1, B_relu2_1, B_relu3_1, B_relu4_1, B_relu5_1
class WTA(torch.autograd.Function):
"""
We can implement our own custom autograd Functions by subclassing
torch.autograd.Function and implementing the forward and backward passes
which operate on Tensors.
"""
@staticmethod
def forward(ctx, input):
"""
In the forward pass we receive a Tensor containing the input and return a
Tensor containing the output. You can cache arbitrary Tensors for use in the
backward pass using the save_for_backward method.
"""
height = input.shape[2]
width = input.shape[3]
activation_max, index_max = torch.topk(input, -1, keepdim=True)
input_zeros = input * 0
output_max_only = torch.where(input == activation_max, input, input_zeros)
mask = (output_max_only > 0).type(torch.float)
ctx.save_for_backward(input, mask)
return output_max_only
@staticmethod
def backward(ctx, grad_output):
"""
In the backward pass we receive a Tensor containing the gradient of the loss
with respect to the output, and we need to compute the gradient of the loss
with respect to the input.
"""
# import pdb
# pdb.set_trace()
input, mask = ctx.saved_tensors
grad_input = grad_output.clone() * mask
return grad_input
class WTA_scale(torch.autograd.Function):
"""
We can implement our own custom autograd Functions by subclassing
torch.autograd.Function and implementing the forward and backward passes
which operate on Tensors.
"""
We can implement our own custom autograd Functions by subclassing
torch.autograd.Function and implementing the forward and backward passes
which operate on Tensors.
"""
@staticmethod
def forward(ctx, input, scale=1e-4):
"""
In the forward pass we receive a Tensor containing the input and return a
Tensor containing the output. You can cache arbitrary Tensors for use in the
backward pass using the save_for_backward method.
"""
In the forward pass we receive a Tensor containing the input and return a
Tensor containing the output. You can cache arbitrary Tensors for use in the
backward pass using the save_for_backward method.
"""
activation_max, index_max = torch.max(input, -1, keepdim=True)
input_scale = input * scale # default: 1e-4
# input_scale = input * scale # default: 1e-4
@@ -377,10 +311,10 @@ class WTA_scale(torch.autograd.Function):
@staticmethod
def backward(ctx, grad_output):
"""
In the backward pass we receive a Tensor containing the gradient of the loss
with respect to the output, and we need to compute the gradient of the loss
with respect to the input.
"""
In the backward pass we receive a Tensor containing the gradient of the loss
with respect to the output, and we need to compute the gradient of the loss
with respect to the input.
"""
# import pdb
# pdb.set_trace()
input, mask = ctx.saved_tensors
@@ -418,9 +352,9 @@ class ResidualBlock(nn.Module):
return out
## NOTE: THIS IS THE FINAL VERSION
class WarpNet(nn.Module):
# input is Al, Bl, channel = 1, range~[0,255]
""" input is Al, Bl, channel = 1, range~[0,255] """
def __init__(self, batch_size):
super(WarpNet, self).__init__()
self.feature_channel = 64
@@ -549,12 +483,9 @@ class WarpNet(nn.Module):
similarity_map = similarity_map.view(batch_size, 1, feature_height, feature_width)
# f can be negative
if WTA_scale_weight == 1:
f_WTA = f
else:
f_WTA = WTA_scale.apply(f, WTA_scale_weight)
f_WTA = f if WTA_scale_weight == 1 else WTA_scale.apply(f, WTA_scale_weight)
f_WTA = f_WTA / temperature
f_div_C = F.softmax(f_WTA.squeeze_(), dim=-1) # 2*1936*1936; softmax along the horizontal line (dim=-1)
f_div_C = F.softmax(f_WTA.squeeze_(), dim=-1) # 2*1936*1936;
# downsample the reference color
B_lab = F.avg_pool2d(B_lab_map, 4)
models/guided_filter.py
-125
View File
@@ -1,125 +0,0 @@
import torch
from torch import nn
from torch.autograd import Variable
from torch.nn import functional as F
# from .box_filter import BoxFilter
def diff_x(input, r):
assert input.dim() == 4
left = input[:, :, r : 2 * r + 1]
middle = input[:, :, 2 * r + 1 :] - input[:, :, : -2 * r - 1]
right = input[:, :, -1:] - input[:, :, -2 * r - 1 : -r - 1]
output = torch.cat([left, middle, right], dim=2)
return output
def diff_y(input, r):
assert input.dim() == 4
left = input[:, :, :, r : 2 * r + 1]
middle = input[:, :, :, 2 * r + 1 :] - input[:, :, :, : -2 * r - 1]
right = input[:, :, :, -1:] - input[:, :, :, -2 * r - 1 : -r - 1]
output = torch.cat([left, middle, right], dim=3)
return output
class BoxFilter(nn.Module):
def __init__(self, r):
super(BoxFilter, self).__init__()
self.r = r
def forward(self, x):
assert x.dim() == 4
return diff_y(diff_x(x.cumsum(dim=2), self.r).cumsum(dim=3), self.r)
class FastGuidedFilter(nn.Module):
def __init__(self, r, eps=1e-8):
super(FastGuidedFilter, self).__init__()
self.r = r
self.eps = eps
self.boxfilter = BoxFilter(r)
def forward(self, lr_x, lr_y, hr_x):
n_lrx, c_lrx, h_lrx, w_lrx = lr_x.size()
n_lry, c_lry, h_lry, w_lry = lr_y.size()
n_hrx, c_hrx, h_hrx, w_hrx = hr_x.size()
assert n_lrx == n_lry and n_lry == n_hrx
assert c_lrx == c_hrx and (c_lrx == 1 or c_lrx == c_lry)
assert h_lrx == h_lry and w_lrx == w_lry
assert h_lrx > 2 * self.r + 1 and w_lrx > 2 * self.r + 1
## N
N = self.boxfilter(Variable(lr_x.data.new().resize_((1, 1, h_lrx, w_lrx)).fill_(1.0)))
## mean_x
mean_x = self.boxfilter(lr_x) / N
## mean_y
mean_y = self.boxfilter(lr_y) / N
## cov_xy
cov_xy = self.boxfilter(lr_x * lr_y) / N - mean_x * mean_y
## var_x
var_x = self.boxfilter(lr_x * lr_x) / N - mean_x * mean_x
## A
A = cov_xy / (var_x + self.eps)
## b
b = mean_y - A * mean_x
## mean_A; mean_b
mean_A = F.upsample(A, (h_hrx, w_hrx), mode="bilinear")
mean_b = F.upsample(b, (h_hrx, w_hrx), mode="bilinear")
return mean_A * hr_x + mean_b
class GuidedFilter(nn.Module):
def __init__(self, r, eps=1e-8):
super(GuidedFilter, self).__init__()
self.r = r
self.eps = eps
self.boxfilter = BoxFilter(r)
def forward(self, x, y):
n_x, c_x, h_x, w_x = x.size()
n_y, c_y, h_y, w_y = y.size()
assert n_x == n_y
assert c_x == 1 or c_x == c_y
assert h_x == h_y and w_x == w_y
assert h_x > 2 * self.r + 1 and w_x > 2 * self.r + 1
# N
N = self.boxfilter(Variable(x.data.new().resize_((1, 1, h_x, w_x)).fill_(1.0)))
# mean_x
mean_x = self.boxfilter(x) / N
# mean_y
mean_y = self.boxfilter(y) / N
# cov_xy
cov_xy = self.boxfilter(x * y) / N - mean_x * mean_y
# var_x
var_x = self.boxfilter(x * x) / N - mean_x * mean_x
# A
A = cov_xy / (var_x + self.eps)
# b
b = mean_y - A * mean_x
# mean_A; mean_b
mean_A = self.boxfilter(A) / N
mean_b = self.boxfilter(b) / N
return mean_A * x + mean_b
models/spectral_normalization.py
+1 -5
View File
@@ -1,9 +1,6 @@
import torch
import torch.nn.functional as F
from torch import Tensor, nn
from torch.autograd import Variable
from torch import nn
from torch.nn import Parameter
from torch.optim.optimizer import Optimizer, required
def l2normalize(v, eps=1e-12):
@@ -29,7 +26,6 @@ class SpectralNorm(nn.Module):
v.data = l2normalize(torch.mv(torch.t(w.view(height, -1).data), u.data))
u.data = l2normalize(torch.mv(w.view(height, -1).data, v.data))
# sigma = torch.dot(u.data, torch.mv(w.view(height,-1).data, v.data))
sigma = u.dot(w.view(height, -1).mv(v))
setattr(self.module, self.name, w / sigma.expand_as(w))
models/vgg19_gray.py
+2 -44
View File
@@ -1,9 +1,7 @@
from collections import OrderedDict, namedtuple
from functools import reduce
import torch
import torch.nn as nn
from torch.autograd import Variable
class LambdaBase(nn.Sequential):
@@ -15,6 +13,7 @@ class LambdaBase(nn.Sequential):
output = []
for module in self._modules.values():
output.append(module(input))
return output if output else input
@@ -126,13 +125,10 @@ model = nn.Sequential( # Sequential,
)
model.load_state_dict(torch.load("models/vgg19_gray.pth"))
# print(model)
model.load_state_dict(torch.load("data/vgg19_gray.pth"))
vgg19_gray_net = torch.nn.Sequential()
for (name, layer) in model._modules.items():
vgg19_gray_net.add_module(layer_names[int(name)], model[int(name)])
# print(vgg19_gray_net)
for param in vgg19_gray_net.parameters():
param.requires_grad = False
@@ -163,8 +159,6 @@ class vgg19_gray(torch.nn.Module):
h_relu4_1 = h
h = self.slice3(h)
h_relu5_1 = h
# vgg_outputs = namedtuple("VggOutputs", ['h_relu3_1', 'h_relu4_1', 'h_relu5_1'])
# out = vgg_outputs(h_relu3_1, h_relu4_1, h_relu5_1)
return h_relu3_1, h_relu4_1, h_relu5_1
@@ -197,40 +191,4 @@ class vgg19_gray_new(torch.nn.Module):
h_relu4_1 = h
h = self.slice3(h)
h_relu5_1 = h
# vgg_outputs = namedtuple("VggOutputs", ['h_relu3_1', 'h_relu4_1', 'h_relu5_1'])
# out = vgg_outputs(h_relu3_1, h_relu4_1, h_relu5_1)
return h_relu2_1, h_relu3_1, h_relu4_1, h_relu5_1
# class vgg19_gray_conv(torch.nn.Module):
# def __init__(self, requires_grad=False):
# super(vgg19_gray_conv, self).__init__()
# vgg_pretrained_features = vgg19_gray_net
# self.slice0 = torch.nn.Sequential()
# self.slice1 = torch.nn.Sequential()
# self.slice2 = torch.nn.Sequential()
# self.slice3 = torch.nn.Sequential()
# for x in range(6):
# self.slice0.add_module(layer_names[x], vgg_pretrained_features[x])
# for x in range(6, 11):
# self.slice1.add_module(layer_names[x], vgg_pretrained_features[x])
# for x in range(11, 20):
# self.slice2.add_module(layer_names[x], vgg_pretrained_features[x])
# for x in range(20, 29):
# self.slice3.add_module(layer_names[x], vgg_pretrained_features[x])
# if not requires_grad:
# for param in self.parameters():
# param.requires_grad = False
# def forward(self, X):
# h = self.slice0(X)
# h_relu2_1 = h
# h = self.slice1(h)
# h_relu3_1 = h
# h = self.slice2(h)
# h_relu4_1 = h
# h = self.slice3(h)
# h_relu5_1 = h
# # vgg_outputs = namedtuple("VggOutputs", ['h_relu3_1', 'h_relu4_1', 'h_relu5_1'])
# # out = vgg_outputs(h_relu3_1, h_relu4_1, h_relu5_1)
# return h_relu2_1, h_relu3_1, h_relu4_1, h_relu5_1