import traceback from pathlib import Path from utils import Path_utils import cv2 from tqdm import tqdm from utils.AlignedPNG import AlignedPNG 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(model_name, model_dir, in_options, sq, cq): 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 ) } ) closing = False while not closing: 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} ) elif obj_op == 'close': closing = True break time.sleep(0.005) model.finalize() 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): super().__init__('Converter') self.converter = converter self.input_path_image_paths = input_path_image_paths self.output_path = output_path self.alignments = alignments self.debug = debug 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(multiprocessing.cpu_count()) 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 } #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): if len (self.input_data) > 0: return self.input_data.pop(0) return None #override def onHostDataReturn (self, data): self.input_data.insert(0, 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'] 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.name 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_landmarks = None a_png = AlignedPNG.load( str(filename_path) ) if a_png is not None: d = a_png.getFaceswapDictData() if d is not None and 'landmarks' in d.keys(): image_landmarks = np.array(d['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 image_landmarks in faces: image = self.converter.convert_face(image, image_landmarks, self.debug) if self.debug: for img in image: cv2.imshow ('Debug convert', img ) cv2.waitKey(0) 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, data, result): self.files_processed += result[0] self.faces_processed += result[1] return 1 #override def get_start_return(self): return self.files_processed, self.faces_processed def main (input_dir, output_dir, aligned_dir, model_dir, model_name, **in_options): print ("Running converter.\r\n") debug = in_options['debug'] try: input_path = Path(input_dir) output_path = Path(output_dir) aligned_path = Path(aligned_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 aligned_path.exists(): print('Aligned directory not found. Please ensure it exists.') return 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=(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 = {} if converter.get_mode() == ConverterBase.MODE_FACE: aligned_path_image_paths = Path_utils.get_image_paths(aligned_path) for filename in tqdm(aligned_path_image_paths, desc= "Collecting alignments" ): a_png = AlignedPNG.load( str(filename) ) if a_png is None: print ( "%s - no embedded data found." % (filename) ) continue d = a_png.getFaceswapDictData() if d is None or d['source_filename'] is None or d['source_rect'] is None or d['source_landmarks'] is None: print ( "%s - no embedded data found." % (filename) ) continue source_filename_stem = Path(d['source_filename']).stem if source_filename_stem not in alignments.keys(): alignments[ source_filename_stem ] = [] alignments[ source_filename_stem ].append ( np.array(d['source_landmarks']) ) 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, debug = debug ).process() model_sq.put ( {'op':'close'} ) model_p.join() ''' 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()