mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-07-15 01:23:44 -07:00
initial
This commit is contained in:
parent
73de93b4f1
commit
6bd5a44264
71 changed files with 8448 additions and 0 deletions
283
mainscripts/Converter.py
Normal file
283
mainscripts/Converter.py
Normal file
|
@ -0,0 +1,283 @@
|
|||
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()
|
||||
|
||||
|
Loading…
Add table
Add a link
Reference in a new issue