Skip to content
Permalink
Browse files

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

Requires 4GB+ VRAM and internet connection for the first run.
  • Loading branch information...
iperov committed Oct 23, 2019
1 parent 0d3b258 commit 734d97d729e49922fe204c898b8bdf1ab3bcda7b
Showing with 184 additions and 41 deletions.
  1. +1 −1 main.py
  2. +97 −16 mainscripts/Sorter.py
  3. +4 −0 mainscripts/dev_misc.py
  4. +5 −4 models/Model_DEV_FUNIT/Model.py
  5. +10 −19 nnlib/FUNIT.py
  6. +64 −0 nnlib/VGGFace.py
  7. +2 −1 nnlib/__init__.py
  8. +1 −0 nnlib/nnlib.py
@@ -112,7 +112,7 @@ def process_sort(arguments):

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):
@@ -1,19 +1,26 @@
import os
import sys
import multiprocessing
import operator
import numpy as np
import cv2
from shutil import copyfile
import sys
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 shutil import copyfile

import cv2
import numpy as np
from numpy import linalg as npla

import imagelib
from facelib import LandmarksProcessor
from joblib import Subprocessor
import multiprocessing
from interact import interact as io
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)

@@ -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)} ")

@@ -42,7 +42,7 @@ def onInitializeOptions(self, is_first_run, ask_override):
#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 @@ def onInitialize(self, batch_size=-1, **in_options):
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 @@ def onInitialize(self, batch_size=-1, **in_options):

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
@@ -162,10 +162,6 @@ def dis_gather_mean(x,l, func=None, acc_func=None):
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 @@ def get_model_filename_list(self):
[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 @@ def ContentEncoderFlow(downs=2, nf=64, n_res_blks=2):
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 @@ def ClassModelEncoderFlow(downs=4, nf=64, latent_dim=64):
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 @@ def func(x):
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 @@ def func(inputs):
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
@@ -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
@@ -1,3 +1,4 @@
from .nnlib import nnlib
from .FUNIT import FUNIT
from .TernausNet import TernausNet
from .TernausNet import TernausNet
from .VGGFace import VGGFace
@@ -63,6 +63,7 @@ class nnlib(object):
BatchNormalization = KL.BatchNormalization
PixelNormalization = nnlib.PixelNormalization
Activation = KL.Activation
LeakyReLU = KL.LeakyReLU
ELU = KL.ELU
ReLU = KL.ReLU

0 comments on commit 734d97d

Please sign in to comment.
You can’t perform that action at this time.