Added depthwise separable convolutions

core/leras/archis/DeepFakeArchi.py

@@ -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.conv1 = nn.SeparableConv2D( ch, ch, kernel_size=kernel_size, padding='SAME')
-                    self.conv2 = nn.Conv2D( 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.conv1 = nn.SeparableConv2D( ch, ch, kernel_size=kernel_size, padding='SAME')
-                    self.conv2 = nn.Conv2D( 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_conv = nn.SeparableConv2D( 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_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.conv1 = nn.SeparableConv2D( ch, ch, kernel_size=kernel_size, padding='SAME')
-                    self.conv2 = nn.Conv2D( 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

facelib/FANExtractor.py

@@ -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 core.leras import nn
-
-from facelib.coord_conv import CoordConvTh
-
 
 """
-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):
+    def __init__ (self, landmarks_3D=False, place_model_on_cpu=False):
-        model_path = Path(__file__).parent / "AWL.pth"
-        
-        nn_pt.initialize()
         
+        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,
+        nn.initialize(data_format="NHWC")
-            bias=False,dilation=1):
+        tf = nn.tf
-            "3x3 convolution with padding"
-            return nn.Conv2d(in_planes, out_planes, kernel_size=3,
-                     stride=strd, padding=padding, bias=bias,
-                     dilation=dilation)
 
-        class BasicBlock(nn.Module):
+        class ConvBlock(nn.ModelBase):
-            expansion = 1
+            def on_build(self, in_planes, out_planes):
-
-            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__()
                 self.in_planes = in_planes
                 self.out_planes = out_planes
 
-                self.bn1 = nn.BatchNorm2d(in_planes)
+                self.bn1 = nn.BatchNorm2D(in_planes)
-                self.conv1 = conv3x3(in_planes, int(out_planes / 2))
+                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.bn2 = nn.BatchNorm2D(out_planes//2)
-                self.conv2 = conv3x3(int(out_planes / 2), int(out_planes / 4),
+                self.conv2 = nn.Conv2D (out_planes/2, out_planes/4, kernel_size=3, strides=1, padding='SAME', use_bias=False )
-                             padding=1, dilation=1)
 
-                self.bn3 = nn.BatchNorm2d(out_planes//4)
+                self.bn3 = nn.BatchNorm2D(out_planes//4)
-                self.conv3 = conv3x3(int(out_planes / 4), int(out_planes / 4),
+                self.conv3 = nn.Conv2D (out_planes/4, out_planes/4, kernel_size=3, strides=1, padding='SAME', use_bias=False )
-                             padding=1, dilation=1)
 
                 if self.in_planes != self.out_planes:
-                    self.downsample = nn.Sequential(
+                    self.down_bn1 = nn.BatchNorm2D(in_planes)
-                        nn.BatchNorm2d(in_planes),
+                    self.down_conv1 = nn.Conv2D (in_planes, out_planes, kernel_size=1, strides=1, padding='VALID', use_bias=False )
-                        nn.ReLU(True),
-                        nn.Conv2d(in_planes, out_planes,
-                                  kernel_size=1, stride=1, bias=False),
-                    )
                 else:
-                    self.downsample = None
+                    self.down_bn1 = None
+                    self.down_conv1 = None
 
-            def forward(self, x):
+            def forward(self, input):
-                residual = x
+                x = input
+                x = self.bn1(x)
+                x = tf.nn.relu(x)
+                x = out1 = self.conv1(x)
 
-                out1 = self.bn1(x)
+                x = self.bn2(x)
-                out1 = F.relu(out1, True)
+                x = tf.nn.relu(x)
-                out1 = self.conv1(out1)
+                x = out2 = self.conv2(x)
 
-                out2 = self.bn2(out1)
+                x = self.bn3(x)
-                out2 = F.relu(out2, True)
+                x = tf.nn.relu(x)
-                out2 = self.conv2(out2)
+                x = out3 = self.conv3(x)
 
-                out3 = self.bn3(out2)
+                x = tf.concat ([out1, out2, out3], axis=-1)
-                out3 = F.relu(out3, True)
-                out3 = self.conv3(out3)
 
-                out3 = torch.cat((out1, out2, out3), 1)
+                if self.in_planes != self.out_planes:
-
+                    downsample = self.down_bn1(input)
-                if self.downsample is not None:
+                    downsample = tf.nn.relu (downsample)
-                    residual = self.downsample(residual)
+                    downsample = self.down_conv1 (downsample)
-
+                    x = x + downsample
-                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)
                 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):
+        class HourGlass (nn.ModelBase):
-                # Upper branch
+            def on_build(self, in_planes, depth):
-                up1 = inp
+                self.b1 = ConvBlock (in_planes, 256)
-                up1 = self._modules['b1_' + str(level)](up1)
+                self.b2 = ConvBlock (in_planes, 256)
 
-                # Lower branch
+                if depth > 1:
-                low1 = F.avg_pool2d(inp, 2, stride=2)
+                    self.b2_plus = HourGlass(256, depth-1)
-                low1 = self._modules['b2_' + str(level)](low1)
-
-                if level > 1:
-                    low2 = self._forward(level - 1, low1)
                 else:
-                    low2 = low1
+                    self.b2_plus = ConvBlock(256, 256)
-                    low2 = self._modules['b2_plus_' + str(level)](low2)
 
-                low3 = low2
+                self.b3 = ConvBlock(256, 256)
-                low3 = self._modules['b3_' + str(level)](low3)
 
-                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):
+                low2 = self.b2_plus(low1)
-                x, last_channel = self.coordconv(x, heatmap)
+                low3 = self.b3(low2)
-                return self._forward(self.depth, x), last_channel
 
+                up2 = nn.upsample2d(low3)
 
-        class FAN (nn.Module):
+                return up1+up2
-            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
 
-                # Base part
+        class FAN (nn.ModelBase):
-                if self.gray_scale:
+            def __init__(self):
-                    self.conv1 = CoordConvTh(x_dim=256, y_dim=256,
+                super().__init__(name='FAN')
-                                     with_r=True, with_boundary=False,
+
-                                     in_channels=3, out_channels=64,
+            def on_build(self):
-                                     kernel_size=7,
+                self.conv1 = nn.Conv2D (3, 64, kernel_size=7, strides=2, padding='SAME')
-                                     stride=2, padding=3)
+                self.bn1 = nn.BatchNorm2D(64)
-                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
+                self.m = []
-                for hg_module in range(self.num_modules):
+                self.top_m = []
-                    if hg_module == 0:
+                self.conv_last = []
-                        first_one = True
+                self.bn_end = []
-                    else:
+                self.l = []
-                        first_one = False
+                self.bl = []
-                    self.add_module('m' + str(hg_module), HourGlass(1, 4, 256,
+                self.al = []
-                                                            first_one))
+                for i in range(4):
-                    self.add_module('top_m_' + str(hg_module), ConvBlock(256, 256))
+                    self.m += [ HourGlass(256, 4) ]
-                    self.add_module('conv_last' + str(hg_module),
+                    self.top_m += [ ConvBlock(256, 256) ]
-                            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))
 
-                    if hg_module < self.num_modules - 1:
+                    self.conv_last += [ nn.Conv2D (256, 256, kernel_size=1, strides=1, padding='VALID') ]
-                        self.add_module(
+                    self.bn_end += [ nn.BatchNorm2D(256) ]
-                            '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.l += [ nn.Conv2D (256, 68, kernel_size=1, strides=1, padding='VALID') ]
 
+                    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, x):
+            def forward(self, inp) :
-                x, _ = self.conv1(x)
+                x, = inp
-                x = F.relu(self.bn1(x), True)
+                x = self.conv1(x)
-                # x = F.relu(self.bn1(self.conv1(x)), True)
+                x = self.bn1(x)
-                x = F.avg_pool2d(self.conv2(x), 2, stride=2)
+                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 = []
+                previous = x
-                tmp_out = None
+                for i in range(4):
-                for i in range(self.num_modules):
+                    ll = self.m[i] (previous)
-                    hg, boundary_channel = self._modules['m' + str(i)](previous,
+                    ll = self.top_m[i] (ll)
-                                                               tmp_out)
+                    ll = self.conv_last[i] (ll)
-
+                    ll = self.bn_end[i] (ll)
-                    ll = hg
+                    ll = tf.nn.relu(ll)
-                    ll = self._modules['top_m_' + str(i)](ll)
+                    tmp_out = self.l[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
-
                     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:
+        e = None
-                        ll = self._modules['bl' + str(i)](ll)
+        if place_model_on_cpu:
-                        tmp_out_ = self._modules['al' + str(i)](tmp_out)
+            e = tf.device("/CPU:0")
-                        previous = previous + ll + tmp_out_
 
-                return outputs, boundary_channels
+        if e is not None: e.__enter__()
-        
+        self.model = FAN()
-        def load_model(model, pretrained_path, load_to_cpu):
+        self.model.load_weights(str(model_path))
-            if load_to_cpu:
+        if e is not None: e.__exit__(None,None,None)
-                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)
-
-        self.model.eval()
-        
-        self.device = torch.device("cpu" if place_model_on_cpu else "cuda")
-        self.model = self.model.to(self.device)
 
+        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
+                predicted = []
-                for img in images:
+                for i in range( len(images) ):
-                    img = ToTensor()(img)
+                    predicted += [ self.model.run ( [ images[i][None,...] ]  )[0] ]
-                    img = img.to(self.device)
 
-                    outputs, boundary_channels = self.model(img[None,...])
+                predicted = np.stack(predicted)
 
-                    pred_heatmap = outputs[-1][:, :-1, :, :][i].detach().cpu()
+                for i, pred in enumerate(predicted):
-                
+                    ptss += [ self.get_pts_from_predict ( pred, centers[i], scale) ]
-                    pred_landmarks, _ = self.get_pts_from_predict ( pred_heatmap.unsqueeze(0), centers[i], scale)
+                pts_img = np.mean ( np.array(ptss), 0 )
-                    i += 1
-                
-                    pred_landmarks = pred_landmarks.squeeze().numpy()
-                    pred_landmarks = convert_98_to_68(pred_landmarks)
-                
-                    landmarks += [pred_landmarks]
 
+                landmarks.append (pts_img)
             except:
                 landmarks.append (None)
 
@@ -369,7 +230,7 @@ class FANExtractor(object):
 
         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):
+    def crop(self, image, center, scale, resolution=256.0):
-        _pt = np.ones(3)
+        ul = self.transform([1, 1], center, scale, resolution).astype( np.int )
-        _pt[0] = point[0]
+        br = self.transform([resolution, resolution], center, scale, resolution).astype( np.int )
-        _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 )
 
         if image.ndim > 2:
             newDim = np.array([br[1] - ul[1], br[0] - ul[0], image.shape[2]], dtype=np.int32)
@@ -465,31 +262,20 @@ class FANExtractor(object):
         newImg = cv2.resize(newImg, dsize=(int(resolution), int(resolution)), interpolation=cv2.INTER_LINEAR)
         return newImg
 
+    def get_pts_from_predict(self, a, center, scale):
+        a_ch, a_h, a_w = a.shape
 
-    def get_pts_from_predict(self, hm, center=None, scale=None, rot=None):
+        b = a.reshape ( (a_ch, a_h*a_w) )
-        max, idx = torch.max(
+        c = b.argmax(1).reshape ( (a_ch, 1) ).repeat(2, axis=1).astype(np.float)
-        hm.view(hm.size(0), hm.size(1), hm.size(2) * hm.size(3)), 2)
+        c[:,0] %= a_w
-        idx += 1
+        c[:,1] = np.apply_along_axis ( lambda x: np.floor(x / a_w), 0, c[:,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)
 
-        for i in range(preds.size(0)):
+        for i in range(a_ch):
-            for j in range(preds.size(1)):
+            pX, pY = int(c[i,0]), int(c[i,1])
-                hm_ = hm[i, j, :]
+            if pX > 0 and pX < 63 and pY > 0 and pY < 63:
-                pX, pY = int(preds[i, j, 0]) - 1, int(preds[i, j, 1]) - 1
+                diff = np.array ( [a[i,pY,pX+1]-a[i,pY,pX-1], a[i,pY+1,pX]-a[i,pY-1,pX]] )
-                if pX > 0 and pX < 63 and pY > 0 and pY < 63:
+                c[i] += np.sign(diff)*0.25
-                    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))
 
-        preds.add_(-0.5)
+        c += 0.5
 
-        preds_orig = torch.zeros(preds.size())
+        return np.array( [ self.transform (c[i], center, scale, a_w) for i in range(a_ch) ] )
-        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))
-
-        return preds, preds_orig

facelib/__init__.py

@@ -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

requirements-colab.txt

@@ -7,5 +7,3 @@ scikit-image==0.14.2
 scipy==1.4.1
 colorama
 tensorflow-gpu==1.13.2
-pytorch
-torchvision

requirements-cuda.txt

@@ -8,6 +8,4 @@ scipy==1.4.1
 colorama
 labelme==4.2.9
 tensorflow-gpu==1.13.2
-pytorch
-torchvision
 pyqt5