1st Working version

lib/faces_detect.py

@@ -7,7 +7,7 @@ def detect_faces(frame, model="hog", verbose=False, rotation=0):
     landmarks = _raw_face_landmarks(frame, face_locations)
 
     for ((y, right, bottom, x), landmarks) in zip(face_locations, landmarks):
-        yield DetectedFace(frame[y: bottom, x: right], rotation, x, right - x, y, bottom - y, landmarks)
+        yield DetectedFace(frame[y: bottom, x: right], rotation, x, right - x, y, bottom - y, [(p.x, p.y) for p in landmarks.parts()])
 
 # Copy/Paste (mostly) from private method in face_recognition
 predictor_68_point_model = face_recognition_models.pose_predictor_model_location()
@@ -22,14 +22,15 @@ def _css_to_rect(css):
 # end of Copy/Paste
 
 class DetectedFace(object):
-    def __init__(self, image=None, r=0, x=None, w=None, y=None, h=None, landmarks=None, landmarksXY=None):
+    def __init__(self, image=None, r=0, x=None, w=None, y=None, h=None, landmarksXY=None, cropped=None):
         self.image = image
         self.r = r
         self.x = x
         self.w = w
         self.y = y
         self.h = h
-        self.landmarksXY = [(p.x, p.y) for p in landmarks.parts()]
+        self.cropped = cropped
+        self.landmarksXY = landmarksXY
 
     def landmarksAsXY(self):
         return self.landmarksXY

plugins/Model_DFaker/Model.py

@@ -73,8 +73,9 @@ class Model(AutoEncoder):
     def initModel(self):
         optimizer = Adam( lr=5e-5, beta_1=0.5, beta_2=0.999 )
 
-        print(self.encoder.summary()) 
-        print(self.decoder_A.summary())
+        # print(self.encoder.summary()) 
+        # print(self.decoder_A.summary())
+        # print(self.decoder_B.summary())
 
         x1 = Input( shape=IMAGE_SHAPE )
         x2 = Input( shape=IMAGE_SHAPE )
@@ -84,11 +85,15 @@ def initModel(self):
         self.autoencoder_A = KerasModel( [x1,m1], self.decoder_A( self.encoder(x1) ) )
         self.autoencoder_B = KerasModel( [x2,m2], self.decoder_B( self.encoder(x2) ) )
 
-        #autoencoder_A = multi_gpu_model( autoencoder_A ,2)
-        #autoencoder_B = multi_gpu_model( autoencoder_B ,2)
+        # print(self.autoencoder_A.summary())
+        # print(self.autoencoder_B.summary())
 
-        o1,om1  = self.decoder_A( self.encoder(x1))
-        o2,om2  = self.decoder_B( self.encoder(x2))
+        if self.gpus > 1:
+            self.autoencoder_A = multi_gpu_model( self.autoencoder_A ,self.gpus)
+            self.autoencoder_B = multi_gpu_model( self.autoencoder_B ,self.gpus)
+
+        # o1,om1  = self.decoder_A( self.encoder(x1))
+        # o2,om2  = self.decoder_B( self.encoder(x2))
 
         DSSIM = DSSIMObjective()
         self.autoencoder_A.compile( optimizer=optimizer, loss=[ penalized_loss(m1, DSSIM),'mse'] )

plugins/Model_DFaker/Trainer.py

@@ -1,28 +1,29 @@
-
+import cv2
 import time
 import numpy
 import random
 
+from lib.utils import BackgroundGenerator
 from lib.training_data import TrainingDataGenerator, stack_images
 from .image_augmentation import random_transform, random_warp_src_dest
 from .utils import load_images_aligned
 
 class DFTrainingDataGenerator(TrainingDataGenerator):
-    def __init__(self, random_transform_args, coverage, scale, zoom):
-        super().__init__(random_transform_args, coverage, scale, zoom)
+    def __init__(self, random_transform_args, coverage):
+        super().__init__(random_transform_args, coverage)
 
     def minibatchAB(self, images, srcPoints, dstPoints, batch_size ):
         batch = BackgroundGenerator(self.minibatch(images, srcPoints, dstPoints, batch_size), 1)
         return batch.iterator()
 
-    def minibatch(self, data, srcPoints, dstPoints, batchsize):
+    def minibatch(self, images, srcPoints, dstPoints, batch_size):
         epoch = 0
         while True:
             epoch+=1
             indices = numpy.random.choice(range(0,images.shape[0]),size=batch_size,replace=False)
             for i,index in enumerate(indices):
                 image = images[index]
-                image = random_transform( image, **random_transform_args )
+                image = random_transform( image, **self.random_transform_args )
 
                 closest = ( numpy.mean(numpy.square(srcPoints[index]-dstPoints),axis=(1,2)) ).argsort()[:10]
                 closest = numpy.random.choice(closest)
@@ -42,7 +43,7 @@ def minibatch(self, data, srcPoints, dstPoints, batchsize):
                 target_images[i] = target_img
                 mask_images[i]   = mask_image
 
-            yield epoch, warped_img, target_img, mask_image
+            yield epoch, warped_images, target_images, mask_images
 
 class Trainer():
     random_transform_args = {
@@ -56,7 +57,7 @@ def __init__(self, model, fn_A, fn_B, batch_size, *args):
         self.batch_size = batch_size
         self.model = model
 
-        generator = TrainingDataGenerator(self.random_transform_args, 160)
+        generator = DFTrainingDataGenerator(self.random_transform_args, 160)
 
         minImages = 2000#min(len(fn_A),len(fn_B))*20
 
@@ -81,13 +82,21 @@ def train_one_step(self, iter, viewer):
 
         #omask = numpy.ones((target_A.shape[0],64,64,1),float)
 
-        loss_A = autoencoder_A.train_on_batch([warped_A,mask_A], [target_A,mask_A])
-        loss_B = autoencoder_B.train_on_batch([warped_B,mask_B], [target_B,mask_B])
-        print("[{0}] [#{1:05d}] loss_A: {2:.5f}, loss_B: {3:.5f}".format(time.strftime("%H:%M:%S"), iter, loss_A, loss_B),
+        loss_A = self.model.autoencoder_A.train_on_batch([warped_A,mask_A], [target_A,mask_A])
+        loss_B = self.model.autoencoder_B.train_on_batch([warped_B,mask_B], [target_B,mask_B])
+
+        print("Loss A [{}] Loss B [{}]".format(loss_A,loss_B),
             end='\r')
 
         if viewer is not None:
             viewer(self.show_sample(target_A[0:8,:,:,:3], target_B[0:8,:,:,:3]), "training")
+            viewer(self.show_warped(warped_A[:6],warped_B[:6]), "warped")
+
+    def show_warped(self, warped_A, warped_B):
+        figWarped = numpy.stack([warped_A, warped_B],axis=0 )
+        figWarped = numpy.clip( figWarped * 255, 0, 255 ).astype('uint8')
+        figWarped = stack_images( figWarped )
+        return figWarped
 
     def show_sample(self, test_A, test_B):
         test_A_i = []
@@ -100,19 +109,14 @@ def show_sample(self, test_A, test_B):
         for i in test_B:
             test_B_i.append(cv2.resize(i,(64,64),cv2.INTER_AREA))
         test_B_i = numpy.array(test_B_i).reshape((-1,64,64,3))
-
-        figWarped = numpy.stack([warped_A[:6],warped_B[:6]],axis=0 )
-        figWarped = numpy.clip( figWarped * 255, 0, 255 ).astype('uint8')
-        figWarped = stack_images( figWarped )
-        cv2.imshow( "w", figWarped )
 
         zmask = numpy.zeros((test_A.shape[0],128,128,1),float)
 
-        pred_a_a,pred_a_a_m = autoencoder_A.predict([test_A_i,zmask])
-        pred_b_a,pred_b_a_m = autoencoder_B.predict([test_A_i,zmask])
+        pred_a_a,pred_a_a_m = self.model.autoencoder_A.predict([test_A_i,zmask])
+        pred_b_a,pred_b_a_m = self.model.autoencoder_B.predict([test_A_i,zmask])
 
-        pred_a_b,pred_a_b_m = autoencoder_A.predict([test_B_i,zmask])
-        pred_b_b,pred_b_b_m = autoencoder_B.predict([test_B_i,zmask])
+        pred_a_b,pred_a_b_m = self.model.autoencoder_A.predict([test_B_i,zmask])
+        pred_b_b,pred_b_b_m = self.model.autoencoder_B.predict([test_B_i,zmask])
 
         pred_a_a = pred_a_a[0:18,:,:,:3]
         pred_a_b = pred_a_b[0:18,:,:,:3]

plugins/Model_DFaker/utils.py

@@ -3,8 +3,10 @@
 import os
 import json
 
-from tqdm import tqdm
+from lib.faces_detect import DetectedFace
 
+from tqdm import tqdm
+from lib.aligner import get_align_mat
 #guideShades = numpy.linspace(20,250,68)
 
 def load_images_aligned(image_paths):
@@ -17,14 +19,17 @@ def load_images_aligned(image_paths):
 
 
   pbar = tqdm(alignments)
-  for original,cropped,mat,points in pbar:
-    pbar.set_description('loading '+basePath)
-    cropped = os.path.split(cropped)[1]
-    cropped = os.path.join(basePath,cropped)
-    if cropped in image_paths and os.path.exists(cropped):
-        image, facepoints = get_image( cropped, mat, points )
-        all_images.append( image  )
-        landmarks.append( facepoints )
+  for original in pbar:
+    faces = alignments[original]
+    for rawface in faces:
+        detected_face = DetectedFace(**rawface)
+        cropped = detected_face.cropped
+        points = detected_face.landmarksXY
+        mat = get_align_mat( detected_face )
+        if os.path.exists(cropped):
+            image, facepoints = load_image( cropped, mat, points )
+            all_images.append( image  )
+            landmarks.append( facepoints )
 
   return numpy.array(all_images),numpy.array(landmarks)
 

scripts/extract.py

@@ -128,8 +128,10 @@ def handleImage(self, image, filename):
 
             resized_image = self.extractor.extract(image, face, 256)
             output_file = get_folder(self.output_dir) / Path(filename).stem
-            cv2.imwrite('{}_{}{}'.format(str(output_file), str(idx), Path(filename).suffix), resized_image)
+            output_file = '{}_{}{}'.format(str(output_file), str(idx), Path(filename).suffix)
+            cv2.imwrite(output_file, resized_image)
             f = {
+                "cropped": output_file,
                 "r": face.r,
                 "x": face.x,
                 "w": face.w,