mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-07-07 13:32:09 -07:00
changed help message for pixel loss:
Pixel loss may help to enhance fine details and stabilize face color. Use it only if quality does not improve over time. SAE: previous SAE model will not work with this update. Greatly decreased chance of model collapse. Increased model accuracy. Residual blocks now default and this option has been removed. Improved 'learn mask'. Added masked preview (switch by space key) Converter: fixed rct/lct in seamless mode added mask mode (6) learned*FAN-prd*FAN-dst added mask editor, its created for refining dataset for FANSeg model, and not for production, but you can spend your time and test it in regular fakes with face obstructions
This commit is contained in:
parent
01e98cde8e
commit
5ac7e5d7f1
22 changed files with 715 additions and 387 deletions
|
@ -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,25 +167,27 @@ 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
|
||||
|
||||
|
@ -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) ]
|
||||
|
@ -310,14 +313,19 @@ class ConverterMasked(Converter):
|
|||
if self.masked_hist_match:
|
||||
hist_mask_a *= prd_face_mask_a
|
||||
|
||||
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)
|
||||
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 + 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)
|
||||
|
|
|
@ -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)
|
||||
|
|
104
imagelib/IEPolys.py
Normal file
104
imagelib/IEPolys.py
Normal file
|
@ -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
|
|
@ -21,3 +21,5 @@ from .color_transfer import linear_color_transfer
|
|||
from .DCSCN import DCSCN
|
||||
|
||||
from .common import normalize_channels
|
||||
|
||||
from .IEPolys import IEPolys
|
|
@ -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):
|
||||
h,w,c = shape
|
||||
try:
|
||||
size = shape[1]
|
||||
pil_font = _get_pil_font( localization.get_default_ttf_font_name() , size)
|
||||
text_width, text_height = pil_font.getsize(text)
|
||||
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
|
||||
|
|
|
@ -18,6 +18,8 @@ except:
|
|||
class InteractBase(object):
|
||||
EVENT_LBUTTONDOWN = 1
|
||||
EVENT_LBUTTONUP = 2
|
||||
EVENT_MBUTTONDOWN = 3
|
||||
EVENT_MBUTTONUP = 4
|
||||
EVENT_RBUTTONDOWN = 5
|
||||
EVENT_RBUTTONUP = 6
|
||||
EVENT_MOUSEWHEEL = 10
|
||||
|
@ -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
|
||||
|
|
17
main.py
17
main.py
|
@ -186,14 +186,16 @@ 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()
|
||||
|
@ -201,9 +203,6 @@ if __name__ == "__main__":
|
|||
parser.set_defaults(func=bad_args)
|
||||
|
||||
arguments = parser.parse_args()
|
||||
|
||||
#os.environ['force_plaidML'] = '1'
|
||||
|
||||
arguments.func(arguments)
|
||||
|
||||
print ("Done.")
|
||||
|
|
|
@ -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)
|
376
mainscripts/MaskEditorTool.py
Normal file
376
mainscripts/MaskEditorTool.py
Normal file
|
@ -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()
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
@ -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] )
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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)
|
||||
|
||||
|
|
|
@ -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)
|
||||
|
||||
|
|
|
@ -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)
|
||||
|
||||
|
|
|
@ -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)
|
||||
|
||||
|
|
|
@ -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():
|
||||
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]
|
||||
|
|
|
@ -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,
|
||||
|
|
|
@ -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:
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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
|
|
@ -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 "<PNG length={length} chunks={}>".format(len(self.chunks), **self.__dict__)
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue