diff --git a/converters/ConverterMasked.py b/converters/ConverterMasked.py index f599625..b5d7c6f 100644 --- a/converters/ConverterMasked.py +++ b/converters/ConverterMasked.py @@ -77,11 +77,11 @@ class ConverterMasked(Converter): self.hist_match_threshold = np.clip ( io.input_int("Hist match threshold [0..255] (skip:255) : ", 255), 0, 255) if face_type == FaceType.FULL: - self.mask_mode = np.clip ( io.input_int ("Mask mode: (1) learned, (2) dst, (3) FAN-prd, (4) FAN-dst , (5) FAN-prd&dst (?) help. Default - %d : " % (1) , 1, help_message="If you learned mask, then option 1 should be choosed. 'dst' mask is raw shaky mask from dst aligned images. 'FAN-prd' - using super smooth mask by pretrained FAN-model from predicted face. 'FAN-dst' - using super smooth mask by pretrained FAN-model from dst face. 'FAN-prd&dst' - using multiplied FAN prd and dst mask. "), 1, 5 ) + self.mask_mode = np.clip ( io.input_int ("Mask mode: (1) learned, (2) dst, (3) FAN-prd, (4) FAN-dst , (5) FAN-prd*FAN-dst (6) learned*FAN-prd*FAN-dst (?) help. Default - %d : " % (1) , 1, help_message="If you learned mask, then option 1 should be choosed. 'dst' mask is raw shaky mask from dst aligned images. 'FAN-prd' - using super smooth mask by pretrained FAN-model from predicted face. 'FAN-dst' - using super smooth mask by pretrained FAN-model from dst face. 'FAN-prd*FAN-dst' or 'learned*FAN-prd*FAN-dst' - using multiplied masks."), 1, 6 ) else: self.mask_mode = np.clip ( io.input_int ("Mask mode: (1) learned, (2) dst . Default - %d : " % (1) , 1), 1, 2 ) - if self.mask_mode == 3 or self.mask_mode == 4 or self.mask_mode == 5: + if self.mask_mode >= 3 or self.mask_mode <= 6: self.fan_seg = None if self.mode != 'raw': @@ -117,7 +117,7 @@ class ConverterMasked(Converter): #overridable def on_cli_initialize(self): - if (self.mask_mode == 3 or self.mask_mode == 4 or self.mask_mode == 5) and self.fan_seg == None: + if (self.mask_mode >= 3 and self.mask_mode <= 6) and self.fan_seg == None: self.fan_seg = FANSegmentator(256, FaceType.toString(FaceType.FULL) ) #override @@ -167,26 +167,28 @@ class ConverterMasked(Converter): if self.mask_mode == 2: #dst prd_face_mask_a_0 = cv2.resize (dst_face_mask_a_0, (output_size,output_size), cv2.INTER_CUBIC) - elif self.mask_mode >= 3 and self.mask_mode <= 5: + elif self.mask_mode >= 3 and self.mask_mode <= 6: - if self.mask_mode == 3 or self.mask_mode == 5: #FAN-prd + if self.mask_mode == 3 or self.mask_mode == 5 or self.mask_mode == 6: prd_face_bgr_256 = cv2.resize (prd_face_bgr, (256,256) ) prd_face_bgr_256_mask = self.fan_seg.extract_from_bgr( prd_face_bgr_256[np.newaxis,...] ) [0] FAN_prd_face_mask_a_0 = cv2.resize (prd_face_bgr_256_mask, (output_size,output_size), cv2.INTER_CUBIC) - if self.mask_mode == 4 or self.mask_mode == 5: #FAN-dst + if self.mask_mode == 4 or self.mask_mode == 5 or self.mask_mode == 6: face_256_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, 256, face_type=FaceType.FULL) dst_face_256_bgr = cv2.warpAffine(img_bgr, face_256_mat, (256, 256), flags=cv2.INTER_LANCZOS4 ) dst_face_256_mask = self.fan_seg.extract_from_bgr( dst_face_256_bgr[np.newaxis,...] ) [0] FAN_dst_face_mask_a_0 = cv2.resize (dst_face_256_mask, (output_size,output_size), cv2.INTER_CUBIC) - if self.mask_mode == 3: + if self.mask_mode == 3: #FAN-prd prd_face_mask_a_0 = FAN_prd_face_mask_a_0 - elif self.mask_mode == 4: + elif self.mask_mode == 4: #FAN-dst prd_face_mask_a_0 = FAN_dst_face_mask_a_0 elif self.mask_mode == 5: prd_face_mask_a_0 = FAN_prd_face_mask_a_0 * FAN_dst_face_mask_a_0 - + elif self.mask_mode == 6: + prd_face_mask_a_0 = prd_face_mask_a_0 * FAN_prd_face_mask_a_0 * FAN_dst_face_mask_a_0 + prd_face_mask_a_0[ prd_face_mask_a_0 < 0.001 ] = 0.0 prd_face_mask_a = prd_face_mask_a_0[...,np.newaxis] @@ -269,11 +271,12 @@ class ConverterMasked(Converter): img_mask_blurry_aaa = cv2.blur(img_mask_blurry_aaa, (blur, blur) ) img_mask_blurry_aaa = np.clip( img_mask_blurry_aaa, 0, 1.0 ) + face_mask_blurry_aaa = cv2.warpAffine( img_mask_blurry_aaa, face_mat, (output_size, output_size) ) if debug: debugs += [img_mask_blurry_aaa.copy()] - if self.color_transfer_mode is not None: + if 'seamless' not in self.mode and self.color_transfer_mode is not None: if self.color_transfer_mode == 'rct': if debug: debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] @@ -309,15 +312,20 @@ class ConverterMasked(Converter): if self.masked_hist_match: hist_mask_a *= prd_face_mask_a + + white = (1.0-hist_mask_a)* np.ones ( prd_face_bgr.shape[:2] + (1,) , dtype=np.float32) - hist_match_1 = prd_face_bgr*hist_mask_a + (1.0-hist_mask_a)* np.ones ( prd_face_bgr.shape[:2] + (1,) , dtype=np.float32) + hist_match_1 = prd_face_bgr*hist_mask_a + white hist_match_1[ hist_match_1 > 1.0 ] = 1.0 - hist_match_2 = dst_face_bgr*hist_mask_a + (1.0-hist_mask_a)* np.ones ( prd_face_bgr.shape[:2] + (1,) , dtype=np.float32) + hist_match_2 = dst_face_bgr*hist_mask_a + white hist_match_2[ hist_match_1 > 1.0 ] = 1.0 prd_face_bgr = imagelib.color_hist_match(hist_match_1, hist_match_2, self.hist_match_threshold ) - + + #if self.masked_hist_match: + # prd_face_bgr -= white + if self.mode == 'hist-match-bw': prd_face_bgr = prd_face_bgr.astype(dtype=np.float32) @@ -364,6 +372,35 @@ class ConverterMasked(Converter): out_img = np.clip( img_bgr*(1-img_mask_blurry_aaa) + (out_img*img_mask_blurry_aaa) , 0, 1.0 ) + if 'seamless' in self.mode and self.color_transfer_mode is not None: + out_face_bgr = cv2.warpAffine( out_img, face_mat, (output_size, output_size) ) + + if self.color_transfer_mode == 'rct': + if debug: + debugs += [ np.clip( cv2.warpAffine( out_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] + + new_out_face_bgr = imagelib.reinhard_color_transfer ( np.clip( (out_face_bgr*255).astype(np.uint8), 0, 255), + np.clip( (dst_face_bgr*255).astype(np.uint8), 0, 255), + source_mask=face_mask_blurry_aaa, target_mask=face_mask_blurry_aaa) + new_out_face_bgr = np.clip( new_out_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0) + + if debug: + debugs += [ np.clip( cv2.warpAffine( new_out_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] + + + elif self.color_transfer_mode == 'lct': + if debug: + debugs += [ np.clip( cv2.warpAffine( out_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] + + new_out_face_bgr = imagelib.linear_color_transfer (out_face_bgr, dst_face_bgr) + new_out_face_bgr = np.clip( new_out_face_bgr, 0.0, 1.0) + + if debug: + debugs += [ np.clip( cv2.warpAffine( new_out_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] + + new_out = cv2.warpAffine( new_out_face_bgr, face_mat, img_size, img_bgr.copy(), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ) + out_img = np.clip( img_bgr*(1-img_mask_blurry_aaa) + (new_out*img_mask_blurry_aaa) , 0, 1.0 ) + if self.mode == 'seamless-hist-match': out_face_bgr = cv2.warpAffine( out_img, face_mat, (output_size, output_size) ) new_out_face_bgr = imagelib.color_hist_match(out_face_bgr, dst_face_bgr, self.hist_match_threshold) diff --git a/facelib/LandmarksProcessor.py b/facelib/LandmarksProcessor.py index 1ec4607..379da80 100644 --- a/facelib/LandmarksProcessor.py +++ b/facelib/LandmarksProcessor.py @@ -4,6 +4,7 @@ import numpy as np from enum import IntEnum import mathlib import imagelib +from imagelib import IEPolys from mathlib.umeyama import umeyama from facelib import FaceType import math @@ -153,7 +154,7 @@ def transform_points(points, mat, invert=False): return points -def get_image_hull_mask (image_shape, image_landmarks): +def get_image_hull_mask (image_shape, image_landmarks, ie_polys=None): if len(image_landmarks) != 68: raise Exception('get_image_hull_mask works only with 68 landmarks') int_lmrks = np.array(image_landmarks, dtype=np.int) @@ -198,6 +199,9 @@ def get_image_hull_mask (image_shape, image_landmarks): #nose cv2.fillConvexPoly( hull_mask, cv2.convexHull(int_lmrks[27:36]), (1,) ) + if ie_polys is not None: + ie_polys.overlay_mask(hull_mask) + return hull_mask def get_image_eye_mask (image_shape, image_landmarks): @@ -211,11 +215,6 @@ def get_image_eye_mask (image_shape, image_landmarks): return hull_mask -def get_image_hull_mask_3D (image_shape, image_landmarks): - result = get_image_hull_mask(image_shape, image_landmarks) - - return np.repeat ( result, (3,), -1 ) - def blur_image_hull_mask (hull_mask): maxregion = np.argwhere(hull_mask==1.0) @@ -235,9 +234,6 @@ def blur_image_hull_mask (hull_mask): return hull_mask -def get_blurred_image_hull_mask(image_shape, image_landmarks): - return blur_image_hull_mask ( get_image_hull_mask(image_shape, image_landmarks) ) - mirror_idxs = [ [0,16], [1,15], @@ -282,7 +278,7 @@ def mirror_landmarks (landmarks, val): result[:,0] = val - result[:,0] - 1 return result -def draw_landmarks (image, image_landmarks, color=(0,255,0), transparent_mask=False): +def draw_landmarks (image, image_landmarks, color=(0,255,0), transparent_mask=False, ie_polys=None): if len(image_landmarks) != 68: raise Exception('get_image_eye_mask works only with 68 landmarks') @@ -310,11 +306,11 @@ def draw_landmarks (image, image_landmarks, color=(0,255,0), transparent_mask=Fa cv2.circle(image, (x, y), 2, color, lineType=cv2.LINE_AA) if transparent_mask: - mask = get_image_hull_mask (image.shape, image_landmarks) + mask = get_image_hull_mask (image.shape, image_landmarks, ie_polys) image[...] = ( image * (1-mask) + image * mask / 2 )[...] -def draw_rect_landmarks (image, rect, image_landmarks, face_size, face_type, transparent_mask=False, landmarks_color=(0,255,0) ): - draw_landmarks(image, image_landmarks, color=landmarks_color, transparent_mask=transparent_mask) +def draw_rect_landmarks (image, rect, image_landmarks, face_size, face_type, transparent_mask=False, ie_polys=None, landmarks_color=(0,255,0) ): + draw_landmarks(image, image_landmarks, color=landmarks_color, transparent_mask=transparent_mask, ie_polys=ie_polys) imagelib.draw_rect (image, rect, (255,0,0), 2 ) image_to_face_mat = get_transform_mat (image_landmarks, face_size, face_type) diff --git a/imagelib/IEPolys.py b/imagelib/IEPolys.py new file mode 100644 index 0000000..820a229 --- /dev/null +++ b/imagelib/IEPolys.py @@ -0,0 +1,104 @@ +import numpy as np +import cv2 + +class IEPolysPoints: + def __init__(self, IEPolys_parent, type): + self.parent = IEPolys_parent + self.type = type + self.points = np.empty( (0,2), dtype=np.int32 ) + self.n_max = self.n = 0 + + def add(self,x,y): + self.points = np.append(self.points[0:self.n], [ (x,y) ], axis=0) + self.n_max = self.n = self.n + 1 + self.parent.dirty = True + + def n_dec(self): + self.n = max(0, self.n-1) + self.parent.dirty = True + return self.n + + def n_inc(self): + self.n = min(len(self.points), self.n+1) + self.parent.dirty = True + return self.n + + def n_clip(self): + self.points = self.points[0:self.n] + self.n_max = self.n + + def cur_point(self): + return self.points[self.n-1] + + def points_to_n(self): + return self.points[0:self.n] + + def set_points(self, points): + self.points = np.array(points) + self.n_max = self.n = len(points) + self.parent.dirty = True + +class IEPolys: + def __init__(self): + self.list = [] + self.n_max = self.n = 0 + self.dirty = True + + def add(self, type): + self.list = self.list[0:self.n] + self.list.append ( IEPolysPoints(self, type) ) + self.n_max = self.n = self.n + 1 + self.dirty = True + + def n_dec(self): + self.n = max(0, self.n-1) + self.dirty = True + return self.n + + def n_inc(self): + self.n = min(len(self.list), self.n+1) + self.dirty = True + return self.n + + def n_list(self): + return self.list[self.n-1] + + def n_clip(self): + self.list = self.list[0:self.n] + self.n_max = self.n + if self.n > 0: + self.list[-1].n_clip() + + def __iter__(self): + for n in range(self.n): + yield self.list[n] + + def switch_dirty(self): + d = self.dirty + self.dirty = False + return d + + def overlay_mask(self, mask): + h,w,c = mask.shape + white = (1,)*c + black = (0,)*c + for n in range(self.n): + poly = self.list[n] + if poly.n > 0: + cv2.fillPoly(mask, [poly.points_to_n()], white if poly.type == 1 else black ) + + def dump(self): + result = [] + for n in range(self.n): + l = self.list[n] + result += [ (l.type, l.points_to_n().tolist() ) ] + return result + + @staticmethod + def load(ie_polys=None): + obj = IEPolys() + if ie_polys is not None: + for (type, points) in ie_polys: + obj.add(type) + obj.n_list().set_points(points) + return obj \ No newline at end of file diff --git a/imagelib/__init__.py b/imagelib/__init__.py index ab2ed81..a749a28 100644 --- a/imagelib/__init__.py +++ b/imagelib/__init__.py @@ -20,4 +20,6 @@ from .color_transfer import linear_color_transfer from .DCSCN import DCSCN -from .common import normalize_channels \ No newline at end of file +from .common import normalize_channels + +from .IEPolys import IEPolys \ No newline at end of file diff --git a/imagelib/text.py b/imagelib/text.py index 0a486af..8c31e2e 100644 --- a/imagelib/text.py +++ b/imagelib/text.py @@ -15,23 +15,24 @@ def _get_pil_font (font, size): return ImageFont.load_default() def get_text_image( shape, text, color=(1,1,1), border=0.2, font=None): - try: - size = shape[1] - pil_font = _get_pil_font( localization.get_default_ttf_font_name() , size) - text_width, text_height = pil_font.getsize(text) + h,w,c = shape + try: + pil_font = _get_pil_font( localization.get_default_ttf_font_name() , h-2) - canvas = Image.new('RGB', shape[0:2], (0,0,0) ) + canvas = Image.new('RGB', (w,h) , (0,0,0) ) draw = ImageDraw.Draw(canvas) offset = ( 0, 0) draw.text(offset, text, font=pil_font, fill=tuple((np.array(color)*255).astype(np.int)) ) result = np.asarray(canvas) / 255 - if shape[2] != 3: - result = np.concatenate ( (result, np.ones ( (shape[1],) + (shape[0],) + (shape[2]-3,)) ), axis=2 ) - + + if c > 3: + result = np.concatenate ( (result, np.ones ((h,w,c-3)) ), axis=-1 ) + elif c < 3: + result = result[...,0:c] return result except: - return np.zeros ( (shape[1], shape[0], shape[2]), dtype=np.float32 ) + return np.zeros ( (h,w,c) ) def draw_text( image, rect, text, color=(1,1,1), border=0.2, font=None): h,w,c = image.shape diff --git a/interact/interact.py b/interact/interact.py index 6255639..7e5c5e9 100644 --- a/interact/interact.py +++ b/interact/interact.py @@ -18,8 +18,10 @@ except: class InteractBase(object): EVENT_LBUTTONDOWN = 1 EVENT_LBUTTONUP = 2 + EVENT_MBUTTONDOWN = 3 + EVENT_MBUTTONUP = 4 EVENT_RBUTTONDOWN = 5 - EVENT_RBUTTONUP = 6 + EVENT_RBUTTONUP = 6 EVENT_MOUSEWHEEL = 10 def __init__(self): @@ -263,6 +265,8 @@ class InteractDesktop(InteractBase): elif event == cv2.EVENT_LBUTTONUP: ev = InteractBase.EVENT_LBUTTONUP elif event == cv2.EVENT_RBUTTONDOWN: ev = InteractBase.EVENT_RBUTTONDOWN elif event == cv2.EVENT_RBUTTONUP: ev = InteractBase.EVENT_RBUTTONUP + elif event == cv2.EVENT_MBUTTONDOWN: ev = InteractBase.EVENT_MBUTTONDOWN + elif event == cv2.EVENT_MBUTTONUP: ev = InteractBase.EVENT_MBUTTONUP elif event == cv2.EVENT_MOUSEWHEEL: ev = InteractBase.EVENT_MOUSEWHEEL else: ev = 0 diff --git a/main.py b/main.py index fcf4477..f7d77ae 100644 --- a/main.py +++ b/main.py @@ -186,24 +186,23 @@ if __name__ == "__main__": p.add_argument('--lossless', action="store_true", dest="lossless", default=False, help="PNG codec.") p.set_defaults(func=process_videoed_video_from_sequence) - def process_labelingtool(arguments): - from mainscripts import LabelingTool - LabelingTool.main (arguments.input_dir, arguments.output_dir) + def process_labelingtool_edit_mask(arguments): + from mainscripts import MaskEditorTool + MaskEditorTool.mask_editor_main (arguments.input_dir, arguments.confirmed_dir, arguments.skipped_dir) - p = subparsers.add_parser( "labelingtool", help="Labeling tool.") + labeling_parser = subparsers.add_parser( "labelingtool", help="Labeling tool.").add_subparsers() + p = labeling_parser.add_parser ( "edit_mask", help="") p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory of aligned faces.") - p.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the labeled faces will be stored.") - p.set_defaults(func=process_labelingtool) - + p.add_argument('--confirmed-dir', required=True, action=fixPathAction, dest="confirmed_dir", help="This is where the labeled faces will be stored.") + p.add_argument('--skipped-dir', required=True, action=fixPathAction, dest="skipped_dir", help="This is where the labeled faces will be stored.") + p.set_defaults(func=process_labelingtool_edit_mask) + def bad_args(arguments): parser.print_help() exit(0) parser.set_defaults(func=bad_args) arguments = parser.parse_args() - - #os.environ['force_plaidML'] = '1' - arguments.func(arguments) print ("Done.") diff --git a/mainscripts/LabelingTool_unfinished.py b/mainscripts/LabelingTool_unfinished.py deleted file mode 100644 index 60067c2..0000000 --- a/mainscripts/LabelingTool_unfinished.py +++ /dev/null @@ -1,287 +0,0 @@ -import traceback -import os -import sys -import time -import numpy as np -import numpy.linalg as npl - -import cv2 -from pathlib import Path -from interact import interact as io -from utils.cv2_utils import * -from utils import Path_utils -from utils.DFLPNG import DFLPNG -from utils.DFLJPG import DFLJPG -from facelib import LandmarksProcessor - -def main(input_dir, output_dir): - input_path = Path(input_dir) - output_path = Path(output_dir) - - if not input_path.exists(): - raise ValueError('Input directory not found. Please ensure it exists.') - - if not output_path.exists(): - output_path.mkdir(parents=True) - - wnd_name = "Labeling tool" - io.named_window (wnd_name) - io.capture_mouse(wnd_name) - io.capture_keys(wnd_name) - - #for filename in io.progress_bar_generator (Path_utils.get_image_paths(input_path), desc="Labeling"): - for filename in Path_utils.get_image_paths(input_path): - 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 None: - io.log_err ("%s is not a dfl image file" % (filepath.name) ) - continue - - lmrks = dflimg.get_landmarks() - lmrks_list = lmrks.tolist() - orig_img = cv2_imread(str(filepath)) - h,w,c = orig_img.shape - - mask_orig = LandmarksProcessor.get_image_hull_mask( orig_img.shape, lmrks).astype(np.uint8)[:,:,0] - ero_dil_rate = w // 8 - mask_ero = cv2.erode (mask_orig, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(ero_dil_rate,ero_dil_rate)), iterations = 1 ) - mask_dil = cv2.dilate(mask_orig, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(ero_dil_rate,ero_dil_rate)), iterations = 1 ) - - - #mask_bg = np.zeros(orig_img.shape[:2],np.uint8) - mask_bg = 1-mask_dil - mask_bgp = np.ones(orig_img.shape[:2],np.uint8) #default - all background possible - mask_fg = np.zeros(orig_img.shape[:2],np.uint8) - mask_fgp = np.zeros(orig_img.shape[:2],np.uint8) - - img = orig_img.copy() - - l_thick=2 - - def draw_4_lines (masks_out, pts, thickness=1): - fgp,fg,bg,bgp = masks_out - h,w = fg.shape - - fgp_pts = [] - fg_pts = np.array([ pts[i:i+2] for i in range(len(pts)-1)]) - bg_pts = [] - bgp_pts = [] - - for i in range(len(fg_pts)): - a, b = line = fg_pts[i] - - ba = b-a - v = ba / npl.norm(ba) - - ccpv = np.array([v[1],-v[0]]) - cpv = np.array([-v[1],v[0]]) - step = 1 / max(np.abs(cpv)) - - fgp_pts.append ( np.clip (line + ccpv * step * thickness, 0, w-1 ).astype(np.int) ) - bg_pts.append ( np.clip (line + cpv * step * thickness, 0, w-1 ).astype(np.int) ) - bgp_pts.append ( np.clip (line + cpv * step * thickness * 2, 0, w-1 ).astype(np.int) ) - - fgp_pts = np.array(fgp_pts) - bg_pts = np.array(bg_pts) - bgp_pts = np.array(bgp_pts) - - cv2.polylines(fgp, fgp_pts, False, (1,), thickness=thickness) - cv2.polylines(fg, fg_pts, False, (1,), thickness=thickness) - cv2.polylines(bg, bg_pts, False, (1,), thickness=thickness) - cv2.polylines(bgp, bgp_pts, False, (1,), thickness=thickness) - - def draw_lines ( masks_steps, pts, thickness=1): - lines = np.array([ pts[i:i+2] for i in range(len(pts)-1)]) - - - for mask, step in masks_steps: - h,w = mask.shape - - mask_lines = [] - for i in range(len(lines)): - a, b = line = lines[i] - ba = b-a - ba_len = npl.norm(ba) - if ba_len != 0: - v = ba / ba_len - pv = np.array([-v[1],v[0]]) - pv_inv_max = 1 / max(np.abs(pv)) - mask_lines.append ( np.clip (line + pv * pv_inv_max * thickness * step, 0, w-1 ).astype(np.int) ) - else: - mask_lines.append ( np.array(line, dtype=np.int) ) - cv2.polylines(mask, mask_lines, False, (1,), thickness=thickness) - - def draw_fill_convex( mask_out, pts, scale=1.0 ): - hull = cv2.convexHull(np.array(pts)) - - if scale !=1.0: - pts_count = hull.shape[0] - - sum_x = np.sum(hull[:, 0, 0]) - sum_y = np.sum(hull[:, 0, 1]) - - hull_center = np.array([sum_x/pts_count, sum_y/pts_count]) - hull = hull_center+(hull-hull_center)*scale - hull = hull.astype(pts.dtype) - cv2.fillConvexPoly( mask_out, hull, (1,) ) - - def get_gc_mask_bgr(gc_mask): - h, w = gc_mask.shape - bgr = np.zeros( (h,w,3), dtype=np.uint8 ) - - bgr [ gc_mask == 0 ] = (0,0,0) - bgr [ gc_mask == 1 ] = (255,255,255) - bgr [ gc_mask == 2 ] = (0,0,255) #RED - bgr [ gc_mask == 3 ] = (0,255,0) #GREEN - return bgr - - def get_gc_mask_result(gc_mask): - return np.where((gc_mask==1) + (gc_mask==3),1,0).astype(np.int) - - #convex inner of right chin to end of right eyebrow - #draw_fill_convex ( mask_fgp, lmrks_list[8:17]+lmrks_list[26:27] ) - - #convex inner of start right chin to right eyebrow - #draw_fill_convex ( mask_fgp, lmrks_list[8:9]+lmrks_list[22:27] ) - - #convex inner of nose - draw_fill_convex ( mask_fgp, lmrks[27:36] ) - - #convex inner of nose half - draw_fill_convex ( mask_fg, lmrks[27:36], scale=0.5 ) - - - #left corner of mouth to left corner of nose - #draw_lines ( [ (mask_fg,0), ], lmrks_list[49:50]+lmrks_list[32:33], l_thick) - - #convex inner: right corner of nose to centers of eyebrows - #draw_fill_convex ( mask_fgp, lmrks_list[35:36]+lmrks_list[19:20]+lmrks_list[24:25]) - - #right corner of mouth to right corner of nose - #draw_lines ( [ (mask_fg,0), ], lmrks_list[54:55]+lmrks_list[35:36], l_thick) - - #left eye - #draw_fill_convex ( mask_fg, lmrks_list[36:40] ) - #right eye - #draw_fill_convex ( mask_fg, lmrks_list[42:48] ) - - #right chin - draw_lines ( [ (mask_bg,0), (mask_fg,-1), ], lmrks[8:17], l_thick) - - #left eyebrow center to right eyeprow center - draw_lines ( [ (mask_bg,-1), (mask_fg,0), ], lmrks_list[19:20] + lmrks_list[24:25], l_thick) - # #draw_lines ( [ (mask_bg,-1), (mask_fg,0), ], lmrks_list[24:25] + lmrks_list[19:17:-1], l_thick) - - #half right eyebrow to end of right chin - draw_lines ( [ (mask_bg,-1), (mask_fg,0), ], lmrks_list[24:27] + lmrks_list[16:17], l_thick) - - #import code - #code.interact(local=dict(globals(), **locals())) - - #compose mask layers - gc_mask = np.zeros(orig_img.shape[:2],np.uint8) - gc_mask [ mask_bgp==1 ] = 2 - gc_mask [ mask_fgp==1 ] = 3 - gc_mask [ mask_bg==1 ] = 0 - gc_mask [ mask_fg==1 ] = 1 - - gc_bgr_before = get_gc_mask_bgr (gc_mask) - - - - #io.show_image (wnd_name, gc_mask ) - - ##points, hierarcy = cv2.findContours(original_mask,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE) - ##gc_mask = ( (1-erode_mask)*2 + erode_mask )# * dilate_mask - #gc_mask = (1-erode_mask)*2 + erode_mask - #cv2.addWeighted( - #gc_mask = mask_0_27 + (1-mask_0_27)*2 - # - ##import code - ##code.interact(local=dict(globals(), **locals())) - # - #rect = (1,1,img.shape[1]-2,img.shape[0]-2) - # - # - cv2.grabCut(img,gc_mask,None,np.zeros((1,65),np.float64),np.zeros((1,65),np.float64),5, cv2.GC_INIT_WITH_MASK) - - gc_bgr = get_gc_mask_bgr (gc_mask) - gc_mask_result = get_gc_mask_result(gc_mask) - gc_mask_result_1 = gc_mask_result[:,:,np.newaxis] - - #import code - #code.interact(local=dict(globals(), **locals())) - orig_img_gc_layers_masked = (0.5*orig_img + 0.5*gc_bgr).astype(np.uint8) - orig_img_gc_before_layers_masked = (0.5*orig_img + 0.5*gc_bgr_before).astype(np.uint8) - - - - pink_bg = np.full ( orig_img.shape, (255,0,255), dtype=np.uint8 ) - - - orig_img_result = orig_img * gc_mask_result_1 - orig_img_result_pinked = orig_img_result + pink_bg * (1-gc_mask_result_1) - - #io.show_image (wnd_name, blended_img) - - ##gc_mask, bgdModel, fgdModel = - # - #mask2 = np.where((gc_mask==1) + (gc_mask==3),255,0).astype('uint8')[:,:,np.newaxis] - #mask2 = np.repeat(mask2, (3,), -1) - # - ##mask2 = np.where(gc_mask!=0,255,0).astype('uint8') - #blended_img = orig_img #-\ - # #0.3 * np.full(original_img.shape, (50,50,50)) * (1-mask_0_27)[:,:,np.newaxis] - # #0.3 * np.full(original_img.shape, (50,50,50)) * (1-dilate_mask)[:,:,np.newaxis] +\ - # #0.3 * np.full(original_img.shape, (50,50,50)) * (erode_mask)[:,:,np.newaxis] - #blended_img = np.clip(blended_img, 0, 255).astype(np.uint8) - ##import code - ##code.interact(local=dict(globals(), **locals())) - orig_img_lmrked = orig_img.copy() - LandmarksProcessor.draw_landmarks(orig_img_lmrked, lmrks, transparent_mask=True) - - screen = np.concatenate ([orig_img_gc_before_layers_masked, - orig_img_gc_layers_masked, - orig_img, - orig_img_lmrked, - orig_img_result_pinked, - orig_img_result, - ], axis=1) - - io.show_image (wnd_name, screen.astype(np.uint8) ) - - - while True: - io.process_messages() - - for (x,y,ev,flags) in io.get_mouse_events(wnd_name): - pass - #print (x,y,ev,flags) - - key_events = [ ev for ev, in io.get_key_events(wnd_name) ] - for key in key_events: - if key == ord('1'): - pass - if key == ord('2'): - pass - if key == ord('3'): - pass - - if ord(' ') in key_events: - break - - import code - code.interact(local=dict(globals(), **locals())) - - - - -#original_mask = np.ones(original_img.shape[:2],np.uint8)*2 -#cv2.drawContours(original_mask, points, -1, (1,), 1) diff --git a/mainscripts/MaskEditorTool.py b/mainscripts/MaskEditorTool.py new file mode 100644 index 0000000..f8ad4cc --- /dev/null +++ b/mainscripts/MaskEditorTool.py @@ -0,0 +1,376 @@ +import os +import sys +import time +import traceback +from pathlib import Path + +import cv2 +import numpy as np +import numpy.linalg as npl + +import imagelib +from facelib import LandmarksProcessor +from imagelib import IEPolys +from interact import interact as io +from utils import Path_utils +from utils.cv2_utils import * +from utils.DFLJPG import DFLJPG +from utils.DFLPNG import DFLPNG + +class MaskEditor: + STATE_NONE=0 + STATE_MASKING=1 + + def __init__(self, img, mask=None, ie_polys=None, get_status_lines_func=None): + self.img = imagelib.normalize_channels (img,3) + h, w, c = img.shape + ph, pw = h // 4, w // 4 + + if mask is not None: + self.mask = imagelib.normalize_channels (mask,3) + else: + self.mask = np.zeros ( (h,w,3) ) + self.get_status_lines_func = get_status_lines_func + + self.state_prop = self.STATE_NONE + + self.w, self.h = w, h + self.pw, self.ph = pw, ph + self.pwh = np.array([self.pw, self.ph]) + self.pwh2 = np.array([self.pw*2, self.ph*2]) + self.sw, self.sh = w+pw*2, h+ph*2 + + if ie_polys is None: + ie_polys = IEPolys() + self.ie_polys = ie_polys + + self.polys_mask = None + + self.screen_status_block = None + self.screen_status_block_dirty = True + + def set_state(self, state): + self.state = state + + @property + def state(self): + return self.state_prop + + @state.setter + def state(self, value): + self.state_prop = value + if value == self.STATE_MASKING: + self.ie_polys.dirty = True + + def get_mask(self): + if self.ie_polys.switch_dirty(): + self.screen_status_block_dirty = True + self.ie_mask = img = self.mask.copy() + + self.ie_polys.overlay_mask(img) + + return img + return self.ie_mask + + def get_screen_overlay(self): + img = np.zeros ( (self.sh, self.sw, 3) ) + + if self.state == self.STATE_MASKING: + mouse_xy = self.mouse_xy.copy() + self.pwh + + l = self.ie_polys.n_list() + if l.n > 0: + p = l.cur_point().copy() + self.pwh + color = (0,1,0) if l.type == 1 else (0,0,1) + cv2.line(img, tuple(p), tuple(mouse_xy), color ) + + return img + + def undo_to_begin_point(self): + while not self.undo_point(): + pass + + def undo_point(self): + if self.state == self.STATE_NONE: + if self.ie_polys.n > 0: + self.state = self.STATE_MASKING + + if self.state == self.STATE_MASKING: + if self.ie_polys.n_list().n_dec() == 0 and \ + self.ie_polys.n_dec() == 0: + self.state = self.STATE_NONE + else: + return False + + return True + + def redo_to_end_point(self): + while not self.redo_point(): + pass + + def redo_point(self): + if self.state == self.STATE_NONE: + if self.ie_polys.n_max > 0: + self.state = self.STATE_MASKING + if self.ie_polys.n == 0: + self.ie_polys.n_inc() + + if self.state == self.STATE_MASKING: + while True: + l = self.ie_polys.n_list() + if l.n_inc() == l.n_max: + if self.ie_polys.n == self.ie_polys.n_max: + break + self.ie_polys.n_inc() + else: + return False + + return True + + def combine_screens(self, screens): + + screens_len = len(screens) + + new_screens = [] + for screen, padded_overlay in screens: + screen_img = np.zeros( (self.sh, self.sw, 3), dtype=np.float32 ) + + screen = imagelib.normalize_channels (screen, 3) + h,w,c = screen.shape + + screen_img[self.ph:-self.ph, self.pw:-self.pw, :] = screen + + if padded_overlay is not None: + screen_img = screen_img + padded_overlay + + screen_img = np.clip(screen_img*255, 0, 255).astype(np.uint8) + new_screens.append(screen_img) + + return np.concatenate (new_screens, axis=1) + + def get_screen_status_block(self, w, c): + if self.screen_status_block_dirty: + self.screen_status_block_dirty = False + lines = [ + 'Polys current/max = %d/%d' % (self.ie_polys.n, self.ie_polys.n_max), + ] + if self.get_status_lines_func is not None: + lines += self.get_status_lines_func() + + lines_count = len(lines) + + + h_line = 21 + h = lines_count * h_line + img = np.ones ( (h,w,c) ) * 0.1 + + for i in range(lines_count): + img[ i*h_line:(i+1)*h_line, 0:w] += \ + imagelib.get_text_image ( (h_line,w,c), lines[i], color=[0.8]*c ) + + self.screen_status_block = np.clip(img*255, 0, 255).astype(np.uint8) + + return self.screen_status_block + + def set_screen_status_block_dirty(self): + self.screen_status_block_dirty = True + + def make_screen(self): + + + + screen_overlay = self.get_screen_overlay() + final_mask = self.get_mask() + + masked_img = self.img*final_mask*0.5 + self.img*(1-final_mask) + + pink = np.full ( (self.h, self.w, 3), (1,0,1) ) + pink_masked_img = self.img*final_mask + pink*(1-final_mask) + + screens = [ (self.img, None), + (masked_img, screen_overlay), + (pink_masked_img, screen_overlay), + ] + screens = self.combine_screens(screens) + + status_img = self.get_screen_status_block( screens.shape[1], screens.shape[2] ) + + result = np.concatenate ( [screens, status_img], axis=0 ) + + return result + + def mask_finish(self, n_clip=True): + if self.state == self.STATE_MASKING: + if self.ie_polys.n_list().n <= 2: + self.ie_polys.n_dec() + self.state = self.STATE_NONE + if n_clip: + self.ie_polys.n_clip() + + def set_mouse_pos(self,x,y): + mouse_x = x % (self.sw) - self.pw + mouse_y = y % (self.sh) - self.ph + self.mouse_xy = np.array( [mouse_x, mouse_y] ) + self.mouse_x, self.mouse_y = self.mouse_xy + + def mask_point(self, type): + if self.state == self.STATE_MASKING and \ + self.ie_polys.n_list().type != type: + self.mask_finish() + + elif self.state == self.STATE_NONE: + self.state = self.STATE_MASKING + self.ie_polys.add(type) + + if self.state == self.STATE_MASKING: + self.ie_polys.n_list().add (self.mouse_x, self.mouse_y) + + def get_ie_polys(self): + return self.ie_polys + +def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None): + input_path = Path(input_dir) + + confirmed_path = Path(confirmed_dir) + skipped_path = Path(skipped_dir) + + if not input_path.exists(): + raise ValueError('Input directory not found. Please ensure it exists.') + + if not confirmed_path.exists(): + confirmed_path.mkdir(parents=True) + + if not skipped_path.exists(): + skipped_path.mkdir(parents=True) + + wnd_name = "MaskEditor tool" + io.named_window (wnd_name) + io.capture_mouse(wnd_name) + io.capture_keys(wnd_name) + + image_paths = [ Path(x) for x in Path_utils.get_image_paths(input_path)] + done_paths = [] + + image_paths_total = len(image_paths) + + + + + is_exit = False + while not is_exit: + + if len(image_paths) > 0: + filepath = image_paths.pop(0) + else: + filepath = None + + if filepath is not None: + if filepath.suffix == '.png': + dflimg = DFLPNG.load( str(filepath) ) + elif filepath.suffix == '.jpg': + dflimg = DFLJPG.load ( str(filepath) ) + else: + dflimg = None + + if dflimg is None: + io.log_err ("%s is not a dfl image file" % (filepath.name) ) + continue + + lmrks = dflimg.get_landmarks() + ie_polys = dflimg.get_ie_polys() + + img = cv2_imread(str(filepath)) / 255.0 + mask = LandmarksProcessor.get_image_hull_mask( img.shape, lmrks) + else: + img = np.zeros ( (256,256,3) ) + mask = np.ones ( (256,256,3) ) + ie_polys = None + + def get_status_lines_func(): + return ['Progress: %d / %d . Current file: %s' % (len(done_paths), image_paths_total, str(filepath.name) if filepath is not None else "end" ), + '[Left mouse button] - mark include mask.', + '[Right mouse button] - mark exclude mask.', + '[Middle mouse button] - finish current poly.', + '[Mouse wheel] - undo/redo poly or point. [+ctrl] - undo to begin/redo to end', + '[q] - prev image. [w] - skip and move to %s. [e] - save and move to %s. ' % (skipped_path.name, confirmed_path.name), + '[z] - prev image. [x] - skip. [c] - save. ', + '[esc] - quit' + ] + ed = MaskEditor(img, mask, ie_polys, get_status_lines_func) + + next = False + while not next: + io.process_messages(0.005) + + for (x,y,ev,flags) in io.get_mouse_events(wnd_name): + ed.set_mouse_pos(x, y) + if filepath is not None: + if ev == io.EVENT_LBUTTONDOWN: + ed.mask_point(1) + elif ev == io.EVENT_RBUTTONDOWN: + ed.mask_point(0) + elif ev == io.EVENT_MBUTTONDOWN: + ed.mask_finish() + elif ev == io.EVENT_MOUSEWHEEL: + if flags & 0x80000000 != 0: + if flags & 0x8 != 0: + ed.undo_to_begin_point() + else: + ed.undo_point() + else: + if flags & 0x8 != 0: + ed.redo_to_end_point() + else: + ed.redo_point() + + key_events = [ ev for ev, in io.get_key_events(wnd_name) ] + for key in key_events: + if key == ord('q') or key == ord('z'): + if len(done_paths) > 0: + image_paths.insert(0, filepath) + filepath = done_paths.pop(-1) + + if filepath.parent != input_path: + new_filename_path = input_path / filepath.name + filepath.rename ( new_filename_path ) + image_paths.insert(0, new_filename_path) + else: + image_paths.insert(0, filepath) + + next = True + break + elif filepath is not None and ( key == ord('e') or key == ord('c') ): + ed.mask_finish() + dflimg.embed_and_set (str(filepath), ie_polys=ed.get_ie_polys() ) + + if key == ord('e'): + new_filename_path = confirmed_path / filepath.name + filepath.rename(new_filename_path) + done_paths += [new_filename_path] + else: + done_paths += [filepath] + + next = True + break + + elif filepath is not None and ( key == ord('w') or key == ord('x') ): + if key == ord('w'): + new_filename_path = skipped_path / filepath.name + filepath.rename(new_filename_path) + done_paths += [new_filename_path] + else: + done_paths += [filepath] + + next = True + break + elif key == 27: #esc + is_exit = True + next = True + break + screen = ed.make_screen() + + io.show_image (wnd_name, screen ) + io.process_messages(0.005) + + io.destroy_all_windows() + diff --git a/mainscripts/Trainer.py b/mainscripts/Trainer.py index 3f86f38..5bbdaf5 100644 --- a/mainscripts/Trainer.py +++ b/mainscripts/Trainer.py @@ -253,7 +253,7 @@ def main(args, device_args): for i in range(0, len(head_lines)): t = i*head_line_height b = (i+1)*head_line_height - head[t:b, 0:w] += imagelib.get_text_image ( (w,head_line_height,c) , head_lines[i], color=[0.8]*c ) + head[t:b, 0:w] += imagelib.get_text_image ( (head_line_height,w,c) , head_lines[i], color=[0.8]*c ) final = head diff --git a/mainscripts/Util.py b/mainscripts/Util.py index 93f3a4e..a7c4a28 100644 --- a/mainscripts/Util.py +++ b/mainscripts/Util.py @@ -27,6 +27,7 @@ def convert_png_to_jpg_file (filepath): DFLJPG.embed_data( new_filepath, face_type=dfl_dict.get('face_type', None), landmarks=dfl_dict.get('landmarks', None), + ie_polys=dfl_dict.get('ie_polys', None), source_filename=dfl_dict.get('source_filename', None), source_rect=dfl_dict.get('source_rect', None), source_landmarks=dfl_dict.get('source_landmarks', None) ) @@ -63,7 +64,7 @@ def add_landmarks_debug_images(input_path): if img is not None: face_landmarks = dflimg.get_landmarks() - LandmarksProcessor.draw_landmarks(img, face_landmarks, transparent_mask=True) + LandmarksProcessor.draw_landmarks(img, face_landmarks, transparent_mask=True, ie_polys=dflimg.get_ie_polys() ) output_file = '{}{}'.format( str(Path(str(input_path)) / filepath.stem), '_debug.jpg') cv2_imwrite(output_file, img, [int(cv2.IMWRITE_JPEG_QUALITY), 50] ) diff --git a/models/ModelBase.py b/models/ModelBase.py index 23b951b..f39c0e8 100644 --- a/models/ModelBase.py +++ b/models/ModelBase.py @@ -496,5 +496,5 @@ class ModelBase(object): lh_text = 'Iter: %d' % (iter) if iter != 0 else '' - lh_img[last_line_t:last_line_b, 0:w] += imagelib.get_text_image ( (w,last_line_b-last_line_t,c), lh_text, color=[0.8]*c ) + lh_img[last_line_t:last_line_b, 0:w] += imagelib.get_text_image ( (last_line_b-last_line_t,w,c), lh_text, color=[0.8]*c ) return lh_img diff --git a/models/Model_DF/Model.py b/models/Model_DF/Model.py index e6f2cc4..502e823 100644 --- a/models/Model_DF/Model.py +++ b/models/Model_DF/Model.py @@ -12,7 +12,7 @@ class Model(ModelBase): def onInitializeOptions(self, is_first_run, ask_override): if is_first_run or ask_override: def_pixel_loss = self.options.get('pixel_loss', False) - self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 20k iters to enhance fine details and decrease face jitter.") + self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Pixel loss may help to enhance fine details and stabilize face color. Use it only if quality does not improve over time.") else: self.options['pixel_loss'] = self.options.get('pixel_loss', False) diff --git a/models/Model_H128/Model.py b/models/Model_H128/Model.py index 9a9c444..fd20899 100644 --- a/models/Model_H128/Model.py +++ b/models/Model_H128/Model.py @@ -20,7 +20,7 @@ class Model(ModelBase): if is_first_run or ask_override: def_pixel_loss = self.options.get('pixel_loss', False) - self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 20k iters to enhance fine details and decrease face jitter.") + self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Pixel loss may help to enhance fine details and stabilize face color. Use it only if quality does not improve over time.") else: self.options['pixel_loss'] = self.options.get('pixel_loss', False) diff --git a/models/Model_H64/Model.py b/models/Model_H64/Model.py index 2ed75ff..dacab9e 100644 --- a/models/Model_H64/Model.py +++ b/models/Model_H64/Model.py @@ -20,7 +20,7 @@ class Model(ModelBase): if is_first_run or ask_override: def_pixel_loss = self.options.get('pixel_loss', False) - self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 20k iters to enhance fine details and decrease face jitter.") + self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Pixel loss may help to enhance fine details and stabilize face color. Use it only if quality does not improve over time.") else: self.options['pixel_loss'] = self.options.get('pixel_loss', False) diff --git a/models/Model_LIAEF128/Model.py b/models/Model_LIAEF128/Model.py index 53119a0..a7ad5ed 100644 --- a/models/Model_LIAEF128/Model.py +++ b/models/Model_LIAEF128/Model.py @@ -12,7 +12,7 @@ class Model(ModelBase): def onInitializeOptions(self, is_first_run, ask_override): if is_first_run or ask_override: def_pixel_loss = self.options.get('pixel_loss', False) - self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 20k iters to enhance fine details and decrease face jitter.") + self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Pixel loss may help to enhance fine details and stabilize face color. Use it only if quality does not improve over time.") else: self.options['pixel_loss'] = self.options.get('pixel_loss', False) diff --git a/models/Model_SAE/Model.py b/models/Model_SAE/Model.py index 357b2ba..05276e6 100644 --- a/models/Model_SAE/Model.py +++ b/models/Model_SAE/Model.py @@ -63,13 +63,11 @@ class SAEModel(ModelBase): self.options['e_ch_dims'] = np.clip ( io.input_int("Encoder dims per channel (21-85 ?:help skip:%d) : " % (default_e_ch_dims) , default_e_ch_dims, help_message="More encoder dims help to recognize more facial features, but require more VRAM. You can fine-tune model size to fit your GPU." ), 21, 85 ) default_d_ch_dims = self.options['e_ch_dims'] // 2 self.options['d_ch_dims'] = np.clip ( io.input_int("Decoder dims per channel (10-85 ?:help skip:%d) : " % (default_d_ch_dims) , default_d_ch_dims, help_message="More decoder dims help to get better details, but require more VRAM. You can fine-tune model size to fit your GPU." ), 10, 85 ) - self.options['d_residual_blocks'] = io.input_bool ("Add residual blocks to decoder? (y/n, ?:help skip:n) : ", False, help_message="These blocks help to get better details, but require more computing time.") self.options['remove_gray_border'] = io.input_bool ("Remove gray border? (y/n, ?:help skip:n) : ", False, help_message="Removes gray border of predicted face, but requires more computing resources.") else: self.options['ae_dims'] = self.options.get('ae_dims', default_ae_dims) self.options['e_ch_dims'] = self.options.get('e_ch_dims', default_e_ch_dims) self.options['d_ch_dims'] = self.options.get('d_ch_dims', default_d_ch_dims) - self.options['d_residual_blocks'] = self.options.get('d_residual_blocks', False) self.options['remove_gray_border'] = self.options.get('remove_gray_border', False) if is_first_run: @@ -81,7 +79,7 @@ class SAEModel(ModelBase): default_bg_style_power = 0.0 if is_first_run or ask_override: def_pixel_loss = self.options.get('pixel_loss', False) - self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: %s ) : " % (yn_str[def_pixel_loss]), def_pixel_loss, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 60k iters to enhance fine details. Warning: this option may cause collapse the model, make a backup of Model folder before apply it.") + self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: %s ) : " % (yn_str[def_pixel_loss]), def_pixel_loss, help_message="Pixel loss may help to enhance fine details and stabilize face color. Use it only if quality does not improve over time.") default_face_style_power = default_face_style_power if is_first_run else self.options.get('face_style_power', default_face_style_power) self.options['face_style_power'] = np.clip ( io.input_number("Face style power ( 0.0 .. 100.0 ?:help skip:%.2f) : " % (default_face_style_power), default_face_style_power, @@ -105,7 +103,7 @@ class SAEModel(ModelBase): ae_dims = self.options['ae_dims'] e_ch_dims = self.options['e_ch_dims'] d_ch_dims = self.options['d_ch_dims'] - d_residual_blocks = self.options['d_residual_blocks'] + d_residual_blocks = True bgr_shape = (resolution, resolution, 3) mask_shape = (resolution, resolution, 1) @@ -127,7 +125,9 @@ class SAEModel(ModelBase): target_dstm_ar = [ Input ( ( mask_shape[0] // (2**i) ,)*2 + (mask_shape[-1],) ) for i in range(ms_count-1, -1, -1)] padding = 'reflect' if self.options['remove_gray_border'] else 'zero' - common_flow_kwargs = { 'padding': padding } + common_flow_kwargs = { 'padding': padding, + 'norm': 'bn', + 'act':'prelu' } weights_to_load = [] if self.options['archi'] == 'liae': @@ -302,7 +302,7 @@ class SAEModel(ModelBase): self.src_dst_mask_train = K.function (feed,[src_mask_loss, dst_mask_loss], self.src_dst_mask_opt.get_updates(src_mask_loss+dst_mask_loss, src_dst_mask_loss_train_weights) ) if self.options['learn_mask']: - self.AE_view = K.function ([warped_src, warped_dst], [pred_src_src[-1], pred_dst_dst[-1], pred_src_dst[-1], pred_src_dstm[-1]]) + self.AE_view = K.function ([warped_src, warped_dst], [pred_src_src[-1], pred_dst_dst[-1], pred_dst_dstm[-1], pred_src_dst[-1], pred_src_dstm[-1]]) else: self.AE_view = K.function ([warped_src, warped_dst], [pred_src_src[-1], pred_dst_dst[-1], pred_src_dst[-1] ] ) @@ -310,7 +310,7 @@ class SAEModel(ModelBase): else: self.load_weights_safe(weights_to_load) if self.options['learn_mask']: - self.AE_convert = K.function ([warped_dst],[ pred_src_dst[-1], pred_src_dstm[-1] ]) + self.AE_convert = K.function ([warped_dst],[ pred_src_dst[-1], pred_dst_dstm[-1], pred_src_dstm[-1] ]) else: self.AE_convert = K.function ([warped_dst],[ pred_src_dst[-1] ]) @@ -391,24 +391,34 @@ class SAEModel(ModelBase): test_B_m = sample[1][2][0:4] if self.options['learn_mask']: - S, D, SS, DD, SD, SDM = [ np.clip(x, 0.0, 1.0) for x in ([test_A,test_B] + self.AE_view ([test_A, test_B]) ) ] - SDM, = [ np.repeat (x, (3,), -1) for x in [SDM] ] + S, D, SS, DD, DDM, SD, SDM = [ np.clip(x, 0.0, 1.0) for x in ([test_A,test_B] + self.AE_view ([test_A, test_B]) ) ] + DDM, SDM, = [ np.repeat (x, (3,), -1) for x in [DDM, SDM] ] else: S, D, SS, DD, SD, = [ np.clip(x, 0.0, 1.0) for x in ([test_A,test_B] + self.AE_view ([test_A, test_B]) ) ] + result = [] st = [] for i in range(0, len(test_A)): ar = S[i], SS[i], D[i], DD[i], SD[i] - #if self.options['learn_mask']: - # ar += (SDM[i],) st.append ( np.concatenate ( ar, axis=1) ) - - return [ ('SAE', np.concatenate (st, axis=0 )), ] + + result += [ ('SAE', np.concatenate (st, axis=0 )), ] + + if self.options['learn_mask']: + st_m = [] + for i in range(0, len(test_A)): + ar = S[i], SS[i], D[i], DD[i]*DDM[i], SD[i]*(DDM[i]*SDM[i]) + st_m.append ( np.concatenate ( ar, axis=1) ) + + result += [ ('SAE masked', np.concatenate (st_m, axis=0 )), ] + + return result def predictor_func (self, face): if self.options['learn_mask']: - bgr, mask = self.AE_convert ([face[np.newaxis,...]]) - return bgr[0], mask[0][...,0] + bgr, mask_dst_dstm, mask_src_dstm = self.AE_convert ([face[np.newaxis,...]]) + mask = mask_dst_dstm[0] * mask_src_dstm[0] + return bgr[0], mask[...,0] else: bgr, = self.AE_convert ([face[np.newaxis,...]]) return bgr[0] @@ -440,23 +450,39 @@ class SAEModel(ModelBase): def initialize_nn_functions(): exec (nnlib.import_all(), locals(), globals()) - def BatchNorm(): - return BatchNormalization(axis=-1) + def NormPass(x): + return x + + def Norm(norm=''): + if norm == 'bn': + return BatchNormalization(axis=-1) + else: + return NormPass + + def Act(act='', lrelu_alpha=0.1): + if act == 'prelu': + return PReLU() + else: + return LeakyReLU(alpha=lrelu_alpha) class ResidualBlock(object): - def __init__(self, filters, kernel_size=3, padding='zero', use_reflection_padding=False): + def __init__(self, filters, kernel_size=3, padding='zero', use_reflection_padding=False, norm='', act='', **kwargs): self.filters = filters self.kernel_size = kernel_size self.padding = padding + self.norm = norm + self.act = act def __call__(self, inp): - var_x = inp - var_x = Conv2D(self.filters, kernel_size=self.kernel_size, padding=self.padding)(var_x) - var_x = LeakyReLU(alpha=0.2)(var_x) - var_x = Conv2D(self.filters, kernel_size=self.kernel_size, padding=self.padding)(var_x) - var_x = Add()([var_x, inp]) - var_x = LeakyReLU(alpha=0.2)(var_x) - return var_x + x = inp + x = Conv2D(self.filters, kernel_size=self.kernel_size, padding=self.padding)(x) + x = Act(self.act, lrelu_alpha=0.2)(x) + x = Norm(self.norm)(x) + x = Conv2D(self.filters, kernel_size=self.kernel_size, padding=self.padding)(x) + x = Add()([x, inp]) + x = Act(self.act, lrelu_alpha=0.2)(x) + x = Norm(self.norm)(x) + return x SAEModel.ResidualBlock = ResidualBlock def ResidualBlock_pre (**base_kwargs): @@ -466,9 +492,9 @@ class SAEModel(ModelBase): return func SAEModel.ResidualBlock_pre = ResidualBlock_pre - def downscale (dim, padding='zero'): + def downscale (dim, padding='zero', norm='', act='', **kwargs): def func(x): - return LeakyReLU(0.1)(Conv2D(dim, kernel_size=5, strides=2, padding=padding)(x)) + return Norm(norm)( Act(act) (Conv2D(dim, kernel_size=5, strides=2, padding=padding)(x)) ) return func SAEModel.downscale = downscale @@ -479,9 +505,9 @@ class SAEModel(ModelBase): return func SAEModel.downscale_pre = downscale_pre - def upscale (dim, padding='zero'): + def upscale (dim, padding='zero', norm='', act='', **kwargs): def func(x): - return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, kernel_size=3, strides=1, padding=padding)(x))) + return SubpixelUpscaler()(Norm(norm)(Act(act)(Conv2D(dim * 4, kernel_size=3, strides=1, padding=padding)(x)))) return func SAEModel.upscale = upscale @@ -492,7 +518,7 @@ class SAEModel(ModelBase): return func SAEModel.upscale_pre = upscale_pre - def to_bgr (output_nc, padding='zero'): + def to_bgr (output_nc, padding='zero', **kwargs): def func(x): return Conv2D(output_nc, kernel_size=5, padding=padding, activation='sigmoid')(x) return func @@ -506,10 +532,10 @@ class SAEModel(ModelBase): SAEModel.to_bgr_pre = to_bgr_pre @staticmethod - def LIAEEncFlow(resolution, ch_dims, padding='zero', **kwargs): + def LIAEEncFlow(resolution, ch_dims, **kwargs): exec (nnlib.import_all(), locals(), globals()) - upscale = SAEModel.upscale_pre(padding=padding) - downscale = SAEModel.downscale_pre(padding=padding) + upscale = SAEModel.upscale_pre(**kwargs) + downscale = SAEModel.downscale_pre(**kwargs) def func(input): dims = K.int_shape(input)[-1]*ch_dims @@ -525,9 +551,9 @@ class SAEModel(ModelBase): return func @staticmethod - def LIAEInterFlow(resolution, ae_dims=256, padding='zero', **kwargs): + def LIAEInterFlow(resolution, ae_dims=256, **kwargs): exec (nnlib.import_all(), locals(), globals()) - upscale = SAEModel.upscale_pre(padding=padding) + upscale = SAEModel.upscale_pre(**kwargs) lowest_dense_res=resolution // 16 def func(input): @@ -540,12 +566,12 @@ class SAEModel(ModelBase): return func @staticmethod - def LIAEDecFlow(output_nc,ch_dims, multiscale_count=1, add_residual_blocks=False, padding='zero', **kwargs): + def LIAEDecFlow(output_nc,ch_dims, multiscale_count=1, add_residual_blocks=False, padding='zero', norm='', **kwargs): exec (nnlib.import_all(), locals(), globals()) - upscale = SAEModel.upscale_pre(padding=padding) - to_bgr = SAEModel.to_bgr_pre(padding=padding) + upscale = SAEModel.upscale_pre(**kwargs) + to_bgr = SAEModel.to_bgr_pre(**kwargs) dims = output_nc * ch_dims - ResidualBlock = SAEModel.ResidualBlock_pre(padding=padding) + ResidualBlock = SAEModel.ResidualBlock_pre(**kwargs) def func(input): x = input[0] diff --git a/samplelib/Sample.py b/samplelib/Sample.py index f9953ed..1c311e2 100644 --- a/samplelib/Sample.py +++ b/samplelib/Sample.py @@ -16,24 +16,26 @@ class SampleType(IntEnum): QTY = 5 class Sample(object): - def __init__(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, pitch=None, yaw=None, mirror=None, close_target_list=None): + def __init__(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch=None, yaw=None, mirror=None, close_target_list=None): self.sample_type = sample_type if sample_type is not None else SampleType.IMAGE self.filename = filename self.face_type = face_type self.shape = shape self.landmarks = np.array(landmarks) if landmarks is not None else None + self.ie_polys = ie_polys self.pitch = pitch self.yaw = yaw self.mirror = mirror self.close_target_list = close_target_list - def copy_and_set(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, pitch=None, yaw=None, mirror=None, close_target_list=None): + def copy_and_set(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch=None, yaw=None, mirror=None, close_target_list=None): return Sample( sample_type=sample_type if sample_type is not None else self.sample_type, filename=filename if filename is not None else self.filename, face_type=face_type if face_type is not None else self.face_type, shape=shape if shape is not None else self.shape, landmarks=landmarks if landmarks is not None else self.landmarks.copy(), + ie_polys=ie_polys if ie_polys is not None else self.ie_polys, pitch=pitch if pitch is not None else self.pitch, yaw=yaw if yaw is not None else self.yaw, mirror=mirror if mirror is not None else self.mirror, diff --git a/samplelib/SampleLoader.py b/samplelib/SampleLoader.py index 083d81c..8184409 100644 --- a/samplelib/SampleLoader.py +++ b/samplelib/SampleLoader.py @@ -75,6 +75,7 @@ class SampleLoader: face_type=FaceType.fromString (dflimg.get_face_type()), shape=dflimg.get_shape(), landmarks=dflimg.get_landmarks(), + ie_polys=dflimg.get_ie_polys(), pitch=pitch, yaw=yaw) ) except: diff --git a/samplelib/SampleProcessor.py b/samplelib/SampleProcessor.py index 9e5ddff..242a7ca 100644 --- a/samplelib/SampleProcessor.py +++ b/samplelib/SampleProcessor.py @@ -152,7 +152,7 @@ class SampleProcessor(object): if is_face_sample: if face_mask_type == 1: - img = np.concatenate( (img, LandmarksProcessor.get_image_hull_mask (img.shape, cur_sample.landmarks) ), -1 ) + img = np.concatenate( (img, LandmarksProcessor.get_image_hull_mask (img.shape, cur_sample.landmarks, cur_sample.ie_polys) ), -1 ) elif face_mask_type == 2: mask = LandmarksProcessor.get_image_eye_mask (img.shape, cur_sample.landmarks) mask = np.expand_dims (cv2.blur (mask, ( w // 32, w // 32 ) ), -1) diff --git a/utils/DFLJPG.py b/utils/DFLJPG.py index 362cf5d..7aa31b8 100644 --- a/utils/DFLJPG.py +++ b/utils/DFLJPG.py @@ -2,6 +2,7 @@ import struct import pickle import numpy as np from facelib import FaceType +from imagelib import IEPolys from utils.struct_utils import * class DFLJPG(object): @@ -18,7 +19,7 @@ class DFLJPG(object): with open(filename, "rb") as f: data = f.read() except: - raise FileNotFoundError(data) + raise FileNotFoundError(filename) try: inst = DFLJPG() @@ -150,6 +151,7 @@ class DFLJPG(object): @staticmethod def embed_data(filename, face_type=None, landmarks=None, + ie_polys=None, source_filename=None, source_rect=None, source_landmarks=None, @@ -160,6 +162,7 @@ class DFLJPG(object): inst.setDFLDictData ({ 'face_type': face_type, 'landmarks': landmarks, + 'ie_polys' : ie_polys.dump() if ie_polys is not None else None, 'source_filename': source_filename, 'source_rect': source_rect, 'source_landmarks': source_landmarks, @@ -172,6 +175,29 @@ class DFLJPG(object): except: raise Exception( 'cannot save %s' % (filename) ) + def embed_and_set(self, filename, face_type=None, + landmarks=None, + ie_polys=None, + source_filename=None, + source_rect=None, + source_landmarks=None, + image_to_face_mat=None + ): + if face_type is None: face_type = self.get_face_type() + if landmarks is None: landmarks = self.get_landmarks() + if ie_polys is None: ie_polys = self.get_ie_polys() + if source_filename is None: source_filename = self.get_source_filename() + if source_rect is None: source_rect = self.get_source_rect() + if source_landmarks is None: source_landmarks = self.get_source_landmarks() + if image_to_face_mat is None: image_to_face_mat = self.get_image_to_face_mat() + DFLJPG.embed_data (filename, face_type=face_type, + landmarks=landmarks, + ie_polys=ie_polys, + source_filename=source_filename, + source_rect=source_rect, + source_landmarks=source_landmarks, + image_to_face_mat=image_to_face_mat) + def dump(self): data = b"" @@ -222,6 +248,12 @@ class DFLJPG(object): def get_face_type(self): return self.dfl_dict['face_type'] def get_landmarks(self): return np.array ( self.dfl_dict['landmarks'] ) + def get_ie_polys(self): return IEPolys.load(self.dfl_dict.get('ie_polys',None)) def get_source_filename(self): return self.dfl_dict['source_filename'] def get_source_rect(self): return self.dfl_dict['source_rect'] def get_source_landmarks(self): return np.array ( self.dfl_dict['source_landmarks'] ) + def get_image_to_face_mat(self): + mat = self.dfl_dict.get ('image_to_face_mat', None) + if mat is not None: + return np.array (mat) + return None \ No newline at end of file diff --git a/utils/DFLPNG.py b/utils/DFLPNG.py index 5cee5ce..79ec9c1 100644 --- a/utils/DFLPNG.py +++ b/utils/DFLPNG.py @@ -6,6 +6,7 @@ import zlib import pickle import numpy as np from facelib import FaceType +from imagelib import IEPolys class Chunk(object): def __init__(self, name=None, data=None): @@ -226,7 +227,7 @@ class DFLPNG(object): with open(filename, "rb") as f: data = f.read() except: - raise FileNotFoundError(data) + raise FileNotFoundError(filename) inst = DFLPNG() inst.data = data @@ -267,18 +268,22 @@ class DFLPNG(object): @staticmethod def embed_data(filename, face_type=None, landmarks=None, + ie_polys=None, source_filename=None, source_rect=None, - source_landmarks=None + source_landmarks=None, + image_to_face_mat=None ): inst = DFLPNG.load_raw (filename) inst.setDFLDictData ({ 'face_type': face_type, 'landmarks': landmarks, + 'ie_polys' : ie_polys.dump() if ie_polys is not None else None, 'source_filename': source_filename, 'source_rect': source_rect, - 'source_landmarks': source_landmarks + 'source_landmarks': source_landmarks, + 'image_to_face_mat':image_to_face_mat }) try: @@ -287,6 +292,29 @@ class DFLPNG(object): except: raise Exception( 'cannot save %s' % (filename) ) + def embed_and_set(self, filename, face_type=None, + landmarks=None, + ie_polys=None, + source_filename=None, + source_rect=None, + source_landmarks=None, + image_to_face_mat=None + ): + if face_type is None: face_type = self.get_face_type() + if landmarks is None: landmarks = self.get_landmarks() + if ie_polys is None: ie_polys = self.get_ie_polys() + if source_filename is None: source_filename = self.get_source_filename() + if source_rect is None: source_rect = self.get_source_rect() + if source_landmarks is None: source_landmarks = self.get_source_landmarks() + if image_to_face_mat is None: image_to_face_mat = self.get_image_to_face_mat() + DFLPNG.embed_data (filename, face_type=face_type, + landmarks=landmarks, + ie_polys=ie_polys, + source_filename=source_filename, + source_rect=source_rect, + source_landmarks=source_landmarks, + image_to_face_mat=image_to_face_mat) + def dump(self): data = PNG_HEADER for chunk in self.chunks: @@ -326,9 +354,15 @@ class DFLPNG(object): def get_face_type(self): return self.fcwp_dict['face_type'] def get_landmarks(self): return np.array ( self.fcwp_dict['landmarks'] ) + def get_ie_polys(self): return IEPolys.load(self.fcwp_dict.get('ie_polys',None)) def get_source_filename(self): return self.fcwp_dict['source_filename'] def get_source_rect(self): return self.fcwp_dict['source_rect'] def get_source_landmarks(self): return np.array ( self.fcwp_dict['source_landmarks'] ) + def get_image_to_face_mat(self): + mat = self.fcwp_dict.get ('image_to_face_mat', None) + if mat is not None: + return np.array (mat) + return None def __str__(self): return "".format(len(self.chunks), **self.__dict__)