github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-07-06 13:02:15 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

added 'sort by vggface': sorting by face similarity using VGGFace model.

Browse source 
Requires 4GB+ VRAM and internet connection for the first run.
This commit is contained in:
Colombo 2019-10-23 15:06:39 +04:00
parent 0d3b25812d
commit 734d97d729
8 changed files with 186 additions and 43 deletions
Download patch file Download diff file Expand all files Collapse all files

2
main.py
View file

@ -112,7 +112,7 @@ if __name__ == "__main__":
p = subparsers.add_parser( "sort", help="Sort faces in a directory.")
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('--by', required=True, dest="sort_by_method", choices=("blur", "face", "face-dissim", "face-yaw", "face-pitch", "hist", "hist-dissim", "brightness", "hue", "black", "origname", "oneface", "final", "final-no-blur", "test"), help="Method of sorting. 'origname' sort by original filename to recover original sequence." )
p.add_argument('--by', required=True, dest="sort_by_method", choices=("blur", "face", "face-dissim", "face-yaw", "face-pitch", "hist", "hist-dissim", "brightness", "hue", "black", "origname", "oneface", "final", "final-no-blur", "vggface", "test"), help="Method of sorting. 'origname' sort by original filename to recover original sequence." )
p.set_defaults (func=process_sort)
def process_util(arguments):

117
mainscripts/Sorter.py
View file

@ -1,19 +1,26 @@
import os
import sys
import operator
import numpy as np
import cv2
from shutil import copyfile
from pathlib import Path
from utils import Path_utils
from utils.DFLPNG import DFLPNG
from utils.DFLJPG import DFLJPG
from utils.cv2_utils import *
from facelib import LandmarksProcessor
from joblib import Subprocessor
import multiprocessing
from interact import interact as io
import operator
import sys
from pathlib import Path
from shutil import copyfile
import cv2
import numpy as np
from numpy import linalg as npla
import imagelib
from facelib import LandmarksProcessor
from functools import cmp_to_key
from imagelib import estimate_sharpness
from interact import interact as io
from joblib import Subprocessor
from nnlib import VGGFace
from utils import Path_utils
from utils.cv2_utils import *
from utils.DFLJPG import DFLJPG
from utils.DFLPNG import DFLPNG
class BlurEstimatorSubprocessor(Subprocessor):
class Cli(Subprocessor.Cli):
@ -772,24 +779,97 @@ def sort_final(input_path, include_by_blur=True):
for pg in range(pitch_grads):
img_list = pitch_sample_list[pg]
if img_list is None:
continue
continue
final_img_list += [ img_list.pop(0) ]
if len(img_list) == 0:
pitch_sample_list[pg] = None
pitch_sample_list[pg] = None
n -= 1
if n == 0:
break
if n_prev == n:
break
if n_prev == n:
break
for pg in range(pitch_grads):
img_list = pitch_sample_list[pg]
if img_list is None:
continue
continue
trash_img_list += img_list
return final_img_list, trash_img_list
def sort_by_vggface(input_path):
io.log_info ("Sorting by face similarity using VGGFace model...")
model = VGGFace()
final_img_list = []
trash_img_list = []
image_paths = Path_utils.get_image_paths(input_path)
img_list = [ (x,) for x in image_paths ]
img_list_len = len(img_list)
img_list_range = [*range(img_list_len)]
feats = [None]*img_list_len
for i in io.progress_bar_generator(img_list_range, "Loading"):
img = cv2_imread( img_list[i][0] ).astype(np.float32)
img = imagelib.normalize_channels (img, 3)
img = cv2.resize (img, (224,224) )
img = img[..., ::-1]
img[..., 0] -= 93.5940
img[..., 1] -= 104.7624
img[..., 2] -= 129.1863
feats[i] = model.predict( img[None,...] )[0]
tmp = np.zeros( (img_list_len,) )
float_inf = float("inf")
for i in io.progress_bar_generator ( range(img_list_len-1), "Sorting" ):
i_feat = feats[i]
for j in img_list_range:
tmp[j] = npla.norm(i_feat-feats[j]) if j >= i+1 else float_inf
idx = np.argmin(tmp)
img_list[i+1], img_list[idx] = img_list[idx], img_list[i+1]
feats[i+1], feats[idx] = feats[idx], feats[i+1]
return img_list, trash_img_list
"""
img_list_len = len(img_list)
for i in io.progress_bar_generator ( range(img_list_len-1), "Sorting" ):
a = []
i_1 = img_list[i][1]
for j in range(i+1, img_list_len):
a.append ( [ j, np.linalg.norm(i_1-img_list[j][1]) ] )
x = sorted(a, key=operator.itemgetter(1) )[0][0]
saved = img_list[i+1]
img_list[i+1] = img_list[x]
img_list[x] = saved
q = np.array ( [ x[1] for x in img_list ] )
for i in io.progress_bar_generator ( range(img_list_len-1), "Sorting" ):
a = np.linalg.norm( q[i] - q[i+1:], axis=1 )
a = i+1+np.argmin(a)
saved = img_list[i+1]
img_list[i+1] = img_list[a]
img_list[a] = saved
saved = q[i+1]
q[i+1] = q[a]
q[a] = saved
"""
def final_process(input_path, img_list, trash_img_list):
if len(trash_img_list) != 0:
parent_input_path = input_path.parent
@ -851,6 +931,7 @@ def main (input_path, sort_by_method):
elif sort_by_method == 'black': img_list = sort_by_black (input_path)
elif sort_by_method == 'origname': img_list, trash_img_list = sort_by_origname (input_path)
elif sort_by_method == 'oneface': img_list, trash_img_list = sort_by_oneface_in_image (input_path)
elif sort_by_method == 'vggface': img_list, trash_img_list = sort_by_vggface (input_path)
elif sort_by_method == 'final': img_list, trash_img_list = sort_final (input_path)
elif sort_by_method == 'final-no-blur': img_list, trash_img_list = sort_final (input_path, include_by_blur=False)

4
mainscripts/dev_misc.py
View file

@ -37,6 +37,10 @@ def extract_vggface2_dataset(input_dir, device_args={} ):
cur_input_path = input_path / dir_name
cur_output_path = output_path / dir_name
l = len(Path_utils.get_image_paths(cur_input_path))
if l < 250 or l > 350:
continue
io.log_info (f"Processing: {str(cur_input_path)} ")

9
models/Model_DEV_FUNIT/Model.py
View file

@ -42,7 +42,7 @@ class FUNITModel(ModelBase):
#override
def onInitialize(self, batch_size=-1, **in_options):
exec(nnlib.code_import_all, locals(), globals())
self.set_vram_batch_requirements({4:16})
self.set_vram_batch_requirements({4:16,11:24})
resolution = self.options['resolution']
face_type = FaceType.FULL if self.options['face_type'] == 'f' else FaceType.HALF
@ -75,7 +75,8 @@ class FUNITModel(ModelBase):
face_type = t.FACE_TYPE_FULL if self.options['face_type'] == 'f' else t.FACE_TYPE_HALF
output_sample_types=[ {'types': (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'resolution':128, 'normalize_tanh':True} ]
output_sample_types1=[ {'types': (t.IMG_SOURCE, face_type, t.MODE_BGR), 'resolution':128, 'normalize_tanh':True} ]
self.set_training_data_generators ([
SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(random_flip=True),
@ -87,11 +88,11 @@ class FUNITModel(ModelBase):
SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(random_flip=True),
output_sample_types=output_sample_types, person_id_mode=True ),
output_sample_types=output_sample_types1, person_id_mode=True ),
SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(random_flip=True),
output_sample_types=output_sample_types, person_id_mode=True ),
output_sample_types=output_sample_types1, person_id_mode=True ),
])
#override

29
nnlib/FUNIT.py
View file

@ -162,10 +162,6 @@ class FUNIT(object):
for w in weights_list:
K.set_value( w, K.get_value(initer(K.int_shape(w))) )
#if not self.is_first_run():
# self.load_weights_safe(self.get_model_filename_list())
if load_weights_locally:
pass
@ -188,9 +184,6 @@ class FUNIT(object):
[self.D_opt, 'D_opt.h5'],
]
#def save_weights(self):
# self.model.save_weights (str(self.weights_path))
def train(self, xa,la,xb,lb):
D_loss, = self.D_train ([xa,la,xb,lb])
G_loss, = self.G_train ([xa,la,xb,lb])
@ -209,17 +202,17 @@ class FUNIT(object):
def ResBlock(dim):
def func(input):
x = input
x = Conv2D(dim, 3, strides=1, padding='valid')(ZeroPadding2D(1)(x))
x = Conv2D(dim, 3, strides=1, padding='same')(x)
x = InstanceNormalization()(x)
x = ReLU()(x)
x = Conv2D(dim, 3, strides=1, padding='valid')(ZeroPadding2D(1)(x))
x = Conv2D(dim, 3, strides=1, padding='same')(x)
x = InstanceNormalization()(x)
return Add()([x,input])
return func
def func(x):
x = Conv2D (nf, kernel_size=7, strides=1, padding='valid')(ZeroPadding2D(3)(x))
x = Conv2D (nf, kernel_size=7, strides=1, padding='same')(x)
x = InstanceNormalization()(x)
x = ReLU()(x)
for i in range(downs):
@ -237,11 +230,11 @@ class FUNIT(object):
exec (nnlib.import_all(), locals(), globals())
def func(x):
x = Conv2D (nf, kernel_size=7, strides=1, padding='valid', activation='relu')(ZeroPadding2D(3)(x))
x = Conv2D (nf, kernel_size=7, strides=1, padding='same', activation='relu')(x)
for i in range(downs):
x = Conv2D (nf * min ( 4, 2**(i+1) ), kernel_size=4, strides=2, padding='valid', activation='relu')(ZeroPadding2D(1)(x))
x = GlobalAveragePooling2D()(x)
x = Dense(nf)(x)
x = Dense(latent_dim)(x)
return x
return func
@ -250,16 +243,14 @@ class FUNIT(object):
def DecoderFlow(ups, n_res_blks=2, mlp_blks=2, subpixel_decoder=False ):
exec (nnlib.import_all(), locals(), globals())
def ResBlock(dim):
def func(input):
inp, mlp = input
x = inp
x = Conv2D(dim, 3, strides=1, padding='valid')(ZeroPadding2D(1)(x))
x = Conv2D(dim, 3, strides=1, padding='same')(x)
x = FUNITAdain(kernel_initializer='he_normal')([x,mlp])
x = ReLU()(x)
x = Conv2D(dim, 3, strides=1, padding='valid')(ZeroPadding2D(1)(x))
x = Conv2D(dim, 3, strides=1, padding='same')(x)
x = FUNITAdain(kernel_initializer='he_normal')([x,mlp])
return Add()([x,inp])
return func
@ -280,16 +271,16 @@ class FUNIT(object):
for i in range(ups):
if subpixel_decoder:
x = Conv2D (4* (nf // 2**(i+1)), kernel_size=3, strides=1, padding='valid')(ZeroPadding2D(1)(x))
x = Conv2D (4* (nf // 2**(i+1)), kernel_size=3, strides=1, padding='same')(x)
x = SubpixelUpscaler()(x)
else:
x = UpSampling2D()(x)
x = Conv2D (nf // 2**(i+1), kernel_size=5, strides=1, padding='valid')(ZeroPadding2D(2)(x))
x = Conv2D (nf // 2**(i+1), kernel_size=5, strides=1, padding='same')(x)
x = InstanceNormalization()(x)
x = ReLU()(x)
rgb = Conv2D (3, kernel_size=7, strides=1, padding='valid', activation='tanh')(ZeroPadding2D(3)(x))
rgb = Conv2D (3, kernel_size=7, strides=1, padding='same', activation='tanh')(x)
return rgb
return func

64
nnlib/VGGFace.py Normal file
View file

@ -0,0 +1,64 @@
from nnlib import nnlib
def VGGFace():
exec(nnlib.import_all(), locals(), globals())
img_input = Input(shape=(224,224,3) )
# Block 1
x = Conv2D(64, (3, 3), activation='relu', padding='same', name='conv1_1')(
img_input)
x = Conv2D(64, (3, 3), activation='relu', padding='same', name='conv1_2')(x)
x = MaxPooling2D((2, 2), strides=(2, 2), name='pool1')(x)
# Block 2
x = Conv2D(128, (3, 3), activation='relu', padding='same', name='conv2_1')(
x)
x = Conv2D(128, (3, 3), activation='relu', padding='same', name='conv2_2')(
x)
x = MaxPooling2D((2, 2), strides=(2, 2), name='pool2')(x)
# Block 3
x = Conv2D(256, (3, 3), activation='relu', padding='same', name='conv3_1')(
x)
x = Conv2D(256, (3, 3), activation='relu', padding='same', name='conv3_2')(
x)
x = Conv2D(256, (3, 3), activation='relu', padding='same', name='conv3_3')(
x)
x = MaxPooling2D((2, 2), strides=(2, 2), name='pool3')(x)
# Block 4
x = Conv2D(512, (3, 3), activation='relu', padding='same', name='conv4_1')(
x)
x = Conv2D(512, (3, 3), activation='relu', padding='same', name='conv4_2')(
x)
x = Conv2D(512, (3, 3), activation='relu', padding='same', name='conv4_3')(
x)
x = MaxPooling2D((2, 2), strides=(2, 2), name='pool4')(x)
# Block 5
x = Conv2D(512, (3, 3), activation='relu', padding='same', name='conv5_1')(
x)
x = Conv2D(512, (3, 3), activation='relu', padding='same', name='conv5_2')(
x)
x = Conv2D(512, (3, 3), activation='relu', padding='same', name='conv5_3')(
x)
x = MaxPooling2D((2, 2), strides=(2, 2), name='pool5')(x)
# Classification block
x = Flatten(name='flatten')(x)
x = Dense(4096, name='fc6')(x)
x = Activation('relu', name='fc6/relu')(x)
x = Dense(4096, name='fc7')(x)
x = Activation('relu', name='fc7/relu')(x)
x = Dense(2622, name='fc8')(x)
x = Activation('softmax', name='fc8/softmax')(x)
model = Model(img_input, x, name='vggface_vgg16')
weights_path = keras.utils.data_utils.get_file('rcmalli_vggface_tf_vgg16.h5',
'https://github.com/rcmalli/keras-vggface/releases/download/v2.0/rcmalli_vggface_tf_vgg16.h5')
model.load_weights(weights_path, by_name=True)
return model

3
nnlib/__init__.py
View file

@ -1,3 +1,4 @@
from .nnlib import nnlib
from .FUNIT import FUNIT
from .TernausNet import TernausNet
from .TernausNet import TernausNet
from .VGGFace import VGGFace

1
nnlib/nnlib.py
View file

@ -63,6 +63,7 @@ UpSampling2D = KL.UpSampling2D
BatchNormalization = KL.BatchNormalization
PixelNormalization = nnlib.PixelNormalization
Activation = KL.Activation
LeakyReLU = KL.LeakyReLU
ELU = KL.ELU
ReLU = KL.ReLU