plugins/plugin_dmdnet.py

from chain_img_processor import ChainImgProcessor, ChainImgPlugin
import os
from PIL import Image
from numpy import asarray

import torch
import torch.nn as nn
import torch.nn.functional as F
import scipy.io as sio
import numpy as np
import torch.nn.utils.spectral_norm as SpectralNorm
from torchvision.ops import roi_align

from math import sqrt
import os

import cv2 
import os
from torchvision.transforms.functional import normalize
import copy
import threading

modname = os.path.basename(__file__)[:-3] # calculating modname

oDMDNet = None
device = None

THREAD_LOCK_DMDNET = threading.Lock()



# start function
def start(core:ChainImgProcessor):
    manifest = { # plugin settings
        "name": "DMDNet", # name
        "version": "1.0", # version

        "default_options": {},
        "img_processor": {
            "dmdnet": DMDNETPlugin
        }
    }
    return manifest

def start_with_options(core:ChainImgProcessor, manifest:dict):
    pass


class DMDNETPlugin(ChainImgPlugin):

    # 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 init_plugin(self):
        global create
        
        if oDMDNet == None:
            create(self.device)


    def process(self, frame, params:dict):
        if "face_detected" in params:
            if not params["face_detected"]:
                return frame

        temp_frame = copy.copy(frame)
        if "processed_faces" in params:
            for face in params["processed_faces"]:
                start_x, start_y, end_x, end_y = map(int, face['bbox'])
                # padding_x = int((end_x - start_x) * 0.5)
                # padding_y = int((end_y - start_y) * 0.5)
                padding_x = 0
                padding_y = 0

                start_x = max(0, start_x - padding_x)
                start_y = max(0, start_y - padding_y)
                end_x = max(0, end_x + padding_x)
                end_y = max(0, end_y + padding_y)
                temp_face = temp_frame[start_y:end_y, start_x:end_x]
                if temp_face.size:
                    temp_face = self.enhance_face(temp_face, face)
                    temp_face = cv2.resize(temp_face, (end_x - start_x,end_y - start_y), interpolation = cv2.INTER_LANCZOS4)
                    temp_frame[start_y:end_y, start_x:end_x] = temp_face

        temp_frame = Image.blend(Image.fromarray(frame), Image.fromarray(temp_frame), params["blend_ratio"])
        return asarray(temp_frame)


    def enhance_face(self, clip, face):
        global device

        lm106 = face.landmark_2d_106
        lq_landmarks = asarray(self.landmarks106_to_68(lm106))
        lq = read_img_tensor(clip, False)

        LQLocs = get_component_location(lq_landmarks)
        # generic
        SpMem256Para, SpMem128Para, SpMem64Para = None, None, None

        with torch.no_grad():
            with THREAD_LOCK_DMDNET:
                try:
                    GenericResult, SpecificResult = oDMDNet(lq = lq.to(device), 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 clip
        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
        enhanced_img =  np.hstack((check_lq, save_generic))
        temp_frame =  save_generic.astype("uint8")
        # temp_frame = save_generic.astype("uint8")
        return temp_frame
    

def create(devicename):
    global device, oDMDNet

    test = "cuda" if torch.cuda.is_available() else "cpu"
    device = torch.device(devicename)
    oDMDNet = DMDNet().to(device)
    weights = torch.load('./models/DMDNet.pth') 
    oDMDNet.load_state_dict(weights, strict=True)

    oDMDNet.eval()
    num_params = 0
    for param in oDMDNet.parameters():
        num_params += param.numel()

    # print('{:>8s} : {}'.format('Using device', device))
    # print('{:>8s} : {:.2f}M'.format('Model params', num_params/1e6))



def read_img_tensor(Img=None, return_landmark=True): #rgb -1~1 
#    Img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)  # BGR or G
    if Img.ndim == 2:
        Img = cv2.cvtColor(Img, cv2.COLOR_GRAY2RGB)  # GGG
    else:
        Img = cv2.cvtColor(Img, cv2.COLOR_BGR2RGB)  # RGB
    
    if Img.shape[0] < 512 or Img.shape[1] < 512:
        Img = cv2.resize(Img, (512,512), interpolation = cv2.INTER_AREA)
    # ImgForLands = Img.copy()

    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