DeepFaceLab/mainscripts/Converter.py

import sys
import os
import traceback
from pathlib import Path
from utils import Path_utils
import cv2
from tqdm import tqdm
from utils.DFLPNG import DFLPNG
from utils.DFLJPG import DFLJPG
from utils.cv2_utils import *
from utils import image_utils
import shutil
import numpy as np
import time
import multiprocessing
from models import ConverterBase

class model_process_predictor(object):
    def __init__(self, sq, cq, lock):
        self.sq = sq
        self.cq = cq
        self.lock = lock
        
    def __call__(self, face):
        self.lock.acquire()
        
        self.sq.put ( {'op': 'predict', 'face' : face} )
        while True:
            if not self.cq.empty():
                obj = self.cq.get()
                obj_op = obj['op']
                if obj_op == 'predict_result':
                    self.lock.release()
                    return obj['result']
            time.sleep(0.005)
        
def model_process(stdin_fd, model_name, model_dir, in_options, sq, cq):
    sys.stdin = os.fdopen(stdin_fd)
    
    try:    
        model_path = Path(model_dir)
        
        import models 
        model = models.import_model(model_name)(model_path, **in_options)
        converter = model.get_converter(**in_options)
        converter.dummy_predict()
        
        cq.put ( {'op':'init', 'converter' : converter.copy_and_set_predictor( None ) } )

        while True:
            while not sq.empty():
                obj = sq.get()
                obj_op = obj['op']
                if obj_op == 'predict':
                    result = converter.predictor ( obj['face'] )
                    cq.put ( {'op':'predict_result', 'result':result} )
            time.sleep(0.005)        
    except Exception as e:
        print ( 'Error: %s' % (str(e)))
        traceback.print_exc()
    
    
    
from utils.SubprocessorBase import SubprocessorBase
class ConvertSubprocessor(SubprocessorBase):

    #override
    def __init__(self, converter, input_path_image_paths, output_path, alignments, debug = False, **in_options): 
        super().__init__('Converter', 86400 if debug == True else 60)    
        self.converter = converter
        self.input_path_image_paths = input_path_image_paths
        self.output_path = output_path
        self.alignments = alignments
        self.debug = debug
        self.in_options = in_options
        self.input_data = self.input_path_image_paths
        self.files_processed = 0
        self.faces_processed = 0
        
    #override
    def process_info_generator(self):
        r = [0] if self.debug else range( min(multiprocessing.cpu_count(), 6) )
        
        

        for i in r:
            yield 'CPU%d' % (i), {}, {'device_idx': i,
                                      'device_name': 'CPU%d' % (i), 
                                      'converter' : self.converter, 
                                      'output_dir' : str(self.output_path), 
                                      'alignments' : self.alignments,
                                      'debug': self.debug,
                                      'in_options': self.in_options
                                      }
     
    #override
    def get_no_process_started_message(self):
        return 'Unable to start CPU processes.'
        
    #override
    def onHostGetProgressBarDesc(self):
        return "Converting"
        
    #override
    def onHostGetProgressBarLen(self):
        return len (self.input_data)
        
    #override
    def onHostGetData(self, host_dict):
        if len (self.input_data) > 0:
            return self.input_data.pop(0)            
        return None
    
    #override
    def onHostDataReturn (self, host_dict, data):
        self.input_data.insert(0, data)   
        
    #overridable
    def onClientGetDataName (self, data):
        #return string identificator of your data
        return data
        
    #override
    def onClientInitialize(self, client_dict):
        print ('Running on %s.' % (client_dict['device_name']) )
        self.device_idx  = client_dict['device_idx']
        self.device_name = client_dict['device_name']
        self.converter   = client_dict['converter']
        self.output_path = Path(client_dict['output_dir']) if 'output_dir' in client_dict.keys() else None        
        self.alignments  = client_dict['alignments']
        self.debug       = client_dict['debug']
        
        from nnlib import nnlib          
        #model process ate all GPU mem,
        #so we cannot use GPU for any TF operations in converter processes (for example image_utils.TFLabConverter)
        #therefore forcing active_DeviceConfig to CPU only
        nnlib.active_DeviceConfig = nnlib.DeviceConfig (cpu_only=True)
        
        return None

    #override
    def onClientFinalize(self):
        pass
        
    #override
    def onClientProcessData(self, data):
        filename_path = Path(data)

        files_processed = 1
        faces_processed = 0
            
        output_filename_path = self.output_path / (filename_path.stem + '.png')

        if self.converter.get_mode() == ConverterBase.MODE_FACE and filename_path.stem not in self.alignments.keys():                    
            if not self.debug:
                print ( 'no faces found for %s, copying without faces' % (filename_path.name) )
                shutil.copy ( str(filename_path), str(output_filename_path) )
        else:
            image = (cv2_imread(str(filename_path)) / 255.0).astype(np.float32)

            if self.converter.get_mode() == ConverterBase.MODE_IMAGE:
                image = self.converter.convert_image(image, None, self.debug)
                if self.debug:
                    for img in image:
                        cv2.imshow ('Debug convert', img )
                        cv2.waitKey(0)
                faces_processed = 1
            elif self.converter.get_mode() == ConverterBase.MODE_IMAGE_WITH_LANDMARKS:
                if filename_path.suffix == '.png':
                    dflimg = DFLPNG.load( str(filename_path), throw_on_no_embedded_data=True )
                elif filename_path.suffix == '.jpg':
                    dflimg = DFLJPG.load ( str(filename_path), throw_on_no_embedded_data=True )
                else:
                    raise Exception ("%s is not a dfl image file" % (filename_path.name) ) 
            
                image_landmarks = dflimg.get_landmarks()
                        
                image = self.converter.convert_image(image, image_landmarks, self.debug)
                if self.debug:
                    for img in image:
                        cv2.imshow ('Debug convert', img )
                        cv2.waitKey(0)
                faces_processed = 1
            elif self.converter.get_mode() == ConverterBase.MODE_FACE:
                faces = self.alignments[filename_path.stem]
                for face_num, image_landmarks in enumerate(faces):
                    try:
                        if self.debug:
                            print ( '\nConverting face_num [%d] in file [%s]' % (face_num, filename_path) )
                        
                        image = self.converter.convert_face(image, image_landmarks, self.debug)     
                        if self.debug:
                            for img in image:
                                cv2.imshow ('Debug convert', (img*255).astype(np.uint8) )
                                cv2.waitKey(0)
                    except Exception as e:
                        print ( 'Error while converting face_num [%d] in file [%s]: %s' % (face_num, filename_path, str(e)) )
                        traceback.print_exc()
                faces_processed = len(faces)
                    
            if not self.debug:
                cv2_imwrite (str(output_filename_path), (image*255).astype(np.uint8) )
            
            
        return (files_processed, faces_processed)
        
    #override
    def onHostResult (self, host_dict, data, result):
        self.files_processed += result[0]
        self.faces_processed += result[1]    
        return 1
             
    #override
    def onFinalizeAndGetResult(self):
        return self.files_processed, self.faces_processed
        
def main (input_dir, output_dir, model_dir, model_name, aligned_dir=None, **in_options):
    print ("Running converter.\r\n")
    
    debug = in_options['debug']
    
    try:
        input_path = Path(input_dir)
        output_path = Path(output_dir)        
        model_path = Path(model_dir)
        
        if not input_path.exists():
            print('Input directory not found. Please ensure it exists.')
            return

        if output_path.exists():
            for filename in Path_utils.get_image_paths(output_path):
                Path(filename).unlink()
        else:
            output_path.mkdir(parents=True, exist_ok=True)
            
        
            
        if not model_path.exists():
            print('Model directory not found. Please ensure it exists.')
            return
   
        model_sq = multiprocessing.Queue()
        model_cq = multiprocessing.Queue()
        model_lock = multiprocessing.Lock()
        model_p = multiprocessing.Process(target=model_process, args=( sys.stdin.fileno(), model_name, model_dir, in_options, model_sq, model_cq))
        model_p.start()
        
        while True:
            if not model_cq.empty():
                obj = model_cq.get()
                obj_op = obj['op']
                if obj_op == 'init':
                    converter = obj['converter']
                    break

        alignments = None
        
        if converter.get_mode() == ConverterBase.MODE_FACE:            
            if aligned_dir is None:
                print('Aligned directory not found. Please ensure it exists.')
                return 
                
            aligned_path = Path(aligned_dir)
            if not aligned_path.exists():
                print('Aligned directory not found. Please ensure it exists.')
                return 
                
            alignments = {}
            
            aligned_path_image_paths = Path_utils.get_image_paths(aligned_path)
            for filepath in tqdm(aligned_path_image_paths, desc="Collecting alignments", ascii=True ):
                filepath = Path(filepath)
                
                if filepath.suffix == '.png':
                    dflimg = DFLPNG.load( str(filepath), print_on_no_embedded_data=True )
                elif filepath.suffix == '.jpg':
                    dflimg = DFLJPG.load ( str(filepath), print_on_no_embedded_data=True )
                else:
                    print ("%s is not a dfl image file" % (filepath.name) ) 
                
                source_filename_stem = Path( dflimg.get_source_filename() ).stem
                if source_filename_stem not in alignments.keys():
                    alignments[ source_filename_stem ] = []

                alignments[ source_filename_stem ].append (dflimg.get_source_landmarks())
        
            
            #interpolate landmarks
            #from facelib import LandmarksProcessor
            #from facelib import FaceType
            #a = sorted(alignments.keys())
            #a_len = len(a)
            #
            #box_pts = 3            
            #box = np.ones(box_pts)/box_pts
            #for i in range( a_len ):
            #    if i >= box_pts and i <= a_len-box_pts-1:
            #        af0 = alignments[ a[i] ][0] ##first face
            #        m0 = LandmarksProcessor.get_transform_mat (af0, 256, face_type=FaceType.FULL)      
            #    
            #        points = []
            #        
            #        for j in range(-box_pts, box_pts+1):
            #            af = alignments[ a[i+j] ][0] ##first face
            #            m = LandmarksProcessor.get_transform_mat (af, 256, face_type=FaceType.FULL)                    
            #            p = LandmarksProcessor.transform_points (af, m)
            #            points.append (p)
            #    
            #        points = np.array(points)
            #        points_len = len(points)
            #        t_points = np.transpose(points, [1,0,2])
            #        
            #        p1 = np.array ( [ int(np.convolve(x[:,0], box, mode='same')[points_len//2]) for x in t_points ] )
            #        p2 = np.array ( [ int(np.convolve(x[:,1], box, mode='same')[points_len//2]) for x in t_points ] )
            #        
            #        new_points = np.concatenate( [np.expand_dims(p1,-1),np.expand_dims(p2,-1)], -1 )
            #        
            #        alignments[ a[i] ][0]  = LandmarksProcessor.transform_points (new_points, m0, True).astype(np.int32)
            
        files_processed, faces_processed = ConvertSubprocessor ( 
                    converter              = converter.copy_and_set_predictor( model_process_predictor(model_sq,model_cq,model_lock) ), 
                    input_path_image_paths = Path_utils.get_image_paths(input_path), 
                    output_path            = output_path,
                    alignments             = alignments,  
                    **in_options ).process()

        model_p.terminate()
        
        '''            
        if model_name == 'AVATAR':
            output_path_image_paths = Path_utils.get_image_paths(output_path)
            
            last_ok_frame = -1
            for filename in output_path_image_paths:
                filename_path = Path(filename)
                stem = Path(filename).stem
                try:
                    frame = int(stem)
                except:
                    raise Exception ('Aligned avatars must be created from indexed sequence files.')
                    
                if frame-last_ok_frame > 1:
                    start = last_ok_frame + 1
                    end = frame - 1
                    
                    print ("Filling gaps: [%d...%d]" % (start, end) )
                    for i in range (start, end+1):                    
                        shutil.copy ( str(filename), str( output_path / ('%.5d%s' % (i, filename_path.suffix ))  ) )
                    
                last_ok_frame = frame
        '''
        
    except Exception as e:
        print ( 'Error: %s' % (str(e)))
        traceback.print_exc()