mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-08-19 21:13:20 -07:00
moving some files
This commit is contained in:
Colombo
parent
2b264da86b
commit
e013cb0f6b
14 changed files with 303 additions and 57 deletions
|
|
@ -16,7 +16,9 @@ import numpy.linalg as npla
|
|||
import imagelib
|
||||
from converters import (ConverterConfig, ConvertFaceAvatar, ConvertMasked,
|
||||
FrameInfo)
|
||||
from facelib import FaceType, FANSegmentator, LandmarksProcessor
|
||||
from facelib import FaceType, LandmarksProcessor
|
||||
from nnlib import TernausNet
|
||||
|
||||
from interact import interact as io
|
||||
from joblib import SubprocessFunctionCaller, Subprocessor
|
||||
from utils import Path_utils
|
||||
|
|
@ -114,13 +116,25 @@ class ConvertSubprocessor(Subprocessor):
|
|||
def fanseg_extract(face_type, *args, **kwargs):
|
||||
fanseg = self.fanseg_by_face_type.get(face_type, None)
|
||||
if self.fanseg_by_face_type.get(face_type, None) is None:
|
||||
fanseg = FANSegmentator( self.fanseg_input_size , FaceType.toString( face_type ) )
|
||||
fanseg = TernausNet("FANSeg", self.fanseg_input_size , FaceType.toString( face_type ) )
|
||||
self.fanseg_by_face_type[face_type] = fanseg
|
||||
|
||||
return fanseg.extract(*args, **kwargs)
|
||||
|
||||
self.fanseg_extract_func = fanseg_extract
|
||||
|
||||
self.fanchq_by_face_type = {}
|
||||
self.fanchq_input_size = 256
|
||||
def fanchq_extract(face_type, *args, **kwargs):
|
||||
fanchq = self.fanchq_by_face_type.get(face_type, None)
|
||||
if self.fanchq_by_face_type.get(face_type, None) is None:
|
||||
fanchq = TernausNet("FANCHQ", self.fanchq_input_size , FaceType.toString( face_type ) )
|
||||
self.fanchq_by_face_type[face_type] = fanchq
|
||||
|
||||
return fanchq.extract(*args, **kwargs)
|
||||
|
||||
self.fanchq_extract_func = fanchq_extract
|
||||
|
||||
import ebsynth
|
||||
def ebs_ct(*args, **kwargs):
|
||||
return ebsynth.color_transfer(*args, **kwargs)
|
||||
|
|
@ -161,6 +175,8 @@ class ConvertSubprocessor(Subprocessor):
|
|||
if cfg.type == ConverterConfig.TYPE_MASKED:
|
||||
cfg.fanseg_input_size = self.fanseg_input_size
|
||||
cfg.fanseg_extract_func = self.fanseg_extract_func
|
||||
cfg.fanchq_input_size = self.fanchq_input_size
|
||||
cfg.fanchq_extract_func = self.fanchq_extract_func
|
||||
|
||||
try:
|
||||
final_img = ConvertMasked (self.predictor_func, self.predictor_input_shape, cfg, frame_info)
|
||||
|
|
|
|||
|
|
@ -14,10 +14,10 @@ import numpy as np
|
|||
import facelib
|
||||
import imagelib
|
||||
import mathlib
|
||||
from facelib import FaceType, FANSegmentator, LandmarksProcessor
|
||||
from facelib import FaceType, LandmarksProcessor
|
||||
from interact import interact as io
|
||||
from joblib import Subprocessor
|
||||
from nnlib import nnlib
|
||||
from nnlib import TernausNet, nnlib
|
||||
from utils import Path_utils
|
||||
from utils.cv2_utils import *
|
||||
from utils.DFLJPG import DFLJPG
|
||||
|
|
@ -95,7 +95,7 @@ class ExtractSubprocessor(Subprocessor):
|
|||
|
||||
elif self.type == 'fanseg':
|
||||
nnlib.import_all (device_config)
|
||||
self.e = facelib.FANSegmentator(256, FaceType.toString(FaceType.FULL) )
|
||||
self.e = TernausNet(256, FaceType.toString(FaceType.FULL) )
|
||||
self.e.__enter__()
|
||||
|
||||
elif self.type == 'final':
|
||||
|
|
@ -279,7 +279,6 @@ class ExtractSubprocessor(Subprocessor):
|
|||
fanseg_mask = self.e.extract( image / 255.0 )
|
||||
src_dflimg.embed_and_set( filename_path_str,
|
||||
fanseg_mask=fanseg_mask,
|
||||
#fanseg_mask_ver=FANSegmentator.VERSION,
|
||||
)
|
||||
|
||||
#overridable
|
||||
|
|
@ -889,4 +888,3 @@ def extract_umd_csv(input_file_csv,
|
|||
io.log_info ('Images found: %d' % (images_found) )
|
||||
io.log_info ('Faces detected: %d' % (faces_detected) )
|
||||
io.log_info ('-------------------------')
|
||||
|
||||
|
|
|
|||
|
|
@ -1,9 +1,20 @@
|
|||
from . import Extractor
|
||||
from . import Sorter
|
||||
from pathlib import Path
|
||||
from utils import Path_utils
|
||||
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)
|
||||
|
|
@ -47,4 +58,169 @@ def extract_vggface2_dataset(input_dir, device_args={} ):
|
|||
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)} ")
|
||||
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()))
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue