mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-07-06 04:52:13 -07:00
moving some files
This commit is contained in:
Colombo
parent
2b264da86b
commit
e013cb0f6b
14 changed files with 303 additions and 57 deletions
|
|
@ -56,14 +56,14 @@ def ConvertMaskedFace (predictor_func, predictor_input_shape, cfg, frame_info, i
|
|||
|
||||
if cfg.mask_mode == 2: #dst
|
||||
prd_face_mask_a_0 = cv2.resize (dst_face_mask_a_0, (output_size,output_size), cv2.INTER_CUBIC)
|
||||
elif cfg.mask_mode >= 3 and cfg.mask_mode <= 7:
|
||||
elif cfg.mask_mode >= 3 and cfg.mask_mode <= 8:
|
||||
|
||||
if cfg.mask_mode == 3 or cfg.mask_mode == 5 or cfg.mask_mode == 6:
|
||||
prd_face_fanseg_bgr = cv2.resize (prd_face_bgr, (cfg.fanseg_input_size,)*2 )
|
||||
prd_face_fanseg_mask = cfg.fanseg_extract_func(FaceType.FULL, prd_face_fanseg_bgr)
|
||||
FAN_prd_face_mask_a_0 = cv2.resize ( prd_face_fanseg_mask, (output_size, output_size), cv2.INTER_CUBIC)
|
||||
|
||||
if cfg.mask_mode >= 4 or cfg.mask_mode <= 7:
|
||||
if cfg.mask_mode >= 4 and cfg.mask_mode <= 7:
|
||||
|
||||
full_face_fanseg_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, cfg.fanseg_input_size, face_type=FaceType.FULL)
|
||||
dst_face_fanseg_bgr = cv2.warpAffine(img_bgr, full_face_fanseg_mat, (cfg.fanseg_input_size,)*2, flags=cv2.INTER_CUBIC )
|
||||
|
|
@ -80,9 +80,24 @@ def ConvertMaskedFace (predictor_func, predictor_input_shape, cfg, frame_info, i
|
|||
m = cv2.getAffineTransform(b, fanseg_rect_corner_pts)
|
||||
FAN_dst_face_mask_a_0 = cv2.warpAffine(dst_face_fanseg_mask, m, (cfg.fanseg_input_size,)*2, flags=cv2.INTER_CUBIC )
|
||||
FAN_dst_face_mask_a_0 = cv2.resize (FAN_dst_face_mask_a_0, (output_size,output_size), cv2.INTER_CUBIC)
|
||||
#else:
|
||||
# raise ValueError ("cfg.face_type unsupported")
|
||||
"""
|
||||
if cfg.mask_mode == 8: #FANCHQ-dst
|
||||
full_face_fanchq_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, cfg.fanchq_input_size, face_type=FaceType.FULL)
|
||||
dst_face_fanchq_bgr = cv2.warpAffine(img_bgr, full_face_fanchq_mat, (cfg.fanchq_input_size,)*2, flags=cv2.INTER_CUBIC )
|
||||
dst_face_fanchq_mask = cfg.fanchq_extract_func( FaceType.FULL, dst_face_fanchq_bgr )
|
||||
|
||||
if cfg.face_type == FaceType.FULL:
|
||||
FANCHQ_dst_face_mask_a_0 = cv2.resize (dst_face_fanchq_mask, (output_size,output_size), cv2.INTER_CUBIC)
|
||||
else:
|
||||
face_fanchq_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, cfg.fanchq_input_size, face_type=cfg.face_type)
|
||||
|
||||
fanchq_rect_corner_pts = np.array ( [ [0,0], [cfg.fanchq_input_size-1,0], [0,cfg.fanchq_input_size-1] ], dtype=np.float32 )
|
||||
a = LandmarksProcessor.transform_points (fanchq_rect_corner_pts, face_fanchq_mat, invert=True )
|
||||
b = LandmarksProcessor.transform_points (a, full_face_fanchq_mat )
|
||||
m = cv2.getAffineTransform(b, fanchq_rect_corner_pts)
|
||||
FAN_dst_face_mask_a_0 = cv2.warpAffine(dst_face_fanchq_mask, m, (cfg.fanchq_input_size,)*2, flags=cv2.INTER_CUBIC )
|
||||
FAN_dst_face_mask_a_0 = cv2.resize (FAN_dst_face_mask_a_0, (output_size,output_size), cv2.INTER_CUBIC)
|
||||
"""
|
||||
if cfg.mask_mode == 3: #FAN-prd
|
||||
prd_face_mask_a_0 = FAN_prd_face_mask_a_0
|
||||
elif cfg.mask_mode == 4: #FAN-dst
|
||||
|
|
@ -93,6 +108,8 @@ def ConvertMaskedFace (predictor_func, predictor_input_shape, cfg, frame_info, i
|
|||
prd_face_mask_a_0 = prd_face_mask_a_0 * FAN_prd_face_mask_a_0 * FAN_dst_face_mask_a_0
|
||||
elif cfg.mask_mode == 7:
|
||||
prd_face_mask_a_0 = prd_face_mask_a_0 * FAN_dst_face_mask_a_0
|
||||
#elif cfg.mask_mode == 8: #FANCHQ-dst
|
||||
# prd_face_mask_a_0 = FANCHQ_dst_face_mask_a_0
|
||||
|
||||
prd_face_mask_a_0[ prd_face_mask_a_0 < 0.001 ] = 0.0
|
||||
|
||||
|
|
|
|||
|
|
@ -21,6 +21,9 @@ class ConverterConfig(object):
|
|||
self.blursharpen_func = None
|
||||
self.fanseg_input_size = None
|
||||
self.fanseg_extract_func = None
|
||||
|
||||
self.fanchq_input_size = None
|
||||
self.fanchq_extract_func = None
|
||||
self.ebs_ct_func = None
|
||||
|
||||
self.super_res_dict = {0:"None", 1:'RankSRGAN'}
|
||||
|
|
|
|||
|
|
@ -3,5 +3,4 @@ from .DLIBExtractor import DLIBExtractor
|
|||
from .MTCExtractor import MTCExtractor
|
||||
from .S3FDExtractor import S3FDExtractor
|
||||
from .FANExtractor import FANExtractor
|
||||
from .FANSegmentator import FANSegmentator
|
||||
from .PoseEstimator import PoseEstimator
|
||||
10
main.py
10
main.py
|
|
@ -78,6 +78,16 @@ if __name__ == "__main__":
|
|||
p.add_argument('--multi-gpu', action="store_true", dest="multi_gpu", default=False, help="Enables multi GPU.")
|
||||
p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Extract on CPU.")
|
||||
p.set_defaults (func=process_dev_extract_umd_csv)
|
||||
|
||||
|
||||
def process_dev_apply_celebamaskhq(arguments):
|
||||
os_utils.set_process_lowest_prio()
|
||||
from mainscripts import dev_misc
|
||||
dev_misc.apply_celebamaskhq( arguments.input_dir )
|
||||
|
||||
p = subparsers.add_parser( "dev_apply_celebamaskhq", help="")
|
||||
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
|
||||
p.set_defaults (func=process_dev_apply_celebamaskhq)
|
||||
"""
|
||||
def process_extract_fanseg(arguments):
|
||||
os_utils.set_process_lowest_prio()
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -48,3 +59,168 @@ def extract_vggface2_dataset(input_dir, device_args={} ):
|
|||
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()))
|
||||
|
|
|
|||
|
|
@ -1,9 +1,8 @@
|
|||
import numpy as np
|
||||
|
||||
from nnlib import nnlib
|
||||
from nnlib import nnlib, TernausNet
|
||||
from models import ModelBase
|
||||
from facelib import FaceType
|
||||
from facelib import FANSegmentator
|
||||
from samplelib import *
|
||||
from interact import interact as io
|
||||
|
||||
|
|
@ -29,13 +28,14 @@ class Model(ModelBase):
|
|||
#override
|
||||
def onInitialize(self):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
self.set_vram_batch_requirements( {1.5:4} )
|
||||
self.set_vram_batch_requirements( {1.5:4, 11:48} )
|
||||
|
||||
self.resolution = 256
|
||||
self.face_type = FaceType.FULL if self.options['face_type'] == 'f' else FaceType.HALF
|
||||
|
||||
|
||||
self.fan_seg = FANSegmentator(self.resolution,
|
||||
model_name = 'FANSeg'
|
||||
model_name = 'FANCHQ'
|
||||
self.fan_seg = TernausNet(model_name, self.resolution,
|
||||
FaceType.toString(self.face_type),
|
||||
load_weights=not self.is_first_run(),
|
||||
weights_file_root=self.get_model_root_path(),
|
||||
|
|
@ -48,13 +48,13 @@ class Model(ModelBase):
|
|||
self.set_training_data_generators ([
|
||||
SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
sample_process_options=SampleProcessor.Options(random_flip=True),
|
||||
output_sample_types=[ { 'types': (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR_SHUFFLE), 'resolution' : self.resolution, 'motion_blur':(25, 5), 'border_replicate':False },
|
||||
output_sample_types=[ { 'types': (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR), 'resolution' : self.resolution, 'motion_blur':(25, 5), 'gaussian_blur':(25,5), 'border_replicate':False, 'random_hsv_shift' : True },
|
||||
{ 'types': (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_M), 'resolution': self.resolution },
|
||||
]),
|
||||
|
||||
SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
sample_process_options=SampleProcessor.Options(random_flip=True ),
|
||||
output_sample_types=[ { 'types': (t.IMG_TRANSFORMED , face_type, t.MODE_BGR_SHUFFLE), 'resolution' : self.resolution},
|
||||
output_sample_types=[ { 'types': (t.IMG_TRANSFORMED , face_type, t.MODE_BGR), 'resolution' : self.resolution, 'random_hsv_shift' : True},
|
||||
])
|
||||
])
|
||||
|
||||
|
|
@ -73,11 +73,14 @@ class Model(ModelBase):
|
|||
#override
|
||||
def onGetPreview(self, sample):
|
||||
test_A = sample[0][0][0:4] #first 4 samples
|
||||
test_Am = sample[0][1][0:4] #first 4 samples
|
||||
test_B = sample[1][0][0:4] #first 4 samples
|
||||
|
||||
|
||||
mAA = self.fan_seg.extract(test_A)
|
||||
mBB = self.fan_seg.extract(test_B)
|
||||
|
||||
test_Am = np.repeat ( test_Am, (3,), -1)
|
||||
mAA = np.repeat ( mAA, (3,), -1)
|
||||
mBB = np.repeat ( mBB, (3,), -1)
|
||||
|
||||
|
|
@ -85,6 +88,7 @@ class Model(ModelBase):
|
|||
for i in range(0, len(test_A)):
|
||||
st.append ( np.concatenate ( (
|
||||
test_A[i,:,:,0:3],
|
||||
test_Am[i],
|
||||
mAA[i],
|
||||
test_A[i,:,:,0:3]*mAA[i],
|
||||
), axis=1) )
|
||||
|
|
|
|||
|
|
@ -10,8 +10,6 @@ from interact import interact as io
|
|||
from nnlib import nnlib
|
||||
|
||||
"""
|
||||
FANSegmentator is designed to exclude obstructions from faces such as hair, fingers, etc.
|
||||
|
||||
Dataset used to train located in official DFL mega.nz folder
|
||||
https://mega.nz/#F!b9MzCK4B!zEAG9txu7uaRUjXz9PtBqg
|
||||
|
||||
|
|
@ -19,19 +17,19 @@ using https://github.com/ternaus/TernausNet
|
|||
TernausNet: U-Net with VGG11 Encoder Pre-Trained on ImageNet for Image Segmentation
|
||||
"""
|
||||
|
||||
class FANSegmentator(object):
|
||||
class TernausNet(object):
|
||||
VERSION = 1
|
||||
def __init__ (self, resolution, face_type_str, load_weights=True, weights_file_root=None, training=False):
|
||||
def __init__ (self, name, resolution, face_type_str, load_weights=True, weights_file_root=None, training=False):
|
||||
exec( nnlib.import_all(), locals(), globals() )
|
||||
|
||||
self.model = FANSegmentator.BuildModel(resolution, ngf=64)
|
||||
self.model = TernausNet.BuildModel(resolution, ngf=64)
|
||||
|
||||
if weights_file_root is not None:
|
||||
weights_file_root = Path(weights_file_root)
|
||||
else:
|
||||
weights_file_root = Path(__file__).parent
|
||||
|
||||
self.weights_path = weights_file_root / ('FANSeg_%d_%s.h5' % (resolution, face_type_str) )
|
||||
self.weights_path = weights_file_root / ('%s_%d_%s.h5' % (name, resolution, face_type_str) )
|
||||
|
||||
if load_weights:
|
||||
self.model.load_weights (str(self.weights_path))
|
||||
|
|
@ -59,7 +57,6 @@ class FANSegmentator(object):
|
|||
real_t = Input ( (resolution, resolution, 1) )
|
||||
out_t = self.model(inp_t)
|
||||
|
||||
#loss = K.mean(10*K.square(out_t-real_t))
|
||||
loss = K.mean(10*K.binary_crossentropy(real_t,out_t) )
|
||||
|
||||
out_t_diff1 = out_t[:, 1:, :, :] - out_t[:, :-1, :, :]
|
||||
|
|
@ -69,7 +66,7 @@ class FANSegmentator(object):
|
|||
|
||||
opt = Adam(lr=0.0001, beta_1=0.5, beta_2=0.999, tf_cpu_mode=2)
|
||||
|
||||
self.train_func = K.function ( [inp_t, real_t], [K.mean(loss)], opt.get_updates( [loss,total_var_loss], self.model.trainable_weights) )
|
||||
self.train_func = K.function ( [inp_t, real_t], [K.mean(loss)], opt.get_updates( [loss], self.model.trainable_weights) )
|
||||
|
||||
|
||||
def __enter__(self):
|
||||
|
|
@ -103,7 +100,7 @@ class FANSegmentator(object):
|
|||
exec( nnlib.import_all(), locals(), globals() )
|
||||
inp = Input ( (resolution,resolution,3) )
|
||||
x = inp
|
||||
x = FANSegmentator.Flow(ngf=ngf)(x)
|
||||
x = TernausNet.Flow(ngf=ngf)(x)
|
||||
model = Model(inp,x)
|
||||
return model
|
||||
|
||||
|
|
@ -115,7 +112,7 @@ class FANSegmentator(object):
|
|||
x = input
|
||||
|
||||
x0 = x = Conv2D(ngf, kernel_size=3, strides=1, padding='same', activation='relu', name='features.0')(x)
|
||||
x = BlurPool(filt_size=3)(x) #x = MaxPooling2D()(x)
|
||||
x = BlurPool(filt_size=3)(x)
|
||||
|
||||
x1 = x = Conv2D(ngf*2, kernel_size=3, strides=1, padding='same', activation='relu', name='features.3')(x)
|
||||
x = BlurPool(filt_size=3)(x)
|
||||
|
|
@ -1,2 +1,3 @@
|
|||
from .nnlib import nnlib
|
||||
from .FUNIT import FUNIT
|
||||
from .TernausNet import TernausNet
|
||||
|
|
@ -72,7 +72,7 @@ class SampleProcessor(object):
|
|||
MODE_GGG = 42 #3xGrayscale
|
||||
MODE_M = 43 #mask only
|
||||
MODE_BGR_SHUFFLE = 44 #BGR shuffle
|
||||
MODE_BGR_RANDOM_HUE_SHIFT = 45
|
||||
MODE_BGR_RANDOM_HSV_SHIFT = 45
|
||||
MODE_END = 50
|
||||
|
||||
class Options(object):
|
||||
|
|
@ -111,8 +111,6 @@ class SampleProcessor(object):
|
|||
|
||||
sample_rnd_seed = np.random.randint(0x80000000)
|
||||
|
||||
|
||||
|
||||
outputs = []
|
||||
for opts in output_sample_types:
|
||||
|
||||
|
|
@ -124,6 +122,9 @@ class SampleProcessor(object):
|
|||
normalize_std_dev = opts.get('normalize_std_dev', False)
|
||||
normalize_vgg = opts.get('normalize_vgg', False)
|
||||
motion_blur = opts.get('motion_blur', None)
|
||||
gaussian_blur = opts.get('gaussian_blur', None)
|
||||
|
||||
random_hsv_shift = opts.get('random_hsv_shift', None)
|
||||
apply_ct = opts.get('apply_ct', False)
|
||||
normalize_tanh = opts.get('normalize_tanh', False)
|
||||
|
||||
|
|
@ -205,6 +206,13 @@ class SampleProcessor(object):
|
|||
if np.random.randint(100) < chance:
|
||||
img = imagelib.LinearMotionBlur (img, np.random.randint( mb_max_size )+1, np.random.randint(360) )
|
||||
|
||||
if gaussian_blur is not None:
|
||||
chance, kernel_max_size = gaussian_blur
|
||||
chance = np.clip(chance, 0, 100)
|
||||
|
||||
if np.random.randint(100) < chance:
|
||||
img = cv2.GaussianBlur(img, ( np.random.randint( kernel_max_size )*2+1 ,) *2 , 0)
|
||||
|
||||
if is_face_sample and target_face_type != SPTF.NONE:
|
||||
target_ft = SampleProcessor.SPTF_FACETYPE_TO_FACETYPE[target_face_type]
|
||||
if target_ft > sample.face_type:
|
||||
|
|
@ -232,7 +240,7 @@ class SampleProcessor(object):
|
|||
start_y = rnd_state.randint(sub_size+1)
|
||||
img = img[start_y:start_y+sub_size,start_x:start_x+sub_size,:]
|
||||
|
||||
img = np.clip(img, 0, 1)
|
||||
img = np.clip(img, 0, 1).astype(np.float32)
|
||||
img_bgr = img[...,0:3]
|
||||
img_mask = img[...,3:4]
|
||||
|
||||
|
|
@ -245,6 +253,18 @@ class SampleProcessor(object):
|
|||
img_bgr = imagelib.linear_color_transfer (img_bgr, ct_sample_bgr_resized)
|
||||
img_bgr = np.clip( img_bgr, 0.0, 1.0)
|
||||
|
||||
if random_hsv_shift:
|
||||
rnd_state = np.random.RandomState (sample_rnd_seed)
|
||||
hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV)
|
||||
h, s, v = cv2.split(hsv)
|
||||
|
||||
h = (h + rnd_state.randint(360) ) % 360
|
||||
s = np.clip ( s + rnd_state.random()-0.5, 0, 1 )
|
||||
v = np.clip ( v + rnd_state.random()-0.5, 0, 1 )
|
||||
hsv = cv2.merge([h, s, v])
|
||||
|
||||
img_bgr = np.clip( cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR) , 0, 1 )
|
||||
|
||||
if normalize_std_dev:
|
||||
img_bgr = (img_bgr - img_bgr.mean( (0,1)) ) / img_bgr.std( (0,1) )
|
||||
elif normalize_vgg:
|
||||
|
|
@ -258,15 +278,6 @@ class SampleProcessor(object):
|
|||
elif mode_type == SPTF.MODE_BGR_SHUFFLE:
|
||||
rnd_state = np.random.RandomState (sample_rnd_seed)
|
||||
img = np.take (img_bgr, rnd_state.permutation(img_bgr.shape[-1]), axis=-1)
|
||||
elif mode_type == SPTF.MODE_BGR_RANDOM_HUE_SHIFT:
|
||||
rnd_state = np.random.RandomState (sample_rnd_seed)
|
||||
hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV)
|
||||
h, s, v = cv2.split(hsv)
|
||||
|
||||
h = (h + rnd_state.randint(360) ) % 360
|
||||
hsv = cv2.merge([h, s, v])
|
||||
|
||||
img = np.clip( cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR) , 0, 1 )
|
||||
elif mode_type == SPTF.MODE_G:
|
||||
img = np.concatenate ( (np.expand_dims(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY),-1),img_mask) , -1 )
|
||||
elif mode_type == SPTF.MODE_GGG:
|
||||
|
|
|
|||
|
|
@ -3,12 +3,26 @@ from os import scandir
|
|||
|
||||
image_extensions = [".jpg", ".jpeg", ".png", ".tif", ".tiff"]
|
||||
|
||||
def get_image_paths(dir_path, image_extensions=image_extensions):
|
||||
def scantree(path):
|
||||
"""Recursively yield DirEntry objects for given directory."""
|
||||
for entry in scandir(path):
|
||||
if entry.is_dir(follow_symlinks=False):
|
||||
yield from scantree(entry.path) # see below for Python 2.x
|
||||
else:
|
||||
yield entry
|
||||
|
||||
def get_image_paths(dir_path, image_extensions=image_extensions, subdirs=False):
|
||||
dir_path = Path (dir_path)
|
||||
|
||||
result = []
|
||||
if dir_path.exists():
|
||||
for x in list(scandir(str(dir_path))):
|
||||
|
||||
if subdirs:
|
||||
gen = scantree(str(dir_path))
|
||||
else:
|
||||
gen = scandir(str(dir_path))
|
||||
|
||||
for x in list(gen):
|
||||
if any([x.name.lower().endswith(ext) for ext in image_extensions]):
|
||||
result.append(x.path)
|
||||
return sorted(result)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue