596 lines
18 KiB
Python
596 lines
18 KiB
Python
"""
|
|
from photo_wct.py of https://github.com/NVIDIA/FastPhotoStyle
|
|
Copyright (C) 2018 NVIDIA Corporation.
|
|
Licensed under the CC BY-NC-SA 4.0
|
|
"""
|
|
|
|
import datetime
|
|
import os
|
|
|
|
import imageio
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
from PIL import Image
|
|
from torchvision import transforms
|
|
|
|
|
|
def svd(feat, iden=False, device="cuda:0"):
|
|
size = feat.size()
|
|
mean = torch.mean(feat, 1)
|
|
mean = mean.unsqueeze(1).expand_as(feat)
|
|
_feat = feat.clone()
|
|
_feat -= mean
|
|
if size[1] > 1:
|
|
conv = torch.mm(_feat, _feat.t()).div(size[1] - 1)
|
|
else:
|
|
conv = torch.mm(_feat, _feat.t())
|
|
if iden:
|
|
conv += torch.eye(size[0]).to(device)
|
|
# pdb.set_trace()
|
|
u, e, v = torch.svd(conv, some=False)
|
|
return u, e, v
|
|
|
|
|
|
def get_squeeze_feat(feat):
|
|
_feat = feat.squeeze(0)
|
|
size = _feat.size(0)
|
|
return _feat.view(size, -1).clone()
|
|
|
|
|
|
def get_rank(singular_values, dim, eps=0.00001):
|
|
r = dim
|
|
for i in range(dim - 1, -1, -1):
|
|
if singular_values[i] >= eps:
|
|
r = i + 1
|
|
break
|
|
return r
|
|
|
|
|
|
def covsqrt_mean(feature, inverse=False, tolerance=1e-14):
|
|
# I referenced the default svd tolerance value in matlab.
|
|
|
|
b, c, h, w = feature.size()
|
|
|
|
mean = torch.mean(feature.view(b, c, -1), dim=2, keepdim=True)
|
|
zeromean = feature.view(b, c, -1) - mean
|
|
cov = torch.bmm(zeromean, zeromean.transpose(1, 2))
|
|
|
|
eps_matrix_ = (torch.ones_like(cov) * 1e-8).cuda()
|
|
# evals, evects = torch.symeig(cov, eigenvectors=True)
|
|
evals, evects = torch.symeig(cov + eps_matrix_, eigenvectors=True)
|
|
|
|
p = 0.5
|
|
if inverse:
|
|
p *= -1
|
|
|
|
covsqrt = []
|
|
for i in range(b):
|
|
k = 0
|
|
for j in range(c):
|
|
if evals[i][j] > tolerance:
|
|
k = j
|
|
break
|
|
covsqrt.append(
|
|
torch.mm(
|
|
evects[i][:, k:],
|
|
torch.mm(evals[i][k:].pow(p).diag_embed(), evects[i][:, k:].t()),
|
|
).unsqueeze(0)
|
|
)
|
|
covsqrt = torch.cat(covsqrt, dim=0)
|
|
|
|
return covsqrt, mean
|
|
|
|
|
|
def whitening(feature):
|
|
b, c, h, w = feature.size()
|
|
|
|
inv_covsqrt, mean = covsqrt_mean(feature, inverse=True)
|
|
|
|
normalized_feature = torch.matmul(inv_covsqrt, feature.view(b, c, -1) - mean)
|
|
|
|
return normalized_feature.view(b, c, h, w)
|
|
|
|
|
|
def whitening_edit(feature):
|
|
b, c, h, w = feature.size()
|
|
cont_feat = get_squeeze_feat(feature)
|
|
# cont_min = cont_feat.min()
|
|
# cont_max = cont_feat.max()
|
|
cont_mean = torch.mean(cont_feat, 1).unsqueeze(1).expand_as(cont_feat)
|
|
cont_feat -= cont_mean
|
|
|
|
_, c_e, c_v = svd(cont_feat, iden=True)
|
|
k_c = get_rank(c_e, cont_feat.size()[0])
|
|
c_d = (c_e[0:k_c]).pow(-0.5)
|
|
step1 = torch.mm(c_v[:, 0:k_c], torch.diag(c_d))
|
|
step2 = torch.mm(step1, (c_v[:, 0:k_c].t()))
|
|
normalized_feature = torch.mm(step2, cont_feat)
|
|
|
|
# inv_covsqrt, mean = covsqrt_mean(feature, inverse=True)
|
|
|
|
# normalized_feature = torch.matmul(inv_covsqrt, feature.view(b, c, -1)-mean)
|
|
|
|
return normalized_feature.view(b, c, h, w)
|
|
|
|
|
|
def coloring(feature, target):
|
|
b, c, h, w = feature.size()
|
|
|
|
covsqrt, mean = covsqrt_mean(target)
|
|
|
|
colored_feature = torch.matmul(covsqrt, feature.view(b, c, -1)) + mean
|
|
|
|
return colored_feature.view(b, c, h, w)
|
|
|
|
|
|
def SwitchWhiten2d(x):
|
|
N, C, H, W = x.size()
|
|
|
|
in_data = x.view(N, C, -1)
|
|
|
|
eye = in_data.data.new().resize_(C, C)
|
|
eye = torch.nn.init.eye_(eye).view(1, C, C).expand(N, C, C)
|
|
|
|
# calculate other statistics
|
|
mean_in = in_data.mean(-1, keepdim=True)
|
|
x_in = in_data - mean_in
|
|
# (N x g) x C x C
|
|
cov_in = torch.bmm(x_in, torch.transpose(x_in, 1, 2)).div(H * W)
|
|
|
|
mean = mean_in
|
|
cov = cov_in + 1e-5 * eye
|
|
|
|
# perform whitening using Newton's iteration
|
|
Ng, c, _ = cov.size()
|
|
P = torch.eye(c).to(cov).expand(Ng, c, c)
|
|
|
|
rTr = (cov * P).sum((1, 2), keepdim=True).reciprocal_()
|
|
cov_N = cov * rTr
|
|
for k in range(5):
|
|
P = torch.baddbmm(1.5, P, -0.5, torch.matrix_power(P, 3), cov_N)
|
|
|
|
wm = P.mul_(rTr.sqrt())
|
|
x_hat = torch.bmm(wm, in_data - mean)
|
|
|
|
return x_hat, wm, mean
|
|
|
|
|
|
def wct_core(cont_feat, styl_feat, weight=1, registers=None, device="cuda:0"):
|
|
cont_feat = get_squeeze_feat(cont_feat)
|
|
cont_min = cont_feat.min()
|
|
cont_max = cont_feat.max()
|
|
cont_mean = torch.mean(cont_feat, 1).unsqueeze(1).expand_as(cont_feat)
|
|
cont_feat -= cont_mean
|
|
|
|
if not registers:
|
|
_, c_e, c_v = svd(cont_feat, iden=True, device=device)
|
|
|
|
styl_feat = get_squeeze_feat(styl_feat)
|
|
s_mean = torch.mean(styl_feat, 1)
|
|
_, s_e, s_v = svd(styl_feat, iden=True, device=device)
|
|
k_s = get_rank(s_e, styl_feat.size()[0])
|
|
s_d = (s_e[0:k_s]).pow(0.5)
|
|
EDE = torch.mm(
|
|
torch.mm(s_v[:, 0:k_s], torch.diag(s_d) * weight), (s_v[:, 0:k_s].t())
|
|
)
|
|
|
|
if registers is not None:
|
|
registers["EDE"] = EDE
|
|
registers["s_mean"] = s_mean
|
|
registers["c_v"] = c_v
|
|
registers["c_e"] = c_e
|
|
else:
|
|
EDE = registers["EDE"]
|
|
s_mean = registers["s_mean"]
|
|
_, c_e, c_v = svd(cont_feat, iden=True, device=device)
|
|
|
|
k_c = get_rank(c_e, cont_feat.size()[0])
|
|
c_d = (c_e[0:k_c]).pow(-0.5)
|
|
# TODO could be more fast
|
|
step1 = torch.mm(c_v[:, 0:k_c], torch.diag(c_d))
|
|
step2 = torch.mm(step1, (c_v[:, 0:k_c].t()))
|
|
whiten_cF = torch.mm(step2, cont_feat)
|
|
|
|
targetFeature = torch.mm(EDE, whiten_cF)
|
|
targetFeature = targetFeature + s_mean.unsqueeze(1).expand_as(targetFeature)
|
|
targetFeature.clamp_(cont_min, cont_max)
|
|
|
|
return targetFeature
|
|
|
|
|
|
def Bw_wct_core(content_feat, style_feat, weight=1, registers=None, device="cpu"):
|
|
N, C, H, W = content_feat.size()
|
|
cont_min = content_feat.min().item()
|
|
cont_max = content_feat.max().item()
|
|
|
|
whiten_cF, _, _ = SwitchWhiten2d(content_feat)
|
|
_, wm_s, s_mean = SwitchWhiten2d(style_feat)
|
|
|
|
targetFeature = torch.bmm(torch.inverse(wm_s), whiten_cF)
|
|
targetFeature = targetFeature.view(N, C, H, W)
|
|
targetFeature = targetFeature + s_mean.unsqueeze(2).expand_as(targetFeature)
|
|
targetFeature.clamp_(cont_min, cont_max)
|
|
|
|
return targetFeature
|
|
|
|
|
|
def wct_core_segment(
|
|
content_feat,
|
|
style_feat,
|
|
content_segment,
|
|
style_segment,
|
|
label_set,
|
|
label_indicator,
|
|
weight=1,
|
|
registers=None,
|
|
device="cpu",
|
|
):
|
|
def resize(feat, target):
|
|
size = (target.size(2), target.size(1))
|
|
if len(feat.shape) == 2:
|
|
return np.asarray(Image.fromarray(feat).resize(size, Image.NEAREST))
|
|
else:
|
|
return np.asarray(
|
|
Image.fromarray(feat, mode="RGB").resize(size, Image.NEAREST)
|
|
)
|
|
|
|
def get_index(feat, label):
|
|
mask = np.where(feat.reshape(feat.shape[0] * feat.shape[1]) == label)
|
|
if mask[0].size <= 0:
|
|
return None
|
|
return torch.LongTensor(mask[0]).cuda()
|
|
|
|
squeeze_content_feat = content_feat.squeeze(0)
|
|
squeeze_style_feat = style_feat.squeeze(0)
|
|
|
|
content_feat_view = squeeze_content_feat.view(
|
|
squeeze_content_feat.size(0), -1
|
|
).clone()
|
|
style_feat_view = squeeze_style_feat.view(squeeze_style_feat.size(0), -1).clone()
|
|
|
|
resized_content_segment = resize(content_segment, squeeze_content_feat)
|
|
resized_style_segment = resize(style_segment, squeeze_style_feat)
|
|
|
|
target_feature = content_feat_view.clone()
|
|
for label in label_set:
|
|
if not label_indicator[label]:
|
|
continue
|
|
content_index = get_index(resized_content_segment, label)
|
|
style_index = get_index(resized_style_segment, label)
|
|
if content_index is None or style_index is None:
|
|
continue
|
|
masked_content_feat = torch.index_select(content_feat_view, 1, content_index)
|
|
masked_style_feat = torch.index_select(style_feat_view, 1, style_index)
|
|
# _target_feature = Bw_wct_core(masked_content_feat, masked_style_feat, device=device)
|
|
_target_feature = wct_core(
|
|
masked_content_feat, masked_style_feat, weight, registers, device=device
|
|
)
|
|
if torch.__version__ >= "0.4.0":
|
|
# XXX reported bug in the original repository
|
|
new_target_feature = torch.transpose(target_feature, 1, 0)
|
|
new_target_feature.index_copy_(
|
|
0, content_index, torch.transpose(_target_feature, 1, 0)
|
|
)
|
|
target_feature = torch.transpose(new_target_feature, 1, 0)
|
|
else:
|
|
target_feature.index_copy_(1, content_index, _target_feature)
|
|
return target_feature
|
|
|
|
|
|
def feature_wct(
|
|
content_feat,
|
|
style_feat,
|
|
content_segment=None,
|
|
style_segment=None,
|
|
label_set=None,
|
|
label_indicator=None,
|
|
weight=1,
|
|
registers=None,
|
|
alpha=1,
|
|
device="cuda:0",
|
|
):
|
|
if label_set is not None:
|
|
target_feature = wct_core_segment(
|
|
content_feat,
|
|
style_feat,
|
|
content_segment,
|
|
style_segment,
|
|
label_set,
|
|
label_indicator,
|
|
weight,
|
|
registers,
|
|
device=device,
|
|
)
|
|
else:
|
|
target_feature = Bw_wct_core(content_feat, style_feat, device=device)
|
|
|
|
target_feature = target_feature.view_as(content_feat)
|
|
target_feature = alpha * target_feature + (1 - alpha) * content_feat
|
|
return target_feature
|
|
|
|
|
|
def feature_wct_simple(content_feat, style_feat, alpha=1):
|
|
target_feature = Bw_wct_core(content_feat, style_feat)
|
|
|
|
target_feature = target_feature.view_as(content_feat)
|
|
target_feature = alpha * target_feature + (1 - alpha) * content_feat
|
|
return target_feature
|
|
|
|
|
|
|
|
|
|
|
|
def init_weights(net):
|
|
for m in net.modules():
|
|
if isinstance(m, nn.Conv2d):
|
|
m.weight.data.normal_(0, 0.02)
|
|
if m.bias is not None:
|
|
m.bias.data.zero_()
|
|
elif isinstance(m, nn.ConvTranspose2d):
|
|
m.weight.data.normal_(0, 0.02)
|
|
if m.bias is not None:
|
|
m.bias.data.zero_()
|
|
elif isinstance(m, nn.Linear):
|
|
m.weight.data.normal_(0, 0.02)
|
|
if m.bias is not None:
|
|
m.bias.data.zero_()
|
|
|
|
|
|
class Timer:
|
|
def __init__(self, msg="Elapsed time: {}", verbose=True):
|
|
self.msg = msg
|
|
self.start_time = None
|
|
self.verbose = verbose
|
|
|
|
def __enter__(self):
|
|
self.start_time = datetime.datetime.now()
|
|
|
|
def __exit__(self, exc_type, exc_value, exc_tb):
|
|
if self.verbose:
|
|
print(self.msg.format(datetime.datetime.now() - self.start_time))
|
|
|
|
|
|
def _normalizer(denormalize=False):
|
|
# set Mean and Std of RGB channels of IMAGENET to use pre-trained VGG net
|
|
MEAN = [0.485, 0.456, 0.406]
|
|
STD = [0.229, 0.224, 0.225]
|
|
|
|
if denormalize:
|
|
MEAN = [-mean / std for mean, std in zip(MEAN, STD)]
|
|
STD = [1 / std for std in STD]
|
|
|
|
return transforms.Normalize(mean=MEAN, std=STD)
|
|
|
|
|
|
def open_image(image_path, image_size=None):
|
|
normalize = _normalizer()
|
|
image = Image.open(image_path)
|
|
_transforms = []
|
|
if image_size is not None:
|
|
image = transforms.Resize(image_size)(image)
|
|
# _transforms.append(transforms.Resize(image_size))
|
|
w, h = image.size
|
|
_transforms.append(transforms.CenterCrop((h // 16 * 16, w // 16 * 16)))
|
|
_transforms.append(transforms.ToTensor())
|
|
_transforms.append(normalize)
|
|
transform = transforms.Compose(_transforms)
|
|
return transform(image).unsqueeze(0)
|
|
|
|
|
|
def change_seg(seg):
|
|
color_dict = {
|
|
(0, 0, 255): 3, # blue
|
|
(0, 255, 0): 2, # green
|
|
(0, 0, 0): 0, # black
|
|
(255, 255, 255): 1, # white
|
|
(255, 0, 0): 4, # red
|
|
(255, 255, 0): 5, # yellow
|
|
(128, 128, 128): 6, # grey
|
|
(0, 255, 255): 7, # lightblue
|
|
(255, 0, 255): 8, # purple
|
|
}
|
|
arr_seg = np.asarray(seg)
|
|
new_seg = np.zeros(arr_seg.shape[:-1])
|
|
for x in range(arr_seg.shape[0]):
|
|
for y in range(arr_seg.shape[1]):
|
|
if tuple(arr_seg[x, y, :]) in color_dict:
|
|
new_seg[x, y] = color_dict[tuple(arr_seg[x, y, :])]
|
|
else:
|
|
min_dist_index = 0
|
|
min_dist = 99999
|
|
for key in color_dict:
|
|
dist = np.sum(np.abs(np.asarray(key) - arr_seg[x, y, :]))
|
|
if dist < min_dist:
|
|
min_dist = dist
|
|
min_dist_index = color_dict[key]
|
|
elif dist == min_dist:
|
|
try:
|
|
min_dist_index = new_seg[x, y - 1, :]
|
|
except Exception:
|
|
pass
|
|
new_seg[x, y] = min_dist_index
|
|
return new_seg.astype(np.uint8)
|
|
|
|
|
|
def load_segment(image_path, image_size=None):
|
|
if not image_path:
|
|
return np.asarray([])
|
|
image = Image.open(image_path)
|
|
if image_size is not None:
|
|
transform = transforms.Resize(image_size, interpolation=Image.NEAREST)
|
|
image = transform(image)
|
|
w, h = image.size
|
|
transform = transforms.CenterCrop((h // 16 * 16, w // 16 * 16))
|
|
image = transform(image)
|
|
if len(np.asarray(image).shape) == 3:
|
|
image = change_seg(image)
|
|
return np.asarray(image)
|
|
|
|
|
|
def compute_label_info(content_segment, style_segment):
|
|
if not content_segment.size or not style_segment.size:
|
|
return None, None
|
|
max_label = np.max(content_segment) + 1
|
|
label_set = np.unique(content_segment)
|
|
label_indicator = np.zeros(max_label)
|
|
for l in label_set: # noqa: E741
|
|
content_mask = np.where(
|
|
content_segment.reshape(content_segment.shape[0] * content_segment.shape[1])
|
|
== l
|
|
)
|
|
style_mask = np.where(
|
|
style_segment.reshape(style_segment.shape[0] * style_segment.shape[1]) == l
|
|
)
|
|
|
|
c_size = content_mask[0].size
|
|
s_size = style_mask[0].size
|
|
if (
|
|
c_size > 10
|
|
and s_size > 10
|
|
and c_size / s_size < 100
|
|
and s_size / c_size < 100
|
|
):
|
|
label_indicator[l] = True
|
|
else:
|
|
label_indicator[l] = False
|
|
return label_set, label_indicator
|
|
|
|
|
|
def mkdir(dname):
|
|
if not os.path.exists(dname):
|
|
os.makedirs(dname)
|
|
else:
|
|
assert os.path.isdir(dname), "alread exists filename {}".format(dname)
|
|
|
|
|
|
def gram_matrix(y):
|
|
(b, ch, h, w) = y.size()
|
|
features = y.view(b, ch, w * h)
|
|
features_t = features.transpose(1, 2)
|
|
gram = features.bmm(features_t) / (ch * h * w)
|
|
return gram
|
|
|
|
|
|
def TVloss(img, tv_weight):
|
|
"""
|
|
Compute total variation loss.
|
|
Inputs:
|
|
- img: PyTorch Variable of shape (1, 3, H, W) holding an input image.
|
|
- tv_weight: Scalar giving the weight w_t to use for the TV loss.
|
|
Returns:
|
|
- loss: PyTorch Variable holding a scalar giving the total variation loss
|
|
for img weighted by tv_weight.
|
|
"""
|
|
# w_variance = torch.sum(torch.pow(img[:, :, :, :-1] - img[:, :, :, 1:], 2))
|
|
# h_variance = torch.sum(torch.pow(img[:, :, :-1, :] - img[:, :, 1:, :], 2))
|
|
# loss = tv_weight * (h_variance + w_variance)
|
|
w_variance = torch.mean(torch.abs(img[:, :, :, :-1] - img[:, :, :, 1:]))
|
|
h_variance = torch.mean(torch.abs(img[:, :, :-1, :] - img[:, :, 1:, :]))
|
|
loss = tv_weight * (h_variance + w_variance)
|
|
return loss
|
|
|
|
|
|
def zeros_like(x):
|
|
return torch.autograd.Variable(torch.zeros_like(x).cuda())
|
|
|
|
|
|
def ones_like(x):
|
|
return torch.autograd.Variable(torch.ones_like(x).cuda())
|
|
|
|
|
|
def denorm_2(x):
|
|
out = (x + 1) / 2
|
|
return out.clamp_(0, 1)
|
|
|
|
|
|
def save_video(video, save_path, type="photo"):
|
|
video = denorm_2(video)
|
|
"""
|
|
vid_lst=[]
|
|
for i in range(0, num_samples*num_samples, num_samples):
|
|
temp_vid = list(video[i:i+num_samples])
|
|
temp_vid = torch.cat(temp_vid, dim=-1)
|
|
vid_lst.append(temp_vid)
|
|
|
|
save_videos = torch.cat(vid_lst, dim=2)
|
|
"""
|
|
|
|
save_videos = video.data.cpu().numpy().transpose(0, 2, 3, 1)
|
|
outputdata = save_videos * 255
|
|
# outputdata = ((save_videos+1)/2) * 255
|
|
outputdata = outputdata.astype(np.uint8)
|
|
dir_path = save_path
|
|
if not os.path.exists(dir_path):
|
|
os.makedirs(dir_path)
|
|
if type == "photo":
|
|
gif_file_path = os.path.join(dir_path, "Photo_StylizedVideo.gif")
|
|
elif type == "art":
|
|
gif_file_path = os.path.join(dir_path, "Art_StylizedVideo.gif")
|
|
else:
|
|
gif_file_path = os.path.join(dir_path, "content_StylizedVideo.gif")
|
|
imageio.mimsave(gif_file_path, outputdata, fps=25)
|
|
|
|
|
|
class GuidedFilter(nn.Module):
|
|
def box_filter(self, x: torch.Tensor, r):
|
|
ch = x.shape[1]
|
|
k = 2 * r + 1
|
|
weight = 1 / ((k) ** 2) # 1/9
|
|
# [c,1,3,3] * 1/9
|
|
box_kernel = torch.ones(
|
|
(ch, 1, k, k), dtype=torch.float32, device=x.device
|
|
).fill_(weight)
|
|
# same padding
|
|
return torch.nn.functional.conv2d(x, box_kernel, padding=r, groups=ch)
|
|
|
|
def forward(self, x: torch.Tensor, y: torch.Tensor, r, eps=1e-2):
|
|
b, c, h, w = x.shape
|
|
device = x.device
|
|
|
|
N = self.box_filter(torch.ones((1, 1, h, w), dtype=x.dtype, device=device), r)
|
|
|
|
mean_x = self.box_filter(x, r) / N
|
|
mean_y = self.box_filter(y, r) / N
|
|
cov_xy = self.box_filter(x * y, r) / N - mean_x * mean_y
|
|
var_x = self.box_filter(x * x, r) / N - mean_x * mean_x
|
|
|
|
A = cov_xy / (var_x + eps)
|
|
b = mean_y - A * mean_x
|
|
|
|
mean_A = self.box_filter(A, r) / N
|
|
mean_b = self.box_filter(b, r) / N
|
|
|
|
output = mean_A * x + mean_b
|
|
return output
|
|
|
|
|
|
def attn_visualization_all(Ic, Is, Fc, Fs, Attn, slide_size=4):
|
|
if Ic.size(3) != Is.size(3):
|
|
Is = torch.nn.functional.interpolate(Is, size=(Ic.size(2), Ic.size(3)))
|
|
B, C, S = Attn.size()
|
|
B, c, c_w, c_h = Ic.size()
|
|
B, c, s_w, s_h = Is.size()
|
|
_, _, f_w, f_h = Fs.size()
|
|
attn = torch.zeros(1, Fs.size(2) * Fs.size(3)).cuda()
|
|
mask = torch.zeros(1, Fc.size(2) * Fc.size(3)).cuda()
|
|
for index in range(0, Fc.size(2) * Fc.size(3), slide_size):
|
|
for idx in range(slide_size):
|
|
start_idx = index + Fc.size(3) * idx
|
|
# start_idx_ref = index + Fs.size(3) * idx
|
|
mask[0][start_idx : start_idx + slide_size] = 1
|
|
attn += torch.sum(Attn[0][start_idx : start_idx + slide_size], dim=0)
|
|
masked = torch.nn.functional.interpolate(
|
|
mask.view(1, 1, Fc.size(2), Fc.size(3)), size=(c_w, c_h), mode="nearest"
|
|
)
|
|
attn = torch.nn.functional.interpolate(
|
|
attn.view(1, 1, Fs.size(2), Fs.size(3)), size=(s_w, s_h), mode="nearest"
|
|
)
|
|
# imsave_no_norm(Is, attn.repeat(1,3,1,1),'style.png')
|
|
return (
|
|
attn.repeat(1, 3, 1, 1),
|
|
Ic * (masked + 0.2),
|
|
Is * (attn * (1 / attn.max()) + 0.2),
|
|
Is * (attn * 0.05 + 0.2),
|
|
)
|