from typing import Any, List, Callable import cv2 import numpy as np import torch import torch.nn as nn import torch.nn.functional as F import torch.nn.utils.spectral_norm as SpectralNorm import threading from torchvision.ops import roi_align from math import sqrt from torchvision.transforms.functional import normalize from roop.typing import Face, Frame, FaceSet THREAD_LOCK_DMDNET = threading.Lock() class Enhance_DMDNet(): model_dmdnet = None torchdevice = None processorname = 'dmdnet' type = 'enhance' def Initialize(self, devicename): if self.model_dmdnet is None: self.model_dmdnet = self.create(devicename) # temp_frame already cropped+aligned, bbox not def Run(self, source_faceset: FaceSet, target_face: Face, temp_frame: Frame) -> Frame: input_size = temp_frame.shape[1] result = self.enhance_face(source_faceset, temp_frame, target_face) scale_factor = int(result.shape[1] / input_size) return result.astype(np.uint8), scale_factor def Release(self): self.model_gfpgan = None # https://stackoverflow.com/a/67174339 def landmarks106_to_68(self, pt106): map106to68=[1,10,12,14,16,3,5,7,0,23,21,19,32,30,28,26,17, 43,48,49,51,50, 102,103,104,105,101, 72,73,74,86,78,79,80,85,84, 35,41,42,39,37,36, 89,95,96,93,91,90, 52,64,63,71,67,68,61,58,59,53,56,55,65,66,62,70,69,57,60,54 ] pt68 = [] for i in range(68): index = map106to68[i] pt68.append(pt106[index]) return pt68 def check_bbox(self, imgs, boxes): boxes = boxes.view(-1, 4, 4) colors = [(0, 255, 0), (0, 255, 0), (255, 255, 0), (255, 0, 0)] i = 0 for img, box in zip(imgs, boxes): img = (img + 1)/2 * 255 img2 = img.permute(1, 2, 0).float().cpu().flip(2).numpy().copy() for idx, point in enumerate(box): cv2.rectangle(img2, (int(point[0]), int(point[1])), (int(point[2]), int(point[3])), color=colors[idx], thickness=2) cv2.imwrite('dmdnet_{:02d}.png'.format(i), img2) i += 1 def trans_points2d(self, pts, M): new_pts = np.zeros(shape=pts.shape, dtype=np.float32) for i in range(pts.shape[0]): pt = pts[i] new_pt = np.array([pt[0], pt[1], 1.0], dtype=np.float32) new_pt = np.dot(M, new_pt) new_pts[i] = new_pt[0:2] return new_pts def enhance_face(self, ref_faceset: FaceSet, temp_frame, face: Face): # preprocess start_x, start_y, end_x, end_y = map(int, face['bbox']) lm106 = face.landmark_2d_106 lq_landmarks = np.asarray(self.landmarks106_to_68(lm106)) if temp_frame.shape[0] != 512 or temp_frame.shape[1] != 512: # scale to 512x512 scale_factor = 512 / temp_frame.shape[1] M = face.matrix * scale_factor lq_landmarks = self.trans_points2d(lq_landmarks, M) temp_frame = cv2.resize(temp_frame, (512,512), interpolation = cv2.INTER_AREA) if temp_frame.ndim == 2: temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_GRAY2RGB) # GGG # else: # temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB) # RGB lq = read_img_tensor(temp_frame) LQLocs = get_component_location(lq_landmarks) # self.check_bbox(lq, LQLocs.unsqueeze(0)) # specific, change 1000 to 1 to activate if len(ref_faceset.faces) > 1: SpecificImgs = [] SpecificLocs = [] for i,face in enumerate(ref_faceset.faces): lm106 = face.landmark_2d_106 lq_landmarks = np.asarray(self.landmarks106_to_68(lm106)) ref_image = ref_faceset.ref_images[i] if ref_image.shape[0] != 512 or ref_image.shape[1] != 512: # scale to 512x512 scale_factor = 512 / ref_image.shape[1] M = face.matrix * scale_factor lq_landmarks = self.trans_points2d(lq_landmarks, M) ref_image = cv2.resize(ref_image, (512,512), interpolation = cv2.INTER_AREA) if ref_image.ndim == 2: temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_GRAY2RGB) # GGG # else: # temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB) # RGB ref_tensor = read_img_tensor(ref_image) ref_locs = get_component_location(lq_landmarks) # self.check_bbox(ref_tensor, ref_locs.unsqueeze(0)) SpecificImgs.append(ref_tensor) SpecificLocs.append(ref_locs.unsqueeze(0)) SpecificImgs = torch.cat(SpecificImgs, dim=0) SpecificLocs = torch.cat(SpecificLocs, dim=0) # check_bbox(SpecificImgs, SpecificLocs) SpMem256, SpMem128, SpMem64 = self.model_dmdnet.generate_specific_dictionary(sp_imgs = SpecificImgs.to(self.torchdevice), sp_locs = SpecificLocs) SpMem256Para = {} SpMem128Para = {} SpMem64Para = {} for k, v in SpMem256.items(): SpMem256Para[k] = v for k, v in SpMem128.items(): SpMem128Para[k] = v for k, v in SpMem64.items(): SpMem64Para[k] = v else: # generic SpMem256Para, SpMem128Para, SpMem64Para = None, None, None with torch.no_grad(): with THREAD_LOCK_DMDNET: try: GenericResult, SpecificResult = self.model_dmdnet(lq = lq.to(self.torchdevice), loc = LQLocs.unsqueeze(0), sp_256 = SpMem256Para, sp_128 = SpMem128Para, sp_64 = SpMem64Para) except Exception as e: print(f'Error {e} there may be something wrong with the detected component locations.') return temp_frame if SpecificResult is not None: save_specific = SpecificResult * 0.5 + 0.5 save_specific = save_specific.squeeze(0).permute(1, 2, 0).flip(2) # RGB->BGR save_specific = np.clip(save_specific.float().cpu().numpy(), 0, 1) * 255.0 temp_frame = save_specific.astype("uint8") if False: save_generic = GenericResult * 0.5 + 0.5 save_generic = save_generic.squeeze(0).permute(1, 2, 0).flip(2) # RGB->BGR save_generic = np.clip(save_generic.float().cpu().numpy(), 0, 1) * 255.0 check_lq = lq * 0.5 + 0.5 check_lq = check_lq.squeeze(0).permute(1, 2, 0).flip(2) # RGB->BGR check_lq = np.clip(check_lq.float().cpu().numpy(), 0, 1) * 255.0 cv2.imwrite('dmdnet_comparison.png', cv2.cvtColor(np.hstack((check_lq, save_generic, save_specific)),cv2.COLOR_RGB2BGR)) else: save_generic = GenericResult * 0.5 + 0.5 save_generic = save_generic.squeeze(0).permute(1, 2, 0).flip(2) # RGB->BGR save_generic = np.clip(save_generic.float().cpu().numpy(), 0, 1) * 255.0 temp_frame = save_generic.astype("uint8") temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_RGB2BGR) # RGB return temp_frame def create(self, devicename): self.torchdevice = torch.device(devicename) model_dmdnet = DMDNet().to(self.torchdevice) weights = torch.load('./models/DMDNet.pth') model_dmdnet.load_state_dict(weights, strict=True) model_dmdnet.eval() num_params = 0 for param in model_dmdnet.parameters(): num_params += param.numel() return model_dmdnet # print('{:>8s} : {}'.format('Using device', device)) # print('{:>8s} : {:.2f}M'.format('Model params', num_params/1e6)) def read_img_tensor(Img=None): #rgb -1~1 Img = Img.transpose((2, 0, 1))/255.0 Img = torch.from_numpy(Img).float() normalize(Img, [0.5,0.5,0.5], [0.5,0.5,0.5], inplace=True) ImgTensor = Img.unsqueeze(0) return ImgTensor def get_component_location(Landmarks, re_read=False): if re_read: ReadLandmark = [] with open(Landmarks,'r') as f: for line in f: tmp = [float(i) for i in line.split(' ') if i != '\n'] ReadLandmark.append(tmp) ReadLandmark = np.array(ReadLandmark) # Landmarks = np.reshape(ReadLandmark, [-1, 2]) # 68*2 Map_LE_B = list(np.hstack((range(17,22), range(36,42)))) Map_RE_B = list(np.hstack((range(22,27), range(42,48)))) Map_LE = list(range(36,42)) Map_RE = list(range(42,48)) Map_NO = list(range(29,36)) Map_MO = list(range(48,68)) Landmarks[Landmarks>504]=504 Landmarks[Landmarks<8]=8 #left eye Mean_LE = np.mean(Landmarks[Map_LE],0) L_LE1 = Mean_LE[1] - np.min(Landmarks[Map_LE_B,1]) L_LE1 = L_LE1 * 1.3 L_LE2 = L_LE1 / 1.9 L_LE_xy = L_LE1 + L_LE2 L_LE_lt = [L_LE_xy/2, L_LE1] L_LE_rb = [L_LE_xy/2, L_LE2] Location_LE = np.hstack((Mean_LE - L_LE_lt + 1, Mean_LE + L_LE_rb)).astype(int) #right eye Mean_RE = np.mean(Landmarks[Map_RE],0) L_RE1 = Mean_RE[1] - np.min(Landmarks[Map_RE_B,1]) L_RE1 = L_RE1 * 1.3 L_RE2 = L_RE1 / 1.9 L_RE_xy = L_RE1 + L_RE2 L_RE_lt = [L_RE_xy/2, L_RE1] L_RE_rb = [L_RE_xy/2, L_RE2] Location_RE = np.hstack((Mean_RE - L_RE_lt + 1, Mean_RE + L_RE_rb)).astype(int) #nose Mean_NO = np.mean(Landmarks[Map_NO],0) L_NO1 =( np.max([Mean_NO[0] - Landmarks[31][0], Landmarks[35][0] - Mean_NO[0]])) * 1.25 L_NO2 = (Landmarks[33][1] - Mean_NO[1]) * 1.1 L_NO_xy = L_NO1 * 2 L_NO_lt = [L_NO_xy/2, L_NO_xy - L_NO2] L_NO_rb = [L_NO_xy/2, L_NO2] Location_NO = np.hstack((Mean_NO - L_NO_lt + 1, Mean_NO + L_NO_rb)).astype(int) #mouth Mean_MO = np.mean(Landmarks[Map_MO],0) L_MO = np.max((np.max(np.max(Landmarks[Map_MO],0) - np.min(Landmarks[Map_MO],0))/2,16)) * 1.1 MO_O = Mean_MO - L_MO + 1 MO_T = Mean_MO + L_MO MO_T[MO_T>510]=510 Location_MO = np.hstack((MO_O, MO_T)).astype(int) return torch.cat([torch.FloatTensor(Location_LE).unsqueeze(0), torch.FloatTensor(Location_RE).unsqueeze(0), torch.FloatTensor(Location_NO).unsqueeze(0), torch.FloatTensor(Location_MO).unsqueeze(0)], dim=0) def calc_mean_std_4D(feat, eps=1e-5): # eps is a small value added to the variance to avoid divide-by-zero. size = feat.size() assert (len(size) == 4) N, C = size[:2] feat_var = feat.view(N, C, -1).var(dim=2) + eps feat_std = feat_var.sqrt().view(N, C, 1, 1) feat_mean = feat.view(N, C, -1).mean(dim=2).view(N, C, 1, 1) return feat_mean, feat_std def adaptive_instance_normalization_4D(content_feat, style_feat): # content_feat is ref feature, style is degradate feature size = content_feat.size() style_mean, style_std = calc_mean_std_4D(style_feat) content_mean, content_std = calc_mean_std_4D(content_feat) normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(size) return normalized_feat * style_std.expand(size) + style_mean.expand(size) def convU(in_channels, out_channels,conv_layer, norm_layer, kernel_size=3, stride=1,dilation=1, bias=True): return nn.Sequential( SpectralNorm(conv_layer(in_channels, out_channels, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=((kernel_size-1)//2)*dilation, bias=bias)), nn.LeakyReLU(0.2), SpectralNorm(conv_layer(out_channels, out_channels, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=((kernel_size-1)//2)*dilation, bias=bias)), ) class MSDilateBlock(nn.Module): def __init__(self, in_channels,conv_layer=nn.Conv2d, norm_layer=nn.BatchNorm2d, kernel_size=3, dilation=[1,1,1,1], bias=True): super(MSDilateBlock, self).__init__() self.conv1 = convU(in_channels, in_channels,conv_layer, norm_layer, kernel_size,dilation=dilation[0], bias=bias) self.conv2 = convU(in_channels, in_channels,conv_layer, norm_layer, kernel_size,dilation=dilation[1], bias=bias) self.conv3 = convU(in_channels, in_channels,conv_layer, norm_layer, kernel_size,dilation=dilation[2], bias=bias) self.conv4 = convU(in_channels, in_channels,conv_layer, norm_layer, kernel_size,dilation=dilation[3], bias=bias) self.convi = SpectralNorm(conv_layer(in_channels*4, in_channels, kernel_size=kernel_size, stride=1, padding=(kernel_size-1)//2, bias=bias)) def forward(self, x): conv1 = self.conv1(x) conv2 = self.conv2(x) conv3 = self.conv3(x) conv4 = self.conv4(x) cat = torch.cat([conv1, conv2, conv3, conv4], 1) out = self.convi(cat) + x return out class AdaptiveInstanceNorm(nn.Module): def __init__(self, in_channel): super().__init__() self.norm = nn.InstanceNorm2d(in_channel) def forward(self, input, style): style_mean, style_std = calc_mean_std_4D(style) out = self.norm(input) size = input.size() out = style_std.expand(size) * out + style_mean.expand(size) return out class NoiseInjection(nn.Module): def __init__(self, channel): super().__init__() self.weight = nn.Parameter(torch.zeros(1, channel, 1, 1)) def forward(self, image, noise): if noise is None: b, c, h, w = image.shape noise = image.new_empty(b, 1, h, w).normal_() return image + self.weight * noise class StyledUpBlock(nn.Module): def __init__(self, in_channel, out_channel, kernel_size=3, padding=1,upsample=False, noise_inject=False): super().__init__() self.noise_inject = noise_inject if upsample: self.conv1 = nn.Sequential( nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False), SpectralNorm(nn.Conv2d(in_channel, out_channel, kernel_size, padding=padding)), nn.LeakyReLU(0.2), ) else: self.conv1 = nn.Sequential( SpectralNorm(nn.Conv2d(in_channel, out_channel, kernel_size, padding=padding)), nn.LeakyReLU(0.2), SpectralNorm(nn.Conv2d(out_channel, out_channel, kernel_size, padding=padding)), ) self.convup = nn.Sequential( nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False), SpectralNorm(nn.Conv2d(out_channel, out_channel, kernel_size, padding=padding)), nn.LeakyReLU(0.2), SpectralNorm(nn.Conv2d(out_channel, out_channel, kernel_size, padding=padding)), ) if self.noise_inject: self.noise1 = NoiseInjection(out_channel) self.lrelu1 = nn.LeakyReLU(0.2) self.ScaleModel1 = nn.Sequential( SpectralNorm(nn.Conv2d(in_channel,out_channel,3, 1, 1)), nn.LeakyReLU(0.2), SpectralNorm(nn.Conv2d(out_channel, out_channel, 3, 1, 1)) ) self.ShiftModel1 = nn.Sequential( SpectralNorm(nn.Conv2d(in_channel,out_channel,3, 1, 1)), nn.LeakyReLU(0.2), SpectralNorm(nn.Conv2d(out_channel, out_channel, 3, 1, 1)), ) def forward(self, input, style): out = self.conv1(input) out = self.lrelu1(out) Shift1 = self.ShiftModel1(style) Scale1 = self.ScaleModel1(style) out = out * Scale1 + Shift1 if self.noise_inject: out = self.noise1(out, noise=None) outup = self.convup(out) return outup #################################################################### ###############Face Dictionary Generator #################################################################### def AttentionBlock(in_channel): return nn.Sequential( SpectralNorm(nn.Conv2d(in_channel, in_channel, 3, 1, 1)), nn.LeakyReLU(0.2), SpectralNorm(nn.Conv2d(in_channel, in_channel, 3, 1, 1)), ) class DilateResBlock(nn.Module): def __init__(self, dim, dilation=[5,3] ): super(DilateResBlock, self).__init__() self.Res = nn.Sequential( SpectralNorm(nn.Conv2d(dim, dim, 3, 1, ((3-1)//2)*dilation[0], dilation[0])), nn.LeakyReLU(0.2), SpectralNorm(nn.Conv2d(dim, dim, 3, 1, ((3-1)//2)*dilation[1], dilation[1])), ) def forward(self, x): out = x + self.Res(x) return out class KeyValue(nn.Module): def __init__(self, indim, keydim, valdim): super(KeyValue, self).__init__() self.Key = nn.Sequential( SpectralNorm(nn.Conv2d(indim, keydim, kernel_size=(3,3), padding=(1,1), stride=1)), nn.LeakyReLU(0.2), SpectralNorm(nn.Conv2d(keydim, keydim, kernel_size=(3,3), padding=(1,1), stride=1)), ) self.Value = nn.Sequential( SpectralNorm(nn.Conv2d(indim, valdim, kernel_size=(3,3), padding=(1,1), stride=1)), nn.LeakyReLU(0.2), SpectralNorm(nn.Conv2d(valdim, valdim, kernel_size=(3,3), padding=(1,1), stride=1)), ) def forward(self, x): return self.Key(x), self.Value(x) class MaskAttention(nn.Module): def __init__(self, indim): super(MaskAttention, self).__init__() self.conv1 = nn.Sequential( SpectralNorm(nn.Conv2d(indim, indim//3, kernel_size=(3,3), padding=(1,1), stride=1)), nn.LeakyReLU(0.2), SpectralNorm(nn.Conv2d(indim//3, indim//3, kernel_size=(3,3), padding=(1,1), stride=1)), ) self.conv2 = nn.Sequential( SpectralNorm(nn.Conv2d(indim, indim//3, kernel_size=(3,3), padding=(1,1), stride=1)), nn.LeakyReLU(0.2), SpectralNorm(nn.Conv2d(indim//3, indim//3, kernel_size=(3,3), padding=(1,1), stride=1)), ) self.conv3 = nn.Sequential( SpectralNorm(nn.Conv2d(indim, indim//3, kernel_size=(3,3), padding=(1,1), stride=1)), nn.LeakyReLU(0.2), SpectralNorm(nn.Conv2d(indim//3, indim//3, kernel_size=(3,3), padding=(1,1), stride=1)), ) self.convCat = nn.Sequential( SpectralNorm(nn.Conv2d(indim//3 * 3, indim, kernel_size=(3,3), padding=(1,1), stride=1)), nn.LeakyReLU(0.2), SpectralNorm(nn.Conv2d(indim, indim, kernel_size=(3,3), padding=(1,1), stride=1)), ) def forward(self, x, y, z): c1 = self.conv1(x) c2 = self.conv2(y) c3 = self.conv3(z) return self.convCat(torch.cat([c1,c2,c3], dim=1)) class Query(nn.Module): def __init__(self, indim, quedim): super(Query, self).__init__() self.Query = nn.Sequential( SpectralNorm(nn.Conv2d(indim, quedim, kernel_size=(3,3), padding=(1,1), stride=1)), nn.LeakyReLU(0.2), SpectralNorm(nn.Conv2d(quedim, quedim, kernel_size=(3,3), padding=(1,1), stride=1)), ) def forward(self, x): return self.Query(x) def roi_align_self(input, location, target_size): return torch.cat([F.interpolate(input[i:i+1,:,location[i,1]:location[i,3],location[i,0]:location[i,2]],(target_size,target_size),mode='bilinear',align_corners=False) for i in range(input.size(0))],0) class FeatureExtractor(nn.Module): def __init__(self, ngf = 64, key_scale = 4):# super().__init__() self.key_scale = 4 self.part_sizes = np.array([80,80,50,110]) # self.feature_sizes = np.array([256,128,64]) # self.conv1 = nn.Sequential( SpectralNorm(nn.Conv2d(3, ngf, 3, 2, 1)), nn.LeakyReLU(0.2), SpectralNorm(nn.Conv2d(ngf, ngf, 3, 1, 1)), ) self.conv2 = nn.Sequential( SpectralNorm(nn.Conv2d(ngf, ngf, 3, 1, 1)), nn.LeakyReLU(0.2), SpectralNorm(nn.Conv2d(ngf, ngf, 3, 1, 1)) ) self.res1 = DilateResBlock(ngf, [5,3]) self.res2 = DilateResBlock(ngf, [5,3]) self.conv3 = nn.Sequential( SpectralNorm(nn.Conv2d(ngf, ngf*2, 3, 2, 1)), nn.LeakyReLU(0.2), SpectralNorm(nn.Conv2d(ngf*2, ngf*2, 3, 1, 1)), ) self.conv4 = nn.Sequential( SpectralNorm(nn.Conv2d(ngf*2, ngf*2, 3, 1, 1)), nn.LeakyReLU(0.2), SpectralNorm(nn.Conv2d(ngf*2, ngf*2, 3, 1, 1)) ) self.res3 = DilateResBlock(ngf*2, [3,1]) self.res4 = DilateResBlock(ngf*2, [3,1]) self.conv5 = nn.Sequential( SpectralNorm(nn.Conv2d(ngf*2, ngf*4, 3, 2, 1)), nn.LeakyReLU(0.2), SpectralNorm(nn.Conv2d(ngf*4, ngf*4, 3, 1, 1)), ) self.conv6 = nn.Sequential( SpectralNorm(nn.Conv2d(ngf*4, ngf*4, 3, 1, 1)), nn.LeakyReLU(0.2), SpectralNorm(nn.Conv2d(ngf*4, ngf*4, 3, 1, 1)) ) self.res5 = DilateResBlock(ngf*4, [1,1]) self.res6 = DilateResBlock(ngf*4, [1,1]) self.LE_256_Q = Query(ngf, ngf // self.key_scale) self.RE_256_Q = Query(ngf, ngf // self.key_scale) self.MO_256_Q = Query(ngf, ngf // self.key_scale) self.LE_128_Q = Query(ngf * 2, ngf * 2 // self.key_scale) self.RE_128_Q = Query(ngf * 2, ngf * 2 // self.key_scale) self.MO_128_Q = Query(ngf * 2, ngf * 2 // self.key_scale) self.LE_64_Q = Query(ngf * 4, ngf * 4 // self.key_scale) self.RE_64_Q = Query(ngf * 4, ngf * 4 // self.key_scale) self.MO_64_Q = Query(ngf * 4, ngf * 4 // self.key_scale) def forward(self, img, locs): le_location = locs[:,0,:].int().cpu().numpy() re_location = locs[:,1,:].int().cpu().numpy() no_location = locs[:,2,:].int().cpu().numpy() mo_location = locs[:,3,:].int().cpu().numpy() f1_0 = self.conv1(img) f1_1 = self.res1(f1_0) f2_0 = self.conv2(f1_1) f2_1 = self.res2(f2_0) f3_0 = self.conv3(f2_1) f3_1 = self.res3(f3_0) f4_0 = self.conv4(f3_1) f4_1 = self.res4(f4_0) f5_0 = self.conv5(f4_1) f5_1 = self.res5(f5_0) f6_0 = self.conv6(f5_1) f6_1 = self.res6(f6_0) ####ROI Align le_part_256 = roi_align_self(f2_1.clone(), le_location//2, self.part_sizes[0]//2) re_part_256 = roi_align_self(f2_1.clone(), re_location//2, self.part_sizes[1]//2) mo_part_256 = roi_align_self(f2_1.clone(), mo_location//2, self.part_sizes[3]//2) le_part_128 = roi_align_self(f4_1.clone(), le_location//4, self.part_sizes[0]//4) re_part_128 = roi_align_self(f4_1.clone(), re_location//4, self.part_sizes[1]//4) mo_part_128 = roi_align_self(f4_1.clone(), mo_location//4, self.part_sizes[3]//4) le_part_64 = roi_align_self(f6_1.clone(), le_location//8, self.part_sizes[0]//8) re_part_64 = roi_align_self(f6_1.clone(), re_location//8, self.part_sizes[1]//8) mo_part_64 = roi_align_self(f6_1.clone(), mo_location//8, self.part_sizes[3]//8) le_256_q = self.LE_256_Q(le_part_256) re_256_q = self.RE_256_Q(re_part_256) mo_256_q = self.MO_256_Q(mo_part_256) le_128_q = self.LE_128_Q(le_part_128) re_128_q = self.RE_128_Q(re_part_128) mo_128_q = self.MO_128_Q(mo_part_128) le_64_q = self.LE_64_Q(le_part_64) re_64_q = self.RE_64_Q(re_part_64) mo_64_q = self.MO_64_Q(mo_part_64) return {'f256': f2_1, 'f128': f4_1, 'f64': f6_1,\ 'le256': le_part_256, 're256': re_part_256, 'mo256': mo_part_256, \ 'le128': le_part_128, 're128': re_part_128, 'mo128': mo_part_128, \ 'le64': le_part_64, 're64': re_part_64, 'mo64': mo_part_64, \ 'le_256_q': le_256_q, 're_256_q': re_256_q, 'mo_256_q': mo_256_q,\ 'le_128_q': le_128_q, 're_128_q': re_128_q, 'mo_128_q': mo_128_q,\ 'le_64_q': le_64_q, 're_64_q': re_64_q, 'mo_64_q': mo_64_q} class DMDNet(nn.Module): def __init__(self, ngf = 64, banks_num = 128): super().__init__() self.part_sizes = np.array([80,80,50,110]) # size for 512 self.feature_sizes = np.array([256,128,64]) # size for 512 self.banks_num = banks_num self.key_scale = 4 self.E_lq = FeatureExtractor(key_scale = self.key_scale) self.E_hq = FeatureExtractor(key_scale = self.key_scale) self.LE_256_KV = KeyValue(ngf, ngf // self.key_scale, ngf) self.RE_256_KV = KeyValue(ngf, ngf // self.key_scale, ngf) self.MO_256_KV = KeyValue(ngf, ngf // self.key_scale, ngf) self.LE_128_KV = KeyValue(ngf * 2 , ngf * 2 // self.key_scale, ngf * 2) self.RE_128_KV = KeyValue(ngf * 2 , ngf * 2 // self.key_scale, ngf * 2) self.MO_128_KV = KeyValue(ngf * 2 , ngf * 2 // self.key_scale, ngf * 2) self.LE_64_KV = KeyValue(ngf * 4 , ngf * 4 // self.key_scale, ngf * 4) self.RE_64_KV = KeyValue(ngf * 4 , ngf * 4 // self.key_scale, ngf * 4) self.MO_64_KV = KeyValue(ngf * 4 , ngf * 4 // self.key_scale, ngf * 4) self.LE_256_Attention = AttentionBlock(64) self.RE_256_Attention = AttentionBlock(64) self.MO_256_Attention = AttentionBlock(64) self.LE_128_Attention = AttentionBlock(128) self.RE_128_Attention = AttentionBlock(128) self.MO_128_Attention = AttentionBlock(128) self.LE_64_Attention = AttentionBlock(256) self.RE_64_Attention = AttentionBlock(256) self.MO_64_Attention = AttentionBlock(256) self.LE_256_Mask = MaskAttention(64) self.RE_256_Mask = MaskAttention(64) self.MO_256_Mask = MaskAttention(64) self.LE_128_Mask = MaskAttention(128) self.RE_128_Mask = MaskAttention(128) self.MO_128_Mask = MaskAttention(128) self.LE_64_Mask = MaskAttention(256) self.RE_64_Mask = MaskAttention(256) self.MO_64_Mask = MaskAttention(256) self.MSDilate = MSDilateBlock(ngf*4, dilation = [4,3,2,1]) self.up1 = StyledUpBlock(ngf*4, ngf*2, noise_inject=False) # self.up2 = StyledUpBlock(ngf*2, ngf, noise_inject=False) # self.up3 = StyledUpBlock(ngf, ngf, noise_inject=False) # self.up4 = nn.Sequential( SpectralNorm(nn.Conv2d(ngf, ngf, 3, 1, 1)), nn.LeakyReLU(0.2), UpResBlock(ngf), UpResBlock(ngf), SpectralNorm(nn.Conv2d(ngf, 3, kernel_size=3, stride=1, padding=1)), nn.Tanh() ) # define generic memory, revise register_buffer to register_parameter for backward update self.register_buffer('le_256_mem_key', torch.randn(128,16,40,40)) self.register_buffer('re_256_mem_key', torch.randn(128,16,40,40)) self.register_buffer('mo_256_mem_key', torch.randn(128,16,55,55)) self.register_buffer('le_256_mem_value', torch.randn(128,64,40,40)) self.register_buffer('re_256_mem_value', torch.randn(128,64,40,40)) self.register_buffer('mo_256_mem_value', torch.randn(128,64,55,55)) self.register_buffer('le_128_mem_key', torch.randn(128,32,20,20)) self.register_buffer('re_128_mem_key', torch.randn(128,32,20,20)) self.register_buffer('mo_128_mem_key', torch.randn(128,32,27,27)) self.register_buffer('le_128_mem_value', torch.randn(128,128,20,20)) self.register_buffer('re_128_mem_value', torch.randn(128,128,20,20)) self.register_buffer('mo_128_mem_value', torch.randn(128,128,27,27)) self.register_buffer('le_64_mem_key', torch.randn(128,64,10,10)) self.register_buffer('re_64_mem_key', torch.randn(128,64,10,10)) self.register_buffer('mo_64_mem_key', torch.randn(128,64,13,13)) self.register_buffer('le_64_mem_value', torch.randn(128,256,10,10)) self.register_buffer('re_64_mem_value', torch.randn(128,256,10,10)) self.register_buffer('mo_64_mem_value', torch.randn(128,256,13,13)) def readMem(self, k, v, q): sim = F.conv2d(q, k) score = F.softmax(sim/sqrt(sim.size(1)), dim=1) #B * S * 1 * 1 6*128 sb,sn,sw,sh = score.size() s_m = score.view(sb, -1).unsqueeze(1)#2*1*M vb,vn,vw,vh = v.size() v_in = v.view(vb, -1).repeat(sb,1,1)#2*M*(c*w*h) mem_out = torch.bmm(s_m, v_in).squeeze(1).view(sb, vn, vw,vh) max_inds = torch.argmax(score, dim=1).squeeze() return mem_out, max_inds def memorize(self, img, locs): fs = self.E_hq(img, locs) LE256_key, LE256_value = self.LE_256_KV(fs['le256']) RE256_key, RE256_value = self.RE_256_KV(fs['re256']) MO256_key, MO256_value = self.MO_256_KV(fs['mo256']) LE128_key, LE128_value = self.LE_128_KV(fs['le128']) RE128_key, RE128_value = self.RE_128_KV(fs['re128']) MO128_key, MO128_value = self.MO_128_KV(fs['mo128']) LE64_key, LE64_value = self.LE_64_KV(fs['le64']) RE64_key, RE64_value = self.RE_64_KV(fs['re64']) MO64_key, MO64_value = self.MO_64_KV(fs['mo64']) Mem256 = {'LE256Key': LE256_key, 'LE256Value': LE256_value, 'RE256Key': RE256_key, 'RE256Value': RE256_value,'MO256Key': MO256_key, 'MO256Value': MO256_value} Mem128 = {'LE128Key': LE128_key, 'LE128Value': LE128_value, 'RE128Key': RE128_key, 'RE128Value': RE128_value,'MO128Key': MO128_key, 'MO128Value': MO128_value} Mem64 = {'LE64Key': LE64_key, 'LE64Value': LE64_value, 'RE64Key': RE64_key, 'RE64Value': RE64_value,'MO64Key': MO64_key, 'MO64Value': MO64_value} FS256 = {'LE256F':fs['le256'], 'RE256F':fs['re256'], 'MO256F':fs['mo256']} FS128 = {'LE128F':fs['le128'], 'RE128F':fs['re128'], 'MO128F':fs['mo128']} FS64 = {'LE64F':fs['le64'], 'RE64F':fs['re64'], 'MO64F':fs['mo64']} return Mem256, Mem128, Mem64 def enhancer(self, fs_in, sp_256=None, sp_128=None, sp_64=None): le_256_q = fs_in['le_256_q'] re_256_q = fs_in['re_256_q'] mo_256_q = fs_in['mo_256_q'] le_128_q = fs_in['le_128_q'] re_128_q = fs_in['re_128_q'] mo_128_q = fs_in['mo_128_q'] le_64_q = fs_in['le_64_q'] re_64_q = fs_in['re_64_q'] mo_64_q = fs_in['mo_64_q'] ####for 256 le_256_mem_g, le_256_inds = self.readMem(self.le_256_mem_key, self.le_256_mem_value, le_256_q) re_256_mem_g, re_256_inds = self.readMem(self.re_256_mem_key, self.re_256_mem_value, re_256_q) mo_256_mem_g, mo_256_inds = self.readMem(self.mo_256_mem_key, self.mo_256_mem_value, mo_256_q) le_128_mem_g, le_128_inds = self.readMem(self.le_128_mem_key, self.le_128_mem_value, le_128_q) re_128_mem_g, re_128_inds = self.readMem(self.re_128_mem_key, self.re_128_mem_value, re_128_q) mo_128_mem_g, mo_128_inds = self.readMem(self.mo_128_mem_key, self.mo_128_mem_value, mo_128_q) le_64_mem_g, le_64_inds = self.readMem(self.le_64_mem_key, self.le_64_mem_value, le_64_q) re_64_mem_g, re_64_inds = self.readMem(self.re_64_mem_key, self.re_64_mem_value, re_64_q) mo_64_mem_g, mo_64_inds = self.readMem(self.mo_64_mem_key, self.mo_64_mem_value, mo_64_q) if sp_256 is not None and sp_128 is not None and sp_64 is not None: le_256_mem_s, _ = self.readMem(sp_256['LE256Key'], sp_256['LE256Value'], le_256_q) re_256_mem_s, _ = self.readMem(sp_256['RE256Key'], sp_256['RE256Value'], re_256_q) mo_256_mem_s, _ = self.readMem(sp_256['MO256Key'], sp_256['MO256Value'], mo_256_q) le_256_mask = self.LE_256_Mask(fs_in['le256'],le_256_mem_s,le_256_mem_g) le_256_mem = le_256_mask*le_256_mem_s + (1-le_256_mask)*le_256_mem_g re_256_mask = self.RE_256_Mask(fs_in['re256'],re_256_mem_s,re_256_mem_g) re_256_mem = re_256_mask*re_256_mem_s + (1-re_256_mask)*re_256_mem_g mo_256_mask = self.MO_256_Mask(fs_in['mo256'],mo_256_mem_s,mo_256_mem_g) mo_256_mem = mo_256_mask*mo_256_mem_s + (1-mo_256_mask)*mo_256_mem_g le_128_mem_s, _ = self.readMem(sp_128['LE128Key'], sp_128['LE128Value'], le_128_q) re_128_mem_s, _ = self.readMem(sp_128['RE128Key'], sp_128['RE128Value'], re_128_q) mo_128_mem_s, _ = self.readMem(sp_128['MO128Key'], sp_128['MO128Value'], mo_128_q) le_128_mask = self.LE_128_Mask(fs_in['le128'],le_128_mem_s,le_128_mem_g) le_128_mem = le_128_mask*le_128_mem_s + (1-le_128_mask)*le_128_mem_g re_128_mask = self.RE_128_Mask(fs_in['re128'],re_128_mem_s,re_128_mem_g) re_128_mem = re_128_mask*re_128_mem_s + (1-re_128_mask)*re_128_mem_g mo_128_mask = self.MO_128_Mask(fs_in['mo128'],mo_128_mem_s,mo_128_mem_g) mo_128_mem = mo_128_mask*mo_128_mem_s + (1-mo_128_mask)*mo_128_mem_g le_64_mem_s, _ = self.readMem(sp_64['LE64Key'], sp_64['LE64Value'], le_64_q) re_64_mem_s, _ = self.readMem(sp_64['RE64Key'], sp_64['RE64Value'], re_64_q) mo_64_mem_s, _ = self.readMem(sp_64['MO64Key'], sp_64['MO64Value'], mo_64_q) le_64_mask = self.LE_64_Mask(fs_in['le64'],le_64_mem_s,le_64_mem_g) le_64_mem = le_64_mask*le_64_mem_s + (1-le_64_mask)*le_64_mem_g re_64_mask = self.RE_64_Mask(fs_in['re64'],re_64_mem_s,re_64_mem_g) re_64_mem = re_64_mask*re_64_mem_s + (1-re_64_mask)*re_64_mem_g mo_64_mask = self.MO_64_Mask(fs_in['mo64'],mo_64_mem_s,mo_64_mem_g) mo_64_mem = mo_64_mask*mo_64_mem_s + (1-mo_64_mask)*mo_64_mem_g else: le_256_mem = le_256_mem_g re_256_mem = re_256_mem_g mo_256_mem = mo_256_mem_g le_128_mem = le_128_mem_g re_128_mem = re_128_mem_g mo_128_mem = mo_128_mem_g le_64_mem = le_64_mem_g re_64_mem = re_64_mem_g mo_64_mem = mo_64_mem_g le_256_mem_norm = adaptive_instance_normalization_4D(le_256_mem, fs_in['le256']) re_256_mem_norm = adaptive_instance_normalization_4D(re_256_mem, fs_in['re256']) mo_256_mem_norm = adaptive_instance_normalization_4D(mo_256_mem, fs_in['mo256']) ####for 128 le_128_mem_norm = adaptive_instance_normalization_4D(le_128_mem, fs_in['le128']) re_128_mem_norm = adaptive_instance_normalization_4D(re_128_mem, fs_in['re128']) mo_128_mem_norm = adaptive_instance_normalization_4D(mo_128_mem, fs_in['mo128']) ####for 64 le_64_mem_norm = adaptive_instance_normalization_4D(le_64_mem, fs_in['le64']) re_64_mem_norm = adaptive_instance_normalization_4D(re_64_mem, fs_in['re64']) mo_64_mem_norm = adaptive_instance_normalization_4D(mo_64_mem, fs_in['mo64']) EnMem256 = {'LE256Norm': le_256_mem_norm, 'RE256Norm': re_256_mem_norm, 'MO256Norm': mo_256_mem_norm} EnMem128 = {'LE128Norm': le_128_mem_norm, 'RE128Norm': re_128_mem_norm, 'MO128Norm': mo_128_mem_norm} EnMem64 = {'LE64Norm': le_64_mem_norm, 'RE64Norm': re_64_mem_norm, 'MO64Norm': mo_64_mem_norm} Ind256 = {'LE': le_256_inds, 'RE': re_256_inds, 'MO': mo_256_inds} Ind128 = {'LE': le_128_inds, 'RE': re_128_inds, 'MO': mo_128_inds} Ind64 = {'LE': le_64_inds, 'RE': re_64_inds, 'MO': mo_64_inds} return EnMem256, EnMem128, EnMem64, Ind256, Ind128, Ind64 def reconstruct(self, fs_in, locs, memstar): le_256_mem_norm, re_256_mem_norm, mo_256_mem_norm = memstar[0]['LE256Norm'], memstar[0]['RE256Norm'], memstar[0]['MO256Norm'] le_128_mem_norm, re_128_mem_norm, mo_128_mem_norm = memstar[1]['LE128Norm'], memstar[1]['RE128Norm'], memstar[1]['MO128Norm'] le_64_mem_norm, re_64_mem_norm, mo_64_mem_norm = memstar[2]['LE64Norm'], memstar[2]['RE64Norm'], memstar[2]['MO64Norm'] le_256_final = self.LE_256_Attention(le_256_mem_norm - fs_in['le256']) * le_256_mem_norm + fs_in['le256'] re_256_final = self.RE_256_Attention(re_256_mem_norm - fs_in['re256']) * re_256_mem_norm + fs_in['re256'] mo_256_final = self.MO_256_Attention(mo_256_mem_norm - fs_in['mo256']) * mo_256_mem_norm + fs_in['mo256'] le_128_final = self.LE_128_Attention(le_128_mem_norm - fs_in['le128']) * le_128_mem_norm + fs_in['le128'] re_128_final = self.RE_128_Attention(re_128_mem_norm - fs_in['re128']) * re_128_mem_norm + fs_in['re128'] mo_128_final = self.MO_128_Attention(mo_128_mem_norm - fs_in['mo128']) * mo_128_mem_norm + fs_in['mo128'] le_64_final = self.LE_64_Attention(le_64_mem_norm - fs_in['le64']) * le_64_mem_norm + fs_in['le64'] re_64_final = self.RE_64_Attention(re_64_mem_norm - fs_in['re64']) * re_64_mem_norm + fs_in['re64'] mo_64_final = self.MO_64_Attention(mo_64_mem_norm - fs_in['mo64']) * mo_64_mem_norm + fs_in['mo64'] le_location = locs[:,0,:] re_location = locs[:,1,:] mo_location = locs[:,3,:] le_location = le_location.cpu().int().numpy() re_location = re_location.cpu().int().numpy() mo_location = mo_location.cpu().int().numpy() up_in_256 = fs_in['f256'].clone()# * 0 up_in_128 = fs_in['f128'].clone()# * 0 up_in_64 = fs_in['f64'].clone()# * 0 for i in range(fs_in['f256'].size(0)): up_in_256[i:i+1,:,le_location[i,1]//2:le_location[i,3]//2,le_location[i,0]//2:le_location[i,2]//2] = F.interpolate(le_256_final[i:i+1,:,:,:].clone(), (le_location[i,3]//2-le_location[i,1]//2,le_location[i,2]//2-le_location[i,0]//2),mode='bilinear',align_corners=False) up_in_256[i:i+1,:,re_location[i,1]//2:re_location[i,3]//2,re_location[i,0]//2:re_location[i,2]//2] = F.interpolate(re_256_final[i:i+1,:,:,:].clone(), (re_location[i,3]//2-re_location[i,1]//2,re_location[i,2]//2-re_location[i,0]//2),mode='bilinear',align_corners=False) up_in_256[i:i+1,:,mo_location[i,1]//2:mo_location[i,3]//2,mo_location[i,0]//2:mo_location[i,2]//2] = F.interpolate(mo_256_final[i:i+1,:,:,:].clone(), (mo_location[i,3]//2-mo_location[i,1]//2,mo_location[i,2]//2-mo_location[i,0]//2),mode='bilinear',align_corners=False) up_in_128[i:i+1,:,le_location[i,1]//4:le_location[i,3]//4,le_location[i,0]//4:le_location[i,2]//4] = F.interpolate(le_128_final[i:i+1,:,:,:].clone(), (le_location[i,3]//4-le_location[i,1]//4,le_location[i,2]//4-le_location[i,0]//4),mode='bilinear',align_corners=False) up_in_128[i:i+1,:,re_location[i,1]//4:re_location[i,3]//4,re_location[i,0]//4:re_location[i,2]//4] = F.interpolate(re_128_final[i:i+1,:,:,:].clone(), (re_location[i,3]//4-re_location[i,1]//4,re_location[i,2]//4-re_location[i,0]//4),mode='bilinear',align_corners=False) up_in_128[i:i+1,:,mo_location[i,1]//4:mo_location[i,3]//4,mo_location[i,0]//4:mo_location[i,2]//4] = F.interpolate(mo_128_final[i:i+1,:,:,:].clone(), (mo_location[i,3]//4-mo_location[i,1]//4,mo_location[i,2]//4-mo_location[i,0]//4),mode='bilinear',align_corners=False) up_in_64[i:i+1,:,le_location[i,1]//8:le_location[i,3]//8,le_location[i,0]//8:le_location[i,2]//8] = F.interpolate(le_64_final[i:i+1,:,:,:].clone(), (le_location[i,3]//8-le_location[i,1]//8,le_location[i,2]//8-le_location[i,0]//8),mode='bilinear',align_corners=False) up_in_64[i:i+1,:,re_location[i,1]//8:re_location[i,3]//8,re_location[i,0]//8:re_location[i,2]//8] = F.interpolate(re_64_final[i:i+1,:,:,:].clone(), (re_location[i,3]//8-re_location[i,1]//8,re_location[i,2]//8-re_location[i,0]//8),mode='bilinear',align_corners=False) up_in_64[i:i+1,:,mo_location[i,1]//8:mo_location[i,3]//8,mo_location[i,0]//8:mo_location[i,2]//8] = F.interpolate(mo_64_final[i:i+1,:,:,:].clone(), (mo_location[i,3]//8-mo_location[i,1]//8,mo_location[i,2]//8-mo_location[i,0]//8),mode='bilinear',align_corners=False) ms_in_64 = self.MSDilate(fs_in['f64'].clone()) fea_up1 = self.up1(ms_in_64, up_in_64) fea_up2 = self.up2(fea_up1, up_in_128) # fea_up3 = self.up3(fea_up2, up_in_256) # output = self.up4(fea_up3) # return output def generate_specific_dictionary(self, sp_imgs=None, sp_locs=None): return self.memorize(sp_imgs, sp_locs) def forward(self, lq=None, loc=None, sp_256 = None, sp_128 = None, sp_64 = None): fs_in = self.E_lq(lq, loc) # low quality images GeMemNorm256, GeMemNorm128, GeMemNorm64, Ind256, Ind128, Ind64 = self.enhancer(fs_in) GeOut = self.reconstruct(fs_in, loc, memstar = [GeMemNorm256, GeMemNorm128, GeMemNorm64]) if sp_256 is not None and sp_128 is not None and sp_64 is not None: GSMemNorm256, GSMemNorm128, GSMemNorm64, _, _, _ = self.enhancer(fs_in, sp_256, sp_128, sp_64) GSOut = self.reconstruct(fs_in, loc, memstar = [GSMemNorm256, GSMemNorm128, GSMemNorm64]) else: GSOut = None return GeOut, GSOut class UpResBlock(nn.Module): def __init__(self, dim, conv_layer = nn.Conv2d, norm_layer = nn.BatchNorm2d): super(UpResBlock, self).__init__() self.Model = nn.Sequential( SpectralNorm(conv_layer(dim, dim, 3, 1, 1)), nn.LeakyReLU(0.2), SpectralNorm(conv_layer(dim, dim, 3, 1, 1)), ) def forward(self, x): out = x + self.Model(x) return out