diff --git a/converters/__init__.py b/converters/__init__.py
deleted file mode 100644
index 8aa1057..0000000
--- a/converters/__init__.py
+++ /dev/null
@@ -1,4 +0,0 @@
-from .FrameInfo import FrameInfo
-from .ConverterConfig import ConverterConfig, ConverterConfigMasked, ConverterConfigFaceAvatar
-from .ConvertMasked import ConvertMasked
-from .ConvertAvatar import ConvertFaceAvatar
diff --git a/utils/cv2_utils.py b/core/cv2ex.py
similarity index 100%
rename from utils/cv2_utils.py
rename to core/cv2ex.py
diff --git a/imagelib/IEPolys.py b/core/imagelib/IEPolys.py
similarity index 100%
rename from imagelib/IEPolys.py
rename to core/imagelib/IEPolys.py
diff --git a/imagelib/__init__.py b/core/imagelib/__init__.py
similarity index 86%
rename from imagelib/__init__.py
rename to core/imagelib/__init__.py
index 7a1ed9b..6c0236f 100644
--- a/imagelib/__init__.py
+++ b/core/imagelib/__init__.py
@@ -13,9 +13,7 @@ from .reduce_colors import reduce_colors
from .color_transfer import color_transfer_mix, color_transfer_sot, color_transfer_mkl, color_transfer_idt, color_hist_match, reinhard_color_transfer, linear_color_transfer, seamless_clone
-from .RankSRGAN import RankSRGAN
-
-from .common import normalize_channels, overlay_alpha_image
+from .common import normalize_channels, cut_odd_image, overlay_alpha_image
from .IEPolys import IEPolys
diff --git a/imagelib/blur.py b/core/imagelib/blur.py
similarity index 100%
rename from imagelib/blur.py
rename to core/imagelib/blur.py
diff --git a/imagelib/color_transfer.py b/core/imagelib/color_transfer.py
similarity index 100%
rename from imagelib/color_transfer.py
rename to core/imagelib/color_transfer.py
diff --git a/imagelib/common.py b/core/imagelib/common.py
similarity index 88%
rename from imagelib/common.py
rename to core/imagelib/common.py
index 2e66441..d73df8b 100644
--- a/imagelib/common.py
+++ b/core/imagelib/common.py
@@ -23,7 +23,14 @@ def normalize_channels(img, target_channels):
c = target_channels
return img
-
+
+def cut_odd_image(img):
+ h, w, c = img.shape
+ wm, hm = w % 2, h % 2
+ if wm + hm != 0:
+ img = img[0:h-hm,0:w-wm,:]
+ return img
+
def overlay_alpha_image(img_target, img_source, xy_offset=(0,0) ):
(h,w,c) = img_source.shape
if c != 4:
diff --git a/imagelib/draw.py b/core/imagelib/draw.py
similarity index 100%
rename from imagelib/draw.py
rename to core/imagelib/draw.py
diff --git a/imagelib/equalize_and_stack_square.py b/core/imagelib/equalize_and_stack_square.py
similarity index 100%
rename from imagelib/equalize_and_stack_square.py
rename to core/imagelib/equalize_and_stack_square.py
diff --git a/imagelib/estimate_sharpness.py b/core/imagelib/estimate_sharpness.py
similarity index 100%
rename from imagelib/estimate_sharpness.py
rename to core/imagelib/estimate_sharpness.py
diff --git a/imagelib/morph.py b/core/imagelib/morph.py
similarity index 100%
rename from imagelib/morph.py
rename to core/imagelib/morph.py
diff --git a/imagelib/reduce_colors.py b/core/imagelib/reduce_colors.py
similarity index 100%
rename from imagelib/reduce_colors.py
rename to core/imagelib/reduce_colors.py
diff --git a/imagelib/text.py b/core/imagelib/text.py
similarity index 100%
rename from imagelib/text.py
rename to core/imagelib/text.py
diff --git a/imagelib/warp.py b/core/imagelib/warp.py
similarity index 88%
rename from imagelib/warp.py
rename to core/imagelib/warp.py
index a3c9490..d5d79b4 100644
--- a/imagelib/warp.py
+++ b/core/imagelib/warp.py
@@ -1,6 +1,6 @@
import numpy as np
import cv2
-from utils import random_utils
+from core import randomex
def gen_warp_params (source, flip, rotation_range=[-10,10], scale_range=[-0.5, 0.5], tx_range=[-0.05, 0.05], ty_range=[-0.05, 0.05], rnd_seed=None ):
h,w,c = source.shape
@@ -26,8 +26,8 @@ def gen_warp_params (source, flip, rotation_range=[-10,10], scale_range=[-0.5, 0
mapx = np.broadcast_to(grid_points, (cell_count, cell_count)).copy()
mapy = mapx.T
- mapx[1:-1,1:-1] = mapx[1:-1,1:-1] + random_utils.random_normal( size=(cell_count-2, cell_count-2) )*(cell_size*0.24)
- mapy[1:-1,1:-1] = mapy[1:-1,1:-1] + random_utils.random_normal( size=(cell_count-2, cell_count-2) )*(cell_size*0.24)
+ mapx[1:-1,1:-1] = mapx[1:-1,1:-1] + randomex.random_normal( size=(cell_count-2, cell_count-2) )*(cell_size*0.24)
+ mapy[1:-1,1:-1] = mapy[1:-1,1:-1] + randomex.random_normal( size=(cell_count-2, cell_count-2) )*(cell_size*0.24)
half_cell_size = cell_size // 2
diff --git a/interact/__init__.py b/core/interact/__init__.py
similarity index 100%
rename from interact/__init__.py
rename to core/interact/__init__.py
diff --git a/interact/interact.py b/core/interact/interact.py
similarity index 74%
rename from interact/interact.py
rename to core/interact/interact.py
index dce7c37..46fd450 100644
--- a/interact/interact.py
+++ b/core/interact/interact.py
@@ -4,6 +4,7 @@ import sys
import time
import types
+import colorama
import cv2
from tqdm import tqdm
@@ -16,6 +17,8 @@ try:
except:
is_colab = False
+yn_str = {True:'y',False:'n'}
+
class InteractBase(object):
EVENT_LBUTTONDOWN = 1
EVENT_LBUTTONUP = 2
@@ -186,13 +189,35 @@ class InteractBase(object):
ar = self.key_events.get(wnd_name, [])
self.key_events[wnd_name] = []
return ar
+
+ def input(self, s):
+ return input(s)
- def input_number(self, s, default_value, valid_list=None, help_message=None):
+ def input_number(self, s, default_value, valid_list=None, show_default_value=True, add_info=None, help_message=None):
+ if show_default_value and default_value is not None:
+ s = f"[{default_value}] {s}"
+
+ if add_info is not None or \
+ help_message is not None:
+ s += " ("
+
+ if add_info is not None:
+ s += f" {add_info}"
+ if help_message is not None:
+ s += " ?:help"
+
+ if add_info is not None or \
+ help_message is not None:
+ s += " )"
+
+ s += " : "
+
while True:
try:
inp = input(s)
if len(inp) == 0:
- raise ValueError("")
+ result = default_value
+ break
if help_message is not None and inp == '?':
print (help_message)
@@ -200,13 +225,39 @@ class InteractBase(object):
i = float(inp)
if (valid_list is not None) and (i not in valid_list):
- return default_value
- return i
+ result = default_value
+ break
+ result = i
+ break
except:
- print (default_value)
- return default_value
-
- def input_int(self,s, default_value, valid_list=None, help_message=None):
+ result = default_value
+ break
+
+ print(result)
+ return result
+
+ def input_int(self, s, default_value, valid_list=None, add_info=None, show_default_value=True, help_message=None):
+ if show_default_value:
+ if len(s) != 0:
+ s = f"[{default_value}] {s}"
+ else:
+ s = f"[{default_value}]"
+
+ if add_info is not None or \
+ help_message is not None:
+ s += " ("
+
+ if add_info is not None:
+ s += f" {add_info}"
+ if help_message is not None:
+ s += " ?:help"
+
+ if add_info is not None or \
+ help_message is not None:
+ s += " )"
+
+ s += " : "
+
while True:
try:
inp = input(s)
@@ -219,13 +270,23 @@ class InteractBase(object):
i = int(inp)
if (valid_list is not None) and (i not in valid_list):
- return default_value
- return i
+ result = default_value
+ break
+ result = i
+ break
except:
- print (default_value)
- return default_value
+ result = default_value
+ break
+ print (result)
+ return result
- def input_bool(self, s, default_value, help_message=None):
+ def input_bool(self, s, default_value, help_message=None):
+ s = f"[{yn_str[default_value]}] {s} ( y/n"
+
+ if help_message is not None:
+ s += " ?:help"
+ s += " ) : "
+
while True:
try:
inp = input(s)
@@ -236,35 +297,65 @@ class InteractBase(object):
print (help_message)
continue
- return bool ( {"y":True,"n":False,"1":True,"0":False}.get(inp.lower(), default_value) )
+ return bool ( {"y":True,"n":False}.get(inp.lower(), default_value) )
except:
print ( "y" if default_value else "n" )
return default_value
- def input_str(self, s, default_value, valid_list=None, help_message=None):
+ def input_str(self, s, default_value=None, valid_list=None, show_default_value=True, help_message=None):
+ if show_default_value and default_value is not None:
+ s = f"[{default_value}] {s}"
+
+ if valid_list is not None or \
+ help_message is not None:
+ s += " ("
+
+ if valid_list is not None:
+ s += " " + "/".join(valid_list)
+
+ if help_message is not None:
+ s += " ?:help"
+
+ if valid_list is not None or \
+ help_message is not None:
+ s += " )"
+
+ s += " : "
+
+
while True:
try:
inp = input(s)
+
if len(inp) == 0:
- raise ValueError("")
-
+ if default_value is None:
+ print("")
+ return None
+ result = default_value
+ break
+
if help_message is not None and inp == '?':
- print (help_message)
+ print(help_message)
continue
-
+
if valid_list is not None:
if inp.lower() in valid_list:
- return inp.lower()
+ result = inp.lower()
+ break
if inp in valid_list:
- return inp
- return default_value
-
- return inp
+ result = inp
+ break
+ continue
+ result = inp
+ break
except:
- print (default_value)
- return default_value
-
+ result = default_value
+ break
+
+ print(result)
+ return result
+
def input_process(self, stdin_fd, sq, str):
sys.stdin = os.fdopen(stdin_fd)
try:
@@ -292,7 +383,14 @@ class InteractBase(object):
class InteractDesktop(InteractBase):
-
+ def __init__(self):
+ colorama.init()
+ super().__init__()
+
+ def color_red(self):
+ pass
+
+
def is_support_windows(self):
return True
@@ -347,9 +445,20 @@ class InteractDesktop(InteractBase):
ord_key = cv2.waitKey(wait_key_time)
shift_pressed = False
if ord_key != -1:
- if chr(ord_key) >= 'A' and chr(ord_key) <= 'Z':
+ chr_key = chr(ord_key)
+
+ if chr_key >= 'A' and chr_key <= 'Z':
shift_pressed = True
ord_key += 32
+ elif chr_key == '?':
+ shift_pressed = True
+ ord_key = ord('/')
+ elif chr_key == '<':
+ shift_pressed = True
+ ord_key = ord(',')
+ elif chr_key == '>':
+ shift_pressed = True
+ ord_key = ord('.')
else:
if sleep_time != 0:
time.sleep(sleep_time)
diff --git a/joblib/SubprocessFunctionCaller.py b/core/joblib/SubprocessFunctionCaller.py
similarity index 100%
rename from joblib/SubprocessFunctionCaller.py
rename to core/joblib/SubprocessFunctionCaller.py
diff --git a/utils/iter_utils.py b/core/joblib/SubprocessGenerator.py
similarity index 75%
rename from utils/iter_utils.py
rename to core/joblib/SubprocessGenerator.py
index e690e3b..b0d893e 100644
--- a/utils/iter_utils.py
+++ b/core/joblib/SubprocessGenerator.py
@@ -1,28 +1,8 @@
-import threading
import queue as Queue
import multiprocessing
-import time
-
-
-class ThisThreadGenerator(object):
- def __init__(self, generator_func, user_param=None):
- super().__init__()
- self.generator_func = generator_func
- self.user_param = user_param
- self.initialized = False
-
- def __iter__(self):
- return self
-
- def __next__(self):
- if not self.initialized:
- self.initialized = True
- self.generator_func = self.generator_func(self.user_param)
-
- return next(self.generator_func)
class SubprocessGenerator(object):
- def __init__(self, generator_func, user_param=None, prefetch=2, start_now=False):
+ def __init__(self, generator_func, user_param=None, prefetch=2, start_now=True):
super().__init__()
self.prefetch = prefetch
self.generator_func = generator_func
diff --git a/joblib/SubprocessorBase.py b/core/joblib/SubprocessorBase.py
similarity index 99%
rename from joblib/SubprocessorBase.py
rename to core/joblib/SubprocessorBase.py
index a9cbc36..993d5cb 100644
--- a/joblib/SubprocessorBase.py
+++ b/core/joblib/SubprocessorBase.py
@@ -2,7 +2,7 @@ import traceback
import multiprocessing
import time
import sys
-from interact import interact as io
+from core.interact import interact as io
class Subprocessor(object):
@@ -87,7 +87,7 @@ class Subprocessor(object):
c2s.put ( {'op': 'error', 'data' : data} )
#overridable
- def __init__(self, name, SubprocessorCli_class, no_response_time_sec = 0, io_loop_sleep_time=0.005, initialize_subprocesses_in_serial=True):
+ def __init__(self, name, SubprocessorCli_class, no_response_time_sec = 0, io_loop_sleep_time=0.005, initialize_subprocesses_in_serial=False):
if not issubclass(SubprocessorCli_class, Subprocessor.Cli):
raise ValueError("SubprocessorCli_class must be subclass of Subprocessor.Cli")
@@ -130,7 +130,7 @@ class Subprocessor(object):
#overridable
def get_result(self):
#return result that will be returned in func run()
- raise NotImplementedError
+ return None
#overridable
def on_tick(self):
diff --git a/core/joblib/ThisThreadGenerator.py b/core/joblib/ThisThreadGenerator.py
new file mode 100644
index 0000000..e6f77a4
--- /dev/null
+++ b/core/joblib/ThisThreadGenerator.py
@@ -0,0 +1,16 @@
+class ThisThreadGenerator(object):
+ def __init__(self, generator_func, user_param=None):
+ super().__init__()
+ self.generator_func = generator_func
+ self.user_param = user_param
+ self.initialized = False
+
+ def __iter__(self):
+ return self
+
+ def __next__(self):
+ if not self.initialized:
+ self.initialized = True
+ self.generator_func = self.generator_func(self.user_param)
+
+ return next(self.generator_func)
\ No newline at end of file
diff --git a/joblib/__init__.py b/core/joblib/__init__.py
similarity index 50%
rename from joblib/__init__.py
rename to core/joblib/__init__.py
index fbbc20c..68b3101 100644
--- a/joblib/__init__.py
+++ b/core/joblib/__init__.py
@@ -1,2 +1,4 @@
from .SubprocessorBase import Subprocessor
from .SubprocessFunctionCaller import SubprocessFunctionCaller
+from .ThisThreadGenerator import ThisThreadGenerator
+from .SubprocessGenerator import SubprocessGenerator
\ No newline at end of file
diff --git a/core/leras/__init__.py b/core/leras/__init__.py
new file mode 100644
index 0000000..7d9fb2b
--- /dev/null
+++ b/core/leras/__init__.py
@@ -0,0 +1 @@
+from .nn import nn
\ No newline at end of file
diff --git a/core/leras/device.py b/core/leras/device.py
new file mode 100644
index 0000000..e18ea2a
--- /dev/null
+++ b/core/leras/device.py
@@ -0,0 +1,205 @@
+import sys
+import ctypes
+import os
+
+class Device(object):
+ def __init__(self, index, name, total_mem, free_mem, cc=0):
+ self.index = index
+ self.name = name
+ self.cc = cc
+ self.total_mem = total_mem
+ self.total_mem_gb = total_mem / 1024**3
+ self.free_mem = free_mem
+ self.free_mem_gb = free_mem / 1024**3
+
+ def __str__(self):
+ return f"[{self.index}]:[{self.name}][{self.free_mem_gb:.3}/{self.total_mem_gb :.3}]"
+
+class Devices(object):
+ all_devices = None
+
+ def __init__(self, devices):
+ self.devices = devices
+
+ def __len__(self):
+ return len(self.devices)
+
+ def __getitem__(self, key):
+ result = self.devices[key]
+ if isinstance(key, slice):
+ return Devices(result)
+ return result
+
+ def __iter__(self):
+ for device in self.devices:
+ yield device
+
+ def get_best_device(self):
+ result = None
+ idx_mem = 0
+ for device in self.devices:
+ mem = device.total_mem
+ if mem > idx_mem:
+ result = device
+ idx_mem = mem
+ return result
+
+ def get_worst_device(self):
+ result = None
+ idx_mem = sys.maxsize
+ for device in self.devices:
+ mem = device.total_mem
+ if mem < idx_mem:
+ result = device
+ idx_mem = mem
+ return result
+
+ def get_device_by_index(self, idx):
+ for device in self.devices:
+ if device.index == idx:
+ return device
+ return None
+
+ def get_devices_from_index_list(self, idx_list):
+ result = []
+ for device in self.devices:
+ if device.index in idx_list:
+ result += [device]
+ return Devices(result)
+
+ def get_equal_devices(self, device):
+ device_name = device.name
+ result = []
+ for device in self.devices:
+ if device.name == device_name:
+ result.append (device)
+ return Devices(result)
+
+ def get_devices_at_least_mem(self, totalmemsize_gb):
+ result = []
+ for device in self.devices:
+ if device.total_mem >= totalmemsize_gb*(1024**3):
+ result.append (device)
+ return Devices(result)
+
+ @staticmethod
+ def initialize_main_env():
+ min_cc = int(os.environ.get("TF_MIN_REQ_CAP", 35))
+ libnames = ('libcuda.so', 'libcuda.dylib', 'nvcuda.dll')
+ for libname in libnames:
+ try:
+ cuda = ctypes.CDLL(libname)
+ except:
+ continue
+ else:
+ break
+ else:
+ return Devices([])
+
+ nGpus = ctypes.c_int()
+ name = b' ' * 200
+ cc_major = ctypes.c_int()
+ cc_minor = ctypes.c_int()
+ freeMem = ctypes.c_size_t()
+ totalMem = ctypes.c_size_t()
+
+ result = ctypes.c_int()
+ device = ctypes.c_int()
+ context = ctypes.c_void_p()
+ error_str = ctypes.c_char_p()
+
+ devices = []
+
+ if cuda.cuInit(0) == 0 and \
+ cuda.cuDeviceGetCount(ctypes.byref(nGpus)) == 0:
+ for i in range(nGpus.value):
+ if cuda.cuDeviceGet(ctypes.byref(device), i) != 0 or \
+ cuda.cuDeviceGetName(ctypes.c_char_p(name), len(name), device) != 0 or \
+ cuda.cuDeviceComputeCapability(ctypes.byref(cc_major), ctypes.byref(cc_minor), device) != 0:
+ continue
+
+ if cuda.cuCtxCreate_v2(ctypes.byref(context), 0, device) == 0:
+ if cuda.cuMemGetInfo_v2(ctypes.byref(freeMem), ctypes.byref(totalMem)) == 0:
+ cc = cc_major.value * 10 + cc_minor.value
+ if cc >= min_cc:
+ devices.append ( {'name' : name.split(b'\0', 1)[0].decode(),
+ 'total_mem' : totalMem.value,
+ 'free_mem' : freeMem.value,
+ 'cc' : cc
+ })
+ cuda.cuCtxDetach(context)
+
+ os.environ['NN_DEVICES_INITIALIZED'] = '1'
+ os.environ['NN_DEVICES_COUNT'] = str(len(devices))
+ for i, device in enumerate(devices):
+ os.environ[f'NN_DEVICE_{i}_NAME'] = device['name']
+ os.environ[f'NN_DEVICE_{i}_TOTAL_MEM'] = str(device['total_mem'])
+ os.environ[f'NN_DEVICE_{i}_FREE_MEM'] = str(device['free_mem'])
+ os.environ[f'NN_DEVICE_{i}_CC'] = str(device['cc'])
+
+ @staticmethod
+ def getDevices():
+ if Devices.all_devices is None:
+ if int(os.environ.get("NN_DEVICES_INITIALIZED", 0)) != 1:
+ raise Exception("nn devices are not initialized. Run initialize_main_env() in main process.")
+ devices = []
+ for i in range ( int(os.environ['NN_DEVICES_COUNT']) ):
+ devices.append ( Device(index=i,
+ name=os.environ[f'NN_DEVICE_{i}_NAME'],
+ total_mem=int(os.environ[f'NN_DEVICE_{i}_TOTAL_MEM']),
+ free_mem=int(os.environ[f'NN_DEVICE_{i}_FREE_MEM']),
+ cc=int(os.environ[f'NN_DEVICE_{i}_CC']) ))
+ Devices.all_devices = Devices(devices)
+
+ return Devices.all_devices
+
+"""
+if Devices.all_devices is None:
+ min_cc = int(os.environ.get("TF_MIN_REQ_CAP", 35))
+
+ libnames = ('libcuda.so', 'libcuda.dylib', 'nvcuda.dll')
+ for libname in libnames:
+ try:
+ cuda = ctypes.CDLL(libname)
+ except:
+ continue
+ else:
+ break
+ else:
+ return Devices([])
+
+ nGpus = ctypes.c_int()
+ name = b' ' * 200
+ cc_major = ctypes.c_int()
+ cc_minor = ctypes.c_int()
+ freeMem = ctypes.c_size_t()
+ totalMem = ctypes.c_size_t()
+
+ result = ctypes.c_int()
+ device = ctypes.c_int()
+ context = ctypes.c_void_p()
+ error_str = ctypes.c_char_p()
+
+ devices = []
+
+ if cuda.cuInit(0) == 0 and \
+ cuda.cuDeviceGetCount(ctypes.byref(nGpus)) == 0:
+ for i in range(nGpus.value):
+ if cuda.cuDeviceGet(ctypes.byref(device), i) != 0 or \
+ cuda.cuDeviceGetName(ctypes.c_char_p(name), len(name), device) != 0 or \
+ cuda.cuDeviceComputeCapability(ctypes.byref(cc_major), ctypes.byref(cc_minor), device) != 0:
+ continue
+
+ if cuda.cuCtxCreate_v2(ctypes.byref(context), 0, device) == 0:
+ if cuda.cuMemGetInfo_v2(ctypes.byref(freeMem), ctypes.byref(totalMem)) == 0:
+ cc = cc_major.value * 10 + cc_minor.value
+ if cc >= min_cc:
+ devices.append ( Device(index=i,
+ name=name.split(b'\0', 1)[0].decode(),
+ total_mem=totalMem.value,
+ free_mem=freeMem.value,
+ cc=cc) )
+ cuda.cuCtxDetach(context)
+ Devices.all_devices = Devices(devices)
+ return Devices.all_devices
+"""
\ No newline at end of file
diff --git a/core/leras/initializers.py b/core/leras/initializers.py
new file mode 100644
index 0000000..a3294cb
--- /dev/null
+++ b/core/leras/initializers.py
@@ -0,0 +1,52 @@
+import numpy as np
+
+def initialize_initializers(nn):
+ tf = nn.tf
+ from tensorflow.python.ops import init_ops
+
+ class initializers():
+ class ca (init_ops.Initializer):
+ def __init__(self, dtype=None):
+ pass
+
+ def __call__(self, shape, dtype=None, partition_info=None):
+ return tf.zeros( shape, name="_cai_")
+
+ @staticmethod
+ def generate(shape, eps_std=0.05, dtype=np.float32):
+ """
+ Super fast implementation of Convolution Aware Initialization for 4D shapes
+ Convolution Aware Initialization https://arxiv.org/abs/1702.06295
+ """
+ if len(shape) != 4:
+ raise ValueError("only shape with rank 4 supported.")
+
+ row, column, stack_size, filters_size = shape
+
+ fan_in = stack_size * (row * column)
+
+ kernel_shape = (row, column)
+
+ kernel_fft_shape = np.fft.rfft2(np.zeros(kernel_shape)).shape
+
+ basis_size = np.prod(kernel_fft_shape)
+ if basis_size == 1:
+ x = np.random.normal( 0.0, eps_std, (filters_size, stack_size, basis_size) )
+ else:
+ nbb = stack_size // basis_size + 1
+ x = np.random.normal(0.0, 1.0, (filters_size, nbb, basis_size, basis_size))
+ x = x + np.transpose(x, (0,1,3,2) ) * (1-np.eye(basis_size))
+ u, _, v = np.linalg.svd(x)
+ x = np.transpose(u, (0,1,3,2) )
+ x = np.reshape(x, (filters_size, -1, basis_size) )
+ x = x[:,:stack_size,:]
+
+ x = np.reshape(x, ( (filters_size,stack_size,) + kernel_fft_shape ) )
+
+ x = np.fft.irfft2( x, kernel_shape ) \
+ + np.random.normal(0, eps_std, (filters_size,stack_size,)+kernel_shape)
+
+ x = x * np.sqrt( (2/fan_in) / np.var(x) )
+ x = np.transpose( x, (2, 3, 1, 0) )
+ return x.astype(dtype)
+ nn.initializers = initializers
\ No newline at end of file
diff --git a/core/leras/layers.py b/core/leras/layers.py
new file mode 100644
index 0000000..7597ccf
--- /dev/null
+++ b/core/leras/layers.py
@@ -0,0 +1,591 @@
+import pickle
+import types
+from pathlib import Path
+from core import pathex
+from core.interact import interact as io
+import numpy as np
+
+
+def initialize_layers(nn):
+ tf = nn.tf
+
+ class Saveable():
+ def __init__(self, name=None):
+ self.name = name
+
+ #override
+ def get_weights(self):
+ #return tf tensors that should be initialized/loaded/saved
+ pass
+
+ def save_weights(self, filename, force_dtype=None):
+ d = {}
+ weights = self.get_weights()
+
+ if self.name is None:
+ raise Exception("name must be defined.")
+
+ name = self.name
+ for w, w_val in zip(weights, nn.tf_sess.run (weights)):
+ w_name_split = w.name.split('/', 1)
+ if name != w_name_split[0]:
+ raise Exception("weight first name != Saveable.name")
+
+ if force_dtype is not None:
+ w_val = w_val.astype(force_dtype)
+
+ d[ w_name_split[1] ] = w_val
+
+ d_dumped = pickle.dumps (d, 4)
+ pathex.write_bytes_safe ( Path(filename), d_dumped )
+
+ def load_weights(self, filename):
+ """
+ returns True if file exists
+ """
+ filepath = Path(filename)
+ if filepath.exists():
+ result = True
+ d_dumped = filepath.read_bytes()
+ d = pickle.loads(d_dumped)
+ else:
+ return False
+
+ weights = self.get_weights()
+
+ if self.name is None:
+ raise Exception("name must be defined.")
+
+ tuples = []
+ for w in weights:
+ w_name_split = w.name.split('/')
+ if self.name != w_name_split[0]:
+ raise Exception("weight first name != Saveable.name")
+
+ sub_w_name = "/".join(w_name_split[1:])
+
+ w_val = d.get(sub_w_name, None)
+ if w_val is None:
+ io.log_err(f"Weight {w.name} was not loaded from file {filename}")
+ tuples.append ( (w, w.initializer) )
+ else:
+ tuples.append ( (w, w_val) )
+
+ nn.tf_batch_set_value(tuples)
+
+ return True
+
+ def init_weights(self):
+ ops = []
+ tuples = []
+ for w in self.get_weights():
+ initializer = w.initializer
+ for input in initializer.inputs:
+ if "_cai_" in input.name:
+ tuples.append ( (w, nn.initializers.ca.generate(w.shape.as_list(), dtype= w.dtype.as_numpy_dtype) ) )
+ break
+ else:
+ ops.append (initializer)
+
+ nn.tf_sess.run (ops)
+ nn.tf_batch_set_value(tuples)
+ nn.Saveable = Saveable
+
+ class LayerBase():
+ def __init__(self, name=None, **kwargs):
+ self.name = name
+
+ #override
+ def build_weights(self):
+ pass
+
+ #override
+ def get_weights(self):
+ return []
+
+ def set_weights(self, new_weights):
+ weights = self.get_weights()
+ if len(weights) != len(new_weights):
+ raise ValueError ('len of lists mismatch')
+
+ tuples = []
+ for w, new_w in zip(weights, new_weights):
+ if len(w.shape) != new_w.shape:
+ new_w = new_w.reshape(w.shape)
+
+ tuples.append ( (w, new_w) )
+
+ nn.tf_batch_set_value (tuples)
+ nn.LayerBase = LayerBase
+
+ class ModelBase(Saveable):
+ def __init__(self, *args, name=None, **kwargs):
+ super().__init__(name=name)
+ self.layers = []
+ self.built = False
+ self.args = args
+ self.kwargs = kwargs
+ self.run_placeholders = None
+
+ def _build_sub(self, layer, name):
+ if isinstance (layer, list):
+ for i,sublayer in enumerate(layer):
+ self._build_sub(sublayer, f"{name}_{i}")
+ elif isinstance (layer, LayerBase) or \
+ isinstance (layer, ModelBase):
+
+ if layer.name is None:
+ layer.name = name
+
+ if isinstance (layer, LayerBase):
+ with tf.variable_scope(layer.name):
+ layer.build_weights()
+ elif isinstance (layer, ModelBase):
+ layer.build()
+
+ self.layers.append (layer)
+
+ def xor_list(self, lst1, lst2):
+ return [value for value in lst1+lst2 if (value not in lst1) or (value not in lst2) ]
+
+ def build(self):
+ with tf.variable_scope(self.name):
+
+ current_vars = []
+ generator = None
+ while True:
+
+ if generator is None:
+ generator = self.on_build(*self.args, **self.kwargs)
+ if not isinstance(generator, types.GeneratorType):
+ generator = None
+
+ if generator is not None:
+ try:
+ next(generator)
+ except StopIteration:
+ generator = None
+
+ v = vars(self)
+ new_vars = self.xor_list (current_vars, list(v.keys()) )
+
+ for name in new_vars:
+ self._build_sub(v[name],name)
+
+ current_vars += new_vars
+
+ if generator is None:
+ break
+
+ self.built = True
+
+ #override
+ def get_weights(self):
+ if not self.built:
+ self.build()
+
+ weights = []
+ for layer in self.layers:
+ weights += layer.get_weights()
+ return weights
+
+ def get_layers(self):
+ if not self.built:
+ self.build()
+ layers = []
+ for layer in self.layers:
+ if isinstance (layer, LayerBase):
+ layers.append(layer)
+ else:
+ layers += layer.get_layers()
+ return layers
+
+ #override
+ def on_build(self, *args, **kwargs):
+ """
+ init model layers here
+
+ return 'yield' if build is not finished
+ therefore dependency models will be initialized
+ """
+ pass
+
+ #override
+ def forward(self, *args, **kwargs):
+ #flow layers/models/tensors here
+ pass
+
+ def __call__(self, *args, **kwargs):
+ if not self.built:
+ self.build()
+
+ return self.forward(*args, **kwargs)
+
+ def compute_output_shape(self, shapes):
+ if not self.built:
+ self.build()
+
+ not_list = False
+ if not isinstance(shapes, list):
+ not_list = True
+ shapes = [shapes]
+
+ with tf.device('/CPU:0'):
+ # CPU tensors will not impact any performance, only slightly RAM "leakage"
+ phs = []
+ for dtype,sh in shapes:
+ phs += [ tf.placeholder(dtype, sh) ]
+
+ result = self.__call__(phs[0] if not_list else phs)
+
+ if not isinstance(result, list):
+ result = [result]
+
+ result_shapes = []
+
+ for t in result:
+ result_shapes += [ t.shape.as_list() ]
+
+ return result_shapes[0] if not_list else result_shapes
+
+ def build_for_run(self, shapes_list):
+ if not isinstance(shapes_list, list):
+ raise ValueError("shapes_list must be a list.")
+
+ self.run_placeholders = []
+ for dtype,sh in shapes_list:
+ self.run_placeholders.append ( tf.placeholder(dtype, (None,)+sh) )
+
+ self.run_output = self.__call__(self.run_placeholders)
+
+ def run (self, inputs):
+ if self.run_placeholders is None:
+ raise Exception ("Model didn't build for run.")
+
+ if len(inputs) != len(self.run_placeholders):
+ raise ValueError("len(inputs) != self.run_placeholders")
+
+ feed_dict = {}
+ for ph, inp in zip(self.run_placeholders, inputs):
+ feed_dict[ph] = inp
+
+ return nn.tf_sess.run ( self.run_output, feed_dict=feed_dict)
+
+ nn.ModelBase = ModelBase
+
+ class Conv2D(LayerBase):
+ """
+ use_wscale bool enables equalized learning rate, kernel_initializer will be forced to random_normal
+
+
+ """
+ def __init__(self, in_ch, out_ch, kernel_size, strides=1, padding='SAME', dilations=1, use_bias=True, use_wscale=False, kernel_initializer=None, bias_initializer=None, trainable=True, dtype=None, **kwargs ):
+ if not isinstance(strides, int):
+ raise ValueError ("strides must be an int type")
+ if not isinstance(dilations, int):
+ raise ValueError ("dilations must be an int type")
+
+ if isinstance(padding, str):
+ if padding == "SAME":
+ padding = ( (kernel_size - 1) * dilations + 1 ) // 2
+ elif padding == "VALID":
+ padding = 0
+ else:
+ raise ValueError ("Wrong padding type. Should be VALID SAME or INT or 4x INTs")
+
+ if isinstance(padding, int):
+ if padding != 0:
+ padding = [ [0,0], [padding,padding], [padding,padding], [0,0] ]
+ else:
+ padding = None
+
+ self.in_ch = in_ch
+ self.out_ch = out_ch
+ self.kernel_size = kernel_size
+ self.strides = [1,strides,strides,1]
+ self.padding = padding
+ self.dilations = [1,dilations,dilations,1]
+ self.use_bias = use_bias
+ self.use_wscale = use_wscale
+ self.kernel_initializer = None if use_wscale else kernel_initializer
+ self.bias_initializer = bias_initializer
+ self.trainable = trainable
+ if dtype is None:
+ dtype = nn.tf_floatx
+ self.dtype = dtype
+ super().__init__(**kwargs)
+
+ def build_weights(self):
+ kernel_initializer = self.kernel_initializer
+ if kernel_initializer is None:
+ if self.use_wscale:
+ gain = 1.0 if self.kernel_size == 1 else np.sqrt(2)
+ fan_in = self.kernel_size*self.kernel_size*self.in_ch
+ he_std = gain / np.sqrt(fan_in) # He init
+ self.wscale = tf.constant(he_std, dtype=self.dtype )
+ kernel_initializer = tf.initializers.random_normal(0, 1.0, dtype=self.dtype)
+ else:
+ kernel_initializer = tf.initializers.glorot_uniform(dtype=self.dtype)
+
+ self.weight = tf.get_variable("weight", (self.kernel_size,self.kernel_size,self.in_ch,self.out_ch), dtype=self.dtype, initializer=kernel_initializer, trainable=self.trainable )
+
+ if self.use_bias:
+ bias_initializer = self.bias_initializer
+ if bias_initializer is None:
+ bias_initializer = tf.initializers.zeros(dtype=self.dtype)
+
+ self.bias = tf.get_variable("bias", (1,1,1,self.out_ch), dtype=self.dtype, initializer=bias_initializer, trainable=self.trainable )
+
+ def get_weights(self):
+ weights = [self.weight]
+ if self.use_bias:
+ weights += [self.bias]
+ return weights
+
+ def __call__(self, x):
+ weight = self.weight
+ if self.use_wscale:
+ weight = weight * self.wscale
+
+ if self.padding is not None:
+ x = tf.pad (x, self.padding, mode='CONSTANT')
+
+ x = tf.nn.conv2d(x, weight, self.strides, 'VALID', dilations=self.dilations)
+ if self.use_bias:
+ x = x + self.bias
+ return x
+
+ def __str__(self):
+ r = f"{self.__class__.__name__} : in_ch:{self.in_ch} out_ch:{self.out_ch} "
+
+ return r
+ nn.Conv2D = Conv2D
+
+ class Conv2DTranspose(LayerBase):
+ """
+ use_wscale enables weight scale (equalized learning rate)
+ kernel_initializer will be forced to random_normal
+ """
+ def __init__(self, in_ch, out_ch, kernel_size, strides=2, padding='SAME', use_bias=True, use_wscale=False, kernel_initializer=None, bias_initializer=None, trainable=True, dtype=None, **kwargs ):
+ if not isinstance(strides, int):
+ raise ValueError ("strides must be an int type")
+ self.in_ch = in_ch
+ self.out_ch = out_ch
+ self.kernel_size = kernel_size
+ self.strides = strides
+ self.padding = padding
+ self.use_bias = use_bias
+ self.use_wscale = use_wscale
+ self.kernel_initializer = None if use_wscale else kernel_initializer
+ self.bias_initializer = bias_initializer
+ self.trainable = trainable
+ if dtype is None:
+ dtype = nn.tf_floatx
+ self.dtype = dtype
+ super().__init__(**kwargs)
+
+ def build_weights(self):
+ kernel_initializer = self.kernel_initializer
+ if kernel_initializer is None:
+ if self.use_wscale:
+ gain = 1.0 if self.kernel_size == 1 else np.sqrt(2)
+ fan_in = self.kernel_size*self.kernel_size*self.in_ch
+ he_std = gain / np.sqrt(fan_in) # He init
+ self.wscale = tf.constant(he_std, dtype=self.dtype )
+ kernel_initializer = tf.initializers.random_normal(0, 1.0, dtype=self.dtype)
+ else:
+ kernel_initializer = tf.initializers.glorot_uniform(dtype=self.dtype)
+
+ self.weight = tf.get_variable("weight", (self.kernel_size,self.kernel_size,self.out_ch,self.in_ch), dtype=self.dtype, initializer=kernel_initializer, trainable=self.trainable )
+
+ if self.use_bias:
+ bias_initializer = self.bias_initializer
+ if bias_initializer is None:
+ bias_initializer = tf.initializers.zeros(dtype=self.dtype)
+ self.bias = tf.get_variable("bias", (1,1,1,self.out_ch), dtype=self.dtype, initializer=bias_initializer, trainable=self.trainable )
+
+ def get_weights(self):
+ weights = [self.weight]
+ if self.use_bias:
+ weights += [self.bias]
+ return weights
+
+ def __call__(self, x):
+ shape = x.shape
+
+ h,w,c = shape[1], shape[2], shape[3]
+
+ output_shape = tf.stack ( (tf.shape(x)[0],
+ self.deconv_length(w, self.strides, self.kernel_size, self.padding),
+ self.deconv_length(h, self.strides, self.kernel_size, self.padding),
+ self.out_ch) )
+
+ weight = self.weight
+ if self.use_wscale:
+ weight = weight * self.wscale
+
+ x = tf.nn.conv2d_transpose(x, weight, output_shape, [1,self.strides,self.strides,1], padding=self.padding)
+
+ if self.use_bias:
+ x = x + self.bias
+ return x
+
+ def __str__(self):
+ r = f"{self.__class__.__name__} : in_ch:{self.in_ch} out_ch:{self.out_ch} "
+
+ return r
+
+ def deconv_length(self, dim_size, stride_size, kernel_size, padding):
+ assert padding in {'SAME', 'VALID', 'FULL'}
+ if dim_size is None:
+ return None
+ if padding == 'VALID':
+ dim_size = dim_size * stride_size + max(kernel_size - stride_size, 0)
+ elif padding == 'FULL':
+ dim_size = dim_size * stride_size - (stride_size + kernel_size - 2)
+ elif padding == 'SAME':
+ dim_size = dim_size * stride_size
+ return dim_size
+ nn.Conv2DTranspose = Conv2DTranspose
+
+ class BlurPool(LayerBase):
+ def __init__(self, filt_size=3, stride=2, **kwargs ):
+ self.strides = [1,stride,stride,1]
+ self.filt_size = filt_size
+ self.padding = [ [0,0],
+ [ int(1.*(filt_size-1)/2), int(np.ceil(1.*(filt_size-1)/2)) ],
+ [ int(1.*(filt_size-1)/2), int(np.ceil(1.*(filt_size-1)/2)) ],
+ [0,0] ]
+ if(self.filt_size==1):
+ a = np.array([1.,])
+ elif(self.filt_size==2):
+ a = np.array([1., 1.])
+ elif(self.filt_size==3):
+ a = np.array([1., 2., 1.])
+ elif(self.filt_size==4):
+ a = np.array([1., 3., 3., 1.])
+ elif(self.filt_size==5):
+ a = np.array([1., 4., 6., 4., 1.])
+ elif(self.filt_size==6):
+ a = np.array([1., 5., 10., 10., 5., 1.])
+ elif(self.filt_size==7):
+ a = np.array([1., 6., 15., 20., 15., 6., 1.])
+
+ a = a[:,None]*a[None,:]
+ a = a / np.sum(a)
+ a = a[:,:,None,None]
+ self.a = a
+ super().__init__(**kwargs)
+
+ def build_weights(self):
+ self.k = tf.constant (self.a, dtype=nn.tf_floatx )
+
+ def __call__(self, x):
+ k = tf.tile (self.k, (1,1,x.shape[-1],1) )
+ x = tf.pad(x, self.padding )
+ x = tf.nn.depthwise_conv2d(x, k, self.strides, 'VALID')
+ return x
+ nn.BlurPool = BlurPool
+
+ class Dense(LayerBase):
+ def __init__(self, in_ch, out_ch, use_bias=True, use_wscale=False, maxout_ch=0, kernel_initializer=None, bias_initializer=None, trainable=True, dtype=None, **kwargs ):
+ """
+ use_wscale enables weight scale (equalized learning rate)
+ kernel_initializer will be forced to random_normal
+
+ maxout_ch https://link.springer.com/article/10.1186/s40537-019-0233-0
+ typical 2-4 if you want to enable DenseMaxout behaviour
+ """
+ self.in_ch = in_ch
+ self.out_ch = out_ch
+ self.use_bias = use_bias
+ self.use_wscale = use_wscale
+ self.maxout_ch = maxout_ch
+ self.kernel_initializer = kernel_initializer
+ self.bias_initializer = bias_initializer
+ self.trainable = trainable
+ if dtype is None:
+ dtype = tf.float32
+ self.dtype = dtype
+ super().__init__(**kwargs)
+
+ def build_weights(self):
+ if self.maxout_ch > 1:
+ weight_shape = (self.in_ch,self.out_ch*self.maxout_ch)
+ else:
+ weight_shape = (self.in_ch,self.out_ch)
+
+ kernel_initializer = self.kernel_initializer
+ if kernel_initializer is None:
+ if self.use_wscale:
+ gain = 1.0
+ fan_in = np.prod( weight_shape[:-1] )
+ he_std = gain / np.sqrt(fan_in) # He init
+ self.wscale = tf.constant(he_std, dtype=self.dtype )
+ kernel_initializer = tf.initializers.random_normal(0, 1.0, dtype=self.dtype)
+ else:
+ kernel_initializer = tf.initializers.glorot_uniform(dtype=self.dtype)
+
+ self.weight = tf.get_variable("weight", weight_shape, dtype=self.dtype, initializer=kernel_initializer, trainable=self.trainable )
+
+ if self.use_bias:
+ bias_initializer = self.bias_initializer
+ if bias_initializer is None:
+ bias_initializer = tf.initializers.zeros(dtype=self.dtype)
+ self.bias = tf.get_variable("bias", (1,self.out_ch), dtype=self.dtype, initializer=bias_initializer, trainable=self.trainable )
+
+ def get_weights(self):
+ weights = [self.weight]
+ if self.use_bias:
+ weights += [self.bias]
+ return weights
+
+ def __call__(self, x):
+ weight = self.weight
+ if self.use_wscale:
+ weight = weight * self.wscale
+
+ x = tf.matmul(x, weight)
+
+ if self.maxout_ch > 1:
+ x = tf.reshape (x, (-1, self.out_ch, self.maxout_ch) )
+ x = tf.reduce_max(x, axis=-1)
+
+ if self.use_bias:
+ x = x + self.bias
+
+ return x
+ nn.Dense = Dense
+
+ class BatchNorm2D(LayerBase):
+ """
+ currently not for training
+ """
+ def __init__(self, dim, eps=1e-05, momentum=0.1, dtype=None, **kwargs ):
+ self.dim = dim
+ self.eps = eps
+ self.momentum = momentum
+ if dtype is None:
+ dtype = nn.tf_floatx
+ self.dtype = dtype
+
+ self.shape = (1,1,1,dim)
+
+ super().__init__(**kwargs)
+
+ def build_weights(self):
+ self.weight = tf.get_variable("weight", self.shape, dtype=self.dtype, initializer=tf.initializers.ones() )
+ self.bias = tf.get_variable("bias", self.shape, dtype=self.dtype, initializer=tf.initializers.zeros() )
+ self.running_mean = tf.get_variable("running_mean", self.shape, dtype=self.dtype, initializer=tf.initializers.zeros(), trainable=False )
+ self.running_var = tf.get_variable("running_var", self.shape, dtype=self.dtype, initializer=tf.initializers.zeros(), trainable=False )
+
+ def get_weights(self):
+ return [self.weight, self.bias, self.running_mean, self.running_var]
+
+ def __call__(self, x):
+ x = (x - self.running_mean) / tf.sqrt( self.running_var + self.eps )
+ x *= self.weight
+ x += self.bias
+ return x
+
+ nn.BatchNorm2D = BatchNorm2D
\ No newline at end of file
diff --git a/core/leras/nn.py b/core/leras/nn.py
new file mode 100644
index 0000000..280bcdb
--- /dev/null
+++ b/core/leras/nn.py
@@ -0,0 +1,256 @@
+"""
+Leras.
+
+like lighter keras.
+This is my lightweight neural network library written from scratch
+based on pure tensorflow without keras.
+
+Provides:
++ full freedom of tensorflow operations without keras model's restrictions
++ easy model operations like in PyTorch, but in graph mode (no eager execution)
++ convenient and understandable logic
+
+Reasons why we cannot import tensorflow or any tensorflow.sub modules right here:
+1) change env variables based on DeviceConfig before import tensorflow
+2) multiprocesses will import tensorflow every spawn
+"""
+
+import os
+import sys
+from pathlib import Path
+from core.interact import interact as io
+from .device import Devices
+
+class nn():
+ current_DeviceConfig = None
+
+ tf = None
+ tf_sess = None
+ tf_sess_config = None
+
+ # Tensor ops
+ tf_get_value = None
+ tf_batch_set_value = None
+ tf_gradients = None
+ tf_average_gv_list = None
+ tf_average_tensor_list = None
+ tf_dot = None
+ tf_gelu = None
+ tf_upsample2d = None
+ tf_upsample2d_bilinear = None
+ tf_flatten = None
+ tf_random_binomial = None
+ tf_gaussian_blur = None
+ tf_style_loss = None
+ tf_dssim = None
+
+ # Layers
+ Saveable = None
+ LayerBase = None
+ ModelBase = None
+ Conv2D = None
+ Conv2DTranspose = None
+ BlurPool = None
+ Dense = None
+ BatchNorm2D = None
+
+ # Initializers
+ initializers = None
+
+ # Optimizers
+ TFBaseOptimizer = None
+ TFRMSpropOptimizer = None
+
+ @staticmethod
+ def initialize(device_config=None):
+ if nn.tf is None:
+ if device_config is None:
+ device_config = nn.getCurrentDeviceConfig()
+ else:
+ nn.setCurrentDeviceConfig(device_config)
+
+ if 'CUDA_VISIBLE_DEVICES' in os.environ.keys():
+ os.environ.pop('CUDA_VISIBLE_DEVICES')
+
+ os.environ['CUDA_CACHE_MAXSIZE'] = '536870912' #512Mb (32mb default)
+
+ first_run = False
+
+ if sys.platform[0:3] == 'win':
+ devices_str = ""
+ for device in device_config.devices:
+ devices_str += "_" + device.name.replace(' ','_')
+
+ compute_cache_path = Path(os.environ['APPDATA']) / 'NVIDIA' / ('ComputeCache' + devices_str)
+ if not compute_cache_path.exists():
+ first_run = True
+ os.environ['CUDA_CACHE_PATH'] = str(compute_cache_path)
+
+ os.environ['TF_MIN_GPU_MULTIPROCESSOR_COUNT'] = '2'
+ os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' # tf log errors only
+
+ import warnings
+ warnings.simplefilter(action='ignore', category=FutureWarning)
+
+ if first_run:
+ io.log_info("Caching GPU kernels...")
+
+ import tensorflow as tf
+ nn.tf = tf
+
+ if device_config.cpu_only:
+ config = tf.ConfigProto(device_count={'GPU': 0})
+ else:
+ config = tf.ConfigProto()
+ config.gpu_options.visible_device_list = ','.join([str(device.index) for device in device_config.devices])
+
+ config.gpu_options.force_gpu_compatible = True
+ config.gpu_options.allow_growth = True
+ nn.tf_sess_config = config
+
+ nn.tf_floatx = nn.tf.float32 #nn.tf.float16 if device_config.use_fp16 else nn.tf.float32
+ nn.np_floatx = nn.tf_floatx.as_numpy_dtype
+
+ from .tensor_ops import initialize_tensor_ops
+ from .layers import initialize_layers
+ from .initializers import initialize_initializers
+ from .optimizers import initialize_optimizers
+
+ initialize_tensor_ops(nn)
+ initialize_layers(nn)
+ initialize_initializers(nn)
+ initialize_optimizers(nn)
+
+ if nn.tf_sess is None:
+ nn.tf_sess = tf.Session(config=nn.tf_sess_config)
+
+ @staticmethod
+ def initialize_main_env():
+ Devices.initialize_main_env()
+
+ @staticmethod
+ def getCurrentDeviceConfig():
+ if nn.current_DeviceConfig is None:
+ nn.current_DeviceConfig = DeviceConfig.BestGPU()
+ return nn.current_DeviceConfig
+
+ @staticmethod
+ def setCurrentDeviceConfig(device_config):
+ nn.current_DeviceConfig = device_config
+
+ @staticmethod
+ def tf_reset_session():
+ if nn.tf is not None:
+ if nn.tf_sess is not None:
+ nn.tf.reset_default_graph()
+ nn.tf_sess.close()
+ nn.tf_sess = nn.tf.Session(config=nn.tf_sess_config)
+
+ @staticmethod
+ def tf_close_session():
+ if nn.tf_sess is not None:
+ nn.tf.reset_default_graph()
+ nn.tf_sess.close()
+ nn.tf_sess = None
+
+
+ @staticmethod
+ def ask_choose_device_idxs(choose_only_one=False, allow_cpu=True, suggest_best_multi_gpu=False, suggest_all_gpu=False, return_device_config=False):
+ devices = Devices.getDevices()
+ if len(devices) == 0:
+ return []
+
+ all_devices_indexes = [device.index for device in devices]
+
+ if choose_only_one:
+ suggest_best_multi_gpu = False
+ suggest_all_gpu = False
+
+ if suggest_all_gpu:
+ best_device_indexes = all_devices_indexes
+ elif suggest_best_multi_gpu:
+ best_device_indexes = [device.index for device in devices.get_equal_devices(devices.get_best_device()) ]
+ else:
+ best_device_indexes = [ devices.get_best_device().index ]
+ best_device_indexes = ",".join([str(x) for x in best_device_indexes])
+
+ io.log_info ("")
+ if choose_only_one:
+ io.log_info ("Choose one GPU idx.")
+ else:
+ io.log_info ("Choose one or several GPU idxs (separated by comma).")
+ io.log_info ("")
+
+ if allow_cpu:
+ io.log_info ("[CPU] : CPU")
+ for device in devices:
+ io.log_info (f" [{device.index}] : {device.name}")
+
+ io.log_info ("")
+
+ while True:
+ try:
+ if choose_only_one:
+ choosed_idxs = io.input_str("Which GPU index to choose?", best_device_indexes)
+ else:
+ choosed_idxs = io.input_str("Which GPU indexes to choose?", best_device_indexes)
+
+ if allow_cpu and choosed_idxs.lower() == "cpu":
+ choosed_idxs = []
+ break
+
+ choosed_idxs = [ int(x) for x in choosed_idxs.split(',') ]
+
+ if choose_only_one:
+ if len(choosed_idxs) == 1:
+ break
+ else:
+ if all( [idx in all_devices_indexes for idx in choosed_idxs] ):
+ break
+ except:
+ pass
+ io.log_info ("")
+
+ if return_device_config:
+ return nn.DeviceConfig.GPUIndexes(choosed_idxs)
+ else:
+ return choosed_idxs
+
+ class DeviceConfig():
+ def __init__ (self, devices=None):
+ devices = devices or []
+
+ if not isinstance(devices, Devices):
+ devices = Devices(devices)
+
+ self.devices = devices
+ self.cpu_only = len(devices) == 0
+
+ @staticmethod
+ def BestGPU():
+ devices = Devices.getDevices()
+ if len(devices) == 0:
+ return nn.DeviceConfig.CPU()
+
+ return nn.DeviceConfig([devices.get_best_device()])
+
+ @staticmethod
+ def WorstGPU():
+ devices = Devices.getDevices()
+ if len(devices) == 0:
+ return nn.DeviceConfig.CPU()
+
+ return nn.DeviceConfig([devices.get_worst_device()])
+
+ @staticmethod
+ def GPUIndexes(indexes):
+ if len(indexes) != 0:
+ devices = Devices.getDevices().get_devices_from_index_list(indexes)
+ else:
+ devices = []
+
+ return nn.DeviceConfig(devices)
+
+ @staticmethod
+ def CPU():
+ return nn.DeviceConfig([])
diff --git a/core/leras/optimizers.py b/core/leras/optimizers.py
new file mode 100644
index 0000000..0a6d477
--- /dev/null
+++ b/core/leras/optimizers.py
@@ -0,0 +1,108 @@
+def initialize_optimizers(nn):
+ tf = nn.tf
+ from tensorflow.python.ops import state_ops, control_flow_ops
+
+ class TFBaseOptimizer(nn.Saveable):
+ def __init__(self, name=None):
+ super().__init__(name=name)
+
+ def tf_clip_norm(self, g, c, n):
+ """Clip the gradient `g` if the L2 norm `n` exceeds `c`.
+ # Arguments
+ g: Tensor, the gradient tensor
+ c: float >= 0. Gradients will be clipped
+ when their L2 norm exceeds this value.
+ n: Tensor, actual norm of `g`.
+ # Returns
+ Tensor, the gradient clipped if required.
+ """
+ if c <= 0: # if clipnorm == 0 no need to add ops to the graph
+ return g
+
+ condition = n >= c
+ then_expression = tf.scalar_mul(c / n, g)
+ else_expression = g
+
+ # saving the shape to avoid converting sparse tensor to dense
+ if isinstance(then_expression, tf.Tensor):
+ g_shape = copy.copy(then_expression.get_shape())
+ elif isinstance(then_expression, tf.IndexedSlices):
+ g_shape = copy.copy(then_expression.dense_shape)
+ if condition.dtype != tf.bool:
+ condition = tf.cast(condition, 'bool')
+ g = tf.cond(condition,
+ lambda: then_expression,
+ lambda: else_expression)
+ if isinstance(then_expression, tf.Tensor):
+ g.set_shape(g_shape)
+ elif isinstance(then_expression, tf.IndexedSlices):
+ g._dense_shape = g_shape
+
+ return g
+ nn.TFBaseOptimizer = TFBaseOptimizer
+
+ class TFRMSpropOptimizer(TFBaseOptimizer):
+ def __init__(self, lr=0.001, rho=0.9, lr_dropout=1.0, epsilon=1e-7, clipnorm=0.0, name=None):
+ super().__init__(name=name)
+
+ if name is None:
+ raise ValueError('name must be defined.')
+
+ self.lr_dropout = lr_dropout
+ self.clipnorm = clipnorm
+
+ with tf.device('/CPU:0') :
+ with tf.variable_scope(self.name):
+ self.lr = tf.Variable (lr, name="lr")
+ self.rho = tf.Variable (rho, name="rho")
+ self.epsilon = tf.Variable (epsilon, name="epsilon")
+ self.iterations = tf.Variable(0, dtype=tf.int64, name='iters')
+
+ self.accumulators = []
+ self.accumulator_counter = 0
+ self.accumulators_dict = {}
+ self.lr_rnds_dict = {}
+
+ def get_weights(self):
+ return [self.lr, self.rho, self.epsilon, self.iterations] + self.accumulators
+
+ def initialize_variables(self, trainable_weights, vars_on_cpu=True):
+ # Initialize here all trainable variables used in training
+ e = tf.device('/CPU:0') if vars_on_cpu else None
+ if e: e.__enter__()
+ with tf.variable_scope(self.name):
+ accumulators = [ tf.get_variable ( f'acc_{i+self.accumulator_counter}', v.shape, initializer=tf.initializers.constant(0.0), trainable=False)
+ for (i, v ) in enumerate(trainable_weights) ]
+
+ self.accumulators_dict.update ( { v.name : acc for v,acc in zip(trainable_weights,accumulators) } )
+ self.accumulators += accumulators
+ self.accumulator_counter += len(trainable_weights)
+
+ if self.lr_dropout != 1.0:
+ lr_rnds = [ nn.tf_random_binomial( v.shape, p=self.lr_dropout) for v in trainable_weights ]
+ self.lr_rnds_dict.update ( { v.name : rnd for v,rnd in zip(trainable_weights,lr_rnds) } )
+ if e: e.__exit__(None, None, None)
+
+ def get_update_op(self, grads_vars):
+ updates = []
+ lr = self.lr
+ if self.clipnorm > 0.0:
+ norm = tf.sqrt( sum([tf.reduce_sum(tf.square(g)) for g,v in grads_vars]))
+ updates += [ state_ops.assign_add( self.iterations, 1) ]
+ for i, (g,v) in enumerate(grads_vars):
+ if self.clipnorm > 0.0:
+ g = self.tf_clip_norm(g, self.clipnorm, norm)
+
+ a = self.accumulators_dict[v.name]
+ new_a = self.rho * a + (1. - self.rho) * tf.square(g)
+ v_diff = - lr * g / (tf.sqrt(new_a) + self.epsilon)
+ if self.lr_dropout != 1.0:
+ lr_rnd = self.lr_rnds_dict[v.name]
+ v_diff *= lr_rnd
+ new_v = v + v_diff
+
+ updates.append (state_ops.assign(a, new_a))
+ updates.append (state_ops.assign(v, new_v))
+
+ return control_flow_ops.group ( *updates, name=self.name+'_updates')
+ nn.TFRMSpropOptimizer = TFRMSpropOptimizer
\ No newline at end of file
diff --git a/core/leras/tensor_ops.py b/core/leras/tensor_ops.py
new file mode 100644
index 0000000..52aa5ab
--- /dev/null
+++ b/core/leras/tensor_ops.py
@@ -0,0 +1,295 @@
+import numpy as np
+
+def initialize_tensor_ops(nn):
+ tf = nn.tf
+ from tensorflow.python.ops import array_ops, random_ops, math_ops, sparse_ops, gradients
+ from tensorflow.python.framework import sparse_tensor
+
+ def tf_get_value(tensor):
+ return nn.tf_sess.run (tensor)
+ nn.tf_get_value = tf_get_value
+
+
+ def tf_batch_set_value(tuples):
+ if len(tuples) != 0:
+ with nn.tf.device('/CPU:0'):
+ assign_ops = []
+ feed_dict = {}
+
+ for x, value in tuples:
+ if isinstance(value, nn.tf.Operation):
+ assign_ops.append(value)
+ else:
+ value = np.asarray(value, dtype=x.dtype.as_numpy_dtype)
+ assign_placeholder = nn.tf.placeholder( x.dtype.base_dtype, shape=[None]*value.ndim )
+ assign_op = nn.tf.assign (x, assign_placeholder )
+ assign_ops.append(assign_op)
+ feed_dict[assign_placeholder] = value
+
+ nn.tf_sess.run(assign_ops, feed_dict=feed_dict)
+ nn.tf_batch_set_value = tf_batch_set_value
+
+
+ def tf_gradients ( loss, vars ):
+ grads = gradients.gradients(loss, vars, colocate_gradients_with_ops=True )
+ #todo none gradient for var
+ return [*zip(grads,vars)]
+ nn.tf_gradients = tf_gradients
+
+ def tf_average_gv_list(grad_var_list, tf_device_string=None):
+ e = tf.device(tf_device_string) if tf_device_string is not None else None
+ if e is not None: e.__enter__()
+ result = []
+ for i, (gv) in enumerate(grad_var_list):
+ for j,(g,v) in enumerate(gv):
+ g = tf.expand_dims(g, 0)
+ if i == 0:
+ result += [ [[g], v] ]
+ else:
+ result[j][0] += [g]
+
+ for i,(gs,v) in enumerate(result):
+ result[i] = ( tf.reduce_mean( tf.concat (gs, 0), 0 ), v )
+ if e is not None: e.__exit__(None,None,None)
+ return result
+ nn.tf_average_gv_list = tf_average_gv_list
+
+ def tf_average_tensor_list(tensors_list, tf_device_string=None):
+ e = tf.device(tf_device_string) if tf_device_string is not None else None
+ if e is not None: e.__enter__()
+ result = tf.reduce_mean(tf.concat ([tf.expand_dims(t, 0) for t in tensors_list], 0), 0)
+ if e is not None: e.__exit__(None,None,None)
+ return result
+ nn.tf_average_tensor_list = tf_average_tensor_list
+
+ def tf_dot(x, y):
+ if x.shape.ndims > 2 or y.shape.ndims > 2:
+ x_shape = []
+ for i, s in zip( x.shape.as_list(), array_ops.unstack(array_ops.shape(x))):
+ if i is not None:
+ x_shape.append(i)
+ else:
+ x_shape.append(s)
+ x_shape = tuple(x_shape)
+ y_shape = []
+ for i, s in zip( y.shape.as_list(), array_ops.unstack(array_ops.shape(y))):
+ if i is not None:
+ y_shape.append(i)
+ else:
+ y_shape.append(s)
+ y_shape = tuple(y_shape)
+ y_permute_dim = list(range(y.shape.ndims))
+ y_permute_dim = [y_permute_dim.pop(-2)] + y_permute_dim
+ xt = array_ops.reshape(x, [-1, x_shape[-1]])
+ yt = array_ops.reshape(array_ops.transpose(y, perm=y_permute_dim), [y_shape[-2], -1])
+
+ import code
+ code.interact(local=dict(globals(), **locals()))
+ return array_ops.reshape(math_ops.matmul(xt, yt), x_shape[:-1] + y_shape[:-2] + y_shape[-1:])
+ if isinstance(x, sparse_tensor.SparseTensor):
+ out = sparse_ops.sparse_tensor_dense_matmul(x, y)
+ else:
+ out = math_ops.matmul(x, y)
+ return out
+ nn.tf_dot = tf_dot
+
+ def tf_gelu(x):
+ cdf = 0.5 * (1.0 + tf.nn.tanh((np.sqrt(2 / np.pi) * (x + 0.044715 * tf.pow(x, 3)))))
+ return x * cdf
+ nn.tf_gelu = tf_gelu
+
+ def tf_upsample2d(x, size=2):
+ return tf.image.resize_nearest_neighbor(x, (x.shape[1]*size, x.shape[2]*size) )
+ nn.tf_upsample2d = tf_upsample2d
+
+ def tf_upsample2d_bilinear(x, size=2):
+ return tf.image.resize_images(x, (x.shape[1]*size, x.shape[2]*size) )
+ nn.tf_upsample2d_bilinear = tf_upsample2d_bilinear
+
+ def tf_flatten(x, dynamic_dims=False):
+ """
+ dynamic_dims allows to flatten without knowing size on input dims
+ """
+ if dynamic_dims:
+ sh = tf.shape(x)
+ return tf.reshape (x, (sh[0], tf.reduce_prod(sh[1:]) ) )
+ else:
+ return tf.reshape (x, (-1, np.prod(x.shape[1:])) )
+
+ nn.tf_flatten = tf_flatten
+
+ def tf_random_binomial(shape, p=0.0, dtype=None, seed=None):
+ if dtype is None:
+ dtype=tf.float32
+
+ if seed is None:
+ seed = np.random.randint(10e6)
+ return array_ops.where(
+ random_ops.random_uniform(shape, dtype=tf.float16, seed=seed) < p,
+ array_ops.ones(shape, dtype=dtype), array_ops.zeros(shape, dtype=dtype))
+ nn.tf_random_binomial = tf_random_binomial
+
+ def tf_gaussian_blur(input, radius=2.0):
+ def gaussian(x, mu, sigma):
+ return np.exp(-(float(x) - float(mu)) ** 2 / (2 * sigma ** 2))
+
+ def make_kernel(sigma):
+ kernel_size = max(3, int(2 * 2 * sigma + 1))
+ mean = np.floor(0.5 * kernel_size)
+ kernel_1d = np.array([gaussian(x, mean, sigma) for x in range(kernel_size)])
+ np_kernel = np.outer(kernel_1d, kernel_1d).astype(np.float32)
+ kernel = np_kernel / np.sum(np_kernel)
+ return kernel
+
+ gauss_kernel = make_kernel(radius)
+ gauss_kernel = gauss_kernel[:, :,np.newaxis, np.newaxis]
+ kernel_size = gauss_kernel.shape[0]
+
+ inputs = [ input[:,:,:,i:i+1] for i in range( input.shape[-1] ) ]
+
+ outputs = []
+ for i in range(len(inputs)):
+ x = inputs[i]
+ if kernel_size != 0:
+ padding = kernel_size//2
+ x = tf.pad (x, [ [0,0], [padding,padding], [padding,padding], [0,0] ] )
+
+ outputs += [ tf.nn.conv2d(x, tf.constant(gauss_kernel, dtype=nn.tf_floatx ) , strides=[1,1,1,1], padding="VALID") ]
+
+ return tf.concat (outputs, axis=-1)
+ nn.tf_gaussian_blur = tf_gaussian_blur
+
+ def tf_style_loss(target, style, gaussian_blur_radius=0.0, loss_weight=1.0, step_size=1):
+ def sd(content, style, loss_weight):
+ content_nc = content.shape[-1]
+ style_nc = style.shape[-1]
+ if content_nc != style_nc:
+ raise Exception("style_loss() content_nc != style_nc")
+
+ axes = [1,2]
+ c_mean, c_var = tf.nn.moments(content, axes=axes, keep_dims=True)
+ s_mean, s_var = tf.nn.moments(style, axes=axes, keep_dims=True)
+ c_std, s_std = tf.sqrt(c_var + 1e-5), tf.sqrt(s_var + 1e-5)
+
+ mean_loss = tf.reduce_sum(tf.square(c_mean-s_mean), axis=[1,2,3])
+ std_loss = tf.reduce_sum(tf.square(c_std-s_std), axis=[1,2,3])
+
+ return (mean_loss + std_loss) * ( loss_weight / content_nc.value )
+
+ if gaussian_blur_radius > 0.0:
+ target = tf_gaussian_blur(target, gaussian_blur_radius)
+ style = tf_gaussian_blur(style, gaussian_blur_radius)
+
+ return sd( target, style, loss_weight=loss_weight )
+
+ nn.tf_style_loss = tf_style_loss
+
+ def tf_dssim(img1,img2, max_val, filter_size=11, filter_sigma=1.5, k1=0.01, k2=0.03):
+
+ ch = img2.shape[-1]
+
+ def _fspecial_gauss(size, sigma):
+ #Function to mimic the 'fspecial' gaussian MATLAB function.
+ coords = np.arange(0, size, dtype=nn.np_floatx)
+ coords -= (size - 1 ) / 2.0
+ g = coords**2
+ g *= ( -0.5 / (sigma**2) )
+ g = np.reshape (g, (1,-1)) + np.reshape(g, (-1,1) )
+ g = tf.constant ( np.reshape (g, (1,-1)), dtype=nn.tf_floatx )
+ g = tf.nn.softmax(g)
+ g = tf.reshape (g, (size, size, 1, 1))
+ g = tf.tile (g, (1,1,ch,1))
+ return g
+
+ kernel = _fspecial_gauss(filter_size,filter_sigma)
+
+ def reducer(x):
+ return tf.nn.depthwise_conv2d(x, kernel, strides=[1,1,1,1], padding='VALID')
+
+ c1 = (k1 * max_val) ** 2
+ c2 = (k2 * max_val) ** 2
+
+ mean0 = reducer(img1)
+ mean1 = reducer(img2)
+ num0 = mean0 * mean1 * 2.0
+ den0 = tf.square(mean0) + tf.square(mean1)
+ luminance = (num0 + c1) / (den0 + c1)
+
+ num1 = reducer(img1 * img2) * 2.0
+ den1 = reducer(tf.square(img1) + tf.square(img2))
+ c2 *= 1.0 #compensation factor
+ cs = (num1 - num0 + c2) / (den1 - den0 + c2)
+
+ ssim_val = tf.reduce_mean(luminance * cs, axis=(-3, -2) )
+ return(1.0 - ssim_val ) / 2.0
+ nn.tf_dssim = tf_dssim
+
+ def tf_rgb_to_lab(srgb):
+ srgb_pixels = tf.reshape(srgb, [-1, 3])
+ linear_mask = tf.cast(srgb_pixels <= 0.04045, dtype=tf.float32)
+ exponential_mask = tf.cast(srgb_pixels > 0.04045, dtype=tf.float32)
+ rgb_pixels = (srgb_pixels / 12.92 * linear_mask) + (((srgb_pixels + 0.055) / 1.055) ** 2.4) * exponential_mask
+ rgb_to_xyz = tf.constant([
+ # X Y Z
+ [0.412453, 0.212671, 0.019334], # R
+ [0.357580, 0.715160, 0.119193], # G
+ [0.180423, 0.072169, 0.950227], # B
+ ])
+ xyz_pixels = tf.matmul(rgb_pixels, rgb_to_xyz)
+
+ xyz_normalized_pixels = tf.multiply(xyz_pixels, [1/0.950456, 1.0, 1/1.088754])
+
+ epsilon = 6/29
+ linear_mask = tf.cast(xyz_normalized_pixels <= (epsilon**3), dtype=tf.float32)
+ exponential_mask = tf.cast(xyz_normalized_pixels > (epsilon**3), dtype=tf.float32)
+ fxfyfz_pixels = (xyz_normalized_pixels / (3 * epsilon**2) + 4/29) * linear_mask + (xyz_normalized_pixels ** (1/3)) * exponential_mask
+
+ fxfyfz_to_lab = tf.constant([
+ # l a b
+ [ 0.0, 500.0, 0.0], # fx
+ [116.0, -500.0, 200.0], # fy
+ [ 0.0, 0.0, -200.0], # fz
+ ])
+ lab_pixels = tf.matmul(fxfyfz_pixels, fxfyfz_to_lab) + tf.constant([-16.0, 0.0, 0.0])
+ return tf.reshape(lab_pixels, tf.shape(srgb))
+ nn.tf_rgb_to_lab = tf_rgb_to_lab
+
+ def tf_suppress_lower_mean(t, eps=0.00001):
+ if t.shape.ndims != 1:
+ raise ValueError("tf_suppress_lower_mean: t rank must be 1")
+ t_mean_eps = tf.reduce_mean(t) - eps
+ q = tf.clip_by_value(t, t_mean_eps, tf.reduce_max(t) )
+ q = tf.clip_by_value(q-t_mean_eps, 0, eps)
+ q = q * (t/eps)
+ return q
+"""
+class GeLU(KL.Layer):
+ Gaussian Error Linear Unit.
+ A smoother version of ReLU generally used
+ in the BERT or BERT architecture based models.
+ Original paper: https://arxiv.org/abs/1606.08415
+ Input shape:
+ Arbitrary. Use the keyword argument `input_shape`
+ (tuple of integers, does not include the samples axis)
+ when using this layer as the first layer in a model.
+ Output shape:
+ Same shape as the input.
+
+ def __init__(self, approximate=True, **kwargs):
+ super(GeLU, self).__init__(**kwargs)
+ self.approximate = approximate
+ self.supports_masking = True
+
+ def call(self, inputs):
+ cdf = 0.5 * (1.0 + K.tanh((np.sqrt(2 / np.pi) * (inputs + 0.044715 * K.pow(inputs, 3)))))
+ return inputs * cdf
+
+ def get_config(self):
+ config = {'approximate': self.approximate}
+ base_config = super(GeLU, self).get_config()
+ return dict(list(base_config.items()) + list(config.items()))
+
+ def compute_output_shape(self, input_shape):
+ return input_shape
+ nn.GeLU = GeLU
+"""
\ No newline at end of file
diff --git a/mathlib/__init__.py b/core/mathlib/__init__.py
similarity index 100%
rename from mathlib/__init__.py
rename to core/mathlib/__init__.py
diff --git a/mathlib/umeyama.py b/core/mathlib/umeyama.py
similarity index 100%
rename from mathlib/umeyama.py
rename to core/mathlib/umeyama.py
diff --git a/utils/mp_utils.py b/core/mplib/__init__.py
similarity index 100%
rename from utils/mp_utils.py
rename to core/mplib/__init__.py
diff --git a/utils/os_utils.py b/core/osex.py
similarity index 100%
rename from utils/os_utils.py
rename to core/osex.py
diff --git a/utils/Path_utils.py b/core/pathex.py
similarity index 88%
rename from utils/Path_utils.py
rename to core/pathex.py
index c609572..5c93eed 100644
--- a/utils/Path_utils.py
+++ b/core/pathex.py
@@ -3,6 +3,16 @@ from os import scandir
image_extensions = [".jpg", ".jpeg", ".png", ".tif", ".tiff"]
+def write_bytes_safe(p, bytes_data):
+ """
+ writes to .tmp first and then rename to target filename
+ """
+ p_tmp = p.parent / (p.name + '.tmp')
+ p_tmp.write_bytes(bytes_data)
+ if p.exists():
+ p.unlink()
+ p_tmp.rename (p)
+
def scantree(path):
"""Recursively yield DirEntry objects for given directory."""
for entry in scandir(path):
@@ -46,7 +56,7 @@ def get_file_paths(dir_path):
dir_path = Path (dir_path)
if dir_path.exists():
- return sorted([ x.path for x in list(scandir(str(dir_path))) if x.is_file() ])
+ return [ Path(x) for x in sorted([ x.path for x in list(scandir(str(dir_path))) if x.is_file() ]) ]
else:
return []
diff --git a/utils/random_utils.py b/core/randomex.py
similarity index 100%
rename from utils/random_utils.py
rename to core/randomex.py
diff --git a/utils/std_utils.py b/core/stdex.py
similarity index 100%
rename from utils/std_utils.py
rename to core/stdex.py
diff --git a/utils/struct_utils.py b/core/structex.py
similarity index 100%
rename from utils/struct_utils.py
rename to core/structex.py
diff --git a/doc/DeepFaceLab is working.png b/doc/DeepFaceLab is working.png
new file mode 100644
index 0000000..4d86d36
Binary files /dev/null and b/doc/DeepFaceLab is working.png differ
diff --git a/doc/doc_build_and_repository_info.md b/doc/doc_build_and_repository_info.md
deleted file mode 100644
index 89a740e..0000000
--- a/doc/doc_build_and_repository_info.md
+++ /dev/null
@@ -1,5 +0,0 @@
-#### **CPU mode**
-
-It is possible to run from script for all stages using the `--cpu-only` flag. To run from script, install the separate dependencies for CPU mode using `pip -r requirements-cpu.txt`.
-
-Please note that extraction and training will take much long without a GPU and performance will greatly suffer without one. In particular, do not use DLIB extractor in CPU mode, it's too slow to run without a GPU. Train only on 64px resolution models like H64 or SAE (with low settings) and the lightweight encoder.
\ No newline at end of file
diff --git a/doc/doc_ready_to_work_facesets.md b/doc/doc_ready_to_work_facesets.md
deleted file mode 100644
index 75f3dee..0000000
--- a/doc/doc_ready_to_work_facesets.md
+++ /dev/null
@@ -1,11 +0,0 @@
-### **Example Face Sets**:
-
-Faces sets for the following have been pre-extracted,
-
-- Nicolas Cage
-- Steve Jobs
-- Putin
-- Elon Musk
-- Harrison Ford
-
-[Download from Google drive](https://drive.google.com/open?id=1LwMdfTxdOaNAHt_sGV76aQVn7XPseXJB)
diff --git a/doc/doc_windows_desktop_app.md b/doc/doc_windows_desktop_app.md
deleted file mode 100644
index e6c83ec..0000000
--- a/doc/doc_windows_desktop_app.md
+++ /dev/null
@@ -1,27 +0,0 @@
-### **Prebuilt Windows Releases**
-
-Windows builds with all dependencies included are released regularly. Only the NVIDIA GeForce display driver needs to be installed. Prebuilt DeepFaceLab, including GPU and CPU versions, can be downloaded from
-
-[Google drive](https://drive.google.com/open?id=1BCFK_L7lPNwMbEQ_kFPqPpDdFEOd_Dci)
-
-if the download qouta is exceeded, add the file to your own google drive and download from it
-
-[Torrent](https://rutracker.org/forum/viewtopic.php?t=5558863)
-
-Available builds:
-
-* DeepFaceLab_CUDA - for NVIDIA cards
-
-* DeepFaceLab_OpenCL - for NVIDIA/AMD/IntelHD cards
-
-Important: you don't need to install CUDA !
-
-#### Video tutorials using prebuilt windows app
-
-* [Basic workflow](https://www.youtube.com/watch?v=K98nTNjXkq8)
-
-* [Basic workflow (thanks @derpfakes)](https://www.youtube.com/watch?v=cVcyghhmQSA)
-
-* [How To Make DeepFakes With DeepFaceLab - An Amatuer's Guide](https://www.youtube.com/watch?v=wBax7_UWXvc)
-
-* [Manual re-extract poorly aligned frames](https://www.youtube.com/watch?v=7z1ykVVCHhM)
\ No newline at end of file
diff --git a/doc/example_faceset.jpg b/doc/example_faceset.jpg
deleted file mode 100644
index 6ee914f..0000000
Binary files a/doc/example_faceset.jpg and /dev/null differ
diff --git a/doc/gallery/1.jpg b/doc/gallery/1.jpg
deleted file mode 100644
index bb4895b..0000000
Binary files a/doc/gallery/1.jpg and /dev/null differ
diff --git a/doc/gallery/2.jpg b/doc/gallery/2.jpg
deleted file mode 100644
index b8e6138..0000000
Binary files a/doc/gallery/2.jpg and /dev/null differ
diff --git a/doc/gallery/doc_gallery.md b/doc/gallery/doc_gallery.md
deleted file mode 100644
index 5ba780f..0000000
--- a/doc/gallery/doc_gallery.md
+++ /dev/null
@@ -1,3 +0,0 @@
-
-
-
\ No newline at end of file
diff --git a/doc/logo_cuda.jpg b/doc/logo_cuda.jpg
deleted file mode 100644
index 472571e..0000000
Binary files a/doc/logo_cuda.jpg and /dev/null differ
diff --git a/doc/logo_cuda.png b/doc/logo_cuda.png
new file mode 100644
index 0000000..0b928a6
Binary files /dev/null and b/doc/logo_cuda.png differ
diff --git a/doc/logo_keras.jpg b/doc/logo_keras.jpg
deleted file mode 100644
index 6082a5b..0000000
Binary files a/doc/logo_keras.jpg and /dev/null differ
diff --git a/doc/logo_opencl.jpg b/doc/logo_opencl.jpg
deleted file mode 100644
index 36ab395..0000000
Binary files a/doc/logo_opencl.jpg and /dev/null differ
diff --git a/doc/logo_plaidml.jpg b/doc/logo_plaidml.jpg
deleted file mode 100644
index c206915..0000000
Binary files a/doc/logo_plaidml.jpg and /dev/null differ
diff --git a/doc/logo_tensorflow.jpg b/doc/logo_tensorflow.jpg
deleted file mode 100644
index 19f44c9..0000000
Binary files a/doc/logo_tensorflow.jpg and /dev/null differ
diff --git a/doc/logo_tensorflow.png b/doc/logo_tensorflow.png
new file mode 100644
index 0000000..b06cdd6
Binary files /dev/null and b/doc/logo_tensorflow.png differ
diff --git a/doc/manual_en_google_translated.docx b/doc/manual_en_google_translated.docx
deleted file mode 100644
index ac64a7d..0000000
Binary files a/doc/manual_en_google_translated.docx and /dev/null differ
diff --git a/doc/manual_en_google_translated.pdf b/doc/manual_en_google_translated.pdf
deleted file mode 100644
index d6fade9..0000000
Binary files a/doc/manual_en_google_translated.pdf and /dev/null differ
diff --git a/doc/manual_extractor_0.jpg b/doc/manual_extractor_0.jpg
deleted file mode 100644
index b88c02c..0000000
Binary files a/doc/manual_extractor_0.jpg and /dev/null differ
diff --git a/doc/manual_ru.pdf b/doc/manual_ru.pdf
deleted file mode 100644
index 2c27e9f..0000000
Binary files a/doc/manual_ru.pdf and /dev/null differ
diff --git a/doc/manual_ru_source.docx b/doc/manual_ru_source.docx
deleted file mode 100644
index f37a425..0000000
Binary files a/doc/manual_ru_source.docx and /dev/null differ
diff --git a/ebsynth/__init__.py b/ebsynth/__init__.py
deleted file mode 100644
index ce31c32..0000000
--- a/ebsynth/__init__.py
+++ /dev/null
@@ -1 +0,0 @@
-from .ebsynth import color_transfer
\ No newline at end of file
diff --git a/ebsynth/ebsynth.dll b/ebsynth/ebsynth.dll
deleted file mode 100644
index f6c085e..0000000
Binary files a/ebsynth/ebsynth.dll and /dev/null differ
diff --git a/ebsynth/ebsynth.py b/ebsynth/ebsynth.py
deleted file mode 100644
index ec2ec08..0000000
--- a/ebsynth/ebsynth.py
+++ /dev/null
@@ -1,201 +0,0 @@
-import os
-import sys
-from ctypes import *
-from pathlib import Path
-
-import cv2
-import numpy as np
-
-libebsynth = None
-cached_buffer = {}
-
-EBSYNTH_BACKEND_CPU = 0x0001
-EBSYNTH_BACKEND_CUDA = 0x0002
-EBSYNTH_BACKEND_AUTO = 0x0000
-EBSYNTH_MAX_STYLE_CHANNELS = 8
-EBSYNTH_MAX_GUIDE_CHANNELS = 24
-EBSYNTH_VOTEMODE_PLAIN = 0x0001 # weight = 1
-EBSYNTH_VOTEMODE_WEIGHTED = 0x0002 # weight = 1/(1+error)
-
-
-def _normalize_img_shape (img):
- img_len = len(img.shape)
- if img_len == 2:
- sh, sw = img.shape
- sc = 0
- elif img_len == 3:
- sh, sw, sc = img.shape
-
- if sc == 0:
- sc = 1
- img = img [...,np.newaxis]
- return img
-
-def run (img_style, guides,
- patch_size=5,
- num_pyramid_levels=-1,
- num_search_vote_iters = 6,
- num_patch_match_iters = 4,
- stop_threshold = 5,
- uniformity_weight = 3500.0,
- extraPass3x3 = False,
- ):
- if patch_size < 3:
- raise ValueError ("patch_size is too small")
- if patch_size % 2 == 0:
- raise ValueError ("patch_size must be an odd number")
- if len(guides) == 0:
- raise ValueError ("at least one guide must be specified")
-
- global libebsynth
- if libebsynth is None:
- if sys.platform[0:3] == 'win':
- libebsynth_path = str ( Path(__file__).parent / 'ebsynth.dll' )
- libebsynth = CDLL(libebsynth_path)
- else:
- #todo: implement for linux
- pass
-
- if libebsynth is not None:
- libebsynth.ebsynthRun.argtypes = ( \
- c_int,
- c_int,
- c_int,
- c_int,
- c_int,
- c_void_p,
- c_void_p,
- c_int,
- c_int,
- c_void_p,
- c_void_p,
- POINTER(c_float),
- POINTER(c_float),
- c_float,
- c_int,
- c_int,
- c_int,
- POINTER(c_int),
- POINTER(c_int),
- POINTER(c_int),
- c_int,
- c_void_p,
- c_void_p
- )
-
- if libebsynth is None:
- return img_style
-
- img_style = _normalize_img_shape (img_style)
- sh, sw, sc = img_style.shape
- t_h, t_w, t_c = 0,0,0
-
- if sc > EBSYNTH_MAX_STYLE_CHANNELS:
- raise ValueError (f"error: too many style channels {sc}, maximum number is {EBSYNTH_MAX_STYLE_CHANNELS}")
-
- guides_source = []
- guides_target = []
- guides_weights = []
-
- for i in range(len(guides)):
- source_guide, target_guide, guide_weight = guides[i]
- source_guide = _normalize_img_shape(source_guide)
- target_guide = _normalize_img_shape(target_guide)
- s_h, s_w, s_c = source_guide.shape
- nt_h, nt_w, nt_c = target_guide.shape
-
- if s_h != sh or s_w != sw:
- raise ValueError ("guide source and style resolution must match style resolution.")
-
- if t_c == 0:
- t_h, t_w, t_c = nt_h, nt_w, nt_c
- elif nt_h != t_h or nt_w != t_w:
- raise ValueError ("guides target resolutions must be equal")
-
- if s_c != nt_c:
- raise ValueError ("guide source and target channels must match exactly.")
-
- guides_source.append (source_guide)
- guides_target.append (target_guide)
-
- guides_weights += [ guide_weight / s_c ] * s_c
-
- guides_source = np.concatenate ( guides_source, axis=-1)
- guides_target = np.concatenate ( guides_target, axis=-1)
- guides_weights = (c_float*len(guides_weights) ) ( *guides_weights )
-
- styleWeight = 1.0
- style_weights = [ styleWeight / sc for i in range(sc) ]
- style_weights = (c_float*sc) ( *style_weights )
-
-
- maxPyramidLevels = 0
- for level in range(32,-1,-1):
- if min( min(sh, t_h)*pow(2.0, -level), \
- min(sw, t_w)*pow(2.0, -level) ) >= (2*patch_size+1):
- maxPyramidLevels = level+1
- break
-
- if num_pyramid_levels == -1:
- num_pyramid_levels = maxPyramidLevels
- num_pyramid_levels = min(num_pyramid_levels, maxPyramidLevels)
-
- num_search_vote_iters_per_level = (c_int*num_pyramid_levels) ( *[num_search_vote_iters]*num_pyramid_levels )
- num_patch_match_iters_per_level = (c_int*num_pyramid_levels) ( *[num_patch_match_iters]*num_pyramid_levels )
- stop_threshold_per_level = (c_int*num_pyramid_levels) ( *[stop_threshold]*num_pyramid_levels )
-
- buffer = cached_buffer.get ( (t_h,t_w,sc), None )
- if buffer is None:
- buffer = create_string_buffer (t_h*t_w*sc)
- cached_buffer[(t_h,t_w,sc)] = buffer
-
- libebsynth.ebsynthRun (EBSYNTH_BACKEND_CPU, #backend
- sc, #numStyleChannels
- guides_source.shape[-1], #numGuideChannels
- sw, #sourceWidth
- sh, #sourceHeight
- img_style.tobytes(), #sourceStyleData (width * height * numStyleChannels) bytes, scan-line order
- guides_source.tobytes(), #sourceGuideData (width * height * numGuideChannels) bytes, scan-line order
- t_w, #targetWidth
- t_h, #targetHeight
- guides_target.tobytes(), #targetGuideData (width * height * numGuideChannels) bytes, scan-line order
- None, #targetModulationData (width * height * numGuideChannels) bytes, scan-line order; pass NULL to switch off the modulation
- style_weights, #styleWeights (numStyleChannels) floats
- guides_weights, #guideWeights (numGuideChannels) floats
- uniformity_weight, #uniformityWeight reasonable values are between 500-15000, 3500 is a good default
- patch_size, #patchSize odd sizes only, use 5 for 5x5 patch, 7 for 7x7, etc.
- EBSYNTH_VOTEMODE_PLAIN, #voteMode use VOTEMODE_WEIGHTED for sharper result
- num_pyramid_levels, #numPyramidLevels
-
- num_search_vote_iters_per_level, #numSearchVoteItersPerLevel how many search/vote iters to perform at each level (array of ints, coarse first, fine last)
- num_patch_match_iters_per_level, #numPatchMatchItersPerLevel how many Patch-Match iters to perform at each level (array of ints, coarse first, fine last)
- stop_threshold_per_level, #stopThresholdPerLevel stop improving pixel when its change since last iteration falls under this threshold
- 1 if extraPass3x3 else 0, #extraPass3x3 perform additional polishing pass with 3x3 patches at the finest level, use 0 to disable
- None, #outputNnfData (width * height * 2) ints, scan-line order; pass NULL to ignore
- buffer #outputImageData (width * height * numStyleChannels) bytes, scan-line order
- )
-
- return np.frombuffer(buffer, dtype=np.uint8).reshape ( (t_h,t_w,sc) ).copy()
-
-#transfer color from source to target
-def color_transfer(img_source, img_target):
- guides = [( cv2.cvtColor(img_source, cv2.COLOR_BGR2GRAY),
- cv2.cvtColor(img_target, cv2.COLOR_BGR2GRAY),
- 1 ) ]
-
- h,w,c = img_source.shape
- result = []
- for i in range(c):
- result += [
- run( img_source[...,i:i+1] , guides=guides,
- patch_size=11,
- num_pyramid_levels=40,
- num_search_vote_iters = 6,
- num_patch_match_iters = 4,
- stop_threshold = 5,
- uniformity_weight=500.0,
- extraPass3x3=True,
- )
- ]
-
- return np.concatenate( result, axis=-1 )
diff --git a/facelib/DLIBExtractor.py b/facelib/DLIBExtractor.py
deleted file mode 100644
index a91164d..0000000
--- a/facelib/DLIBExtractor.py
+++ /dev/null
@@ -1,40 +0,0 @@
-import numpy as np
-import os
-import cv2
-
-from pathlib import Path
-
-class DLIBExtractor(object):
- def __init__(self, dlib):
- self.scale_to = 1850
- #3100 eats ~1.687GB VRAM on 2GB 730 desktop card, but >4Gb on 6GB card,
- #but 3100 doesnt work on 2GB 850M notebook card, I cant understand this behaviour
- #1850 works on 2GB 850M notebook card, works faster than 3100, produces good result
- self.dlib = dlib
-
- def __enter__(self):
- self.dlib_cnn_face_detector = self.dlib.cnn_face_detection_model_v1( str(Path(__file__).parent / "mmod_human_face_detector.dat") )
- self.dlib_cnn_face_detector ( np.zeros ( (self.scale_to, self.scale_to, 3), dtype=np.uint8), 0 )
- return self
-
- def __exit__(self, exc_type=None, exc_value=None, traceback=None):
- del self.dlib_cnn_face_detector
- return False #pass exception between __enter__ and __exit__ to outter level
-
- def extract_from_bgr (self, input_image):
- input_image = input_image[:,:,::-1].copy()
- (h, w, ch) = input_image.shape
-
- detected_faces = []
- input_scale = self.scale_to / (w if w > h else h)
- input_image = cv2.resize (input_image, ( int(w*input_scale), int(h*input_scale) ), interpolation=cv2.INTER_LINEAR)
- detected_faces = self.dlib_cnn_face_detector(input_image, 0)
-
- result = []
- for d_rect in detected_faces:
- if type(d_rect) == self.dlib.mmod_rectangle:
- d_rect = d_rect.rect
- left, top, right, bottom = d_rect.left(), d_rect.top(), d_rect.right(), d_rect.bottom()
- result.append ( (int(left/input_scale), int(top/input_scale), int(right/input_scale), int(bottom/input_scale)) )
-
- return result
diff --git a/facelib/2DFAN-4.h5 b/facelib/FAN.npy
similarity index 98%
rename from facelib/2DFAN-4.h5
rename to facelib/FAN.npy
index ac782ad..ba55f15 100644
Binary files a/facelib/2DFAN-4.h5 and b/facelib/FAN.npy differ
diff --git a/facelib/FANExtractor.py b/facelib/FANExtractor.py
index 9fd9b02..3429172 100644
--- a/facelib/FANExtractor.py
+++ b/facelib/FANExtractor.py
@@ -7,29 +7,159 @@ import numpy as np
from numpy import linalg as npla
from facelib import FaceType, LandmarksProcessor
-from nnlib import nnlib
+from core.leras import nn
"""
ported from https://github.com/1adrianb/face-alignment
"""
class FANExtractor(object):
- def __init__ (self):
- pass
+ def __init__ (self, place_model_on_cpu=False):
+ model_path = Path(__file__).parent / "FAN.npy"
+ if not model_path.exists():
+ raise Exception("Unable to load FANExtractor model")
- def __enter__(self):
- keras_model_path = Path(__file__).parent / "2DFAN-4.h5"
- if not keras_model_path.exists():
- return None
+ nn.initialize()
+ tf = nn.tf
- exec( nnlib.import_all(), locals(), globals() )
- self.model = FANExtractor.BuildModel()
- self.model.load_weights(str(keras_model_path))
+ class ConvBlock(nn.ModelBase):
+ def on_build(self, in_planes, out_planes):
+ self.in_planes = in_planes
+ self.out_planes = out_planes
- return self
+ 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 )
- def __exit__(self, exc_type=None, exc_value=None, traceback=None):
- del self.model
- return False #pass exception between __enter__ and __exit__ to outter level
+ 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 = 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.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.down_bn1 = None
+ self.down_conv1 = None
+
+ def forward(self, input):
+ x = input
+ x = self.bn1(x)
+ x = tf.nn.relu(x)
+ x = out1 = self.conv1(x)
+
+ x = self.bn2(x)
+ x = tf.nn.relu(x)
+ x = out2 = self.conv2(x)
+
+ x = self.bn3(x)
+ x = tf.nn.relu(x)
+ x = out3 = self.conv3(x)
+ x = tf.concat ([out1, out2, out3], axis=-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:
+ x = x + input
+
+ return x
+
+ class HourGlass (nn.ModelBase):
+ def on_build(self, in_planes, depth):
+ self.b1 = ConvBlock (in_planes, 256)
+ self.b2 = ConvBlock (in_planes, 256)
+
+ if depth > 1:
+ self.b2_plus = HourGlass(256, depth-1)
+ else:
+ self.b2_plus = ConvBlock(256, 256)
+
+ self.b3 = ConvBlock(256, 256)
+
+ def forward(self, input):
+ up1 = self.b1(input)
+
+ low1 = tf.nn.avg_pool(input, [1,2,2,1], [1,2,2,1], 'VALID')
+ low1 = self.b2 (low1)
+
+ low2 = self.b2_plus(low1)
+ low3 = self.b3(low2)
+
+ up2 = nn.tf_upsample2d(low3)
+
+ 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)
+
+ self.conv2 = ConvBlock(64, 128)
+ self.conv3 = ConvBlock(128, 128)
+ self.conv4 = ConvBlock(128, 256)
+
+ 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) ]
+
+ 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') ]
+
+ 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)
+
+ outputs = []
+ 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)
+ if i < 4 - 1:
+ ll = self.bl[i](ll)
+ previous = previous + ll + self.al[i](tmp_out)
+ return outputs[-1]
+
+ e = None
+ if place_model_on_cpu:
+ e = tf.device("/CPU:0")
+
+ 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, (256,256,3) ) ])
def extract (self, input_image, rects, second_pass_extractor=None, is_bgr=True, multi_sample=False):
if len(rects) == 0:
@@ -63,13 +193,13 @@ class FANExtractor(object):
images += [ self.crop(input_image, c, scale) ]
images = np.stack (images)
- images = images.astype(np.float32) / 255.0
+ images = images.astype(np.float32) / 255.0
predicted = []
for i in range( len(images) ):
- predicted += [ self.model.predict ( images[i][None,...] ).transpose (0,3,1,2)[0] ]
+ predicted += [ self.model.run ( [ images[i][None,...] ] ).transpose (0,3,1,2)[0] ]
- predicted = np.stack(predicted)
+ predicted = np.stack(predicted)
for i, pred in enumerate(predicted):
ptss += [ self.get_pts_from_predict ( pred, centers[i], scale) ]
@@ -144,81 +274,3 @@ class FANExtractor(object):
c += 0.5
return np.array( [ self.transform (c[i], center, scale, a_w) for i in range(a_ch) ] )
-
- @staticmethod
- def BuildModel():
- def ConvBlock(out_planes, input):
- in_planes = K.int_shape(input)[-1]
- x = input
- x = BatchNormalization(momentum=0.1, epsilon=1e-05)(x)
- x = ReLU() (x)
- x = out1 = Conv2D( int(out_planes/2), kernel_size=3, strides=1, padding='valid', use_bias = False) (ZeroPadding2D(1)(x))
-
- x = BatchNormalization(momentum=0.1, epsilon=1e-05)(x)
- x = ReLU() (x)
- x = out2 = Conv2D( int(out_planes/4), kernel_size=3, strides=1, padding='valid', use_bias = False) (ZeroPadding2D(1)(x))
-
- x = BatchNormalization(momentum=0.1, epsilon=1e-05)(x)
- x = ReLU() (x)
- x = out3 = Conv2D( int(out_planes/4), kernel_size=3, strides=1, padding='valid', use_bias = False) (ZeroPadding2D(1)(x))
-
- x = Concatenate()([out1, out2, out3])
-
- if in_planes != out_planes:
- downsample = BatchNormalization(momentum=0.1, epsilon=1e-05)(input)
- downsample = ReLU() (downsample)
- downsample = Conv2D( out_planes, kernel_size=1, strides=1, padding='valid', use_bias = False) (downsample)
- x = Add ()([x, downsample])
- else:
- x = Add ()([x, input])
-
-
- return x
-
- def HourGlass (depth, input):
- up1 = ConvBlock(256, input)
-
- low1 = AveragePooling2D (pool_size=2, strides=2, padding='valid' )(input)
- low1 = ConvBlock (256, low1)
-
- if depth > 1:
- low2 = HourGlass (depth-1, low1)
- else:
- low2 = ConvBlock(256, low1)
-
- low3 = ConvBlock(256, low2)
-
- up2 = UpSampling2D(size=2) (low3)
- return Add() ( [up1, up2] )
-
- FAN_Input = Input ( (256, 256, 3) )
-
- x = FAN_Input
-
- x = Conv2D (64, kernel_size=7, strides=2, padding='valid')(ZeroPadding2D(3)(x))
- x = BatchNormalization(momentum=0.1, epsilon=1e-05)(x)
- x = ReLU()(x)
-
- x = ConvBlock (128, x)
- x = AveragePooling2D (pool_size=2, strides=2, padding='valid') (x)
- x = ConvBlock (128, x)
- x = ConvBlock (256, x)
-
- outputs = []
- previous = x
- for i in range(4):
- ll = HourGlass (4, previous)
- ll = ConvBlock (256, ll)
-
- ll = Conv2D(256, kernel_size=1, strides=1, padding='valid') (ll)
- ll = BatchNormalization(momentum=0.1, epsilon=1e-05)(ll)
- ll = ReLU() (ll)
-
- tmp_out = Conv2D(68, kernel_size=1, strides=1, padding='valid') (ll)
- outputs.append(tmp_out)
-
- if i < 4 - 1:
- ll = Conv2D(256, kernel_size=1, strides=1, padding='valid') (ll)
- previous = Add() ( [previous, ll, KL.Conv2D(256, kernel_size=1, strides=1, padding='valid') (tmp_out) ] )
-
- return Model(FAN_Input, outputs[-1] )
diff --git a/nnlib/FANSeg_256_full_face.h5 b/facelib/FANSeg_256_full_face.npy
similarity index 99%
rename from nnlib/FANSeg_256_full_face.h5
rename to facelib/FANSeg_256_full_face.npy
index 6886504..53a6664 100644
Binary files a/nnlib/FANSeg_256_full_face.h5 and b/facelib/FANSeg_256_full_face.npy differ
diff --git a/facelib/FaceEnhancer.h5 b/facelib/FaceEnhancer.npy
similarity index 99%
rename from facelib/FaceEnhancer.h5
rename to facelib/FaceEnhancer.npy
index 201105b..1890f42 100644
Binary files a/facelib/FaceEnhancer.h5 and b/facelib/FaceEnhancer.npy differ
diff --git a/facelib/FaceEnhancer.py b/facelib/FaceEnhancer.py
index c3b2016..88c4da4 100644
--- a/facelib/FaceEnhancer.py
+++ b/facelib/FaceEnhancer.py
@@ -4,151 +4,321 @@ from pathlib import Path
import cv2
import numpy as np
-
+from core.leras import nn
class FaceEnhancer(object):
"""
x4 face enhancer
"""
- def __init__(self):
- from nnlib import nnlib
- exec( nnlib.import_all(), locals(), globals() )
+ def __init__(self, place_model_on_cpu=False):
+ nn.initialize()
+ tf = nn.tf
- model_path = Path(__file__).parent / "FaceEnhancer.h5"
+ class FaceEnhancer (nn.ModelBase):
+ def __init__(self, name='FaceEnhancer'):
+ super().__init__(name=name)
+
+ def on_build(self):
+ self.conv1 = nn.Conv2D (3, 64, kernel_size=3, strides=1, padding='SAME')
+
+ self.dense1 = nn.Dense (1, 64, use_bias=False)
+ self.dense2 = nn.Dense (1, 64, use_bias=False)
+
+ self.e0_conv0 = nn.Conv2D (64, 64, kernel_size=3, strides=1, padding='SAME')
+ self.e0_conv1 = nn.Conv2D (64, 64, kernel_size=3, strides=1, padding='SAME')
+
+ self.e1_conv0 = nn.Conv2D (64, 112, kernel_size=3, strides=1, padding='SAME')
+ self.e1_conv1 = nn.Conv2D (112, 112, kernel_size=3, strides=1, padding='SAME')
+
+ self.e2_conv0 = nn.Conv2D (112, 192, kernel_size=3, strides=1, padding='SAME')
+ self.e2_conv1 = nn.Conv2D (192, 192, kernel_size=3, strides=1, padding='SAME')
+
+ self.e3_conv0 = nn.Conv2D (192, 336, kernel_size=3, strides=1, padding='SAME')
+ self.e3_conv1 = nn.Conv2D (336, 336, kernel_size=3, strides=1, padding='SAME')
+
+ self.e4_conv0 = nn.Conv2D (336, 512, kernel_size=3, strides=1, padding='SAME')
+ self.e4_conv1 = nn.Conv2D (512, 512, kernel_size=3, strides=1, padding='SAME')
+
+ self.center_conv0 = nn.Conv2D (512, 512, kernel_size=3, strides=1, padding='SAME')
+ self.center_conv1 = nn.Conv2D (512, 512, kernel_size=3, strides=1, padding='SAME')
+ self.center_conv2 = nn.Conv2D (512, 512, kernel_size=3, strides=1, padding='SAME')
+ self.center_conv3 = nn.Conv2D (512, 512, kernel_size=3, strides=1, padding='SAME')
+
+ self.d4_conv0 = nn.Conv2D (1024, 512, kernel_size=3, strides=1, padding='SAME')
+ self.d4_conv1 = nn.Conv2D (512, 512, kernel_size=3, strides=1, padding='SAME')
+
+ self.d3_conv0 = nn.Conv2D (848, 512, kernel_size=3, strides=1, padding='SAME')
+ self.d3_conv1 = nn.Conv2D (512, 512, kernel_size=3, strides=1, padding='SAME')
+
+ self.d2_conv0 = nn.Conv2D (704, 288, kernel_size=3, strides=1, padding='SAME')
+ self.d2_conv1 = nn.Conv2D (288, 288, kernel_size=3, strides=1, padding='SAME')
+
+ self.d1_conv0 = nn.Conv2D (400, 160, kernel_size=3, strides=1, padding='SAME')
+ self.d1_conv1 = nn.Conv2D (160, 160, kernel_size=3, strides=1, padding='SAME')
+
+ self.d0_conv0 = nn.Conv2D (224, 96, kernel_size=3, strides=1, padding='SAME')
+ self.d0_conv1 = nn.Conv2D (96, 96, kernel_size=3, strides=1, padding='SAME')
+
+ self.out1x_conv0 = nn.Conv2D (96, 48, kernel_size=3, strides=1, padding='SAME')
+ self.out1x_conv1 = nn.Conv2D (48, 3, kernel_size=3, strides=1, padding='SAME')
+
+ self.dec2x_conv0 = nn.Conv2D (96, 96, kernel_size=3, strides=1, padding='SAME')
+ self.dec2x_conv1 = nn.Conv2D (96, 96, kernel_size=3, strides=1, padding='SAME')
+
+ self.out2x_conv0 = nn.Conv2D (96, 48, kernel_size=3, strides=1, padding='SAME')
+ self.out2x_conv1 = nn.Conv2D (48, 3, kernel_size=3, strides=1, padding='SAME')
+
+ self.dec4x_conv0 = nn.Conv2D (96, 72, kernel_size=3, strides=1, padding='SAME')
+ self.dec4x_conv1 = nn.Conv2D (72, 72, kernel_size=3, strides=1, padding='SAME')
+
+ self.out4x_conv0 = nn.Conv2D (72, 36, kernel_size=3, strides=1, padding='SAME')
+ self.out4x_conv1 = nn.Conv2D (36, 3 , kernel_size=3, strides=1, padding='SAME')
+
+ def forward(self, inp):
+ bgr, param, param1 = inp
+
+ x = self.conv1(bgr)
+ a = self.dense1(param)
+ a = tf.reshape(a, (-1,1,1,64) )
+
+ b = self.dense2(param1)
+ b = tf.reshape(b, (-1,1,1,64) )
+
+ x = tf.nn.leaky_relu(x+a+b, 0.1)
+
+ x = tf.nn.leaky_relu(self.e0_conv0(x), 0.1)
+ x = e0 = tf.nn.leaky_relu(self.e0_conv1(x), 0.1)
+
+ x = tf.nn.avg_pool(x, [1,2,2,1], [1,2,2,1], "VALID")
+ x = tf.nn.leaky_relu(self.e1_conv0(x), 0.1)
+ x = e1 = tf.nn.leaky_relu(self.e1_conv1(x), 0.1)
+
+ x = tf.nn.avg_pool(x, [1,2,2,1], [1,2,2,1], "VALID")
+ x = tf.nn.leaky_relu(self.e2_conv0(x), 0.1)
+ x = e2 = tf.nn.leaky_relu(self.e2_conv1(x), 0.1)
+
+ x = tf.nn.avg_pool(x, [1,2,2,1], [1,2,2,1], "VALID")
+ x = tf.nn.leaky_relu(self.e3_conv0(x), 0.1)
+ x = e3 = tf.nn.leaky_relu(self.e3_conv1(x), 0.1)
+
+ x = tf.nn.avg_pool(x, [1,2,2,1], [1,2,2,1], "VALID")
+ x = tf.nn.leaky_relu(self.e4_conv0(x), 0.1)
+ x = e4 = tf.nn.leaky_relu(self.e4_conv1(x), 0.1)
+
+ x = tf.nn.avg_pool(x, [1,2,2,1], [1,2,2,1], "VALID")
+ x = tf.nn.leaky_relu(self.center_conv0(x), 0.1)
+ x = tf.nn.leaky_relu(self.center_conv1(x), 0.1)
+ x = tf.nn.leaky_relu(self.center_conv2(x), 0.1)
+ x = tf.nn.leaky_relu(self.center_conv3(x), 0.1)
+
+ x = tf.concat( [nn.tf_upsample2d_bilinear(x), e4], -1 )
+ x = tf.nn.leaky_relu(self.d4_conv0(x), 0.1)
+ x = tf.nn.leaky_relu(self.d4_conv1(x), 0.1)
+
+ x = tf.concat( [nn.tf_upsample2d_bilinear(x), e3], -1 )
+ x = tf.nn.leaky_relu(self.d3_conv0(x), 0.1)
+ x = tf.nn.leaky_relu(self.d3_conv1(x), 0.1)
+
+ x = tf.concat( [nn.tf_upsample2d_bilinear(x), e2], -1 )
+ x = tf.nn.leaky_relu(self.d2_conv0(x), 0.1)
+ x = tf.nn.leaky_relu(self.d2_conv1(x), 0.1)
+
+ x = tf.concat( [nn.tf_upsample2d_bilinear(x), e1], -1 )
+ x = tf.nn.leaky_relu(self.d1_conv0(x), 0.1)
+ x = tf.nn.leaky_relu(self.d1_conv1(x), 0.1)
+
+ x = tf.concat( [nn.tf_upsample2d_bilinear(x), e0], -1 )
+ x = tf.nn.leaky_relu(self.d0_conv0(x), 0.1)
+ x = d0 = tf.nn.leaky_relu(self.d0_conv1(x), 0.1)
+
+ x = tf.nn.leaky_relu(self.out1x_conv0(x), 0.1)
+ x = self.out1x_conv1(x)
+ out1x = bgr + tf.nn.tanh(x)
+
+ x = d0
+ x = tf.nn.leaky_relu(self.dec2x_conv0(x), 0.1)
+ x = tf.nn.leaky_relu(self.dec2x_conv1(x), 0.1)
+ x = d2x = nn.tf_upsample2d_bilinear(x)
+
+ x = tf.nn.leaky_relu(self.out2x_conv0(x), 0.1)
+ x = self.out2x_conv1(x)
+
+ out2x = nn.tf_upsample2d_bilinear(out1x) + tf.nn.tanh(x)
+
+ x = d2x
+ x = tf.nn.leaky_relu(self.dec4x_conv0(x), 0.1)
+ x = tf.nn.leaky_relu(self.dec4x_conv1(x), 0.1)
+ x = d4x = nn.tf_upsample2d_bilinear(x)
+
+ x = tf.nn.leaky_relu(self.out4x_conv0(x), 0.1)
+ x = self.out4x_conv1(x)
+
+ out4x = nn.tf_upsample2d_bilinear(out2x) + tf.nn.tanh(x)
+
+ return out4x
+
+ model_path = Path(__file__).parent / "FaceEnhancer.npy"
if not model_path.exists():
- return
-
- bgr_inp = Input ( (192,192,3) )
- t_param_inp = Input ( (1,) )
- t_param1_inp = Input ( (1,) )
- x = Conv2D (64, 3, strides=1, padding='same' )(bgr_inp)
-
- a = Dense (64, use_bias=False) ( t_param_inp )
- a = Reshape( (1,1,64) )(a)
- b = Dense (64, use_bias=False ) ( t_param1_inp )
- b = Reshape( (1,1,64) )(b)
- x = Add()([x,a,b])
-
- x = LeakyReLU(0.1)(x)
+ raise Exception("Unable to load FaceEnhancer.npy")
- x = LeakyReLU(0.1)(Conv2D (64, 3, strides=1, padding='same' )(x))
- x = e0 = LeakyReLU(0.1)(Conv2D (64, 3, strides=1, padding='same')(x))
-
- x = AveragePooling2D()(x)
- x = LeakyReLU(0.1)(Conv2D (112, 3, strides=1, padding='same')(x))
- x = e1 = LeakyReLU(0.1)(Conv2D (112, 3, strides=1, padding='same')(x))
-
- x = AveragePooling2D()(x)
- x = LeakyReLU(0.1)(Conv2D (192, 3, strides=1, padding='same')(x))
- x = e2 = LeakyReLU(0.1)(Conv2D (192, 3, strides=1, padding='same')(x))
-
- x = AveragePooling2D()(x)
- x = LeakyReLU(0.1)(Conv2D (336, 3, strides=1, padding='same')(x))
- x = e3 = LeakyReLU(0.1)(Conv2D (336, 3, strides=1, padding='same')(x))
-
- x = AveragePooling2D()(x)
- x = LeakyReLU(0.1)(Conv2D (512, 3, strides=1, padding='same')(x))
- x = e4 = LeakyReLU(0.1)(Conv2D (512, 3, strides=1, padding='same')(x))
-
- x = AveragePooling2D()(x)
- x = LeakyReLU(0.1)(Conv2D (512, 3, strides=1, padding='same')(x))
- x = LeakyReLU(0.1)(Conv2D (512, 3, strides=1, padding='same')(x))
- x = LeakyReLU(0.1)(Conv2D (512, 3, strides=1, padding='same')(x))
- x = LeakyReLU(0.1)(Conv2D (512, 3, strides=1, padding='same')(x))
+ e = tf.device("/CPU:0") if place_model_on_cpu else None
+ if e is not None: e.__enter__()
+ self.model = FaceEnhancer()
+ self.model.load_weights (model_path)
+ if e is not None: e.__exit__(None,None,None)
- x = Concatenate()([ BilinearInterpolation()(x), e4 ])
-
- x = LeakyReLU(0.1)(Conv2D (512, 3, strides=1, padding='same')(x))
- x = LeakyReLU(0.1)(Conv2D (512, 3, strides=1, padding='same')(x))
-
- x = Concatenate()([ BilinearInterpolation()(x), e3 ])
- x = LeakyReLU(0.1)(Conv2D (512, 3, strides=1, padding='same')(x))
- x = LeakyReLU(0.1)(Conv2D (512, 3, strides=1, padding='same')(x))
-
- x = Concatenate()([ BilinearInterpolation()(x), e2 ])
- x = LeakyReLU(0.1)(Conv2D (288, 3, strides=1, padding='same')(x))
- x = LeakyReLU(0.1)(Conv2D (288, 3, strides=1, padding='same')(x))
-
- x = Concatenate()([ BilinearInterpolation()(x), e1 ])
- x = LeakyReLU(0.1)(Conv2D (160, 3, strides=1, padding='same')(x))
- x = LeakyReLU(0.1)(Conv2D (160, 3, strides=1, padding='same')(x))
-
- x = Concatenate()([ BilinearInterpolation()(x), e0 ])
- x = LeakyReLU(0.1)(Conv2D (96, 3, strides=1, padding='same')(x))
- x = d0 = LeakyReLU(0.1)(Conv2D (96, 3, strides=1, padding='same')(x))
-
- x = LeakyReLU(0.1)(Conv2D (48, 3, strides=1, padding='same')(x))
-
- x = Conv2D (3, 3, strides=1, padding='same', activation='tanh')(x)
- out1x = Add()([bgr_inp, x])
-
- x = d0
- x = LeakyReLU(0.1)(Conv2D (96, 3, strides=1, padding='same')(x))
- x = LeakyReLU(0.1)(Conv2D (96, 3, strides=1, padding='same')(x))
- x = d2x = BilinearInterpolation()(x)
-
- x = LeakyReLU(0.1)(Conv2D (48, 3, strides=1, padding='same')(x))
- x = Conv2D (3, 3, strides=1, padding='same', activation='tanh')(x)
-
- out2x = Add()([BilinearInterpolation()(out1x), x])
-
- x = d2x
- x = LeakyReLU(0.1)(Conv2D (72, 3, strides=1, padding='same')(x))
- x = LeakyReLU(0.1)(Conv2D (72, 3, strides=1, padding='same')(x))
- x = d4x = BilinearInterpolation()(x)
-
- x = LeakyReLU(0.1)(Conv2D (36, 3, strides=1, padding='same')(x))
- x = Conv2D (3, 3, strides=1, padding='same', activation='tanh')(x)
- out4x = Add()([BilinearInterpolation()(out2x), x ])
-
- self.model = keras.models.Model ( [bgr_inp,t_param_inp,t_param1_inp], [out4x] )
- self.model.load_weights (str(model_path))
+ self.model.build_for_run ([ (tf.float32, (192,192,3) ),
+ (tf.float32, (1,) ),
+ (tf.float32, (1,) ),
+ ])
def enhance (self, inp_img, is_tanh=False, preserve_size=True):
if not is_tanh:
inp_img = np.clip( inp_img * 2 -1, -1, 1 )
-
+
param = np.array([0.2])
- param1 = np.array([1.0])
+ param1 = np.array([1.0])
up_res = 4
patch_size = 192
patch_size_half = patch_size // 2
-
- h,w,c = inp_img.shape
+
+ ih,iw,ic = inp_img.shape
+ h,w,c = ih,iw,ic
+ th,tw = h*up_res, w*up_res
+
+ t_padding = 0
+ b_padding = 0
+ l_padding = 0
+ r_padding = 0
+
+ if h < patch_size:
+ t_padding = (patch_size-h)//2
+ b_padding = (patch_size-h) - t_padding
+
+ if w < patch_size:
+ l_padding = (patch_size-w)//2
+ r_padding = (patch_size-w) - l_padding
+
+ if t_padding != 0:
+ inp_img = np.concatenate ([ np.zeros ( (t_padding,w,c), dtype=np.float32 ), inp_img ], axis=0 )
+ h,w,c = inp_img.shape
+
+ if b_padding != 0:
+ inp_img = np.concatenate ([ inp_img, np.zeros ( (b_padding,w,c), dtype=np.float32 ) ], axis=0 )
+ h,w,c = inp_img.shape
+
+ if l_padding != 0:
+ inp_img = np.concatenate ([ np.zeros ( (h,l_padding,c), dtype=np.float32 ), inp_img ], axis=1 )
+ h,w,c = inp_img.shape
+
+ if r_padding != 0:
+ inp_img = np.concatenate ([ inp_img, np.zeros ( (h,r_padding,c), dtype=np.float32 ) ], axis=1 )
+ h,w,c = inp_img.shape
+
+
i_max = w-patch_size+1
- j_max = h-patch_size+1
-
+ j_max = h-patch_size+1
+
final_img = np.zeros ( (h*up_res,w*up_res,c), dtype=np.float32 )
final_img_div = np.zeros ( (h*up_res,w*up_res,1), dtype=np.float32 )
-
+
x = np.concatenate ( [ np.linspace (0,1,patch_size_half*up_res), np.linspace (1,0,patch_size_half*up_res) ] )
x,y = np.meshgrid(x,x)
patch_mask = (x*y)[...,None]
-
+
j=0
while j < j_max:
i = 0
- while i < i_max:
- patch_img = inp_img[j:j+patch_size, i:i+patch_size,:]
- x = self.model.predict( [ patch_img[None,...], param, param1 ] )[0]
+ while i < i_max:
+ patch_img = inp_img[j:j+patch_size, i:i+patch_size,:]
+ x = self.model.run( [ patch_img[None,...], [param], [param1] ] )[0]
final_img [j*up_res:(j+patch_size)*up_res, i*up_res:(i+patch_size)*up_res,:] += x*patch_mask
final_img_div[j*up_res:(j+patch_size)*up_res, i*up_res:(i+patch_size)*up_res,:] += patch_mask
if i == i_max-1:
break
- i = min( i+patch_size_half, i_max-1)
+ i = min( i+patch_size_half, i_max-1)
if j == j_max-1:
break
j = min( j+patch_size_half, j_max-1)
-
+
final_img_div[final_img_div==0] = 1.0
final_img /= final_img_div
-
+
+ if t_padding+b_padding+l_padding+r_padding != 0:
+ final_img = final_img [t_padding*up_res:(h-b_padding)*up_res, l_padding*up_res:(w-r_padding)*up_res,:]
+
if preserve_size:
- final_img = cv2.resize (final_img, (w,h), cv2.INTER_LANCZOS4)
-
+ final_img = cv2.resize (final_img, (iw,ih), cv2.INTER_LANCZOS4)
+
if not is_tanh:
final_img = np.clip( final_img/2+0.5, 0, 1 )
-
+
return final_img
+
+
+"""
+
+ def enhance (self, inp_img, is_tanh=False, preserve_size=True):
+ if not is_tanh:
+ inp_img = np.clip( inp_img * 2 -1, -1, 1 )
+
+ param = np.array([0.2])
+ param1 = np.array([1.0])
+ up_res = 4
+ patch_size = 192
+ patch_size_half = patch_size // 2
+
+ h,w,c = inp_img.shape
+
+ th,tw = h*up_res, w*up_res
+
+ preupscale_rate = 1.0
+
+ if h < patch_size or w < patch_size:
+ preupscale_rate = 1.0 / ( max(h,w) / patch_size )
+
+ if preupscale_rate != 1.0:
+ inp_img = cv2.resize (inp_img, ( int(w*preupscale_rate), int(h*preupscale_rate) ), cv2.INTER_LANCZOS4)
+ h,w,c = inp_img.shape
+
+ i_max = w-patch_size+1
+ j_max = h-patch_size+1
+
+ final_img = np.zeros ( (h*up_res,w*up_res,c), dtype=np.float32 )
+ final_img_div = np.zeros ( (h*up_res,w*up_res,1), dtype=np.float32 )
+
+ x = np.concatenate ( [ np.linspace (0,1,patch_size_half*up_res), np.linspace (1,0,patch_size_half*up_res) ] )
+ x,y = np.meshgrid(x,x)
+ patch_mask = (x*y)[...,None]
+
+ j=0
+ while j < j_max:
+ i = 0
+ while i < i_max:
+ patch_img = inp_img[j:j+patch_size, i:i+patch_size,:]
+ x = self.model.run( [ patch_img[None,...], [param], [param1] ] )[0]
+ final_img [j*up_res:(j+patch_size)*up_res, i*up_res:(i+patch_size)*up_res,:] += x*patch_mask
+ final_img_div[j*up_res:(j+patch_size)*up_res, i*up_res:(i+patch_size)*up_res,:] += patch_mask
+ if i == i_max-1:
+ break
+ i = min( i+patch_size_half, i_max-1)
+ if j == j_max-1:
+ break
+ j = min( j+patch_size_half, j_max-1)
+
+ final_img_div[final_img_div==0] = 1.0
+ final_img /= final_img_div
+
+ if preserve_size:
+ final_img = cv2.resize (final_img, (w,h), cv2.INTER_LANCZOS4)
+ else:
+ if preupscale_rate != 1.0:
+ final_img = cv2.resize (final_img, (tw,th), cv2.INTER_LANCZOS4)
+
+ if not is_tanh:
+ final_img = np.clip( final_img/2+0.5, 0, 1 )
+
+ return final_img
+"""
\ No newline at end of file
diff --git a/facelib/LandmarksProcessor.py b/facelib/LandmarksProcessor.py
index 7d9f706..560ae29 100644
--- a/facelib/LandmarksProcessor.py
+++ b/facelib/LandmarksProcessor.py
@@ -6,11 +6,11 @@ import cv2
import numpy as np
import numpy.linalg as npla
-import imagelib
-import mathlib
+from core import imagelib
+from core import mathlib
from facelib import FaceType
-from imagelib import IEPolys
-from mathlib.umeyama import umeyama
+from core.imagelib import IEPolys
+from core.mathlib.umeyama import umeyama
landmarks_2D = np.array([
[ 0.000213256, 0.106454 ], #17
@@ -665,8 +665,10 @@ def calc_face_yaw(landmarks):
r = ( (landmarks[16][0]-landmarks[27][0]) + (landmarks[15][0]-landmarks[28][0]) + (landmarks[14][0]-landmarks[29][0]) ) / 3.0
return float(r-l)
-#returns pitch,yaw,roll [-1...+1]
def estimate_pitch_yaw_roll(aligned_256px_landmarks):
+ """
+ returns pitch,yaw,roll [-pi...+pi]
+ """
shape = (256,256)
focal_length = shape[1]
camera_center = (shape[1] / 2, shape[0] / 2)
@@ -682,7 +684,8 @@ def estimate_pitch_yaw_roll(aligned_256px_landmarks):
np.zeros((4, 1)) )
pitch, yaw, roll = mathlib.rotationMatrixToEulerAngles( cv2.Rodrigues(rotation_vector)[0] )
- pitch = np.clip ( pitch/1.30, -1.0, 1.0 )
- yaw = np.clip ( yaw / 1.11, -1.0, 1.0 )
- roll = np.clip ( roll/3.15, -1.0, 1.0 ) #todo radians
+ pitch = np.clip ( pitch, -math.pi, math.pi )
+ yaw = np.clip ( yaw , -math.pi, math.pi )
+ roll = np.clip ( roll, -math.pi, math.pi )
+
return -pitch, yaw, roll
diff --git a/facelib/MTCExtractor.py b/facelib/MTCExtractor.py
deleted file mode 100644
index c524ab9..0000000
--- a/facelib/MTCExtractor.py
+++ /dev/null
@@ -1,350 +0,0 @@
-import numpy as np
-import os
-import cv2
-
-from pathlib import Path
-from nnlib import nnlib
-
-class MTCExtractor(object):
- def __init__(self):
- self.scale_to = 1920
-
- self.min_face_size = self.scale_to * 0.042
- self.thresh1 = 0.7
- self.thresh2 = 0.85
- self.thresh3 = 0.6
- self.scale_factor = 0.95
-
- exec( nnlib.import_all(), locals(), globals() )
- PNet_Input = Input ( (None, None,3) )
- x = PNet_Input
- x = Conv2D (10, kernel_size=(3,3), strides=(1,1), padding='valid', name="conv1")(x)
- x = PReLU (shared_axes=[1,2], name="PReLU1" )(x)
- x = MaxPooling2D( pool_size=(2,2), strides=(2,2), padding='same' ) (x)
- x = Conv2D (16, kernel_size=(3,3), strides=(1,1), padding='valid', name="conv2")(x)
- x = PReLU (shared_axes=[1,2], name="PReLU2" )(x)
- x = Conv2D (32, kernel_size=(3,3), strides=(1,1), padding='valid', name="conv3")(x)
- x = PReLU (shared_axes=[1,2], name="PReLU3" )(x)
- prob = Conv2D (2, kernel_size=(1,1), strides=(1,1), padding='valid', name="conv41")(x)
- prob = Softmax()(prob)
- x = Conv2D (4, kernel_size=(1,1), strides=(1,1), padding='valid', name="conv42")(x)
-
- PNet_model = Model(PNet_Input, [x,prob] )
- PNet_model.load_weights ( (Path(__file__).parent / 'mtcnn_pnet.h5').__str__() )
-
- RNet_Input = Input ( (24, 24, 3) )
- x = RNet_Input
- x = Conv2D (28, kernel_size=(3,3), strides=(1,1), padding='valid', name="conv1")(x)
- x = PReLU (shared_axes=[1,2], name="prelu1" )(x)
- x = MaxPooling2D( pool_size=(3,3), strides=(2,2), padding='same' ) (x)
- x = Conv2D (48, kernel_size=(3,3), strides=(1,1), padding='valid', name="conv2")(x)
- x = PReLU (shared_axes=[1,2], name="prelu2" )(x)
- x = MaxPooling2D( pool_size=(3,3), strides=(2,2), padding='valid' ) (x)
- x = Conv2D (64, kernel_size=(2,2), strides=(1,1), padding='valid', name="conv3")(x)
- x = PReLU (shared_axes=[1,2], name="prelu3" )(x)
- x = Lambda ( lambda x: K.reshape (x, (-1, np.prod(K.int_shape(x)[1:]),) ), output_shape=(np.prod(K.int_shape(x)[1:]),) ) (x)
- x = Dense (128, name='conv4')(x)
- x = PReLU (name="prelu4" )(x)
- prob = Dense (2, name='conv51')(x)
- prob = Softmax()(prob)
- x = Dense (4, name='conv52')(x)
- RNet_model = Model(RNet_Input, [x,prob] )
- RNet_model.load_weights ( (Path(__file__).parent / 'mtcnn_rnet.h5').__str__() )
-
- ONet_Input = Input ( (48, 48, 3) )
- x = ONet_Input
- x = Conv2D (32, kernel_size=(3,3), strides=(1,1), padding='valid', name="conv1")(x)
- x = PReLU (shared_axes=[1,2], name="prelu1" )(x)
- x = MaxPooling2D( pool_size=(3,3), strides=(2,2), padding='same' ) (x)
- x = Conv2D (64, kernel_size=(3,3), strides=(1,1), padding='valid', name="conv2")(x)
- x = PReLU (shared_axes=[1,2], name="prelu2" )(x)
- x = MaxPooling2D( pool_size=(3,3), strides=(2,2), padding='valid' ) (x)
- x = Conv2D (64, kernel_size=(3,3), strides=(1,1), padding='valid', name="conv3")(x)
- x = PReLU (shared_axes=[1,2], name="prelu3" )(x)
- x = MaxPooling2D( pool_size=(2,2), strides=(2,2), padding='same' ) (x)
- x = Conv2D (128, kernel_size=(2,2), strides=(1,1), padding='valid', name="conv4")(x)
- x = PReLU (shared_axes=[1,2], name="prelu4" )(x)
- x = Lambda ( lambda x: K.reshape (x, (-1, np.prod(K.int_shape(x)[1:]),) ), output_shape=(np.prod(K.int_shape(x)[1:]),) ) (x)
- x = Dense (256, name='conv5')(x)
- x = PReLU (name="prelu5" )(x)
- prob = Dense (2, name='conv61')(x)
- prob = Softmax()(prob)
- x1 = Dense (4, name='conv62')(x)
- x2 = Dense (10, name='conv63')(x)
- ONet_model = Model(ONet_Input, [x1,x2,prob] )
- ONet_model.load_weights ( (Path(__file__).parent / 'mtcnn_onet.h5').__str__() )
-
- self.pnet_fun = K.function ( PNet_model.inputs, PNet_model.outputs )
- self.rnet_fun = K.function ( RNet_model.inputs, RNet_model.outputs )
- self.onet_fun = K.function ( ONet_model.inputs, ONet_model.outputs )
-
- def __enter__(self):
- faces, pnts = detect_face ( np.zeros ( (self.scale_to, self.scale_to, 3)), self.min_face_size, self.pnet_fun, self.rnet_fun, self.onet_fun, [ self.thresh1, self.thresh2, self.thresh3 ], self.scale_factor )
-
- return self
-
- def __exit__(self, exc_type=None, exc_value=None, traceback=None):
- return False #pass exception between __enter__ and __exit__ to outter level
-
- def extract (self, input_image, is_bgr=True):
-
- if is_bgr:
- input_image = input_image[:,:,::-1].copy()
- is_bgr = False
-
- (h, w, ch) = input_image.shape
-
- input_scale = self.scale_to / max(w,h)
- input_image = cv2.resize (input_image, ( int(w*input_scale), int(h*input_scale) ), interpolation=cv2.INTER_LINEAR)
-
- detected_faces, pnts = detect_face ( input_image, self.min_face_size, self.pnet_fun, self.rnet_fun, self.onet_fun, [ self.thresh1, self.thresh2, self.thresh3 ], self.scale_factor )
- detected_faces = [ ( int(face[0]/input_scale), int(face[1]/input_scale), int(face[2]/input_scale), int(face[3]/input_scale)) for face in detected_faces ]
-
- return detected_faces
-
-def detect_face(img, minsize, pnet, rnet, onet, threshold, factor):
- """Detects faces in an image, and returns bounding boxes and points for them.
- img: input image
- minsize: minimum faces' size
- pnet, rnet, onet: caffemodel
- threshold: threshold=[th1, th2, th3], th1-3 are three steps's threshold
- factor: the factor used to create a scaling pyramid of face sizes to detect in the image.
- """
- factor_count=0
- total_boxes=np.empty((0,9))
- points=np.empty(0)
- h=img.shape[0]
- w=img.shape[1]
- minl=np.amin([h, w])
- m=12.0/minsize
- minl=minl*m
- # create scale pyramid
- scales=[]
- while minl>=12:
- scales += [m*np.power(factor, factor_count)]
- minl = minl*factor
- factor_count += 1
- # first stage
- for scale in scales:
- hs=int(np.ceil(h*scale))
- ws=int(np.ceil(w*scale))
- #print ('scale %f %d %d' % (scale, ws,hs))
- im_data = imresample(img, (hs, ws))
- im_data = (im_data-127.5)*0.0078125
- img_x = np.expand_dims(im_data, 0)
- img_y = np.transpose(img_x, (0,2,1,3))
- out = pnet([img_y])
- out0 = np.transpose(out[0], (0,2,1,3))
- out1 = np.transpose(out[1], (0,2,1,3))
-
- boxes, _ = generateBoundingBox(out1[0,:,:,1].copy(), out0[0,:,:,:].copy(), scale, threshold[0])
-
- # inter-scale nms
- pick = nms(boxes.copy(), 0.5, 'Union')
- if boxes.size>0 and pick.size>0:
- boxes = boxes[pick,:]
- total_boxes = np.append(total_boxes, boxes, axis=0)
-
- numbox = total_boxes.shape[0]
- if numbox>0:
- pick = nms(total_boxes.copy(), 0.7, 'Union')
- total_boxes = total_boxes[pick,:]
- regw = total_boxes[:,2]-total_boxes[:,0]
- regh = total_boxes[:,3]-total_boxes[:,1]
- qq1 = total_boxes[:,0]+total_boxes[:,5]*regw
- qq2 = total_boxes[:,1]+total_boxes[:,6]*regh
- qq3 = total_boxes[:,2]+total_boxes[:,7]*regw
- qq4 = total_boxes[:,3]+total_boxes[:,8]*regh
- total_boxes = np.transpose(np.vstack([qq1, qq2, qq3, qq4, total_boxes[:,4]]))
- total_boxes = rerec(total_boxes.copy())
- total_boxes[:,0:4] = np.fix(total_boxes[:,0:4]).astype(np.int32)
- dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph = pad(total_boxes.copy(), w, h)
-
- numbox = total_boxes.shape[0]
- if numbox>0:
- # second stage
- tempimg = np.zeros((24,24,3,numbox))
- for k in range(0,numbox):
- tmp = np.zeros((int(tmph[k]),int(tmpw[k]),3))
- tmp[dy[k]-1:edy[k],dx[k]-1:edx[k],:] = img[y[k]-1:ey[k],x[k]-1:ex[k],:]
- if tmp.shape[0]>0 and tmp.shape[1]>0 or tmp.shape[0]==0 and tmp.shape[1]==0:
- tempimg[:,:,:,k] = imresample(tmp, (24, 24))
- else:
- return np.empty()
- tempimg = (tempimg-127.5)*0.0078125
- tempimg1 = np.transpose(tempimg, (3,1,0,2))
- out = rnet([tempimg1])
- out0 = np.transpose(out[0])
- out1 = np.transpose(out[1])
- score = out1[1,:]
- ipass = np.where(score>threshold[1])
- total_boxes = np.hstack([total_boxes[ipass[0],0:4].copy(), np.expand_dims(score[ipass].copy(),1)])
- mv = out0[:,ipass[0]]
- if total_boxes.shape[0]>0:
- pick = nms(total_boxes, 0.7, 'Union')
- total_boxes = total_boxes[pick,:]
- total_boxes = bbreg(total_boxes.copy(), np.transpose(mv[:,pick]))
- total_boxes = rerec(total_boxes.copy())
-
- numbox = total_boxes.shape[0]
- if numbox>0:
- # third stage
- total_boxes = np.fix(total_boxes).astype(np.int32)
- dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph = pad(total_boxes.copy(), w, h)
- tempimg = np.zeros((48,48,3,numbox))
- for k in range(0,numbox):
- tmp = np.zeros((int(tmph[k]),int(tmpw[k]),3))
- tmp[dy[k]-1:edy[k],dx[k]-1:edx[k],:] = img[y[k]-1:ey[k],x[k]-1:ex[k],:]
- if tmp.shape[0]>0 and tmp.shape[1]>0 or tmp.shape[0]==0 and tmp.shape[1]==0:
- tempimg[:,:,:,k] = imresample(tmp, (48, 48))
- else:
- return np.empty()
- tempimg = (tempimg-127.5)*0.0078125
- tempimg1 = np.transpose(tempimg, (3,1,0,2))
- out = onet([tempimg1])
- out0 = np.transpose(out[0])
- out1 = np.transpose(out[1])
- out2 = np.transpose(out[2])
- score = out2[1,:]
- points = out1
- ipass = np.where(score>threshold[2])
- points = points[:,ipass[0]]
- total_boxes = np.hstack([total_boxes[ipass[0],0:4].copy(), np.expand_dims(score[ipass].copy(),1)])
- mv = out0[:,ipass[0]]
-
- w = total_boxes[:,2]-total_boxes[:,0]+1
- h = total_boxes[:,3]-total_boxes[:,1]+1
- points[0:5,:] = np.tile(w,(5, 1))*points[0:5,:] + np.tile(total_boxes[:,0],(5, 1))-1
- points[5:10,:] = np.tile(h,(5, 1))*points[5:10,:] + np.tile(total_boxes[:,1],(5, 1))-1
- if total_boxes.shape[0]>0:
- total_boxes = bbreg(total_boxes.copy(), np.transpose(mv))
- pick = nms(total_boxes.copy(), 0.7, 'Min')
- total_boxes = total_boxes[pick,:]
- points = points[:,pick]
-
- return total_boxes, points
-
-
-# function [boundingbox] = bbreg(boundingbox,reg)
-def bbreg(boundingbox,reg):
- """Calibrate bounding boxes"""
- if reg.shape[1]==1:
- reg = np.reshape(reg, (reg.shape[2], reg.shape[3]))
-
- w = boundingbox[:,2]-boundingbox[:,0]+1
- h = boundingbox[:,3]-boundingbox[:,1]+1
- b1 = boundingbox[:,0]+reg[:,0]*w
- b2 = boundingbox[:,1]+reg[:,1]*h
- b3 = boundingbox[:,2]+reg[:,2]*w
- b4 = boundingbox[:,3]+reg[:,3]*h
- boundingbox[:,0:4] = np.transpose(np.vstack([b1, b2, b3, b4 ]))
- return boundingbox
-
-def generateBoundingBox(imap, reg, scale, t):
- """Use heatmap to generate bounding boxes"""
- stride=2
- cellsize=12
-
- imap = np.transpose(imap)
- dx1 = np.transpose(reg[:,:,0])
- dy1 = np.transpose(reg[:,:,1])
- dx2 = np.transpose(reg[:,:,2])
- dy2 = np.transpose(reg[:,:,3])
- y, x = np.where(imap >= t)
- if y.shape[0]==1:
- dx1 = np.flipud(dx1)
- dy1 = np.flipud(dy1)
- dx2 = np.flipud(dx2)
- dy2 = np.flipud(dy2)
- score = imap[(y,x)]
- reg = np.transpose(np.vstack([ dx1[(y,x)], dy1[(y,x)], dx2[(y,x)], dy2[(y,x)] ]))
- if reg.size==0:
- reg = np.empty((0,3))
- bb = np.transpose(np.vstack([y,x]))
- q1 = np.fix((stride*bb+1)/scale)
- q2 = np.fix((stride*bb+cellsize-1+1)/scale)
- boundingbox = np.hstack([q1, q2, np.expand_dims(score,1), reg])
- return boundingbox, reg
-
-# function pick = nms(boxes,threshold,type)
-def nms(boxes, threshold, method):
- if boxes.size==0:
- return np.empty((0,3))
- x1 = boxes[:,0]
- y1 = boxes[:,1]
- x2 = boxes[:,2]
- y2 = boxes[:,3]
- s = boxes[:,4]
- area = (x2-x1+1) * (y2-y1+1)
- I = np.argsort(s)
- pick = np.zeros_like(s, dtype=np.int16)
- counter = 0
- while I.size>0:
- i = I[-1]
- pick[counter] = i
- counter += 1
- idx = I[0:-1]
- xx1 = np.maximum(x1[i], x1[idx])
- yy1 = np.maximum(y1[i], y1[idx])
- xx2 = np.minimum(x2[i], x2[idx])
- yy2 = np.minimum(y2[i], y2[idx])
- w = np.maximum(0.0, xx2-xx1+1)
- h = np.maximum(0.0, yy2-yy1+1)
- inter = w * h
- if method == 'Min':
- o = inter / np.minimum(area[i], area[idx])
- else:
- o = inter / (area[i] + area[idx] - inter)
- I = I[np.where(o<=threshold)]
- pick = pick[0:counter]
- return pick
-
-# function [dy edy dx edx y ey x ex tmpw tmph] = pad(total_boxes,w,h)
-def pad(total_boxes, w, h):
- """Compute the padding coordinates (pad the bounding boxes to square)"""
- tmpw = (total_boxes[:,2]-total_boxes[:,0]+1).astype(np.int32)
- tmph = (total_boxes[:,3]-total_boxes[:,1]+1).astype(np.int32)
- numbox = total_boxes.shape[0]
-
- dx = np.ones((numbox), dtype=np.int32)
- dy = np.ones((numbox), dtype=np.int32)
- edx = tmpw.copy().astype(np.int32)
- edy = tmph.copy().astype(np.int32)
-
- x = total_boxes[:,0].copy().astype(np.int32)
- y = total_boxes[:,1].copy().astype(np.int32)
- ex = total_boxes[:,2].copy().astype(np.int32)
- ey = total_boxes[:,3].copy().astype(np.int32)
-
- tmp = np.where(ex>w)
- edx.flat[tmp] = np.expand_dims(-ex[tmp]+w+tmpw[tmp],1)
- ex[tmp] = w
-
- tmp = np.where(ey>h)
- edy.flat[tmp] = np.expand_dims(-ey[tmp]+h+tmph[tmp],1)
- ey[tmp] = h
-
- tmp = np.where(x<1)
- dx.flat[tmp] = np.expand_dims(2-x[tmp],1)
- x[tmp] = 1
-
- tmp = np.where(y<1)
- dy.flat[tmp] = np.expand_dims(2-y[tmp],1)
- y[tmp] = 1
-
- return dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph
-
-# function [bboxA] = rerec(bboxA)
-def rerec(bboxA):
- """Convert bboxA to square."""
- h = bboxA[:,3]-bboxA[:,1]
- w = bboxA[:,2]-bboxA[:,0]
- l = np.maximum(w, h)
- bboxA[:,0] = bboxA[:,0]+w*0.5-l*0.5
- bboxA[:,1] = bboxA[:,1]+h*0.5-l*0.5
- bboxA[:,2:4] = bboxA[:,0:2] + np.transpose(np.tile(l,(2,1)))
- return bboxA
-
-def imresample(img, sz):
- im_data = cv2.resize(img, (sz[1], sz[0]), interpolation=cv2.INTER_LINEAR) #@UndefinedVariable
- return im_data
diff --git a/facelib/PoseEstimator.py b/facelib/PoseEstimator.py
deleted file mode 100644
index 7939d90..0000000
--- a/facelib/PoseEstimator.py
+++ /dev/null
@@ -1,302 +0,0 @@
-import os
-import pickle
-from functools import partial
-from pathlib import Path
-
-import cv2
-import numpy as np
-
-from interact import interact as io
-from nnlib import nnlib
-
-"""
-PoseEstimator estimates pitch, yaw, roll, from FAN aligned face.
-trained on https://www.umdfaces.io
-based on https://arxiv.org/pdf/1901.06778.pdf HYBRID COARSE-FINE CLASSIFICATION FOR HEAD POSE ESTIMATION
-"""
-
-class PoseEstimator(object):
- VERSION = 1
- def __init__ (self, resolution, face_type_str, load_weights=True, weights_file_root=None, training=False):
- exec( nnlib.import_all(), locals(), globals() )
- self.resolution = resolution
-
- self.angles = [60, 45, 30, 10, 2]
- self.alpha_cat_losses = [7,5,3,1,1]
- self.class_nums = [ angle+1 for angle in self.angles ]
- self.encoder, self.decoder, self.model_l = PoseEstimator.BuildModels(resolution, class_nums=self.class_nums)
-
- if weights_file_root is not None:
- weights_file_root = Path(weights_file_root)
- else:
- weights_file_root = Path(__file__).parent
-
- self.encoder_weights_path = weights_file_root / ('PoseEst_%d_%s_enc.h5' % (resolution, face_type_str) )
- self.decoder_weights_path = weights_file_root / ('PoseEst_%d_%s_dec.h5' % (resolution, face_type_str) )
- self.l_weights_path = weights_file_root / ('PoseEst_%d_%s_l.h5' % (resolution, face_type_str) )
-
- self.model_weights_path = weights_file_root / ('PoseEst_%d_%s.h5' % (resolution, face_type_str) )
-
- self.input_bgr_shape = (resolution, resolution, 3)
-
- def ResamplerFunc(input):
- mean_t, logvar_t = input
- return mean_t + K.exp(0.5*logvar_t)*K.random_normal(K.shape(mean_t))
-
- self.BVAEResampler = Lambda ( lambda x: x[0] + K.random_normal(K.shape(x[0])) * K.sqrt(K.exp(0.5*x[1])),
- output_shape=K.int_shape(self.encoder.outputs[0])[1:] )
-
- inp_t = Input (self.input_bgr_shape)
- inp_real_t = Input (self.input_bgr_shape)
- inp_pitch_t = Input ( (1,) )
- inp_yaw_t = Input ( (1,) )
- inp_roll_t = Input ( (1,) )
-
-
- mean_t, logvar_t = self.encoder(inp_t)
-
- latent_t = self.BVAEResampler([mean_t, logvar_t])
-
- if training:
- bgr_t = self.decoder (latent_t)
- pyrs_t = self.model_l(latent_t)
- else:
- self.model = Model(inp_t, self.model_l(latent_t) )
- pyrs_t = self.model(inp_t)
-
- if load_weights:
- if training:
- self.encoder.load_weights (str(self.encoder_weights_path))
- self.decoder.load_weights (str(self.decoder_weights_path))
- self.model_l.load_weights (str(self.l_weights_path))
- else:
- self.model.load_weights (str(self.model_weights_path))
-
- else:
- def gather_Conv2D_layers(models_list):
- conv_weights_list = []
- for model in models_list:
- for layer in model.layers:
- layer_type = type(layer)
- if layer_type == keras.layers.Conv2D:
- conv_weights_list += [layer.weights[0]] #Conv2D kernel_weights
- elif layer_type == keras.engine.training.Model:
- conv_weights_list += gather_Conv2D_layers ([layer])
- return conv_weights_list
-
- CAInitializerMP ( gather_Conv2D_layers( [self.encoder, self.decoder] ) )
-
-
- if training:
- inp_pyrs_t = []
- for class_num in self.class_nums:
- inp_pyrs_t += [ Input ((3,)) ]
-
- pyr_loss = []
-
- for i,class_num in enumerate(self.class_nums):
- a = self.alpha_cat_losses[i]
- pyr_loss += [ a*K.mean( K.square ( inp_pyrs_t[i] - pyrs_t[i]) ) ]
-
- def BVAELoss(beta=4):
- def func(input):
- mean_t, logvar_t = input
- return beta * K.mean ( K.sum( 0.5*(K.exp(logvar_t)+ K.square(mean_t)-logvar_t-1), axis=1) )
- return func
-
- BVAE_loss = BVAELoss()([mean_t, logvar_t])
-
- bgr_loss = K.mean(K.sum(K.abs(inp_real_t-bgr_t), axis=[1,2,3]))
-
- G_loss = BVAE_loss+bgr_loss
- pyr_loss = sum(pyr_loss)
-
-
- self.train = K.function ([inp_t, inp_real_t],
- [ G_loss ], Adam(lr=0.0005, beta_1=0.9, beta_2=0.999).get_updates( G_loss, self.encoder.trainable_weights+self.decoder.trainable_weights ) )
-
- self.train_l = K.function ([inp_t] + inp_pyrs_t,
- [pyr_loss], Adam(lr=0.0001).get_updates( pyr_loss, self.model_l.trainable_weights) )
-
-
- self.view = K.function ([inp_t], [ bgr_t, pyrs_t[0] ] )
-
- def __enter__(self):
- return self
-
- def __exit__(self, exc_type=None, exc_value=None, traceback=None):
- return False #pass exception between __enter__ and __exit__ to outter level
-
- def save_weights(self):
- self.encoder.save_weights (str(self.encoder_weights_path))
- self.decoder.save_weights (str(self.decoder_weights_path))
- self.model_l.save_weights (str(self.l_weights_path))
-
- inp_t = Input (self.input_bgr_shape)
-
- Model(inp_t, self.model_l(self.BVAEResampler(self.encoder(inp_t))) ).save_weights (str(self.model_weights_path))
-
- def train_on_batch(self, warps, imgs, pyr_tanh, skip_bgr_train=False):
- if not skip_bgr_train:
- bgr_loss, = self.train( [warps, imgs] )
- pyr_loss = 0
- else:
- bgr_loss = 0
-
- feed = [imgs]
- for i, (angle, class_num) in enumerate(zip(self.angles, self.class_nums)):
- a = angle / 2
- c = np.round( (pyr_tanh+1) * a ) / a -1 #.astype(K.floatx())
- feed += [c]
-
- pyr_loss, = self.train_l(feed)
-
- return bgr_loss, pyr_loss
-
- def extract (self, input_image, is_input_tanh=False):
- if is_input_tanh:
- raise NotImplemented("is_input_tanh")
-
- input_shape_len = len(input_image.shape)
- if input_shape_len == 3:
- input_image = input_image[np.newaxis,...]
-
- bgr, result, = self.view( [input_image] )
-
-
- #result = np.clip ( result / (self.angles[0] / 2) - 1, 0.0, 1.0 )
-
- if input_shape_len == 3:
- bgr = bgr[0]
- result = result[0]
-
- return bgr, result
-
- @staticmethod
- def BuildModels ( resolution, class_nums, ae_dims=128):
- exec( nnlib.import_all(), locals(), globals() )
-
- x = inp = Input ( (resolution,resolution,3) )
- x = PoseEstimator.EncFlow(ae_dims)(x)
- encoder = Model(inp,x)
-
- x = inp = Input ( K.int_shape(encoder.outputs[0][1:]) )
- x = PoseEstimator.DecFlow(resolution, ae_dims)(x)
- decoder = Model(inp,x)
-
- x = inp = Input ( K.int_shape(encoder.outputs[0][1:]) )
- x = PoseEstimator.LatentFlow(class_nums=class_nums)(x)
- model_l = Model(inp, x )
-
- return encoder, decoder, model_l
-
- @staticmethod
- def EncFlow(ae_dims):
- exec( nnlib.import_all(), locals(), globals() )
-
- def downscale (dim, **kwargs):
- def func(x):
- return ReLU() ( Conv2D(dim, kernel_size=5, strides=2, padding='same')(x))
- return func
-
-
- downscale = partial(downscale)
-
- ed_ch_dims = 128
-
- def func(input):
- x = input
- x = downscale(64)(x)
- x = downscale(128)(x)
- x = downscale(256)(x)
- x = downscale(512)(x)
- x = Flatten()(x)
-
- x = Dense(256)(x)
- x = ReLU()(x)
-
- x = Dense(256)(x)
- x = ReLU()(x)
-
- mean = Dense(ae_dims)(x)
- logvar = Dense(ae_dims)(x)
-
- return mean, logvar
-
- return func
-
- @staticmethod
- def DecFlow(resolution, ae_dims):
- exec( nnlib.import_all(), locals(), globals() )
-
- def upscale (dim, strides=2, **kwargs):
- def func(x):
- return ReLU()( ( Conv2DTranspose(dim, kernel_size=3, strides=strides, padding='same')(x)) )
- return func
-
- def to_bgr (output_nc, **kwargs):
- def func(x):
- return Conv2D(output_nc, kernel_size=5, padding='same', activation='sigmoid')(x)
- return func
-
- upscale = partial(upscale)
- lowest_dense_res = resolution // 16
-
- def func(input):
- x = input
-
- x = Dense(256)(x)
- x = ReLU()(x)
-
- x = Dense(256)(x)
- x = ReLU()(x)
-
- x = Dense( (lowest_dense_res*lowest_dense_res*256) ) (x)
- x = ReLU()(x)
-
- x = Reshape( (lowest_dense_res,lowest_dense_res,256) )(x)
-
- x = upscale(512)(x)
- x = upscale(256)(x)
- x = upscale(128)(x)
- x = upscale(64)(x)
- x = to_bgr(3)(x)
-
- return x
- return func
-
- @staticmethod
- def LatentFlow(class_nums):
- exec( nnlib.import_all(), locals(), globals() )
-
- def func(latent):
- x = latent
-
- x = Dense(1024, activation='relu')(x)
- x = Dropout(0.5)(x)
- x = Dense(1024, activation='relu')(x)
- # x = Dropout(0.5)(x)
- # x = Dense(4096, activation='relu')(x)
-
- output = []
- for class_num in class_nums:
- pyr = Dense(3, activation='tanh')(x)
- output += [pyr]
-
- return output
-
- #y = Dropout(0.5)(y)
- #y = Dense(1024, activation='relu')(y)
- return func
-
-
-# resnet50 = keras.applications.ResNet50(include_top=False, weights=None, input_shape=K.int_shape(x)[1:], pooling='avg')
-# x = resnet50(x)
-# output = []
-# for class_num in class_nums:
-# pitch = Dense(class_num)(x)
-# yaw = Dense(class_num)(x)
-# roll = Dense(class_num)(x)
-# output += [pitch,yaw,roll]
-
-# return output
diff --git a/facelib/S3FD.h5 b/facelib/S3FD.npy
similarity index 84%
rename from facelib/S3FD.h5
rename to facelib/S3FD.npy
index aed0938..66c054f 100644
Binary files a/facelib/S3FD.h5 and b/facelib/S3FD.npy differ
diff --git a/facelib/S3FDExtractor.py b/facelib/S3FDExtractor.py
index adafae9..c58d931 100644
--- a/facelib/S3FDExtractor.py
+++ b/facelib/S3FDExtractor.py
@@ -4,21 +4,171 @@ from pathlib import Path
import cv2
import numpy as np
-from nnlib import nnlib
+from core.leras import nn
class S3FDExtractor(object):
- def __init__(self, do_dummy_predict=False):
- exec( nnlib.import_all(), locals(), globals() )
-
- model_path = Path(__file__).parent / "S3FD.h5"
- if not model_path.exists():
- return None
-
- self.model = nnlib.keras.models.load_model ( str(model_path) )
+ def __init__(self, place_model_on_cpu=False):
+ nn.initialize()
+ tf = nn.tf
- if do_dummy_predict:
- self.extract ( np.zeros( (640,640,3), dtype=np.uint8) )
+ model_path = Path(__file__).parent / "S3FD.npy"
+ if not model_path.exists():
+ raise Exception("Unable to load S3FD.npy")
+ class L2Norm(nn.LayerBase):
+ def __init__(self, n_channels, **kwargs):
+ self.n_channels = n_channels
+ super().__init__(**kwargs)
+
+ def build_weights(self):
+ self.weight = tf.get_variable ("weight", (1, 1, 1, self.n_channels), dtype=nn.tf_floatx, initializer=tf.initializers.ones )
+
+ def get_weights(self):
+ return [self.weight]
+
+ def __call__(self, inputs):
+ x = inputs
+ x = x / (tf.sqrt( tf.reduce_sum( tf.pow(x, 2), axis=-1, keepdims=True ) ) + 1e-10) * self.weight
+ return x
+
+ class S3FD(nn.ModelBase):
+ def __init__(self):
+ super().__init__(name='S3FD')
+
+ def on_build(self):
+ self.minus = tf.constant([104,117,123], dtype=nn.tf_floatx )
+ self.conv1_1 = nn.Conv2D(3, 64, kernel_size=3, strides=1, padding='SAME')
+ self.conv1_2 = nn.Conv2D(64, 64, kernel_size=3, strides=1, padding='SAME')
+
+ self.conv2_1 = nn.Conv2D(64, 128, kernel_size=3, strides=1, padding='SAME')
+ self.conv2_2 = nn.Conv2D(128, 128, kernel_size=3, strides=1, padding='SAME')
+
+ self.conv3_1 = nn.Conv2D(128, 256, kernel_size=3, strides=1, padding='SAME')
+ self.conv3_2 = nn.Conv2D(256, 256, kernel_size=3, strides=1, padding='SAME')
+ self.conv3_3 = nn.Conv2D(256, 256, kernel_size=3, strides=1, padding='SAME')
+
+ self.conv4_1 = nn.Conv2D(256, 512, kernel_size=3, strides=1, padding='SAME')
+ self.conv4_2 = nn.Conv2D(512, 512, kernel_size=3, strides=1, padding='SAME')
+ self.conv4_3 = nn.Conv2D(512, 512, kernel_size=3, strides=1, padding='SAME')
+
+ self.conv5_1 = nn.Conv2D(512, 512, kernel_size=3, strides=1, padding='SAME')
+ self.conv5_2 = nn.Conv2D(512, 512, kernel_size=3, strides=1, padding='SAME')
+ self.conv5_3 = nn.Conv2D(512, 512, kernel_size=3, strides=1, padding='SAME')
+
+ self.fc6 = nn.Conv2D(512, 1024, kernel_size=3, strides=1, padding=3)
+ self.fc7 = nn.Conv2D(1024, 1024, kernel_size=1, strides=1, padding='SAME')
+
+ self.conv6_1 = nn.Conv2D(1024, 256, kernel_size=1, strides=1, padding='SAME')
+ self.conv6_2 = nn.Conv2D(256, 512, kernel_size=3, strides=2, padding='SAME')
+
+ self.conv7_1 = nn.Conv2D(512, 128, kernel_size=1, strides=1, padding='SAME')
+ self.conv7_2 = nn.Conv2D(128, 256, kernel_size=3, strides=2, padding='SAME')
+
+ self.conv3_3_norm = L2Norm(256)
+ self.conv4_3_norm = L2Norm(512)
+ self.conv5_3_norm = L2Norm(512)
+
+
+ self.conv3_3_norm_mbox_conf = nn.Conv2D(256, 4, kernel_size=3, strides=1, padding='SAME')
+ self.conv3_3_norm_mbox_loc = nn.Conv2D(256, 4, kernel_size=3, strides=1, padding='SAME')
+
+ self.conv4_3_norm_mbox_conf = nn.Conv2D(512, 2, kernel_size=3, strides=1, padding='SAME')
+ self.conv4_3_norm_mbox_loc = nn.Conv2D(512, 4, kernel_size=3, strides=1, padding='SAME')
+
+ self.conv5_3_norm_mbox_conf = nn.Conv2D(512, 2, kernel_size=3, strides=1, padding='SAME')
+ self.conv5_3_norm_mbox_loc = nn.Conv2D(512, 4, kernel_size=3, strides=1, padding='SAME')
+
+ self.fc7_mbox_conf = nn.Conv2D(1024, 2, kernel_size=3, strides=1, padding='SAME')
+ self.fc7_mbox_loc = nn.Conv2D(1024, 4, kernel_size=3, strides=1, padding='SAME')
+
+ self.conv6_2_mbox_conf = nn.Conv2D(512, 2, kernel_size=3, strides=1, padding='SAME')
+ self.conv6_2_mbox_loc = nn.Conv2D(512, 4, kernel_size=3, strides=1, padding='SAME')
+
+ self.conv7_2_mbox_conf = nn.Conv2D(256, 2, kernel_size=3, strides=1, padding='SAME')
+ self.conv7_2_mbox_loc = nn.Conv2D(256, 4, kernel_size=3, strides=1, padding='SAME')
+
+ def forward(self, inp):
+ x, = inp
+ x = x - self.minus
+ x = tf.nn.relu(self.conv1_1(x))
+ x = tf.nn.relu(self.conv1_2(x))
+ x = tf.nn.max_pool(x, [1,2,2,1], [1,2,2,1], "VALID")
+
+ x = tf.nn.relu(self.conv2_1(x))
+ x = tf.nn.relu(self.conv2_2(x))
+ x = tf.nn.max_pool(x, [1,2,2,1], [1,2,2,1], "VALID")
+
+ x = tf.nn.relu(self.conv3_1(x))
+ x = tf.nn.relu(self.conv3_2(x))
+ x = tf.nn.relu(self.conv3_3(x))
+ f3_3 = x
+ x = tf.nn.max_pool(x, [1,2,2,1], [1,2,2,1], "VALID")
+
+ x = tf.nn.relu(self.conv4_1(x))
+ x = tf.nn.relu(self.conv4_2(x))
+ x = tf.nn.relu(self.conv4_3(x))
+ f4_3 = x
+ x = tf.nn.max_pool(x, [1,2,2,1], [1,2,2,1], "VALID")
+
+ x = tf.nn.relu(self.conv5_1(x))
+ x = tf.nn.relu(self.conv5_2(x))
+ x = tf.nn.relu(self.conv5_3(x))
+ f5_3 = x
+ x = tf.nn.max_pool(x, [1,2,2,1], [1,2,2,1], "VALID")
+
+ x = tf.nn.relu(self.fc6(x))
+ x = tf.nn.relu(self.fc7(x))
+ ffc7 = x
+
+ x = tf.nn.relu(self.conv6_1(x))
+ x = tf.nn.relu(self.conv6_2(x))
+ f6_2 = x
+
+ x = tf.nn.relu(self.conv7_1(x))
+ x = tf.nn.relu(self.conv7_2(x))
+ f7_2 = x
+
+ f3_3 = self.conv3_3_norm(f3_3)
+ f4_3 = self.conv4_3_norm(f4_3)
+ f5_3 = self.conv5_3_norm(f5_3)
+
+ cls1 = self.conv3_3_norm_mbox_conf(f3_3)
+ reg1 = self.conv3_3_norm_mbox_loc(f3_3)
+
+ cls2 = tf.nn.softmax(self.conv4_3_norm_mbox_conf(f4_3))
+ reg2 = self.conv4_3_norm_mbox_loc(f4_3)
+
+ cls3 = tf.nn.softmax(self.conv5_3_norm_mbox_conf(f5_3))
+ reg3 = self.conv5_3_norm_mbox_loc(f5_3)
+
+ cls4 = tf.nn.softmax(self.fc7_mbox_conf(ffc7))
+ reg4 = self.fc7_mbox_loc(ffc7)
+
+ cls5 = tf.nn.softmax(self.conv6_2_mbox_conf(f6_2))
+ reg5 = self.conv6_2_mbox_loc(f6_2)
+
+ cls6 = tf.nn.softmax(self.conv7_2_mbox_conf(f7_2))
+ reg6 = self.conv7_2_mbox_loc(f7_2)
+
+ # max-out background label
+ bmax = tf.maximum(tf.maximum(cls1[:,:,:,0:1], cls1[:,:,:,1:2]), cls1[:,:,:,2:3])
+
+ cls1 = tf.concat ([bmax, cls1[:,:,:,3:4] ], axis=-1)
+ cls1 = tf.nn.softmax(cls1)
+
+ return [cls1, reg1, cls2, reg2, cls3, reg3, cls4, reg4, cls5, reg5, cls6, reg6]
+
+ e = None
+ if place_model_on_cpu:
+ e = tf.device("/CPU:0")
+
+ if e is not None: e.__enter__()
+ self.model = S3FD()
+ self.model.load_weights (model_path)
+ if e is not None: e.__exit__(None,None,None)
+
+ self.model.build_for_run ([ ( tf.float32, (None,None,3) ) ])
+
def __enter__(self):
return self
@@ -40,7 +190,7 @@ class S3FDExtractor(object):
input_scale = d / scale_to
input_image = cv2.resize (input_image, ( int(w/input_scale), int(h/input_scale) ), interpolation=cv2.INTER_LINEAR)
- olist = self.model.predict( np.expand_dims(input_image,0) )
+ olist = self.model.run ([ input_image[None,...] ] )
detected_faces = []
for ltrb in self.refine (olist):
@@ -75,8 +225,8 @@ class S3FDExtractor(object):
s_d2 = stride / 2
s_m4 = stride * 4
- for hindex, windex in zip(*np.where(ocls > 0.05)):
- score = ocls[hindex, windex]
+ for hindex, windex in zip(*np.where(ocls[...,1] > 0.05)):
+ score = ocls[hindex, windex, 1]
loc = oreg[hindex, windex, :]
priors = np.array([windex * stride + s_d2, hindex * stride + s_d2, s_m4, s_m4])
priors_2p = priors[2:]
diff --git a/facelib/TernausNet.py b/facelib/TernausNet.py
new file mode 100644
index 0000000..a8fed6c
--- /dev/null
+++ b/facelib/TernausNet.py
@@ -0,0 +1,318 @@
+import os
+import pickle
+from functools import partial
+from pathlib import Path
+
+import cv2
+import numpy as np
+
+from core.interact import interact as io
+from core.leras import nn
+
+"""
+Dataset used to train located in official DFL mega.nz folder
+https://mega.nz/#F!b9MzCK4B!zEAG9txu7uaRUjXz9PtBqg
+
+using https://github.com/ternaus/TernausNet
+TernausNet: U-Net with VGG11 Encoder Pre-Trained on ImageNet for Image Segmentation
+"""
+
+class TernausNet(object):
+ VERSION = 1
+ def __init__ (self, name, resolution, face_type_str, load_weights=True, weights_file_root=None, training=False, place_model_on_cpu=False):
+ nn.initialize()
+ tf = nn.tf
+
+ class Ternaus(nn.ModelBase):
+ def on_build(self, in_ch, ch):
+
+ self.features_0 = nn.Conv2D (in_ch, ch, kernel_size=3, padding='SAME')
+ self.blurpool_0 = nn.BlurPool (filt_size=3)
+
+ self.features_3 = nn.Conv2D (ch, ch*2, kernel_size=3, padding='SAME')
+ self.blurpool_3 = nn.BlurPool (filt_size=3)
+
+ self.features_6 = nn.Conv2D (ch*2, ch*4, kernel_size=3, padding='SAME')
+ self.features_8 = nn.Conv2D (ch*4, ch*4, kernel_size=3, padding='SAME')
+ self.blurpool_8 = nn.BlurPool (filt_size=3)
+
+ self.features_11 = nn.Conv2D (ch*4, ch*8, kernel_size=3, padding='SAME')
+ self.features_13 = nn.Conv2D (ch*8, ch*8, kernel_size=3, padding='SAME')
+ self.blurpool_13 = nn.BlurPool (filt_size=3)
+
+ self.features_16 = nn.Conv2D (ch*8, ch*8, kernel_size=3, padding='SAME')
+ self.features_18 = nn.Conv2D (ch*8, ch*8, kernel_size=3, padding='SAME')
+ self.blurpool_18 = nn.BlurPool (filt_size=3)
+
+ self.conv_center = nn.Conv2D (ch*8, ch*8, kernel_size=3, padding='SAME')
+
+ self.conv1_up = nn.Conv2DTranspose (ch*8, ch*4, kernel_size=3, padding='SAME')
+ self.conv1 = nn.Conv2D (ch*12, ch*8, kernel_size=3, padding='SAME')
+
+ self.conv2_up = nn.Conv2DTranspose (ch*8, ch*4, kernel_size=3, padding='SAME')
+ self.conv2 = nn.Conv2D (ch*12, ch*8, kernel_size=3, padding='SAME')
+
+ self.conv3_up = nn.Conv2DTranspose (ch*8, ch*2, kernel_size=3, padding='SAME')
+ self.conv3 = nn.Conv2D (ch*6, ch*4, kernel_size=3, padding='SAME')
+
+ self.conv4_up = nn.Conv2DTranspose (ch*4, ch, kernel_size=3, padding='SAME')
+ self.conv4 = nn.Conv2D (ch*3, ch*2, kernel_size=3, padding='SAME')
+
+ self.conv5_up = nn.Conv2DTranspose (ch*2, ch//2, kernel_size=3, padding='SAME')
+ self.conv5 = nn.Conv2D (ch//2+ch, ch, kernel_size=3, padding='SAME')
+
+ self.out_conv = nn.Conv2D (ch, 1, kernel_size=3, padding='SAME')
+
+ def forward(self, inp):
+ x, = inp
+
+ x = x0 = tf.nn.relu(self.features_0(x))
+ x = self.blurpool_0(x)
+
+ x = x1 = tf.nn.relu(self.features_3(x))
+ x = self.blurpool_3(x)
+
+ x = tf.nn.relu(self.features_6(x))
+ x = x2 = tf.nn.relu(self.features_8(x))
+ x = self.blurpool_8(x)
+
+ x = tf.nn.relu(self.features_11(x))
+ x = x3 = tf.nn.relu(self.features_13(x))
+ x = self.blurpool_13(x)
+
+ x = tf.nn.relu(self.features_16(x))
+ x = x4 = tf.nn.relu(self.features_18(x))
+ x = self.blurpool_18(x)
+
+ x = self.conv_center(x)
+
+ x = tf.nn.relu(self.conv1_up(x))
+ x = tf.concat( [x,x4], -1)
+ x = tf.nn.relu(self.conv1(x))
+
+ x = tf.nn.relu(self.conv2_up(x))
+ x = tf.concat( [x,x3], -1)
+ x = tf.nn.relu(self.conv2(x))
+
+ x = tf.nn.relu(self.conv3_up(x))
+ x = tf.concat( [x,x2], -1)
+ x = tf.nn.relu(self.conv3(x))
+
+ x = tf.nn.relu(self.conv4_up(x))
+ x = tf.concat( [x,x1], -1)
+ x = tf.nn.relu(self.conv4(x))
+
+ x = tf.nn.relu(self.conv5_up(x))
+ x = tf.concat( [x,x0], -1)
+ x = tf.nn.relu(self.conv5(x))
+
+ x = tf.nn.sigmoid(self.out_conv(x))
+ return x
+
+ if weights_file_root is not None:
+ weights_file_root = Path(weights_file_root)
+ else:
+ weights_file_root = Path(__file__).parent
+ self.weights_path = weights_file_root / ('%s_%d_%s.npy' % (name, resolution, face_type_str) )
+
+ e = tf.device('/CPU:0') if place_model_on_cpu else None
+
+ if e is not None: e.__enter__()
+ self.net = Ternaus(3, 64, name='Ternaus')
+ if load_weights:
+ self.net.load_weights (self.weights_path)
+ else:
+ self.net.init_weights()
+ if e is not None: e.__exit__(None,None,None)
+
+ self.net.build_for_run ( [(tf.float32, (resolution,resolution,3))] )
+
+ if training:
+ raise Exception("training not supported yet")
+
+
+ """
+ if training:
+ try:
+ with open( Path(__file__).parent / 'vgg11_enc_weights.npy', 'rb' ) as f:
+ d = pickle.loads (f.read())
+
+ for i in [0,3,6,8,11,13,16,18]:
+ s = 'features.%d' % i
+
+ self.model.get_layer (s).set_weights ( d[s] )
+ except:
+ io.log_err("Unable to load VGG11 pretrained weights from vgg11_enc_weights.npy")
+
+ conv_weights_list = []
+ for layer in self.model.layers:
+ if 'CA.' in layer.name:
+ conv_weights_list += [layer.weights[0]] #Conv2D kernel_weights
+ CAInitializerMP ( conv_weights_list )
+ """
+
+
+ """
+ if training:
+ inp_t = Input ( (resolution, resolution, 3) )
+ real_t = Input ( (resolution, resolution, 1) )
+ out_t = self.model(inp_t)
+
+ loss = K.mean(10*K.binary_crossentropy(real_t,out_t) )
+
+ out_t_diff1 = out_t[:, 1:, :, :] - out_t[:, :-1, :, :]
+ out_t_diff2 = out_t[:, :, 1:, :] - out_t[:, :, :-1, :]
+
+ total_var_loss = K.mean( 0.1*K.abs(out_t_diff1), axis=[1, 2, 3] ) + K.mean( 0.1*K.abs(out_t_diff2), axis=[1, 2, 3] )
+
+ opt = Adam(lr=0.0001, beta_1=0.5, beta_2=0.999, tf_cpu_mode=2)
+
+ self.train_func = K.function ( [inp_t, real_t], [K.mean(loss)], opt.get_updates( [loss], self.model.trainable_weights) )
+ """
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, exc_type=None, exc_value=None, traceback=None):
+ return False #pass exception between __enter__ and __exit__ to outter level
+
+ def save_weights(self):
+ self.net.save_weights (str(self.weights_path))
+
+ def train(self, inp, real):
+ loss, = self.train_func ([inp, real])
+ return loss
+
+ def extract (self, input_image):
+ input_shape_len = len(input_image.shape)
+ if input_shape_len == 3:
+ input_image = input_image[np.newaxis,...]
+
+ result = np.clip ( self.net.run([input_image]), 0, 1.0 )
+ result[result < 0.1] = 0 #get rid of noise
+
+ if input_shape_len == 3:
+ result = result[0]
+
+ return result
+
+
+
+
+
+
+
+"""
+ self.weights_path = weights_file_root / ('%s_%d_%s.h5' % (name, resolution, face_type_str) )
+
+
+ self.net.build()
+
+
+ self.net.features_0.set_weights ( self.model.get_layer('features.0').get_weights() )
+ self.net.features_3.set_weights ( self.model.get_layer('features.3').get_weights() )
+ self.net.features_6.set_weights ( self.model.get_layer('features.6').get_weights() )
+ self.net.features_8.set_weights ( self.model.get_layer('features.8').get_weights() )
+ self.net.features_11.set_weights ( self.model.get_layer('features.11').get_weights() )
+ self.net.features_13.set_weights ( self.model.get_layer('features.13').get_weights() )
+ self.net.features_16.set_weights ( self.model.get_layer('features.16').get_weights() )
+ self.net.features_18.set_weights ( self.model.get_layer('features.18').get_weights() )
+
+ self.net.conv_center.set_weights ( self.model.get_layer('CA.1').get_weights() )
+
+ self.net.conv1_up.set_weights ( self.model.get_layer('CA.2').get_weights() )
+ self.net.conv1.set_weights ( self.model.get_layer('CA.3').get_weights() )
+
+ self.net.conv2_up.set_weights ( self.model.get_layer('CA.4').get_weights() )
+ self.net.conv2.set_weights ( self.model.get_layer('CA.5').get_weights() )
+
+ self.net.conv3_up.set_weights ( self.model.get_layer('CA.6').get_weights() )
+ self.net.conv3.set_weights ( self.model.get_layer('CA.7').get_weights() )
+
+ self.net.conv4_up.set_weights ( self.model.get_layer('CA.8').get_weights() )
+ self.net.conv4.set_weights ( self.model.get_layer('CA.9').get_weights() )
+
+ self.net.conv5_up.set_weights ( self.model.get_layer('CA.10').get_weights() )
+ self.net.conv5.set_weights ( self.model.get_layer('CA.11').get_weights() )
+
+ self.net.out_conv.set_weights ( self.model.get_layer('CA.12').get_weights() )
+
+ self.net.build_for_run ( [ (tf.float32, (resolution,resolution,3)) ])
+ self.net.save_weights (self.weights_path2)
+
+
+ def extract (self, input_image):
+ input_shape_len = len(input_image.shape)
+ if input_shape_len == 3:
+ input_image = input_image[np.newaxis,...]
+
+ result = np.clip ( self.model.predict( [input_image] ), 0, 1.0 )
+ result[result < 0.1] = 0 #get rid of noise
+
+ if input_shape_len == 3:
+ result = result[0]
+
+ return result
+
+
+ @staticmethod
+ def BuildModel ( resolution, ngf=64):
+ exec( nn.initialize(), locals(), globals() )
+ inp = Input ( (resolution,resolution,3) )
+ x = inp
+ x = TernausNet.Flow(ngf=ngf)(x)
+ model = Model(inp,x)
+ return model
+
+ @staticmethod
+ def Flow(ngf=64):
+ exec( nn.initialize(), locals(), globals() )
+
+ def func(input):
+ x = input
+
+ x0 = x = Conv2D(ngf, kernel_size=3, strides=1, padding='same', activation='relu', name='features.0')(x)
+ x = BlurPool(filt_size=3)(x)
+
+ x1 = x = Conv2D(ngf*2, kernel_size=3, strides=1, padding='same', activation='relu', name='features.3')(x)
+ x = BlurPool(filt_size=3)(x)
+
+ x = Conv2D(ngf*4, kernel_size=3, strides=1, padding='same', activation='relu', name='features.6')(x)
+ x2 = x = Conv2D(ngf*4, kernel_size=3, strides=1, padding='same', activation='relu', name='features.8')(x)
+ x = BlurPool(filt_size=3)(x)
+
+ x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.11')(x)
+ x3 = x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.13')(x)
+ x = BlurPool(filt_size=3)(x)
+
+ x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.16')(x)
+ x4 = x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.18')(x)
+ x = BlurPool(filt_size=3)(x)
+
+ x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', name='CA.1')(x)
+
+ x = Conv2DTranspose (ngf*4, 3, strides=2, padding='same', activation='relu', name='CA.2') (x)
+ x = Concatenate(axis=3)([ x, x4])
+ x = Conv2D (ngf*8, 3, strides=1, padding='same', activation='relu', name='CA.3') (x)
+
+ x = Conv2DTranspose (ngf*4, 3, strides=2, padding='same', activation='relu', name='CA.4') (x)
+ x = Concatenate(axis=3)([ x, x3])
+ x = Conv2D (ngf*8, 3, strides=1, padding='same', activation='relu', name='CA.5') (x)
+
+ x = Conv2DTranspose (ngf*2, 3, strides=2, padding='same', activation='relu', name='CA.6') (x)
+ x = Concatenate(axis=3)([ x, x2])
+ x = Conv2D (ngf*4, 3, strides=1, padding='same', activation='relu', name='CA.7') (x)
+
+ x = Conv2DTranspose (ngf, 3, strides=2, padding='same', activation='relu', name='CA.8') (x)
+ x = Concatenate(axis=3)([ x, x1])
+ x = Conv2D (ngf*2, 3, strides=1, padding='same', activation='relu', name='CA.9') (x)
+
+ x = Conv2DTranspose (ngf // 2, 3, strides=2, padding='same', activation='relu', name='CA.10') (x)
+ x = Concatenate(axis=3)([ x, x0])
+ x = Conv2D (ngf, 3, strides=1, padding='same', activation='relu', name='CA.11') (x)
+
+ return Conv2D(1, 3, strides=1, padding='same', activation='sigmoid', name='CA.12')(x)
+
+
+ return func
+"""
diff --git a/facelib/__init__.py b/facelib/__init__.py
index cde2ab5..ca3292e 100644
--- a/facelib/__init__.py
+++ b/facelib/__init__.py
@@ -1,7 +1,5 @@
from .FaceType import FaceType
-from .DLIBExtractor import DLIBExtractor
-from .MTCExtractor import MTCExtractor
from .S3FDExtractor import S3FDExtractor
from .FANExtractor import FANExtractor
-from .PoseEstimator import PoseEstimator
-from .FaceEnhancer import FaceEnhancer
\ No newline at end of file
+from .FaceEnhancer import FaceEnhancer
+from .TernausNet import TernausNet
\ No newline at end of file
diff --git a/facelib/mmod_human_face_detector.dat b/facelib/mmod_human_face_detector.dat
deleted file mode 100644
index f1f73a5..0000000
Binary files a/facelib/mmod_human_face_detector.dat and /dev/null differ
diff --git a/facelib/mtcnn_onet.h5 b/facelib/mtcnn_onet.h5
deleted file mode 100644
index bd615de..0000000
Binary files a/facelib/mtcnn_onet.h5 and /dev/null differ
diff --git a/facelib/mtcnn_pnet.h5 b/facelib/mtcnn_pnet.h5
deleted file mode 100644
index e13f81b..0000000
Binary files a/facelib/mtcnn_pnet.h5 and /dev/null differ
diff --git a/facelib/mtcnn_rnet.h5 b/facelib/mtcnn_rnet.h5
deleted file mode 100644
index 798a807..0000000
Binary files a/facelib/mtcnn_rnet.h5 and /dev/null differ
diff --git a/nnlib/vgg11_enc_weights.npy b/facelib/vgg11_enc_weights.npy
similarity index 100%
rename from nnlib/vgg11_enc_weights.npy
rename to facelib/vgg11_enc_weights.npy
diff --git a/imagelib/RankSRGAN.h5 b/imagelib/RankSRGAN.h5
deleted file mode 100644
index 765ae2e..0000000
Binary files a/imagelib/RankSRGAN.h5 and /dev/null differ
diff --git a/imagelib/RankSRGAN.py b/imagelib/RankSRGAN.py
deleted file mode 100644
index 529dea8..0000000
--- a/imagelib/RankSRGAN.py
+++ /dev/null
@@ -1,109 +0,0 @@
-import numpy as np
-import cv2
-from pathlib import Path
-from nnlib import nnlib
-from interact import interact as io
-
-class RankSRGAN():
- def __init__(self):
- exec( nnlib.import_all(), locals(), globals() )
-
- class PixelShufflerTorch(KL.Layer):
- def __init__(self, size=(2, 2), data_format='channels_last', **kwargs):
- super(PixelShufflerTorch, self).__init__(**kwargs)
- self.data_format = data_format
- self.size = size
-
- def call(self, inputs):
- input_shape = K.shape(inputs)
- if K.int_shape(input_shape)[0] != 4:
- raise ValueError('Inputs should have rank 4; Received input shape:', str(K.int_shape(inputs)))
-
- batch_size, h, w, c = input_shape[0], input_shape[1], input_shape[2], K.int_shape(inputs)[-1]
- rh, rw = self.size
- oh, ow = h * rh, w * rw
- oc = c // (rh * rw)
-
- out = inputs
- out = K.permute_dimensions(out, (0, 3, 1, 2)) #NCHW
-
- out = K.reshape(out, (batch_size, oc, rh, rw, h, w))
- out = K.permute_dimensions(out, (0, 1, 4, 2, 5, 3))
- out = K.reshape(out, (batch_size, oc, oh, ow))
-
- out = K.permute_dimensions(out, (0, 2, 3, 1))
- return out
-
- def compute_output_shape(self, input_shape):
- if len(input_shape) != 4:
- raise ValueError('Inputs should have rank ' + str(4) + '; Received input shape:', str(input_shape))
-
- height = input_shape[1] * self.size[0] if input_shape[1] is not None else None
- width = input_shape[2] * self.size[1] if input_shape[2] is not None else None
- channels = input_shape[3] // self.size[0] // self.size[1]
-
- if channels * self.size[0] * self.size[1] != input_shape[3]:
- raise ValueError('channels of input and size are incompatible')
-
- return (input_shape[0],
- height,
- width,
- channels)
-
- def get_config(self):
- config = {'size': self.size,
- 'data_format': self.data_format}
- base_config = super(PixelShufflerTorch, self).get_config()
-
- return dict(list(base_config.items()) + list(config.items()))
-
- def res_block(inp, name_prefix):
- x = inp
- x = Conv2D (ndf, kernel_size=3, strides=1, padding='same', activation="relu", name=name_prefix+"0")(x)
- x = Conv2D (ndf, kernel_size=3, strides=1, padding='same', name=name_prefix+"2")(x)
- return Add()([inp,x])
-
- ndf = 64
- nb = 16
- inp = Input ( (None, None,3) )
- x = inp
-
- x = x0 = Conv2D (ndf, kernel_size=3, strides=1, padding='same', name="model0")(x)
- for i in range(nb):
- x = res_block(x, "model1%.2d" %i )
- x = Conv2D (ndf, kernel_size=3, strides=1, padding='same', name="model1160")(x)
- x = Add()([x0,x])
-
- x = ReLU() ( PixelShufflerTorch() ( Conv2D (ndf*4, kernel_size=3, strides=1, padding='same', name="model2")(x) ) )
- x = ReLU() ( PixelShufflerTorch() ( Conv2D (ndf*4, kernel_size=3, strides=1, padding='same', name="model5")(x) ) )
-
- x = Conv2D (ndf, kernel_size=3, strides=1, padding='same', activation="relu", name="model8")(x)
- x = Conv2D (3, kernel_size=3, strides=1, padding='same', name="model10")(x)
- self.model = Model(inp, x )
- self.model.load_weights ( Path(__file__).parent / 'RankSRGAN.h5')
-
- def upscale(self, img, scale=2, is_bgr=True, is_float=True):
- if scale not in [2,4]:
- raise ValueError ("RankSRGAN: supported scale are 2 or 4.")
-
- if not is_bgr:
- img = img[...,::-1]
-
- if not is_float:
- img /= 255.0
-
- h, w = img.shape[:2]
- ch = img.shape[2] if len(img.shape) >= 3 else 1
-
- output = self.model.predict([img[None,...]])[0]
-
- if scale == 2:
- output = cv2.resize (output, (w*scale, h*scale), cv2.INTER_CUBIC)
-
- if not is_float:
- output = np.clip (output * 255.0, 0, 255.0)
-
- if not is_bgr:
- output = output[...,::-1]
-
- return output
\ No newline at end of file
diff --git a/main.py b/main.py
index 2166817..f7a78ba 100644
--- a/main.py
+++ b/main.py
@@ -1,14 +1,17 @@
if __name__ == "__main__":
+ from core.leras import nn
+ nn.initialize_main_env()
+
import os
import sys
import time
import argparse
import multiprocessing
multiprocessing.set_start_method("spawn")
- from utils import Path_utils
- from utils import os_utils
+ from core import pathex
+ from core import osex
from pathlib import Path
- from interact import interact as io
+ from core.interact import interact as io
if sys.version_info[0] < 3 or (sys.version_info[0] == 3 and sys.version_info[1] < 6):
raise Exception("This program requires at least Python 3.6")
@@ -21,36 +24,37 @@ if __name__ == "__main__":
subparsers = parser.add_subparsers()
def process_extract(arguments):
- os_utils.set_process_lowest_prio()
+ osex.set_process_lowest_prio()
from mainscripts import Extractor
- Extractor.main( arguments.input_dir,
- arguments.output_dir,
- arguments.debug_dir,
- arguments.detector,
- arguments.manual_fix,
- arguments.manual_output_debug_fix,
- arguments.manual_window_size,
- face_type=arguments.face_type,
- device_args={'cpu_only' : arguments.cpu_only,
- 'multi_gpu' : arguments.multi_gpu,
- }
+ Extractor.main( detector = arguments.detector,
+ input_path = Path(arguments.input_dir),
+ output_path = Path(arguments.output_dir),
+ output_debug = arguments.output_debug,
+ manual_fix = arguments.manual_fix,
+ manual_output_debug_fix = arguments.manual_output_debug_fix,
+ manual_window_size = arguments.manual_window_size,
+ face_type = arguments.face_type,
+ cpu_only = arguments.cpu_only,
+ force_gpu_idxs = [ int(x) for x in arguments.force_gpu_idxs.split(',') ] if arguments.force_gpu_idxs is not None else None,
)
p = subparsers.add_parser( "extract", help="Extract the faces from a pictures.")
+ p.add_argument('--detector', dest="detector", choices=['s3fd','manual'], default=None, help="Type of detector.")
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
p.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the extracted files will be stored.")
- p.add_argument('--debug-dir', action=fixPathAction, dest="debug_dir", help="Writes debug images to this directory.")
+ p.add_argument('--output-debug', action="store_true", dest="output_debug", default=None, help="Writes debug images to 
+
