diff --git a/mainscripts/XSegUtil.py b/mainscripts/XSegUtil.py index c75a14a..96af480 100644 --- a/mainscripts/XSegUtil.py +++ b/mainscripts/XSegUtil.py @@ -32,7 +32,7 @@ def apply_xseg(input_path, model_path): if face_type is None: - face_type = io.input_str ("XSeg model face type", 'same', ['h','mf','f','wf','head','same'], help_message="Specify face type of trained XSeg model. For example if XSeg model trained as WF, but faceset is HEAD, specify WF to apply xseg only on WF part of HEAD. Default is 'same'").lower() + face_type = io.input_str ("XSeg model face type", 'same', ['h','mf','f','wf','head','same', 'custom'], help_message="Specify face type of trained XSeg model. For example if XSeg model trained as WF, but faceset is HEAD, specify WF to apply xseg only on WF part of HEAD. Default is 'same'").lower() if face_type == 'same': face_type = None @@ -41,6 +41,7 @@ def apply_xseg(input_path, model_path): 'mf' : FaceType.MID_FULL, 'f' : FaceType.FULL, 'wf' : FaceType.WHOLE_FACE, + 'custom' : FaceType.CUSTOM, 'head' : FaceType.HEAD}[face_type] io.log_info(f'Applying trained XSeg model to {input_path.name}/ folder.') @@ -69,7 +70,7 @@ def apply_xseg(input_path, model_path): h,w,c = img.shape img_face_type = FaceType.fromString( dflimg.get_face_type() ) - if face_type is not None and img_face_type != face_type: + if face_type is not None and img_face_type != face_type or img_face_type == FaceType.CUSTOM: # custom always goes for eqvivalents lmrks = dflimg.get_source_landmarks() fmat = LandmarksProcessor.get_transform_mat(lmrks, w, face_type) @@ -91,7 +92,7 @@ def apply_xseg(input_path, model_path): mask = xseg.extract(img) - if face_type is not None and img_face_type != face_type: + if face_type is not None and img_face_type != face_type or img_face_type == FaceType.CUSTOM: mask = cv2.resize(mask, (w, w), interpolation=cv2.INTER_LANCZOS4) mask = cv2.warpAffine( mask, mat, (w,w), np.zeros( (h,w,c), dtype=np.float), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4) mask = cv2.resize(mask, (xseg_res, xseg_res), interpolation=cv2.INTER_LANCZOS4) diff --git a/samplelib/SampleGeneratorFaceXSeg.py b/samplelib/SampleGeneratorFaceXSeg.py index 7e38e64..334213d 100644 --- a/samplelib/SampleGeneratorFaceXSeg.py +++ b/samplelib/SampleGeneratorFaceXSeg.py @@ -12,7 +12,7 @@ from core.imagelib import sd from core.cv2ex import * from core.interact import interact as io from core.joblib import Subprocessor, SubprocessGenerator, ThisThreadGenerator -from facelib import LandmarksProcessor +from facelib import LandmarksProcessor, FaceType from samplelib import (SampleGeneratorBase, SampleLoader, SampleProcessor, SampleType) class SampleGeneratorFaceXSeg(SampleGeneratorBase): @@ -96,7 +96,7 @@ class SampleGeneratorFaceXSeg(SampleGeneratorBase): else: raise Exception(f'no mask in sample {sample.filename}') - if face_type == sample.face_type: + if face_type == sample.face_type and sample.face_type != FaceType.CUSTOM: # custom always valid for stuff like for wf custom equivivelnet if w != resolution: img = cv2.resize( img, (resolution, resolution), interpolation=cv2.INTER_LANCZOS4 ) mask = cv2.resize( mask, (resolution, resolution), interpolation=cv2.INTER_LANCZOS4 )