mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-08-22 06:23:20 -07:00
Added depthwise separable convolutions
This commit is contained in:
TalosOfCrete
parent
a36359f480
commit
4e46708303
5 changed files with 161 additions and 383 deletions
|
|
@ -24,7 +24,7 @@ class DeepFakeArchi(nn.ArchiBase):
|
|||
super().__init__(*kwargs)
|
||||
|
||||
def on_build(self, *args, **kwargs ):
|
||||
self.conv1 = nn.Conv2D( self.in_ch,
|
||||
self.conv1 = nn.SeparableConv2D( self.in_ch,
|
||||
self.out_ch // (4 if self.subpixel else 1),
|
||||
kernel_size=self.kernel_size,
|
||||
strides=1 if self.subpixel else 2,
|
||||
|
|
@ -59,7 +59,7 @@ class DeepFakeArchi(nn.ArchiBase):
|
|||
|
||||
class Upscale(nn.ModelBase):
|
||||
def on_build(self, in_ch, out_ch, kernel_size=3 ):
|
||||
self.conv1 = nn.Conv2D( in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME')
|
||||
self.conv1 = nn.SeparableConv2D( in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME')
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
|
|
@ -69,8 +69,8 @@ class DeepFakeArchi(nn.ArchiBase):
|
|||
|
||||
class ResidualBlock(nn.ModelBase):
|
||||
def on_build(self, ch, kernel_size=3 ):
|
||||
self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME')
|
||||
self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME')
|
||||
self.conv1 = nn.SeparableConv2D( ch, ch, kernel_size=kernel_size, padding='SAME')
|
||||
self.conv2 = nn.SeparableConv2D( ch, ch, kernel_size=kernel_size, padding='SAME')
|
||||
|
||||
def forward(self, inp):
|
||||
x = self.conv1(inp)
|
||||
|
|
@ -160,12 +160,12 @@ class DeepFakeArchi(nn.ArchiBase):
|
|||
self.res1 = ResidualBlock(d_ch*4, kernel_size=3)
|
||||
self.res2 = ResidualBlock(d_ch*2, kernel_size=3)
|
||||
|
||||
self.out_conv = nn.Conv2D( d_ch*2, 3, kernel_size=1, padding='SAME')
|
||||
self.out_conv = nn.SeparableConv2D( d_ch*2, 3, kernel_size=1, padding='SAME')
|
||||
|
||||
self.upscalem0 = Upscale(in_ch, d_mask_ch*8, kernel_size=3)
|
||||
self.upscalem1 = Upscale(d_mask_ch*8, d_mask_ch*4, kernel_size=3)
|
||||
self.upscalem2 = Upscale(d_mask_ch*4, d_mask_ch*2, kernel_size=3)
|
||||
self.out_convm = nn.Conv2D( d_mask_ch*2, 1, kernel_size=1, padding='SAME')
|
||||
self.out_convm = nn.SeparableConv2D( d_mask_ch*2, 1, kernel_size=1, padding='SAME')
|
||||
|
||||
def forward(self, inp):
|
||||
z = inp
|
||||
|
|
@ -210,7 +210,7 @@ class DeepFakeArchi(nn.ArchiBase):
|
|||
super().__init__(*kwargs)
|
||||
|
||||
def on_build(self, *args, **kwargs ):
|
||||
self.conv1 = nn.Conv2D( self.in_ch,
|
||||
self.conv1 = nn.SeparableConv2D( self.in_ch,
|
||||
self.out_ch // (4 if self.subpixel else 1),
|
||||
kernel_size=self.kernel_size,
|
||||
strides=1 if self.subpixel else 2,
|
||||
|
|
@ -247,7 +247,7 @@ class DeepFakeArchi(nn.ArchiBase):
|
|||
|
||||
class Upscale(nn.ModelBase):
|
||||
def on_build(self, in_ch, out_ch, kernel_size=3 ):
|
||||
self.conv1 = nn.Conv2D( in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME')
|
||||
self.conv1 = nn.SeparableConv2D( in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME')
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
|
|
@ -257,8 +257,8 @@ class DeepFakeArchi(nn.ArchiBase):
|
|||
|
||||
class ResidualBlock(nn.ModelBase):
|
||||
def on_build(self, ch, kernel_size=3 ):
|
||||
self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME')
|
||||
self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME')
|
||||
self.conv1 = nn.SeparableConv2D( ch, ch, kernel_size=kernel_size, padding='SAME')
|
||||
self.conv2 = nn.SeparableConv2D( ch, ch, kernel_size=kernel_size, padding='SAME')
|
||||
|
||||
def forward(self, inp):
|
||||
x = self.conv1(inp)
|
||||
|
|
@ -313,8 +313,8 @@ class DeepFakeArchi(nn.ArchiBase):
|
|||
self.upscalem2 = Upscale(d_ch, d_ch//2)
|
||||
self.upscalem3 = Upscale(d_ch//2, d_ch//2)
|
||||
|
||||
self.out_conv = nn.Conv2D( d_ch*1, 3, kernel_size=1, padding='SAME')
|
||||
self.out_convm = nn.Conv2D( d_ch//2, 1, kernel_size=1, padding='SAME')
|
||||
self.out_conv = nn.SeparableConv2D( d_ch*1, 3, kernel_size=1, padding='SAME')
|
||||
self.out_convm = nn.SeparableConv2D( d_ch//2, 1, kernel_size=1, padding='SAME')
|
||||
|
||||
def forward(self, inp):
|
||||
z = inp
|
||||
|
|
@ -344,7 +344,7 @@ class DeepFakeArchi(nn.ArchiBase):
|
|||
super().__init__(*kwargs)
|
||||
|
||||
def on_build(self, *args, **kwargs ):
|
||||
self.conv1 = nn.Conv2D( self.in_ch,
|
||||
self.conv1 = nn.SeparableConv2D( self.in_ch,
|
||||
self.out_ch // (4 if self.subpixel else 1),
|
||||
kernel_size=self.kernel_size,
|
||||
strides=1 if self.subpixel else 2,
|
||||
|
|
@ -379,7 +379,7 @@ class DeepFakeArchi(nn.ArchiBase):
|
|||
|
||||
class Upscale(nn.ModelBase):
|
||||
def on_build(self, in_ch, out_ch, kernel_size=3 ):
|
||||
self.conv1 = nn.Conv2D( in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME')
|
||||
self.conv1 = nn.SeparableConv2D( in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME')
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
|
|
@ -389,8 +389,8 @@ class DeepFakeArchi(nn.ArchiBase):
|
|||
|
||||
class ResidualBlock(nn.ModelBase):
|
||||
def on_build(self, ch, kernel_size=3 ):
|
||||
self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME')
|
||||
self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME')
|
||||
self.conv1 = nn.SeparableConv2D( ch, ch, kernel_size=kernel_size, padding='SAME')
|
||||
self.conv2 = nn.SeparableConv2D( ch, ch, kernel_size=kernel_size, padding='SAME')
|
||||
|
||||
def forward(self, inp):
|
||||
x = self.conv1(inp)
|
||||
|
|
@ -443,12 +443,12 @@ class DeepFakeArchi(nn.ArchiBase):
|
|||
self.res1 = ResidualBlock(d_ch*4, kernel_size=3)
|
||||
self.res2 = ResidualBlock(d_ch*2, kernel_size=3)
|
||||
|
||||
self.out_conv = nn.Conv2D( d_ch*2, 3, kernel_size=1, padding='SAME')
|
||||
self.out_conv = nn.SeparableConv2D( d_ch*2, 3, kernel_size=1, padding='SAME')
|
||||
|
||||
self.upscalem0 = Upscale(in_ch, d_mask_ch*8, kernel_size=3)
|
||||
self.upscalem1 = Upscale(d_mask_ch*8, d_mask_ch*4, kernel_size=3)
|
||||
self.upscalem2 = Upscale(d_mask_ch*4, d_mask_ch*2, kernel_size=3)
|
||||
self.out_convm = nn.Conv2D( d_mask_ch*2, 1, kernel_size=1, padding='SAME')
|
||||
self.out_convm = nn.SeparableConv2D( d_mask_ch*2, 1, kernel_size=1, padding='SAME')
|
||||
|
||||
def forward(self, inp):
|
||||
z = inp
|
||||
|
|
|
|||
|
|
@ -2,304 +2,170 @@ import os
|
|||
import traceback
|
||||
from pathlib import Path
|
||||
|
||||
import warnings
|
||||
#warnings.simplefilter(action='ignore', category=FutureWarning)
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from numpy import linalg as npla
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torchvision.transforms import ToTensor
|
||||
|
||||
import math
|
||||
|
||||
from facelib.nn_pt import nn as nn_pt
|
||||
|
||||
from facelib import FaceType, LandmarksProcessor
|
||||
from facelib.LandmarksProcessor import convert_98_to_68
|
||||
|
||||
from facelib.coord_conv import CoordConvTh
|
||||
|
||||
from core.leras import nn
|
||||
|
||||
"""
|
||||
ported from https://github.com/protossw512/AdaptiveWingLoss
|
||||
ported from https://github.com/1adrianb/face-alignment
|
||||
"""
|
||||
class FANExtractor(object):
|
||||
def __init__ (self, place_model_on_cpu=False):
|
||||
model_path = Path(__file__).parent / "AWL.pth"
|
||||
|
||||
nn_pt.initialize()
|
||||
def __init__ (self, landmarks_3D=False, place_model_on_cpu=False):
|
||||
|
||||
model_path = Path(__file__).parent / ( "2DFAN.npy" if not landmarks_3D else "3DFAN.npy")
|
||||
if not model_path.exists():
|
||||
raise Exception("Unable to load AWL.pth")
|
||||
raise Exception("Unable to load FANExtractor model")
|
||||
|
||||
def conv3x3(in_planes, out_planes, strd=1, padding=1,
|
||||
bias=False,dilation=1):
|
||||
"3x3 convolution with padding"
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=3,
|
||||
stride=strd, padding=padding, bias=bias,
|
||||
dilation=dilation)
|
||||
nn.initialize(data_format="NHWC")
|
||||
tf = nn.tf
|
||||
|
||||
class BasicBlock(nn.Module):
|
||||
expansion = 1
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None):
|
||||
super(BasicBlock, self).__init__()
|
||||
self.conv1 = conv3x3(inplanes, planes, stride)
|
||||
# self.bn1 = nn.BatchNorm2d(planes)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.conv2 = conv3x3(planes, planes)
|
||||
# self.bn2 = nn.BatchNorm2d(planes)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
residual = x
|
||||
|
||||
out = self.conv1(x)
|
||||
# out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
# out = self.bn2(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
residual = self.downsample(x)
|
||||
|
||||
out += residual
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class ConvBlock(nn.Module):
|
||||
def __init__(self, in_planes, out_planes):
|
||||
super(ConvBlock, self).__init__()
|
||||
class ConvBlock(nn.ModelBase):
|
||||
def on_build(self, in_planes, out_planes):
|
||||
self.in_planes = in_planes
|
||||
self.out_planes = out_planes
|
||||
|
||||
self.bn1 = nn.BatchNorm2d(in_planes)
|
||||
self.conv1 = conv3x3(in_planes, int(out_planes / 2))
|
||||
self.bn1 = nn.BatchNorm2D(in_planes)
|
||||
self.conv1 = nn.Conv2D (in_planes, out_planes/2, kernel_size=3, strides=1, padding='SAME', use_bias=False )
|
||||
|
||||
self.bn2 = nn.BatchNorm2d(out_planes//2)
|
||||
self.conv2 = conv3x3(int(out_planes / 2), int(out_planes / 4),
|
||||
padding=1, dilation=1)
|
||||
self.bn2 = nn.BatchNorm2D(out_planes//2)
|
||||
self.conv2 = nn.Conv2D (out_planes/2, out_planes/4, kernel_size=3, strides=1, padding='SAME', use_bias=False )
|
||||
|
||||
self.bn3 = nn.BatchNorm2d(out_planes//4)
|
||||
self.conv3 = conv3x3(int(out_planes / 4), int(out_planes / 4),
|
||||
padding=1, dilation=1)
|
||||
self.bn3 = nn.BatchNorm2D(out_planes//4)
|
||||
self.conv3 = nn.Conv2D (out_planes/4, out_planes/4, kernel_size=3, strides=1, padding='SAME', use_bias=False )
|
||||
|
||||
if self.in_planes != self.out_planes:
|
||||
self.downsample = nn.Sequential(
|
||||
nn.BatchNorm2d(in_planes),
|
||||
nn.ReLU(True),
|
||||
nn.Conv2d(in_planes, out_planes,
|
||||
kernel_size=1, stride=1, bias=False),
|
||||
)
|
||||
self.down_bn1 = nn.BatchNorm2D(in_planes)
|
||||
self.down_conv1 = nn.Conv2D (in_planes, out_planes, kernel_size=1, strides=1, padding='VALID', use_bias=False )
|
||||
else:
|
||||
self.downsample = None
|
||||
self.down_bn1 = None
|
||||
self.down_conv1 = None
|
||||
|
||||
def forward(self, x):
|
||||
residual = x
|
||||
def forward(self, input):
|
||||
x = input
|
||||
x = self.bn1(x)
|
||||
x = tf.nn.relu(x)
|
||||
x = out1 = self.conv1(x)
|
||||
|
||||
out1 = self.bn1(x)
|
||||
out1 = F.relu(out1, True)
|
||||
out1 = self.conv1(out1)
|
||||
x = self.bn2(x)
|
||||
x = tf.nn.relu(x)
|
||||
x = out2 = self.conv2(x)
|
||||
|
||||
out2 = self.bn2(out1)
|
||||
out2 = F.relu(out2, True)
|
||||
out2 = self.conv2(out2)
|
||||
x = self.bn3(x)
|
||||
x = tf.nn.relu(x)
|
||||
x = out3 = self.conv3(x)
|
||||
|
||||
out3 = self.bn3(out2)
|
||||
out3 = F.relu(out3, True)
|
||||
out3 = self.conv3(out3)
|
||||
x = tf.concat ([out1, out2, out3], axis=-1)
|
||||
|
||||
out3 = torch.cat((out1, out2, out3), 1)
|
||||
|
||||
if self.downsample is not None:
|
||||
residual = self.downsample(residual)
|
||||
|
||||
out3 += residual
|
||||
|
||||
return out3
|
||||
|
||||
class HourGlass (nn.Module):
|
||||
def __init__(self, num_modules, depth, num_features, first_one=False):
|
||||
super(HourGlass, self).__init__()
|
||||
self.num_modules = num_modules
|
||||
self.depth = depth
|
||||
self.features = num_features
|
||||
self.coordconv = CoordConvTh(x_dim=64, y_dim=64,
|
||||
with_r=True, with_boundary=True,
|
||||
in_channels=256, first_one=first_one,
|
||||
out_channels=256,
|
||||
kernel_size=1,
|
||||
stride=1, padding=0)
|
||||
self._generate_network(self.depth)
|
||||
|
||||
def _generate_network(self, level):
|
||||
self.add_module('b1_' + str(level), ConvBlock(256, 256))
|
||||
|
||||
self.add_module('b2_' + str(level), ConvBlock(256, 256))
|
||||
|
||||
if level > 1:
|
||||
self._generate_network(level - 1)
|
||||
if self.in_planes != self.out_planes:
|
||||
downsample = self.down_bn1(input)
|
||||
downsample = tf.nn.relu (downsample)
|
||||
downsample = self.down_conv1 (downsample)
|
||||
x = x + downsample
|
||||
else:
|
||||
self.add_module('b2_plus_' + str(level), ConvBlock(256, 256))
|
||||
x = x + input
|
||||
|
||||
self.add_module('b3_' + str(level), ConvBlock(256, 256))
|
||||
return x
|
||||
|
||||
def _forward(self, level, inp):
|
||||
# Upper branch
|
||||
up1 = inp
|
||||
up1 = self._modules['b1_' + str(level)](up1)
|
||||
class HourGlass (nn.ModelBase):
|
||||
def on_build(self, in_planes, depth):
|
||||
self.b1 = ConvBlock (in_planes, 256)
|
||||
self.b2 = ConvBlock (in_planes, 256)
|
||||
|
||||
# Lower branch
|
||||
low1 = F.avg_pool2d(inp, 2, stride=2)
|
||||
low1 = self._modules['b2_' + str(level)](low1)
|
||||
|
||||
if level > 1:
|
||||
low2 = self._forward(level - 1, low1)
|
||||
if depth > 1:
|
||||
self.b2_plus = HourGlass(256, depth-1)
|
||||
else:
|
||||
low2 = low1
|
||||
low2 = self._modules['b2_plus_' + str(level)](low2)
|
||||
self.b2_plus = ConvBlock(256, 256)
|
||||
|
||||
low3 = low2
|
||||
low3 = self._modules['b3_' + str(level)](low3)
|
||||
self.b3 = ConvBlock(256, 256)
|
||||
|
||||
up2 = F.upsample(low3, scale_factor=2, mode='nearest')
|
||||
def forward(self, input):
|
||||
up1 = self.b1(input)
|
||||
|
||||
return up1 + up2
|
||||
low1 = tf.nn.avg_pool(input, [1,2,2,1], [1,2,2,1], 'VALID')
|
||||
low1 = self.b2 (low1)
|
||||
|
||||
def forward(self, x, heatmap):
|
||||
x, last_channel = self.coordconv(x, heatmap)
|
||||
return self._forward(self.depth, x), last_channel
|
||||
low2 = self.b2_plus(low1)
|
||||
low3 = self.b3(low2)
|
||||
|
||||
up2 = nn.upsample2d(low3)
|
||||
|
||||
class FAN (nn.Module):
|
||||
def __init__(self, num_modules=1, end_relu=False, gray_scale=False, num_landmarks=68):
|
||||
super(FAN,self).__init__()
|
||||
self.num_modules = num_modules
|
||||
self.gray_scale = gray_scale
|
||||
self.end_relu = end_relu
|
||||
self.num_landmarks = num_landmarks
|
||||
return up1+up2
|
||||
|
||||
class FAN (nn.ModelBase):
|
||||
def __init__(self):
|
||||
super().__init__(name='FAN')
|
||||
|
||||
def on_build(self):
|
||||
self.conv1 = nn.Conv2D (3, 64, kernel_size=7, strides=2, padding='SAME')
|
||||
self.bn1 = nn.BatchNorm2D(64)
|
||||
|
||||
# Base part
|
||||
if self.gray_scale:
|
||||
self.conv1 = CoordConvTh(x_dim=256, y_dim=256,
|
||||
with_r=True, with_boundary=False,
|
||||
in_channels=3, out_channels=64,
|
||||
kernel_size=7,
|
||||
stride=2, padding=3)
|
||||
else:
|
||||
self.conv1 = CoordConvTh(x_dim=256, y_dim=256,
|
||||
with_r=True, with_boundary=False,
|
||||
in_channels=3, out_channels=64,
|
||||
kernel_size=7,
|
||||
stride=2, padding=3)
|
||||
|
||||
self.bn1 = nn.BatchNorm2d(64)
|
||||
self.conv2 = ConvBlock(64, 128)
|
||||
self.conv3 = ConvBlock(128, 128)
|
||||
self.conv4 = ConvBlock(128, 256)
|
||||
|
||||
# Stacking part
|
||||
for hg_module in range(self.num_modules):
|
||||
if hg_module == 0:
|
||||
first_one = True
|
||||
else:
|
||||
first_one = False
|
||||
self.add_module('m' + str(hg_module), HourGlass(1, 4, 256,
|
||||
first_one))
|
||||
self.add_module('top_m_' + str(hg_module), ConvBlock(256, 256))
|
||||
self.add_module('conv_last' + str(hg_module),
|
||||
nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0))
|
||||
self.add_module('bn_end' + str(hg_module), nn.BatchNorm2d(256))
|
||||
self.add_module('l' + str(hg_module), nn.Conv2d(256,
|
||||
num_landmarks+1, kernel_size=1, stride=1, padding=0))
|
||||
self.m = []
|
||||
self.top_m = []
|
||||
self.conv_last = []
|
||||
self.bn_end = []
|
||||
self.l = []
|
||||
self.bl = []
|
||||
self.al = []
|
||||
for i in range(4):
|
||||
self.m += [ HourGlass(256, 4) ]
|
||||
self.top_m += [ ConvBlock(256, 256) ]
|
||||
|
||||
if hg_module < self.num_modules - 1:
|
||||
self.add_module(
|
||||
'bl' + str(hg_module), nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0))
|
||||
self.add_module('al' + str(hg_module), nn.Conv2d(num_landmarks+1,
|
||||
256, kernel_size=1, stride=1, padding=0))
|
||||
|
||||
self.conv_last += [ nn.Conv2D (256, 256, kernel_size=1, strides=1, padding='VALID') ]
|
||||
self.bn_end += [ nn.BatchNorm2D(256) ]
|
||||
|
||||
self.l += [ nn.Conv2D (256, 68, kernel_size=1, strides=1, padding='VALID') ]
|
||||
|
||||
def forward(self, x):
|
||||
x, _ = self.conv1(x)
|
||||
x = F.relu(self.bn1(x), True)
|
||||
# x = F.relu(self.bn1(self.conv1(x)), True)
|
||||
x = F.avg_pool2d(self.conv2(x), 2, stride=2)
|
||||
if i < 4-1:
|
||||
self.bl += [ nn.Conv2D (256, 256, kernel_size=1, strides=1, padding='VALID') ]
|
||||
self.al += [ nn.Conv2D (68, 256, kernel_size=1, strides=1, padding='VALID') ]
|
||||
|
||||
def forward(self, inp) :
|
||||
x, = inp
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = tf.nn.relu(x)
|
||||
|
||||
x = self.conv2(x)
|
||||
x = tf.nn.avg_pool(x, [1,2,2,1], [1,2,2,1], 'VALID')
|
||||
x = self.conv3(x)
|
||||
x = self.conv4(x)
|
||||
|
||||
previous = x
|
||||
|
||||
outputs = []
|
||||
boundary_channels = []
|
||||
tmp_out = None
|
||||
for i in range(self.num_modules):
|
||||
hg, boundary_channel = self._modules['m' + str(i)](previous,
|
||||
tmp_out)
|
||||
|
||||
ll = hg
|
||||
ll = self._modules['top_m_' + str(i)](ll)
|
||||
|
||||
ll = F.relu(self._modules['bn_end' + str(i)]
|
||||
(self._modules['conv_last' + str(i)](ll)), True)
|
||||
|
||||
|
||||
# Predict heatmaps
|
||||
tmp_out = self._modules['l' + str(i)](ll)
|
||||
if self.end_relu:
|
||||
tmp_out = F.relu(tmp_out) # HACK: Added relu
|
||||
|
||||
previous = x
|
||||
for i in range(4):
|
||||
ll = self.m[i] (previous)
|
||||
ll = self.top_m[i] (ll)
|
||||
ll = self.conv_last[i] (ll)
|
||||
ll = self.bn_end[i] (ll)
|
||||
ll = tf.nn.relu(ll)
|
||||
tmp_out = self.l[i](ll)
|
||||
outputs.append(tmp_out)
|
||||
boundary_channels.append(boundary_channel)
|
||||
if i < 4 - 1:
|
||||
ll = self.bl[i](ll)
|
||||
previous = previous + ll + self.al[i](tmp_out)
|
||||
x = outputs[-1]
|
||||
x = tf.transpose(x, (0,3,1,2) )
|
||||
return x
|
||||
|
||||
if i < self.num_modules - 1:
|
||||
ll = self._modules['bl' + str(i)](ll)
|
||||
tmp_out_ = self._modules['al' + str(i)](tmp_out)
|
||||
previous = previous + ll + tmp_out_
|
||||
|
||||
return outputs, boundary_channels
|
||||
|
||||
def load_model(model, pretrained_path, load_to_cpu):
|
||||
if load_to_cpu:
|
||||
checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage)
|
||||
else:
|
||||
device = torch.cuda.current_device()
|
||||
checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage.cuda(device))
|
||||
|
||||
if 'state_dict' not in checkpoint:
|
||||
model.load_state_dict(checkpoint)
|
||||
else:
|
||||
pretrained_weights = checkpoint['state_dict']
|
||||
model_weights = model.state_dict()
|
||||
pretrained_weights = {k: v for k, v in pretrained_weights.items() \
|
||||
if k in model_weights}
|
||||
model_weights.update(pretrained_weights)
|
||||
model.load_state_dict(model_weights)
|
||||
return model
|
||||
|
||||
|
||||
torch.set_grad_enabled(False)
|
||||
|
||||
self.model = FAN(num_modules=4, num_landmarks=98, end_relu=False, gray_scale=False)
|
||||
self.model = load_model(self.model, model_path, place_model_on_cpu)
|
||||
e = None
|
||||
if place_model_on_cpu:
|
||||
e = tf.device("/CPU:0")
|
||||
|
||||
self.model.eval()
|
||||
|
||||
self.device = torch.device("cpu" if place_model_on_cpu else "cuda")
|
||||
self.model = self.model.to(self.device)
|
||||
if e is not None: e.__enter__()
|
||||
self.model = FAN()
|
||||
self.model.load_weights(str(model_path))
|
||||
if e is not None: e.__exit__(None,None,None)
|
||||
|
||||
self.model.build_for_run ([ ( tf.float32, (None,256,256,3) ) ])
|
||||
|
||||
def extract (self, input_image, rects, second_pass_extractor=None, is_bgr=True, multi_sample=False):
|
||||
|
||||
if len(rects) == 0:
|
||||
return []
|
||||
|
||||
|
|
@ -308,14 +174,13 @@ class FANExtractor(object):
|
|||
is_bgr = False
|
||||
|
||||
(h, w, ch) = input_image.shape
|
||||
print("Image shape:", input_image.shape)
|
||||
|
||||
landmarks = []
|
||||
for (left, top, right, bottom) in rects:
|
||||
scale = (right - left + bottom - top) / 195.0
|
||||
print(scale)
|
||||
|
||||
center = np.array( [ (left + right) / 2.0, (top + bottom) / 2.0] )
|
||||
print("Center:", center)
|
||||
centers = [ center ]
|
||||
|
||||
if multi_sample:
|
||||
|
|
@ -326,30 +191,26 @@ class FANExtractor(object):
|
|||
]
|
||||
|
||||
images = []
|
||||
ptss = []
|
||||
|
||||
try:
|
||||
for c in centers:
|
||||
images += [ self.crop_old(input_image, c, scale) ]
|
||||
images += [ self.crop(input_image, c, scale) ]
|
||||
|
||||
images = np.stack (images)
|
||||
images = images.astype(np.float32) / 255.0
|
||||
|
||||
i = 0
|
||||
for img in images:
|
||||
img = ToTensor()(img)
|
||||
img = img.to(self.device)
|
||||
|
||||
outputs, boundary_channels = self.model(img[None,...])
|
||||
|
||||
pred_heatmap = outputs[-1][:, :-1, :, :][i].detach().cpu()
|
||||
|
||||
pred_landmarks, _ = self.get_pts_from_predict ( pred_heatmap.unsqueeze(0), centers[i], scale)
|
||||
i += 1
|
||||
|
||||
pred_landmarks = pred_landmarks.squeeze().numpy()
|
||||
pred_landmarks = convert_98_to_68(pred_landmarks)
|
||||
|
||||
landmarks += [pred_landmarks]
|
||||
|
||||
|
||||
predicted = []
|
||||
for i in range( len(images) ):
|
||||
predicted += [ self.model.run ( [ images[i][None,...] ] )[0] ]
|
||||
|
||||
predicted = np.stack(predicted)
|
||||
|
||||
for i, pred in enumerate(predicted):
|
||||
ptss += [ self.get_pts_from_predict ( pred, centers[i], scale) ]
|
||||
pts_img = np.mean ( np.array(ptss), 0 )
|
||||
|
||||
landmarks.append (pts_img)
|
||||
except:
|
||||
landmarks.append (None)
|
||||
|
||||
|
|
@ -368,8 +229,8 @@ class FANExtractor(object):
|
|||
pass
|
||||
|
||||
return landmarks
|
||||
|
||||
def transform_old(self, point, center, scale, resolution):
|
||||
|
||||
def transform(self, point, center, scale, resolution):
|
||||
pt = np.array ( [point[0], point[1], 1.0] )
|
||||
h = 200.0 * scale
|
||||
m = np.eye(3)
|
||||
|
|
@ -380,73 +241,9 @@ class FANExtractor(object):
|
|||
m = np.linalg.inv(m)
|
||||
return np.matmul (m, pt)[0:2]
|
||||
|
||||
def transform(self, point, center, scale, resolution, rotation=0, invert=False):
|
||||
_pt = np.ones(3)
|
||||
_pt[0] = point[0]
|
||||
_pt[1] = point[1]
|
||||
|
||||
h = 200.0 * scale
|
||||
t = np.eye(3)
|
||||
t[0, 0] = resolution / h
|
||||
t[1, 1] = resolution / h
|
||||
t[0, 2] = resolution * (-center[0] / h + 0.5)
|
||||
t[1, 2] = resolution * (-center[1] / h + 0.5)
|
||||
|
||||
if rotation != 0:
|
||||
rotation = -1*rotation
|
||||
r = np.eye(3)
|
||||
ang = rotation * math.pi / 180.0
|
||||
s = math.sin(ang)
|
||||
c = math.cos(ang)
|
||||
r[0][0] = c
|
||||
r[0][1] = -s
|
||||
r[1][0] = s
|
||||
r[1][1] = c
|
||||
|
||||
t_ = np.eye(3)
|
||||
t_[0][2] = -resolution / 2.0
|
||||
t_[1][2] = -resolution / 2.0
|
||||
t_inv = torch.eye(3)
|
||||
t_inv[0][2] = resolution / 2.0
|
||||
t_inv[1][2] = resolution / 2.0
|
||||
t = reduce(np.matmul, [t_inv, r, t_, t])
|
||||
|
||||
if invert:
|
||||
t = np.linalg.inv(t)
|
||||
new_point = (np.matmul(t, _pt))[0:2]
|
||||
|
||||
return new_point.astype(int)
|
||||
|
||||
def crop(self, image, center, scale, resolution=256, center_shift=0):
|
||||
new_image = cv2.copyMakeBorder(image, center_shift,
|
||||
center_shift,
|
||||
center_shift,
|
||||
center_shift,
|
||||
cv2.BORDER_CONSTANT, value=[0,0,0])
|
||||
if center_shift != 0:
|
||||
center[0] += center_shift
|
||||
center[1] += center_shift
|
||||
length = 200 * scale
|
||||
top = int(center[1] - length // 2)
|
||||
bottom = int(center[1] + length // 2)
|
||||
left = int(center[0] - length // 2)
|
||||
right = int(center[0] + length // 2)
|
||||
y_pad = abs(min(top, new_image.shape[0] - bottom, 0))
|
||||
x_pad = abs(min(left, new_image.shape[1] - right, 0))
|
||||
top, bottom, left, right = top + y_pad, bottom + y_pad, left + x_pad, right + x_pad
|
||||
new_image = cv2.copyMakeBorder(new_image, y_pad,
|
||||
y_pad,
|
||||
x_pad,
|
||||
x_pad,
|
||||
cv2.BORDER_CONSTANT, value=[0,0,0])
|
||||
new_image = new_image[top:bottom, left:right]
|
||||
new_image = cv2.resize(new_image, dsize=(int(resolution), int(resolution)),
|
||||
interpolation=cv2.INTER_LINEAR)
|
||||
return new_image
|
||||
|
||||
def crop_old(self, image, center, scale, resolution=256.0):
|
||||
ul = self.transform_old([1, 1], center, scale, resolution).astype( np.int )
|
||||
br = self.transform_old([resolution, resolution], center, scale, resolution).astype( np.int )
|
||||
def crop(self, image, center, scale, resolution=256.0):
|
||||
ul = self.transform([1, 1], center, scale, resolution).astype( np.int )
|
||||
br = self.transform([resolution, resolution], center, scale, resolution).astype( np.int )
|
||||
|
||||
if image.ndim > 2:
|
||||
newDim = np.array([br[1] - ul[1], br[0] - ul[0], image.shape[2]], dtype=np.int32)
|
||||
|
|
@ -464,32 +261,21 @@ class FANExtractor(object):
|
|||
|
||||
newImg = cv2.resize(newImg, dsize=(int(resolution), int(resolution)), interpolation=cv2.INTER_LINEAR)
|
||||
return newImg
|
||||
|
||||
|
||||
def get_pts_from_predict(self, hm, center=None, scale=None, rot=None):
|
||||
max, idx = torch.max(
|
||||
hm.view(hm.size(0), hm.size(1), hm.size(2) * hm.size(3)), 2)
|
||||
idx += 1
|
||||
preds = idx.view(idx.size(0), idx.size(1), 1).repeat(1, 1, 2).float()
|
||||
preds[..., 0].apply_(lambda x: (x - 1) % hm.size(3) + 1)
|
||||
preds[..., 1].add_(-1).div_(hm.size(2)).floor_().add_(1)
|
||||
def get_pts_from_predict(self, a, center, scale):
|
||||
a_ch, a_h, a_w = a.shape
|
||||
|
||||
for i in range(preds.size(0)):
|
||||
for j in range(preds.size(1)):
|
||||
hm_ = hm[i, j, :]
|
||||
pX, pY = int(preds[i, j, 0]) - 1, int(preds[i, j, 1]) - 1
|
||||
if pX > 0 and pX < 63 and pY > 0 and pY < 63:
|
||||
diff = torch.FloatTensor(
|
||||
[hm_[pY, pX + 1] - hm_[pY, pX - 1],
|
||||
hm_[pY + 1, pX] - hm_[pY - 1, pX]])
|
||||
preds[i, j].add_(diff.sign_().mul_(.25))
|
||||
b = a.reshape ( (a_ch, a_h*a_w) )
|
||||
c = b.argmax(1).reshape ( (a_ch, 1) ).repeat(2, axis=1).astype(np.float)
|
||||
c[:,0] %= a_w
|
||||
c[:,1] = np.apply_along_axis ( lambda x: np.floor(x / a_w), 0, c[:,1] )
|
||||
|
||||
preds.add_(-0.5)
|
||||
for i in range(a_ch):
|
||||
pX, pY = int(c[i,0]), int(c[i,1])
|
||||
if pX > 0 and pX < 63 and pY > 0 and pY < 63:
|
||||
diff = np.array ( [a[i,pY,pX+1]-a[i,pY,pX-1], a[i,pY+1,pX]-a[i,pY-1,pX]] )
|
||||
c[i] += np.sign(diff)*0.25
|
||||
|
||||
preds_orig = torch.zeros(preds.size())
|
||||
if center is not None and scale is not None:
|
||||
for i in range(hm.size(0)):
|
||||
for j in range(hm.size(1)):
|
||||
preds_orig[i, j] = torch.from_numpy(self.transform(preds[i, j], center, scale, hm.size(2), rot if (rot is not None) else (0), True))
|
||||
c += 0.5
|
||||
|
||||
return preds, preds_orig
|
||||
return np.array( [ self.transform (c[i], center, scale, a_w) for i in range(a_ch) ] )
|
||||
|
|
|
|||
|
|
@ -1,9 +1,5 @@
|
|||
from .FaceType import FaceType
|
||||
from .S3FDExtractor import S3FDExtractor
|
||||
from .RetinaFaceExtractor import RetinaFaceExtractor
|
||||
from .nn_pt import nn
|
||||
from .device_pt import Devices
|
||||
from .ResNet50 import resnet50
|
||||
from .FANExtractor import FANExtractor
|
||||
from .FaceEnhancer import FaceEnhancer
|
||||
from .XSegNet import XSegNet
|
||||
|
|
@ -7,5 +7,3 @@ scikit-image==0.14.2
|
|||
scipy==1.4.1
|
||||
colorama
|
||||
tensorflow-gpu==1.13.2
|
||||
pytorch
|
||||
torchvision
|
||||
|
|
|
|||
|
|
@ -8,6 +8,4 @@ scipy==1.4.1
|
|||
colorama
|
||||
labelme==4.2.9
|
||||
tensorflow-gpu==1.13.2
|
||||
pytorch
|
||||
torchvision
|
||||
pyqt5
|
||||
Loading…
Add table
Add a link
Reference in a new issue