mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-07-06 04:52:13 -07:00
226 lines
7.5 KiB
Python
226 lines
7.5 KiB
Python
import multiprocessing
|
|
import shutil
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from facelib import LandmarksProcessor
|
|
from interact import interact as io
|
|
from joblib import Subprocessor
|
|
from utils import Path_utils
|
|
from utils.cv2_utils import *
|
|
from utils.DFLJPG import DFLJPG
|
|
from utils.DFLPNG import DFLPNG
|
|
|
|
from . import Extractor, Sorter
|
|
|
|
|
|
def extract_vggface2_dataset(input_dir, device_args={} ):
|
|
multi_gpu = device_args.get('multi_gpu', False)
|
|
cpu_only = device_args.get('cpu_only', False)
|
|
|
|
input_path = Path(input_dir)
|
|
if not input_path.exists():
|
|
raise ValueError('Input directory not found. Please ensure it exists.')
|
|
|
|
output_path = input_path.parent / (input_path.name + '_out')
|
|
|
|
dir_names = Path_utils.get_all_dir_names(input_path)
|
|
|
|
if not output_path.exists():
|
|
output_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
|
|
|
for dir_name in dir_names:
|
|
|
|
cur_input_path = input_path / dir_name
|
|
cur_output_path = output_path / dir_name
|
|
|
|
io.log_info (f"Processing: {str(cur_input_path)} ")
|
|
|
|
if not cur_output_path.exists():
|
|
cur_output_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
Extractor.main( str(cur_input_path),
|
|
str(cur_output_path),
|
|
detector='s3fd',
|
|
image_size=256,
|
|
face_type='full_face',
|
|
max_faces_from_image=1,
|
|
device_args=device_args )
|
|
|
|
io.log_info (f"Sorting: {str(cur_input_path)} ")
|
|
Sorter.main (input_path=str(cur_output_path), sort_by_method='hist')
|
|
|
|
try:
|
|
io.log_info (f"Removing: {str(cur_input_path)} ")
|
|
shutil.rmtree(cur_input_path)
|
|
except:
|
|
io.log_info (f"unable to remove: {str(cur_input_path)} ")
|
|
|
|
|
|
|
|
class CelebAMASKHQSubprocessor(Subprocessor):
|
|
class Cli(Subprocessor.Cli):
|
|
#override
|
|
def on_initialize(self, client_dict):
|
|
self.masks_files_paths = client_dict['masks_files_paths']
|
|
return None
|
|
|
|
#override
|
|
def process_data(self, data):
|
|
filename = data[0]
|
|
filepath = Path(filename)
|
|
if filepath.suffix == '.png':
|
|
dflimg = DFLPNG.load( str(filepath) )
|
|
elif filepath.suffix == '.jpg':
|
|
dflimg = DFLJPG.load ( str(filepath) )
|
|
else:
|
|
dflimg = None
|
|
|
|
image_to_face_mat = dflimg.get_image_to_face_mat()
|
|
src_filename = dflimg.get_source_filename()
|
|
|
|
img = cv2_imread(filename)
|
|
h,w,c = img.shape
|
|
|
|
fanseg_mask = LandmarksProcessor.get_image_hull_mask(img.shape, dflimg.get_landmarks() )
|
|
|
|
idx_name = '%.5d' % int(src_filename.split('.')[0])
|
|
idx_files = [ x for x in self.masks_files_paths if idx_name in x ]
|
|
|
|
skin_files = [ x for x in idx_files if 'skin' in x ]
|
|
eye_glass_files = [ x for x in idx_files if 'eye_g' in x ]
|
|
|
|
for files, is_invert in [ (skin_files,False),
|
|
(eye_glass_files,True) ]:
|
|
if len(files) > 0:
|
|
mask = cv2_imread(files[0])
|
|
mask = mask[...,0]
|
|
mask[mask == 255] = 1
|
|
mask = mask.astype(np.float32)
|
|
mask = cv2.resize(mask, (1024,1024) )
|
|
mask = cv2.warpAffine(mask, image_to_face_mat, (w, h), cv2.INTER_LANCZOS4)
|
|
|
|
if not is_invert:
|
|
fanseg_mask *= mask[...,None]
|
|
else:
|
|
fanseg_mask *= (1-mask[...,None])
|
|
|
|
dflimg.embed_and_set (filename, fanseg_mask=fanseg_mask)
|
|
return 1
|
|
|
|
#override
|
|
def get_data_name (self, data):
|
|
#return string identificator of your data
|
|
return data[0]
|
|
|
|
#override
|
|
def __init__(self, image_paths, masks_files_paths ):
|
|
self.image_paths = image_paths
|
|
self.masks_files_paths = masks_files_paths
|
|
|
|
self.result = []
|
|
super().__init__('CelebAMASKHQSubprocessor', CelebAMASKHQSubprocessor.Cli, 60)
|
|
|
|
#override
|
|
def process_info_generator(self):
|
|
for i in range(min(multiprocessing.cpu_count(), 8)):
|
|
yield 'CPU%d' % (i), {}, {'masks_files_paths' : self.masks_files_paths }
|
|
|
|
#override
|
|
def on_clients_initialized(self):
|
|
io.progress_bar ("Processing", len (self.image_paths))
|
|
|
|
#override
|
|
def on_clients_finalized(self):
|
|
io.progress_bar_close()
|
|
|
|
#override
|
|
def get_data(self, host_dict):
|
|
if len (self.image_paths) > 0:
|
|
return [self.image_paths.pop(0)]
|
|
return None
|
|
|
|
#override
|
|
def on_data_return (self, host_dict, data):
|
|
self.image_paths.insert(0, data[0])
|
|
|
|
#override
|
|
def on_result (self, host_dict, data, result):
|
|
io.progress_bar_inc(1)
|
|
|
|
#override
|
|
def get_result(self):
|
|
return self.result
|
|
|
|
#unused in end user workflow
|
|
def apply_celebamaskhq(input_dir ):
|
|
|
|
input_path = Path(input_dir)
|
|
|
|
img_path = input_path / 'aligned'
|
|
mask_path = input_path / 'mask'
|
|
|
|
if not img_path.exists():
|
|
raise ValueError(f'{str(img_path)} directory not found. Please ensure it exists.')
|
|
|
|
CelebAMASKHQSubprocessor(Path_utils.get_image_paths(img_path),
|
|
Path_utils.get_image_paths(mask_path, subdirs=True) ).run()
|
|
|
|
return
|
|
|
|
paths_to_extract = []
|
|
for filename in io.progress_bar_generator(Path_utils.get_image_paths(img_path), desc="Processing"):
|
|
filepath = Path(filename)
|
|
if filepath.suffix == '.png':
|
|
dflimg = DFLPNG.load( str(filepath) )
|
|
elif filepath.suffix == '.jpg':
|
|
dflimg = DFLJPG.load ( str(filepath) )
|
|
else:
|
|
dflimg = None
|
|
|
|
if dflimg is not None:
|
|
paths_to_extract.append (filepath)
|
|
|
|
image_to_face_mat = dflimg.get_image_to_face_mat()
|
|
src_filename = dflimg.get_source_filename()
|
|
|
|
#img = cv2_imread(filename)
|
|
h,w,c = dflimg.get_shape()
|
|
|
|
fanseg_mask = LandmarksProcessor.get_image_hull_mask( (h,w,c), dflimg.get_landmarks() )
|
|
|
|
idx_name = '%.5d' % int(src_filename.split('.')[0])
|
|
idx_files = [ x for x in masks_files if idx_name in x ]
|
|
|
|
skin_files = [ x for x in idx_files if 'skin' in x ]
|
|
eye_glass_files = [ x for x in idx_files if 'eye_g' in x ]
|
|
|
|
for files, is_invert in [ (skin_files,False),
|
|
(eye_glass_files,True) ]:
|
|
|
|
if len(files) > 0:
|
|
mask = cv2_imread(files[0])
|
|
mask = mask[...,0]
|
|
mask[mask == 255] = 1
|
|
mask = mask.astype(np.float32)
|
|
mask = cv2.resize(mask, (1024,1024) )
|
|
mask = cv2.warpAffine(mask, image_to_face_mat, (w, h), cv2.INTER_LANCZOS4)
|
|
|
|
if not is_invert:
|
|
fanseg_mask *= mask[...,None]
|
|
else:
|
|
fanseg_mask *= (1-mask[...,None])
|
|
|
|
#cv2.imshow("", (fanseg_mask*255).astype(np.uint8) )
|
|
#cv2.waitKey(0)
|
|
|
|
|
|
dflimg.embed_and_set (filename, fanseg_mask=fanseg_mask)
|
|
|
|
|
|
#import code
|
|
#code.interact(local=dict(globals(), **locals()))
|