diff --git a/converters/Converter.py b/converters/Converter.py index f79f02e..a284419 100644 --- a/converters/Converter.py +++ b/converters/Converter.py @@ -7,25 +7,25 @@ class Converter(object): TYPE_FACE = 0 #calls convert_face TYPE_IMAGE = 1 #calls convert_image without landmarks TYPE_IMAGE_WITH_LANDMARKS = 2 #calls convert_image with landmarks - + #overridable def __init__(self, predictor_func, type): self.predictor_func = predictor_func self.type = type - + #overridable def convert_face (self, img_bgr, img_face_landmarks, debug): - #return float32 image + #return float32 image #if debug , return tuple ( images of any size and channels, ...) return image - + #overridable def convert_image (self, img_bgr, img_landmarks, debug): #img_landmarks not None, if input image is png with embedded data - #return float32 image + #return float32 image #if debug , return tuple ( images of any size and channels, ...) return image - + #overridable def dummy_predict(self): #do dummy predict here @@ -33,8 +33,8 @@ class Converter(object): def copy(self): return copy.copy(self) - + def copy_and_set_predictor(self, predictor_func): result = self.copy() result.predictor_func = predictor_func - return result \ No newline at end of file + return result diff --git a/converters/ConverterImage.py b/converters/ConverterImage.py index 77a4dfc..c507b6c 100644 --- a/converters/ConverterImage.py +++ b/converters/ConverterImage.py @@ -7,7 +7,7 @@ import numpy as np from utils import image_utils ''' -predictor_func: +predictor_func: input: [predictor_input_size, predictor_input_size, BGR] output: [predictor_input_size, predictor_input_size, BGR] ''' @@ -16,18 +16,18 @@ class ConverterImage(Converter): #override def __init__(self, predictor_func, - predictor_input_size=0, + predictor_input_size=0, output_size=0): - + super().__init__(predictor_func, Converter.TYPE_IMAGE) - + self.predictor_input_size = predictor_input_size - self.output_size = output_size - + self.output_size = output_size + #override def dummy_predict(self): self.predictor_func ( np.zeros ( (self.predictor_input_size, self.predictor_input_size,3), dtype=np.float32) ) - + #override def convert_image (self, img_bgr, img_landmarks, debug): img_size = img_bgr.shape[1], img_bgr.shape[0] diff --git a/facelib/DLIBExtractor.py b/facelib/DLIBExtractor.py index fced9f2..a91164d 100644 --- a/facelib/DLIBExtractor.py +++ b/facelib/DLIBExtractor.py @@ -4,36 +4,36 @@ import cv2 from pathlib import Path -class DLIBExtractor(object): +class DLIBExtractor(object): def __init__(self, dlib): - self.scale_to = 1850 - #3100 eats ~1.687GB VRAM on 2GB 730 desktop card, but >4Gb on 6GB card, + self.scale_to = 1850 + #3100 eats ~1.687GB VRAM on 2GB 730 desktop card, but >4Gb on 6GB card, #but 3100 doesnt work on 2GB 850M notebook card, I cant understand this behaviour #1850 works on 2GB 850M notebook card, works faster than 3100, produces good result self.dlib = dlib - def __enter__(self): + def __enter__(self): self.dlib_cnn_face_detector = self.dlib.cnn_face_detection_model_v1( str(Path(__file__).parent / "mmod_human_face_detector.dat") ) - self.dlib_cnn_face_detector ( np.zeros ( (self.scale_to, self.scale_to, 3), dtype=np.uint8), 0 ) + self.dlib_cnn_face_detector ( np.zeros ( (self.scale_to, self.scale_to, 3), dtype=np.uint8), 0 ) return self - + def __exit__(self, exc_type=None, exc_value=None, traceback=None): del self.dlib_cnn_face_detector return False #pass exception between __enter__ and __exit__ to outter level - + def extract_from_bgr (self, input_image): input_image = input_image[:,:,::-1].copy() (h, w, ch) = input_image.shape - detected_faces = [] + detected_faces = [] input_scale = self.scale_to / (w if w > h else h) input_image = cv2.resize (input_image, ( int(w*input_scale), int(h*input_scale) ), interpolation=cv2.INTER_LINEAR) detected_faces = self.dlib_cnn_face_detector(input_image, 0) - result = [] + result = [] for d_rect in detected_faces: if type(d_rect) == self.dlib.mmod_rectangle: - d_rect = d_rect.rect + d_rect = d_rect.rect left, top, right, bottom = d_rect.left(), d_rect.top(), d_rect.right(), d_rect.bottom() result.append ( (int(left/input_scale), int(top/input_scale), int(right/input_scale), int(bottom/input_scale)) ) diff --git a/facelib/FANSegmentator.py b/facelib/FANSegmentator.py index d84f23d..1209c16 100644 --- a/facelib/FANSegmentator.py +++ b/facelib/FANSegmentator.py @@ -8,16 +8,16 @@ from interact import interact as io class FANSegmentator(object): def __init__ (self, resolution, face_type_str, load_weights=True, weights_file_root=None): exec( nnlib.import_all(), locals(), globals() ) - + self.model = FANSegmentator.BuildModel(resolution, ngf=32) - + if weights_file_root: weights_file_root = Path(weights_file_root) else: weights_file_root = Path(__file__).parent - + self.weights_path = weights_file_root / ('FANSeg_%d_%s.h5' % (resolution, face_type_str) ) - + if load_weights: self.model.load_weights (str(self.weights_path)) else: @@ -31,19 +31,19 @@ class FANSegmentator(object): def __enter__(self): return self - + def __exit__(self, exc_type=None, exc_value=None, traceback=None): return False #pass exception between __enter__ and __exit__ to outter level - + def save_weights(self): self.model.save_weights (str(self.weights_path)) - + def train_on_batch(self, inp, outp): return self.model.train_on_batch(inp, outp) - + def extract_from_bgr (self, input_image): return np.clip ( (self.model.predict(input_image) + 1) / 2.0, 0, 1.0 ) - + @staticmethod def BuildModel ( resolution, ngf=64): exec( nnlib.import_all(), locals(), globals() ) @@ -53,7 +53,7 @@ class FANSegmentator(object): x = FANSegmentator.DecFlow(ngf=ngf)(x) model = Model(inp,x) return model - + @staticmethod def EncFlow(ngf=64, num_downs=4): exec( nnlib.import_all(), locals(), globals() ) @@ -65,19 +65,19 @@ class FANSegmentator(object): def downscale (dim): def func(x): return LeakyReLU(0.1)(XNormalization(Conv2D(dim, kernel_size=5, strides=2, padding='same', kernel_initializer=RandomNormal(0, 0.02))(x))) - return func - - def func(input): + return func + + def func(input): x = input - + result = [] for i in range(num_downs): x = downscale ( min(ngf*(2**i), ngf*8) )(x) - result += [x] - + result += [x] + return result return func - + @staticmethod def DecFlow(output_nc=1, ngf=64, activation='tanh'): exec (nnlib.import_all(), locals(), globals()) @@ -85,23 +85,23 @@ class FANSegmentator(object): use_bias = True def XNormalization(x): return InstanceNormalization (axis=3, gamma_initializer=RandomNormal(1., 0.02))(x) - + def Conv2D (filters, kernel_size, strides=(1, 1), padding='valid', data_format=None, dilation_rate=(1, 1), activation=None, use_bias=use_bias, kernel_initializer=RandomNormal(0, 0.02), bias_initializer='zeros', kernel_regularizer=None, bias_regularizer=None, activity_regularizer=None, kernel_constraint=None, bias_constraint=None): return keras.layers.Conv2D( filters=filters, kernel_size=kernel_size, strides=strides, padding=padding, data_format=data_format, dilation_rate=dilation_rate, activation=activation, use_bias=use_bias, kernel_initializer=kernel_initializer, bias_initializer=bias_initializer, kernel_regularizer=kernel_regularizer, bias_regularizer=bias_regularizer, activity_regularizer=activity_regularizer, kernel_constraint=kernel_constraint, bias_constraint=bias_constraint ) - + def upscale (dim): def func(x): return SubpixelUpscaler()( LeakyReLU(0.1)(XNormalization(Conv2D(dim, kernel_size=3, strides=1, padding='same', kernel_initializer=RandomNormal(0, 0.02))(x)))) - return func - + return func + def func(input): input_len = len(input) - + x = input[input_len-1] - for i in range(input_len-1, -1, -1): + for i in range(input_len-1, -1, -1): x = upscale( min(ngf* (2**i) *4, ngf*8 *4 ) )(x) if i != 0: x = Concatenate(axis=3)([ input[i-1] , x]) - + return Conv2D(output_nc, 3, 1, 'same', activation=activation)(x) - return func \ No newline at end of file + return func diff --git a/facelib/FaceType.py b/facelib/FaceType.py index 782f39d..f0d5530 100644 --- a/facelib/FaceType.py +++ b/facelib/FaceType.py @@ -3,7 +3,7 @@ from enum import IntEnum class FaceType(IntEnum): HALF = 0, FULL = 1, - HEAD = 2, + HEAD = 2, AVATAR = 3, #centered nose only MARK_ONLY = 4, #no align at all, just embedded faceinfo QTY = 5 @@ -13,12 +13,12 @@ class FaceType(IntEnum): r = from_string_dict.get (s.lower()) if r is None: raise Exception ('FaceType.fromString value error') - return r - - @staticmethod + return r + + @staticmethod def toString (face_type): return to_string_list[face_type] - + from_string_dict = {'half_face': FaceType.HALF, 'full_face': FaceType.FULL, 'head' : FaceType.HEAD, @@ -29,6 +29,5 @@ to_string_list = [ 'half_face', 'full_face', 'head', 'avatar', - 'mark_only' + 'mark_only' ] - diff --git a/facelib/LandmarksExtractor.py b/facelib/LandmarksExtractor.py index 4ca15c4..b144bcf 100644 --- a/facelib/LandmarksExtractor.py +++ b/facelib/LandmarksExtractor.py @@ -10,38 +10,38 @@ class LandmarksExtractor(object): def __init__ (self, keras): self.keras = keras K = self.keras.backend - - def __enter__(self): + + def __enter__(self): keras_model_path = Path(__file__).parent / "2DFAN-4.h5" if not keras_model_path.exists(): return None - self.keras_model = self.keras.models.load_model (str(keras_model_path)) - + self.keras_model = self.keras.models.load_model (str(keras_model_path)) + return self - + def __exit__(self, exc_type=None, exc_value=None, traceback=None): del self.keras_model return False #pass exception between __enter__ and __exit__ to outter level - + def extract_from_bgr (self, input_image, rects, second_pass_extractor=None): input_image = input_image[:,:,::-1].copy() (h, w, ch) = input_image.shape - + landmarks = [] for (left, top, right, bottom) in rects: try: center = np.array( [ (left + right) / 2.0, (top + bottom) / 2.0] ) #center[1] -= (bottom - top) * 0.12 scale = (right - left + bottom - top) / 195.0 - + image = self.crop(input_image, center, scale).astype(np.float32) image = np.expand_dims(image, 0) - + predicted = self.keras_model.predict (image).transpose (0,3,1,2) - + pts_img = self.get_pts_from_predict ( predicted[-1], center, scale) - pts_img = [ ( int(pt[0]), int(pt[1]) ) for pt in pts_img ] + pts_img = [ ( int(pt[0]), int(pt[1]) ) for pt in pts_img ] landmarks.append ( ( (left, top, right, bottom),pts_img ) ) except Exception as e: landmarks.append ( ( (left, top, right, bottom), None ) ) @@ -52,26 +52,26 @@ class LandmarksExtractor(object): rect, lmrks = landmarks[i] if lmrks is None: continue - + image_to_face_mat = LandmarksProcessor.get_transform_mat (lmrks, 256, FaceType.FULL) face_image = cv2.warpAffine(input_image, image_to_face_mat, (256, 256), cv2.INTER_CUBIC) - + rects2 = second_pass_extractor.extract_from_bgr(face_image) if len(rects2) != 1: #dont do second pass if more than 1 or zero faces detected in cropped image continue - + rect2 = rects2[0] - + lmrks2 = self.extract_from_bgr (face_image, [rect2] )[0][1] - source_lmrks2 = LandmarksProcessor.transform_points (lmrks2, image_to_face_mat, True) + source_lmrks2 = LandmarksProcessor.transform_points (lmrks2, image_to_face_mat, True) landmarks[i] = (rect, source_lmrks2) except: continue - + return landmarks def transform(self, point, center, scale, resolution): - pt = np.array ( [point[0], point[1], 1.0] ) + pt = np.array ( [point[0], point[1], 1.0] ) h = 200.0 * scale m = np.eye(3) m[0,0] = resolution / h @@ -80,11 +80,11 @@ class LandmarksExtractor(object): m[1,2] = resolution * ( -center[1] / h + 0.5 ) m = np.linalg.inv(m) return np.matmul (m, pt)[0:2] - + def crop(self, image, center, scale, resolution=256.0): ul = self.transform([1, 1], center, scale, resolution).astype( np.int ) br = self.transform([resolution, resolution], center, scale, resolution).astype( np.int ) - + if image.ndim > 2: newDim = np.array([br[1] - ul[1], br[0] - ul[0], image.shape[2]], dtype=np.int32) newImg = np.zeros(newDim, dtype=np.uint8) @@ -98,14 +98,14 @@ class LandmarksExtractor(object): oldX = np.array([max(1, ul[0] + 1), min(br[0], wd)], dtype=np.int32) oldY = np.array([max(1, ul[1] + 1), min(br[1], ht)], dtype=np.int32) newImg[newY[0] - 1:newY[1], newX[0] - 1:newX[1] ] = image[oldY[0] - 1:oldY[1], oldX[0] - 1:oldX[1], :] - + newImg = cv2.resize(newImg, dsize=(int(resolution), int(resolution)), interpolation=cv2.INTER_LINEAR) return newImg - + def get_pts_from_predict(self, a, center, scale): - b = a.reshape ( (a.shape[0], a.shape[1]*a.shape[2]) ) + b = a.reshape ( (a.shape[0], a.shape[1]*a.shape[2]) ) c = b.argmax(1).reshape ( (a.shape[0], 1) ).repeat(2, axis=1).astype(np.float) - c[:,0] %= a.shape[2] + c[:,0] %= a.shape[2] c[:,1] = np.apply_along_axis ( lambda x: np.floor(x / a.shape[2]), 0, c[:,1] ) for i in range(a.shape[0]): @@ -113,6 +113,6 @@ class LandmarksExtractor(object): if pX > 0 and pX < 63 and pY > 0 and pY < 63: diff = np.array ( [a[i,pY,pX+1]-a[i,pY,pX-1], a[i,pY+1,pX]-a[i,pY-1,pX]] ) c[i] += np.sign(diff)*0.25 - + c += 0.5 - return [ self.transform (c[i], center, scale, a.shape[2]) for i in range(a.shape[0]) ] \ No newline at end of file + return [ self.transform (c[i], center, scale, a.shape[2]) for i in range(a.shape[0]) ] diff --git a/facelib/LandmarksProcessor.py b/facelib/LandmarksProcessor.py index 88470ff..ccaba46 100644 --- a/facelib/LandmarksProcessor.py +++ b/facelib/LandmarksProcessor.py @@ -36,7 +36,7 @@ landmarks_68_pt = { "mouth": (48,68), "left_eye": (42, 48), "nose": (27, 36), # missed one point "jaw": (0, 17) } - + landmarks_68_3D = np.array( [ [-73.393523 , -29.801432 , 47.667532 ], @@ -107,20 +107,20 @@ landmarks_68_3D = np.array( [ [8.449166 , 30.596216 , -20.671489 ], [0.205322 , 31.408738 , -21.903670 ], [-7.198266 , 30.844876 , -20.328022 ] ], dtype=np.float32) - + def get_transform_mat (image_landmarks, output_size, face_type, scale=1.0): if not isinstance(image_landmarks, np.ndarray): - image_landmarks = np.array (image_landmarks) - + image_landmarks = np.array (image_landmarks) + if face_type == FaceType.AVATAR: centroid = np.mean (image_landmarks, axis=0) - + mat = umeyama(image_landmarks[17:], landmarks_2D, True)[0:2] a, c = mat[0,0], mat[1,0] scale = math.sqrt((a * a) + (c * c)) - + padding = (output_size / 64) * 32 - + mat = np.eye ( 2,3 ) mat[0,2] = -centroid[0] mat[1,2] = -centroid[1] @@ -135,15 +135,15 @@ def get_transform_mat (image_landmarks, output_size, face_type, scale=1.0): padding = (output_size / 64) * 24 else: raise ValueError ('wrong face_type: ', face_type) - + mat = umeyama(image_landmarks[17:], landmarks_2D, True)[0:2] mat = mat * (output_size - 2 * padding) - mat[:,2] += padding + mat[:,2] += padding mat *= (1 / scale) mat[:,2] += -output_size*( ( (1 / scale) - 1.0 ) / 2 ) - + return mat - + def transform_points(points, mat, invert=False): if invert: mat = cv2.invertAffineTransform (mat) @@ -151,68 +151,68 @@ def transform_points(points, mat, invert=False): points = cv2.transform(points, mat, points.shape) points = np.squeeze(points) return points - - -def get_image_hull_mask (image_shape, image_landmarks): + + +def get_image_hull_mask (image_shape, image_landmarks): if len(image_landmarks) != 68: raise Exception('get_image_hull_mask works only with 68 landmarks') - + hull_mask = np.zeros(image_shape[0:2]+(1,),dtype=np.float32) - cv2.fillConvexPoly( hull_mask, cv2.convexHull( - np.concatenate ( (image_landmarks[0:9], + cv2.fillConvexPoly( hull_mask, cv2.convexHull( + np.concatenate ( (image_landmarks[0:9], image_landmarks[17:18]))) , (1,) ) - + cv2.fillConvexPoly( hull_mask, cv2.convexHull( - np.concatenate ( (image_landmarks[8:17], + np.concatenate ( (image_landmarks[8:17], image_landmarks[26:27]))) , (1,) ) - + cv2.fillConvexPoly( hull_mask, cv2.convexHull( - np.concatenate ( (image_landmarks[17:20], + np.concatenate ( (image_landmarks[17:20], image_landmarks[8:9]))) , (1,) ) - + cv2.fillConvexPoly( hull_mask, cv2.convexHull( - np.concatenate ( (image_landmarks[24:27], + np.concatenate ( (image_landmarks[24:27], image_landmarks[8:9]))) , (1,) ) - + cv2.fillConvexPoly( hull_mask, cv2.convexHull( - np.concatenate ( (image_landmarks[19:25], + np.concatenate ( (image_landmarks[19:25], image_landmarks[8:9], ))) , (1,) ) - + cv2.fillConvexPoly( hull_mask, cv2.convexHull( - np.concatenate ( (image_landmarks[17:22], + np.concatenate ( (image_landmarks[17:22], image_landmarks[27:28], image_landmarks[31:36], image_landmarks[8:9] - ))) , (1,) ) - + ))) , (1,) ) + cv2.fillConvexPoly( hull_mask, cv2.convexHull( - np.concatenate ( (image_landmarks[22:27], + np.concatenate ( (image_landmarks[22:27], image_landmarks[27:28], image_landmarks[31:36], image_landmarks[8:9] - ))) , (1,) ) - + ))) , (1,) ) + #nose cv2.fillConvexPoly( hull_mask, cv2.convexHull(image_landmarks[27:36]), (1,) ) - + return hull_mask -def get_image_eye_mask (image_shape, image_landmarks): +def get_image_eye_mask (image_shape, image_landmarks): if len(image_landmarks) != 68: raise Exception('get_image_eye_mask works only with 68 landmarks') - + hull_mask = np.zeros(image_shape[0:2]+(1,),dtype=np.float32) cv2.fillConvexPoly( hull_mask, cv2.convexHull( image_landmarks[36:42]), (1,) ) cv2.fillConvexPoly( hull_mask, cv2.convexHull( image_landmarks[42:48]), (1,) ) - + return hull_mask - -def get_image_hull_mask_3D (image_shape, image_landmarks): + +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): @@ -224,7 +224,7 @@ def blur_image_hull_mask (hull_mask): leny = maxy - miny; masky = int(minx+(lenx//2)) maskx = int(miny+(leny//2)) - lowest_len = min (lenx, leny) + lowest_len = min (lenx, leny) ero = int( lowest_len * 0.085 ) blur = int( lowest_len * 0.10 ) @@ -233,10 +233,10 @@ def blur_image_hull_mask (hull_mask): hull_mask = np.expand_dims (hull_mask,-1) return hull_mask - -def get_blurred_image_hull_mask(image_shape, image_landmarks): + +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], @@ -246,23 +246,23 @@ mirror_idxs = [ [5,11], [6,10], [7,9], - + [17,26], [18,25], [19,24], [20,23], - [21,22], - + [21,22], + [36,45], [37,44], [38,43], [39,42], [40,47], - [41,46], - + [41,46], + [31,35], [32,34], - + [50,52], [49,53], [48,54], @@ -271,28 +271,28 @@ mirror_idxs = [ [67,65], [60,64], [61,63] ] - -def mirror_landmarks (landmarks, val): + +def mirror_landmarks (landmarks, val): result = landmarks.copy() - + for idx in mirror_idxs: result [ idx ] = result [ idx[::-1] ] result[:,0] = val - result[:,0] - 1 return result - + def draw_landmarks (image, image_landmarks, color=(0,255,0), transparent_mask=False): if len(image_landmarks) != 68: - raise Exception('get_image_eye_mask works only with 68 landmarks') - - jaw = image_landmarks[slice(*landmarks_68_pt["jaw"])] + raise Exception('get_image_eye_mask works only with 68 landmarks') + + jaw = image_landmarks[slice(*landmarks_68_pt["jaw"])] right_eyebrow = image_landmarks[slice(*landmarks_68_pt["right_eyebrow"])] left_eyebrow = image_landmarks[slice(*landmarks_68_pt["left_eyebrow"])] - mouth = image_landmarks[slice(*landmarks_68_pt["mouth"])] - right_eye = image_landmarks[slice(*landmarks_68_pt["right_eye"])] - left_eye = image_landmarks[slice(*landmarks_68_pt["left_eye"])] - nose = image_landmarks[slice(*landmarks_68_pt["nose"])] - + mouth = image_landmarks[slice(*landmarks_68_pt["mouth"])] + right_eye = image_landmarks[slice(*landmarks_68_pt["right_eye"])] + left_eye = image_landmarks[slice(*landmarks_68_pt["left_eye"])] + nose = image_landmarks[slice(*landmarks_68_pt["nose"])] + # open shapes cv2.polylines(image, tuple(np.array([v]) for v in ( right_eyebrow, jaw, left_eyebrow, np.concatenate((nose, [nose[-6]])) )), False, color, lineType=cv2.LINE_AA) @@ -303,9 +303,9 @@ def draw_landmarks (image, image_landmarks, color=(0,255,0), transparent_mask=Fa for x, y in np.concatenate((right_eyebrow, left_eyebrow, mouth, right_eye, left_eye, nose), axis=0): cv2.circle(image, (x, y), 1, color, 1, lineType=cv2.LINE_AA) # jaw big circles - for x, y in jaw: + for x, y in jaw: cv2.circle(image, (x, y), 2, color, lineType=cv2.LINE_AA) - + if transparent_mask: mask = get_image_hull_mask (image.shape, image_landmarks) image[...] = ( image * (1-mask) + image * mask / 2 )[...] @@ -314,24 +314,24 @@ def draw_rect_landmarks (image, rect, image_landmarks, face_size, face_type, tra draw_landmarks(image, image_landmarks, color=landmarks_color, transparent_mask=transparent_mask) image_utils.draw_rect (image, rect, (255,0,0), 2 ) - image_to_face_mat = get_transform_mat (image_landmarks, face_size, face_type) + image_to_face_mat = get_transform_mat (image_landmarks, face_size, face_type) points = transform_points ( [ (0,0), (0,face_size-1), (face_size-1, face_size-1), (face_size-1,0) ], image_to_face_mat, True) - image_utils.draw_polygon (image, points, (0,0,255), 2) - + image_utils.draw_polygon (image, points, (0,0,255), 2) + def calc_face_pitch(landmarks): if not isinstance(landmarks, np.ndarray): landmarks = np.array (landmarks) - t = ( (landmarks[6][1]-landmarks[8][1]) + (landmarks[10][1]-landmarks[8][1]) ) / 2.0 + t = ( (landmarks[6][1]-landmarks[8][1]) + (landmarks[10][1]-landmarks[8][1]) ) / 2.0 b = landmarks[8][1] return float(b-t) - + def calc_face_yaw(landmarks): if not isinstance(landmarks, np.ndarray): landmarks = np.array (landmarks) - l = ( (landmarks[27][0]-landmarks[0][0]) + (landmarks[28][0]-landmarks[1][0]) + (landmarks[29][0]-landmarks[2][0]) ) / 3.0 + l = ( (landmarks[27][0]-landmarks[0][0]) + (landmarks[28][0]-landmarks[1][0]) + (landmarks[29][0]-landmarks[2][0]) ) / 3.0 r = ( (landmarks[16][0]-landmarks[27][0]) + (landmarks[15][0]-landmarks[28][0]) + (landmarks[14][0]-landmarks[29][0]) ) / 3.0 return float(r-l) - + #returns pitch,yaw [-1...+1] def estimate_pitch_yaw(aligned_256px_landmarks): shape = (256,256) @@ -347,8 +347,8 @@ def estimate_pitch_yaw(aligned_256px_landmarks): aligned_256px_landmarks.astype(np.float32), camera_matrix, np.zeros((4, 1)) ) - + pitch, yaw, _ = mathlib.rotationMatrixToEulerAngles( cv2.Rodrigues(rotation_vector)[0] ) pitch = np.clip ( pitch*1.25, -1.0, 1.0 ) yaw = np.clip ( yaw*1.25, -1.0, 1.0 ) - return pitch, yaw \ No newline at end of file + return pitch, yaw diff --git a/facelib/MTCExtractor.py b/facelib/MTCExtractor.py index e5f630a..090a3a0 100644 --- a/facelib/MTCExtractor.py +++ b/facelib/MTCExtractor.py @@ -5,10 +5,10 @@ import cv2 from pathlib import Path from nnlib import nnlib -class MTCExtractor(object): +class MTCExtractor(object): def __init__(self): self.scale_to = 1920 - + self.min_face_size = self.scale_to * 0.042 self.thresh1 = 0.7 self.thresh2 = 0.85 @@ -26,12 +26,12 @@ class MTCExtractor(object): x = Conv2D (32, kernel_size=(3,3), strides=(1,1), padding='valid', name="conv3")(x) x = PReLU (shared_axes=[1,2], name="PReLU3" )(x) prob = Conv2D (2, kernel_size=(1,1), strides=(1,1), padding='valid', name="conv41")(x) - prob = Softmax()(prob) + prob = Softmax()(prob) x = Conv2D (4, kernel_size=(1,1), strides=(1,1), padding='valid', name="conv42")(x) - PNet_model = Model(PNet_Input, [x,prob] ) + PNet_model = Model(PNet_Input, [x,prob] ) PNet_model.load_weights ( (Path(__file__).parent / 'mtcnn_pnet.h5').__str__() ) - + RNet_Input = Input ( (24, 24, 3) ) x = RNet_Input x = Conv2D (28, kernel_size=(3,3), strides=(1,1), padding='valid', name="conv1")(x) @@ -39,18 +39,18 @@ class MTCExtractor(object): x = MaxPooling2D( pool_size=(3,3), strides=(2,2), padding='same' ) (x) x = Conv2D (48, kernel_size=(3,3), strides=(1,1), padding='valid', name="conv2")(x) x = PReLU (shared_axes=[1,2], name="prelu2" )(x) - x = MaxPooling2D( pool_size=(3,3), strides=(2,2), padding='valid' ) (x) + x = MaxPooling2D( pool_size=(3,3), strides=(2,2), padding='valid' ) (x) x = Conv2D (64, kernel_size=(2,2), strides=(1,1), padding='valid', name="conv3")(x) x = PReLU (shared_axes=[1,2], name="prelu3" )(x) x = Lambda ( lambda x: K.reshape (x, (-1, np.prod(K.int_shape(x)[1:]),) ), output_shape=(np.prod(K.int_shape(x)[1:]),) ) (x) - x = Dense (128, name='conv4')(x) + x = Dense (128, name='conv4')(x) x = PReLU (name="prelu4" )(x) prob = Dense (2, name='conv51')(x) - prob = Softmax()(prob) - x = Dense (4, name='conv52')(x) - RNet_model = Model(RNet_Input, [x,prob] ) + prob = Softmax()(prob) + x = Dense (4, name='conv52')(x) + RNet_model = Model(RNet_Input, [x,prob] ) RNet_model.load_weights ( (Path(__file__).parent / 'mtcnn_rnet.h5').__str__() ) - + ONet_Input = Input ( (48, 48, 3) ) x = ONet_Input x = Conv2D (32, kernel_size=(3,3), strides=(1,1), padding='valid', name="conv1")(x) @@ -58,20 +58,20 @@ class MTCExtractor(object): x = MaxPooling2D( pool_size=(3,3), strides=(2,2), padding='same' ) (x) x = Conv2D (64, kernel_size=(3,3), strides=(1,1), padding='valid', name="conv2")(x) x = PReLU (shared_axes=[1,2], name="prelu2" )(x) - x = MaxPooling2D( pool_size=(3,3), strides=(2,2), padding='valid' ) (x) + x = MaxPooling2D( pool_size=(3,3), strides=(2,2), padding='valid' ) (x) x = Conv2D (64, kernel_size=(3,3), strides=(1,1), padding='valid', name="conv3")(x) x = PReLU (shared_axes=[1,2], name="prelu3" )(x) - x = MaxPooling2D( pool_size=(2,2), strides=(2,2), padding='same' ) (x) + x = MaxPooling2D( pool_size=(2,2), strides=(2,2), padding='same' ) (x) x = Conv2D (128, kernel_size=(2,2), strides=(1,1), padding='valid', name="conv4")(x) x = PReLU (shared_axes=[1,2], name="prelu4" )(x) - x = Lambda ( lambda x: K.reshape (x, (-1, np.prod(K.int_shape(x)[1:]),) ), output_shape=(np.prod(K.int_shape(x)[1:]),) ) (x) + x = Lambda ( lambda x: K.reshape (x, (-1, np.prod(K.int_shape(x)[1:]),) ), output_shape=(np.prod(K.int_shape(x)[1:]),) ) (x) x = Dense (256, name='conv5')(x) x = PReLU (name="prelu5" )(x) prob = Dense (2, name='conv61')(x) - prob = Softmax()(prob) + prob = Softmax()(prob) x1 = Dense (4, name='conv62')(x) - x2 = Dense (10, name='conv63')(x) - ONet_model = Model(ONet_Input, [x1,x2,prob] ) + x2 = Dense (10, name='conv63')(x) + ONet_model = Model(ONet_Input, [x1,x2,prob] ) ONet_model.load_weights ( (Path(__file__).parent / 'mtcnn_onet.h5').__str__() ) self.pnet_fun = K.function ( PNet_model.inputs, PNet_model.outputs ) @@ -79,13 +79,13 @@ class MTCExtractor(object): self.onet_fun = K.function ( ONet_model.inputs, ONet_model.outputs ) def __enter__(self): - faces, pnts = detect_face ( np.zeros ( (self.scale_to, self.scale_to, 3)), self.min_face_size, self.pnet_fun, self.rnet_fun, self.onet_fun, [ self.thresh1, self.thresh2, self.thresh3 ], self.scale_factor ) - + faces, pnts = detect_face ( np.zeros ( (self.scale_to, self.scale_to, 3)), self.min_face_size, self.pnet_fun, self.rnet_fun, self.onet_fun, [ self.thresh1, self.thresh2, self.thresh3 ], self.scale_factor ) + return self - + def __exit__(self, exc_type=None, exc_value=None, traceback=None): return False #pass exception between __enter__ and __exit__ to outter level - + def extract_from_bgr (self, input_image): input_image = input_image[:,:,::-1].copy() (h, w, ch) = input_image.shape @@ -95,7 +95,7 @@ class MTCExtractor(object): detected_faces, pnts = detect_face ( input_image, self.min_face_size, self.pnet_fun, self.rnet_fun, self.onet_fun, [ self.thresh1, self.thresh2, self.thresh3 ], self.scale_factor ) detected_faces = [ ( int(face[0]/input_scale), int(face[1]/input_scale), int(face[2]/input_scale), int(face[3]/input_scale)) for face in detected_faces ] - + return detected_faces def detect_face(img, minsize, pnet, rnet, onet, threshold, factor): @@ -132,9 +132,9 @@ def detect_face(img, minsize, pnet, rnet, onet, threshold, factor): out = pnet([img_y]) out0 = np.transpose(out[0], (0,2,1,3)) out1 = np.transpose(out[1], (0,2,1,3)) - + boxes, _ = generateBoundingBox(out1[0,:,:,1].copy(), out0[0,:,:,:].copy(), scale, threshold[0]) - + # inter-scale nms pick = nms(boxes.copy(), 0.5, 'Union') if boxes.size>0 and pick.size>0: @@ -217,7 +217,7 @@ def detect_face(img, minsize, pnet, rnet, onet, threshold, factor): pick = nms(total_boxes.copy(), 0.7, 'Min') total_boxes = total_boxes[pick,:] points = points[:,pick] - + return total_boxes, points @@ -235,7 +235,7 @@ def bbreg(boundingbox,reg): b4 = boundingbox[:,3]+reg[:,3]*h boundingbox[:,0:4] = np.transpose(np.vstack([b1, b2, b3, b4 ])) return boundingbox - + def generateBoundingBox(imap, reg, scale, t): """Use heatmap to generate bounding boxes""" stride=2 @@ -261,7 +261,7 @@ def generateBoundingBox(imap, reg, scale, t): q2 = np.fix((stride*bb+cellsize-1+1)/scale) boundingbox = np.hstack([q1, q2, np.expand_dims(score,1), reg]) return boundingbox, reg - + # function pick = nms(boxes,threshold,type) def nms(boxes, threshold, method): if boxes.size==0: @@ -315,7 +315,7 @@ def pad(total_boxes, w, h): tmp = np.where(ex>w) edx.flat[tmp] = np.expand_dims(-ex[tmp]+w+tmpw[tmp],1) ex[tmp] = w - + tmp = np.where(ey>h) edy.flat[tmp] = np.expand_dims(-ey[tmp]+h+tmph[tmp],1) ey[tmp] = h @@ -327,7 +327,7 @@ def pad(total_boxes, w, h): tmp = np.where(y<1) dy.flat[tmp] = np.expand_dims(2-y[tmp],1) y[tmp] = 1 - + return dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph # function [bboxA] = rerec(bboxA) diff --git a/facelib/S3FDExtractor.py b/facelib/S3FDExtractor.py index 2054cee..e4aef82 100644 --- a/facelib/S3FDExtractor.py +++ b/facelib/S3FDExtractor.py @@ -3,35 +3,35 @@ from pathlib import Path import cv2 from nnlib import nnlib -class S3FDExtractor(object): +class S3FDExtractor(object): def __init__(self): exec( nnlib.import_all(), locals(), globals() ) - + model_path = Path(__file__).parent / "S3FD.h5" if not model_path.exists(): return None - self.model = nnlib.keras.models.load_model ( str(model_path) ) + self.model = nnlib.keras.models.load_model ( str(model_path) ) def __enter__(self): return self - + def __exit__(self, exc_type=None, exc_value=None, traceback=None): return False #pass exception between __enter__ and __exit__ to outter level - + def extract_from_bgr (self, input_image): input_image = input_image[:,:,::-1].copy() (h, w, ch) = input_image.shape - + d = max(w, h) scale_to = 640 if d >= 1280 else d / 2 scale_to = max(64, scale_to) - + input_scale = d / scale_to input_image = cv2.resize (input_image, ( int(w/input_scale), int(h/input_scale) ), interpolation=cv2.INTER_LINEAR) olist = self.model.predict( np.expand_dims(input_image,0) ) - + detected_faces = [] for ltrb in self.refine (olist): l,t,r,b = [ x*input_scale for x in ltrb] @@ -42,7 +42,7 @@ class S3FDExtractor(object): detected_faces.append ( [int(x) for x in (l,t,r,b) ] ) return detected_faces - + def refine(self, olist): bboxlist = [] for i, ((ocls,), (oreg,)) in enumerate ( zip ( olist[::2], olist[1::2] ) ): @@ -51,7 +51,7 @@ class S3FDExtractor(object): s_m4 = stride * 4 for hindex, windex in zip(*np.where(ocls > 0.05)): - score = ocls[hindex, windex] + score = ocls[hindex, windex] loc = oreg[hindex, windex, :] priors = np.array([windex * stride + s_d2, hindex * stride + s_d2, s_m4, s_m4]) priors_2p = priors[2:] @@ -61,15 +61,15 @@ class S3FDExtractor(object): box[2:] += box[:2] bboxlist.append([*box, score]) - + bboxlist = np.array(bboxlist) if len(bboxlist) == 0: bboxlist = np.zeros((1, 5)) - + bboxlist = bboxlist[self.refine_nms(bboxlist, 0.3), :] bboxlist = [ x[:-1].astype(np.int) for x in bboxlist if x[-1] >= 0.5] return bboxlist - + def refine_nms(self, dets, thresh): keep = list() if len(dets) == 0: @@ -91,4 +91,4 @@ class S3FDExtractor(object): inds = np.where(ovr <= thresh)[0] order = order[inds + 1] - return keep \ No newline at end of file + return keep diff --git a/imagelib/estimate_sharpness.py b/imagelib/estimate_sharpness.py index 1617acc..01ef0b7 100644 --- a/imagelib/estimate_sharpness.py +++ b/imagelib/estimate_sharpness.py @@ -1,6 +1,6 @@ """ Copyright (c) 2009-2010 Arizona Board of Regents. All Rights Reserved. - Contact: Lina Karam (karam@asu.edu) and Niranjan Narvekar (nnarveka@asu.edu) + Contact: Lina Karam (karam@asu.edu) and Niranjan Narvekar (nnarveka@asu.edu) Image, Video, and Usabilty (IVU) Lab, http://ivulab.asu.edu , Arizona State University This copyright statement may not be removed from any file containing it or from modifications to these files. This copyright notice must also be included in any file or product that is derived from the source files. @@ -267,11 +267,11 @@ def get_block_contrast(block): # type: (numpy.ndarray) -> int return int(np.max(block) - np.min(block)) - -def estimate_sharpness(image): + +def estimate_sharpness(image): height, width = image.shape[:2] - + if image.ndim == 3: image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) - return compute(image) \ No newline at end of file + return compute(image) diff --git a/interact/__init__.py b/interact/__init__.py index c54c814..db40e4f 100644 --- a/interact/__init__.py +++ b/interact/__init__.py @@ -1 +1 @@ -from .interact import interact \ No newline at end of file +from .interact import interact diff --git a/interact/interact.py b/interact/interact.py index 51a55a9..fc2e727 100644 --- a/interact/interact.py +++ b/interact/interact.py @@ -12,28 +12,28 @@ class Interact(object): EVENT_RBUTTONDOWN = 5 EVENT_RBUTTONUP = 6 EVENT_MOUSEWHEEL = 10 - + def __init__(self): self.named_windows = {} self.capture_mouse_windows = {} - self.capture_keys_windows = {} + self.capture_keys_windows = {} self.mouse_events = {} self.key_events = {} self.pg_bar = None - + def log_info(self, msg, end='\n'): print (msg, end=end) - + def log_err(self, msg, end='\n'): - print (msg, end=end) - + print (msg, end=end) + def named_window(self, wnd_name): if wnd_name not in self.named_windows: #we will show window only on first show_image self.named_windows[wnd_name] = 0 - + else: print("named_window: ", wnd_name, " already created.") - + def destroy_all_windows(self): if len( self.named_windows ) != 0: cv2.destroyAllWindows() @@ -42,32 +42,32 @@ class Interact(object): self.capture_keys_windows = {} self.mouse_events = {} self.key_events = {} - + def show_image(self, wnd_name, img): if wnd_name in self.named_windows: if self.named_windows[wnd_name] == 0: self.named_windows[wnd_name] = 1 - cv2.namedWindow(wnd_name) + cv2.namedWindow(wnd_name) if wnd_name in self.capture_mouse_windows: self.capture_mouse(wnd_name) - + cv2.imshow (wnd_name, img) else: print("show_image: named_window ", wnd_name, " not found.") def capture_mouse(self, wnd_name): def onMouse(event, x, y, flags, param): - (inst, wnd_name) = param + (inst, wnd_name) = param if event == cv2.EVENT_LBUTTONDOWN: ev = Interact.EVENT_LBUTTONDOWN elif event == cv2.EVENT_LBUTTONUP: ev = Interact.EVENT_LBUTTONUP elif event == cv2.EVENT_RBUTTONDOWN: ev = Interact.EVENT_RBUTTONDOWN elif event == cv2.EVENT_RBUTTONUP: ev = Interact.EVENT_RBUTTONUP elif event == cv2.EVENT_MOUSEWHEEL: ev = Interact.EVENT_MOUSEWHEEL - + else: ev = 0 - inst.add_mouse_event (wnd_name, x, y, ev, flags) + inst.add_mouse_event (wnd_name, x, y, ev, flags) if wnd_name in self.named_windows: - self.capture_mouse_windows[wnd_name] = True + self.capture_mouse_windows[wnd_name] = True if self.named_windows[wnd_name] == 1: cv2.setMouseCallback(wnd_name, onMouse, (self,wnd_name) ) else: print("capture_mouse: named_window ", wnd_name, " not found.") @@ -95,66 +95,66 @@ class Interact(object): self.pg_bar.close() self.pg_bar = None else: print("progress_bar not set.") - + def progress_bar_generator(self, data, desc, leave=True): for x in tqdm( data, desc=desc, leave=leave, ascii=True ): yield x - + def process_messages(self, sleep_time=0): has_windows = False has_capture_keys = False if len(self.named_windows) != 0: has_windows = True - + if len(self.capture_keys_windows) != 0: has_capture_keys = True - + if has_windows or has_capture_keys: wait_key_time = max(1, int(sleep_time*1000) ) key = cv2.waitKey(wait_key_time) & 0xFF else: if sleep_time != 0: time.sleep(sleep_time) - + if has_capture_keys and key != 255: for wnd_name in self.capture_keys_windows: self.add_key_event (wnd_name, key) - + def wait_any_key(self): cv2.waitKey(0) - + def add_mouse_event(self, wnd_name, x, y, ev, flags): - if wnd_name not in self.mouse_events: + if wnd_name not in self.mouse_events: self.mouse_events[wnd_name] = [] self.mouse_events[wnd_name] += [ (x, y, ev, flags) ] - + def add_key_event(self, wnd_name, key): - if wnd_name not in self.key_events: + if wnd_name not in self.key_events: self.key_events[wnd_name] = [] self.key_events[wnd_name] += [ (key,) ] def get_mouse_events(self, wnd_name): ar = self.mouse_events.get(wnd_name, []) - self.mouse_events[wnd_name] = [] + self.mouse_events[wnd_name] = [] return ar - + def get_key_events(self, wnd_name): ar = self.key_events.get(wnd_name, []) - self.key_events[wnd_name] = [] + self.key_events[wnd_name] = [] return ar - + def input_number(self, s, default_value, valid_list=None, help_message=None): while True: try: inp = input(s) if len(inp) == 0: raise ValueError("") - + if help_message is not None and inp == '?': print (help_message) continue - + i = float(inp) if (valid_list is not None) and (i not in valid_list): return default_value @@ -162,18 +162,18 @@ class Interact(object): except: print (default_value) return default_value - + def input_int(self,s, default_value, valid_list=None, help_message=None): while True: try: inp = input(s) if len(inp) == 0: raise ValueError("") - + if help_message is not None and inp == '?': print (help_message) continue - + i = int(inp) if (valid_list is not None) and (i not in valid_list): return default_value @@ -181,41 +181,41 @@ class Interact(object): except: print (default_value) return default_value - + def input_bool(self, s, default_value, help_message=None): while True: try: inp = input(s) if len(inp) == 0: raise ValueError("") - + if help_message is not None and inp == '?': print (help_message) continue - + return bool ( {"y":True,"n":False,"1":True,"0":False}.get(inp.lower(), default_value) ) except: print ( "y" if default_value else "n" ) return default_value - + def input_str(self, s, default_value, valid_list=None, help_message=None): - while True: + while True: try: inp = input(s) if len(inp) == 0: raise ValueError("") - + if help_message is not None and inp == '?': print (help_message) continue - + if (valid_list is not None) and (inp.lower() not in valid_list): return default_value return inp except: print (default_value) return default_value - + def input_process(self, stdin_fd, sq, str): sys.stdin = os.fdopen(stdin_fd) try: @@ -223,7 +223,7 @@ class Interact(object): sq.put (True) except: sq.put (False) - + def input_in_time (self, str, max_time_sec): sq = multiprocessing.Queue() p = multiprocessing.Process(target=self.input_process, args=( sys.stdin.fileno(), sq, str)) @@ -240,4 +240,4 @@ class Interact(object): sys.stdin = os.fdopen( sys.stdin.fileno() ) return inp -interact = Interact() \ No newline at end of file +interact = Interact() diff --git a/joblib/SubprocessFunctionCaller.py b/joblib/SubprocessFunctionCaller.py index b458120..7b3a528 100644 --- a/joblib/SubprocessFunctionCaller.py +++ b/joblib/SubprocessFunctionCaller.py @@ -7,7 +7,7 @@ class SubprocessFunctionCaller(object): self.s2c = s2c self.c2s = c2s self.lock = lock - + def __call__(self, value): self.lock.acquire() self.c2s.put (value) @@ -17,26 +17,26 @@ class SubprocessFunctionCaller(object): self.lock.release() return obj time.sleep(0.005) - + class HostProcessor(object): def __init__(self, s2c, c2s, func): self.s2c = s2c self.c2s = c2s self.func = func - + def process_messages(self): while not self.c2s.empty(): obj = self.c2s.get() result = self.func (obj) self.s2c.put (result) - + @staticmethod def make_pair( func ): s2c = multiprocessing.Queue() c2s = multiprocessing.Queue() lock = multiprocessing.Lock() - + host_processor = SubprocessFunctionCaller.HostProcessor (s2c, c2s, func) cli_func = SubprocessFunctionCaller.CliFunction (s2c, c2s, lock) - - return host_processor, cli_func \ No newline at end of file + + return host_processor, cli_func diff --git a/joblib/SubprocessorBase.py b/joblib/SubprocessorBase.py index bf8cd86..91276b2 100644 --- a/joblib/SubprocessorBase.py +++ b/joblib/SubprocessorBase.py @@ -3,12 +3,12 @@ import multiprocessing import time import sys from interact import interact as io - + class Subprocessor(object): class SilenceException(Exception): pass - + class Cli(object): def __init__ ( self, client_dict ): self.s2c = multiprocessing.Queue() @@ -16,41 +16,41 @@ class Subprocessor(object): self.p = multiprocessing.Process(target=self._subprocess_run, args=(client_dict,) ) self.p.daemon = True self.p.start() - + self.state = None self.sent_time = None self.sent_data = None self.name = None self.host_dict = None - + def kill(self): self.p.terminate() self.p.join() - + #overridable optional def on_initialize(self, client_dict): #initialize your subprocess here using client_dict pass - + #overridable optional def on_finalize(self): #finalize your subprocess here pass - + #overridable def process_data(self, data): #process 'data' given from host and return result raise NotImplementedError - + #overridable optional def get_data_name (self, data): #return string identificator of your 'data' return "undefined" - + def log_info(self, msg): self.c2s.put ( {'op': 'log_info', 'msg':msg } ) def log_err(self, msg): self.c2s.put ( {'op': 'log_err' , 'msg':msg } ) def progress_bar_inc(self, c): self.c2s.put ( {'op': 'progress_bar_inc' , 'c':c } ) - + def _subprocess_run(self, client_dict): data = None s2c, c2s = self.s2c, self.c2s @@ -65,20 +65,20 @@ class Subprocessor(object): if op == 'data': data = msg['data'] result = self.process_data (data) - c2s.put ( {'op': 'success', 'data' : data, 'result' : result} ) + c2s.put ( {'op': 'success', 'data' : data, 'result' : result} ) data = None elif op == 'close': break time.sleep(0.001) - + self.on_finalize() c2s.put ( {'op': 'finalized'} ) return except Subprocessor.SilenceException as e: pass except Exception as e: - if data is not None: + if data is not None: print ('Exception while process data [%s]: %s' % (self.get_data_name(data), traceback.format_exc()) ) else: print ('Exception: %s' % (traceback.format_exc()) ) @@ -91,10 +91,10 @@ class Subprocessor(object): raise ValueError("SubprocessorCli_class must be subclass of Subprocessor.Cli") self.name = name - self.SubprocessorCli_class = SubprocessorCli_class + self.SubprocessorCli_class = SubprocessorCli_class self.no_response_time_sec = no_response_time_sec - #overridable + #overridable def process_info_generator(self): #yield per process (name, host_dict, client_dict) raise NotImplementedError @@ -103,42 +103,42 @@ class Subprocessor(object): def on_clients_initialized(self): #logic when all subprocesses initialized and ready pass - + #overridable optional def on_clients_finalized(self): #logic when all subprocess finalized pass - - #overridable + + #overridable def get_data(self, host_dict): #return data for processing here raise NotImplementedError - + #overridable def on_data_return (self, host_dict, data): - #you have to place returned 'data' back to your queue + #you have to place returned 'data' back to your queue raise NotImplementedError - + #overridable def on_result (self, host_dict, data, result): #your logic what to do with 'result' of 'data' raise NotImplementedError - + #overridable def get_result(self): #return result that will be returned in func run() raise NotImplementedError - + #overridable def on_tick(self): #tick in main loop pass - + def run(self): self.clis = [] - + #getting info about name of subprocesses, host and client dicts, and spawning them - for name, host_dict, client_dict in self.process_info_generator(): + for name, host_dict, client_dict in self.process_info_generator(): try: cli = self.SubprocessorCli_class(client_dict) cli.state = 1 @@ -146,21 +146,21 @@ class Subprocessor(object): cli.sent_data = None cli.name = name cli.host_dict = host_dict - + self.clis.append (cli) - + except: raise Exception ("Unable to start subprocess %s" % (name)) if len(self.clis) == 0: raise Exception ("Unable to start Subprocessor '%s' " % (self.name)) - + #waiting subprocesses their success(or not) initialization while True: for cli in self.clis[:]: while not cli.c2s.empty(): obj = cli.c2s.get() - op = obj.get('op','') + op = obj.get('op','') if op == 'init_ok': cli.state = 0 elif op == 'log_info': @@ -172,16 +172,16 @@ class Subprocessor(object): self.clis.remove(cli) break if all ([cli.state == 0 for cli in self.clis]): - break + break io.process_messages(0.005) - + if len(self.clis) == 0: raise Exception ( "Unable to start subprocesses." ) - - #ok some processes survived, initialize host logic - + + #ok some processes survived, initialize host logic + self.on_clients_initialized() - + #main loop of data processing while True: for cli in self.clis[:]: @@ -206,10 +206,10 @@ class Subprocessor(object): io.log_err(obj['msg']) elif op == 'progress_bar_inc': io.progress_bar_inc(obj['c']) - + for cli in self.clis[:]: if cli.state == 0: - #free state of subprocess, get some data from get_data + #free state of subprocess, get some data from get_data data = self.get_data(cli.host_dict) if data is not None: #and send it to subprocess @@ -217,7 +217,7 @@ class Subprocessor(object): cli.sent_time = time.time() cli.sent_data = data cli.state = 1 - + elif cli.state == 1: if self.no_response_time_sec != 0 and (time.time() - cli.sent_time) > self.no_response_time_sec: #subprocess busy too long @@ -225,39 +225,39 @@ class Subprocessor(object): self.on_data_return (cli.host_dict, cli.sent_data ) cli.kill() self.clis.remove(cli) - + if all ([cli.state == 0 for cli in self.clis]): #all subprocesses free and no more data available to process, ending loop - break + break io.process_messages(0.005) self.on_tick() - + #gracefully terminating subprocesses for cli in self.clis[:]: cli.s2c.put ( {'op': 'close'} ) cli.sent_time = time.time() - + while True: for cli in self.clis[:]: terminate_it = False while not cli.c2s.empty(): obj = cli.c2s.get() - obj_op = obj['op'] + obj_op = obj['op'] if obj_op == 'finalized': terminate_it = True break - - if self.no_response_time_sec != 0 and (time.time() - cli.sent_time) > self.no_response_time_sec: + + if self.no_response_time_sec != 0 and (time.time() - cli.sent_time) > self.no_response_time_sec: terminate_it = True - + if terminate_it: cli.state = 2 cli.kill() - + if all ([cli.state == 2 for cli in self.clis]): break - + #finalizing host logic and return result self.on_clients_finalized() - + return self.get_result() diff --git a/joblib/__init__.py b/joblib/__init__.py index 28f3313..fbbc20c 100644 --- a/joblib/__init__.py +++ b/joblib/__init__.py @@ -1,2 +1,2 @@ from .SubprocessorBase import Subprocessor -from .SubprocessFunctionCaller import SubprocessFunctionCaller \ No newline at end of file +from .SubprocessFunctionCaller import SubprocessFunctionCaller diff --git a/main.py b/main.py index 8dee298..79e997b 100644 --- a/main.py +++ b/main.py @@ -14,110 +14,110 @@ class fixPathAction(argparse.Action): def __call__(self, parser, namespace, values, option_string=None): setattr(namespace, self.dest, os.path.abspath(os.path.expanduser(values))) -if __name__ == "__main__": +if __name__ == "__main__": multiprocessing.set_start_method("spawn") - - os_utils.set_process_lowest_prio() - parser = argparse.ArgumentParser() + + os_utils.set_process_lowest_prio() + parser = argparse.ArgumentParser() subparsers = parser.add_subparsers() - + def process_extract(arguments): - from mainscripts import Extractor - Extractor.main( arguments.input_dir, - arguments.output_dir, + from mainscripts import Extractor + Extractor.main( arguments.input_dir, + arguments.output_dir, arguments.debug_dir, - arguments.detector, + arguments.detector, arguments.manual_fix, arguments.manual_output_debug_fix, - arguments.manual_window_size, + arguments.manual_window_size, face_type=arguments.face_type, device_args={'cpu_only' : arguments.cpu_only, 'multi_gpu' : arguments.multi_gpu, } ) - + p = subparsers.add_parser( "extract", help="Extract the faces from a pictures.") p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.") p.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the extracted files will be stored.") - p.add_argument('--debug-dir', action=fixPathAction, dest="debug_dir", help="Writes debug images to this directory.") - p.add_argument('--face-type', dest="face_type", choices=['half_face', 'full_face', 'head', 'avatar', 'mark_only'], default='full_face', help="Default 'full_face'. Don't change this option, currently all models uses 'full_face'") + p.add_argument('--debug-dir', action=fixPathAction, dest="debug_dir", help="Writes debug images to this directory.") + p.add_argument('--face-type', dest="face_type", choices=['half_face', 'full_face', 'head', 'avatar', 'mark_only'], default='full_face', help="Default 'full_face'. Don't change this option, currently all models uses 'full_face'") p.add_argument('--detector', dest="detector", choices=['dlib','mt','s3fd','manual'], default='dlib', help="Type of detector. Default 'dlib'. 'mt' (MTCNNv1) - faster, better, almost no jitter, perfect for gathering thousands faces for src-set. It is also good for dst-set, but can generate false faces in frames where main face not recognized! In this case for dst-set use either 'dlib' with '--manual-fix' or '--detector manual'. Manual detector suitable only for dst-set.") p.add_argument('--multi-gpu', action="store_true", dest="multi_gpu", default=False, help="Enables multi GPU.") p.add_argument('--manual-fix', action="store_true", dest="manual_fix", default=False, help="Enables manual extract only frames where faces were not recognized.") p.add_argument('--manual-output-debug-fix', action="store_true", dest="manual_output_debug_fix", default=False, help="Performs manual reextract input-dir frames which were deleted from [output_dir]_debug\ dir.") - p.add_argument('--manual-window-size', type=int, dest="manual_window_size", default=1368, help="Manual fix window size. Default: 1368.") - p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Extract on CPU. Forces to use MT extractor.") + p.add_argument('--manual-window-size', type=int, dest="manual_window_size", default=1368, help="Manual fix window size. Default: 1368.") + p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Extract on CPU. Forces to use MT extractor.") p.set_defaults (func=process_extract) - - def process_sort(arguments): + + def process_sort(arguments): from mainscripts import Sorter Sorter.main (input_path=arguments.input_dir, sort_by_method=arguments.sort_by_method) - - p = subparsers.add_parser( "sort", help="Sort faces in a directory.") + + p = subparsers.add_parser( "sort", help="Sort faces in a directory.") p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.") p.add_argument('--by', required=True, dest="sort_by_method", choices=("blur", "face", "face-dissim", "face-yaw", "face-pitch", "hist", "hist-dissim", "brightness", "hue", "black", "origname", "oneface", "final", "final-no-blur", "test"), help="Method of sorting. 'origname' sort by original filename to recover original sequence." ) p.set_defaults (func=process_sort) - - def process_util(arguments): + + def process_util(arguments): from mainscripts import Util - + if arguments.convert_png_to_jpg: Util.convert_png_to_jpg_folder (input_path=arguments.input_dir) - + if arguments.add_landmarks_debug_images: Util.add_landmarks_debug_images (input_path=arguments.input_dir) if arguments.recover_original_aligned_filename: Util.recover_original_aligned_filename (input_path=arguments.input_dir) - - p = subparsers.add_parser( "util", help="Utilities.") + + p = subparsers.add_parser( "util", help="Utilities.") p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.") p.add_argument('--convert-png-to-jpg', action="store_true", dest="convert_png_to_jpg", default=False, help="Convert DeepFaceLAB PNG files to JPEG.") p.add_argument('--add-landmarks-debug-images', action="store_true", dest="add_landmarks_debug_images", default=False, help="Add landmarks debug image for aligned faces.") p.add_argument('--recover-original-aligned-filename', action="store_true", dest="recover_original_aligned_filename", default=False, help="Recover original aligned filename.") p.set_defaults (func=process_util) - + def process_train(arguments): - args = {'training_data_src_dir' : arguments.training_data_src_dir, - 'training_data_dst_dir' : arguments.training_data_dst_dir, + args = {'training_data_src_dir' : arguments.training_data_src_dir, + 'training_data_dst_dir' : arguments.training_data_dst_dir, 'model_path' : arguments.model_dir, 'model_name' : arguments.model_name, 'no_preview' : arguments.no_preview, - 'debug' : arguments.debug, - } + 'debug' : arguments.debug, + } device_args = {'cpu_only' : arguments.cpu_only, 'force_gpu_idx' : arguments.force_gpu_idx, } - from mainscripts import Trainer + from mainscripts import Trainer Trainer.main(args, device_args) - p = subparsers.add_parser( "train", help="Trainer") + p = subparsers.add_parser( "train", help="Trainer") p.add_argument('--training-data-src-dir', required=True, action=fixPathAction, dest="training_data_src_dir", help="Dir of src-set.") p.add_argument('--training-data-dst-dir', required=True, action=fixPathAction, dest="training_data_dst_dir", help="Dir of dst-set.") p.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Model dir.") p.add_argument('--model', required=True, dest="model_name", choices=Path_utils.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Type of model") p.add_argument('--no-preview', action="store_true", dest="no_preview", default=False, help="Disable preview window.") - p.add_argument('--debug', action="store_true", dest="debug", default=False, help="Debug samples.") + p.add_argument('--debug', action="store_true", dest="debug", default=False, help="Debug samples.") p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Train on CPU.") p.add_argument('--force-gpu-idx', type=int, dest="force_gpu_idx", default=-1, help="Force to choose this GPU idx.") p.set_defaults (func=process_train) - + def process_convert(arguments): - args = {'input_dir' : arguments.input_dir, - 'output_dir' : arguments.output_dir, + args = {'input_dir' : arguments.input_dir, + 'output_dir' : arguments.output_dir, 'aligned_dir' : arguments.aligned_dir, 'model_dir' : arguments.model_dir, 'model_name' : arguments.model_name, - 'debug' : arguments.debug, - } + 'debug' : arguments.debug, + } device_args = {'cpu_only' : arguments.cpu_only, 'force_gpu_idx' : arguments.force_gpu_idx, } from mainscripts import Converter Converter.main (args, device_args) - - p = subparsers.add_parser( "convert", help="Converter") + + p = subparsers.add_parser( "convert", help="Converter") p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.") p.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the converted files will be stored.") p.add_argument('--aligned-dir', action=fixPathAction, dest="aligned_dir", help="Aligned directory. This is where the extracted of dst faces stored. Not used in AVATAR model.") @@ -127,10 +127,10 @@ if __name__ == "__main__": p.add_argument('--force-gpu-idx', type=int, dest="force_gpu_idx", default=-1, help="Force to choose this GPU idx.") p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Convert on CPU.") p.set_defaults(func=process_convert) - + videoed_parser = subparsers.add_parser( "videoed", help="Video processing.").add_subparsers() - - def process_videoed_extract_video(arguments): + + def process_videoed_extract_video(arguments): from mainscripts import VideoEd VideoEd.extract_video (arguments.input_file, arguments.output_dir, arguments.output_ext, arguments.fps) p = videoed_parser.add_parser( "extract-video", help="Extract images from video file.") @@ -139,23 +139,23 @@ if __name__ == "__main__": p.add_argument('--ouptut-ext', dest="output_ext", default='png', help="Image format (extension) of output files.") p.add_argument('--fps', type=int, dest="fps", default=None, help="How many frames of every second of the video will be extracted. 0 - full fps.") p.set_defaults(func=process_videoed_extract_video) - - def process_videoed_cut_video(arguments): + + def process_videoed_cut_video(arguments): from mainscripts import VideoEd - VideoEd.cut_video (arguments.input_file, - arguments.from_time, - arguments.to_time, - arguments.audio_track_id, + VideoEd.cut_video (arguments.input_file, + arguments.from_time, + arguments.to_time, + arguments.audio_track_id, arguments.bitrate) p = videoed_parser.add_parser( "cut-video", help="Cut video file.") p.add_argument('--input-file', required=True, action=fixPathAction, dest="input_file", help="Input file to be processed. Specify .*-extension to find first file.") p.add_argument('--from-time', dest="from_time", default=None, help="From time, for example 00:00:00.000") p.add_argument('--to-time', dest="to_time", default=None, help="To time, for example 00:00:00.000") p.add_argument('--audio-track-id', type=int, dest="audio_track_id", default=None, help="Specify audio track id.") - p.add_argument('--bitrate', type=int, dest="bitrate", default=None, help="Bitrate of output file in Megabits.") + p.add_argument('--bitrate', type=int, dest="bitrate", default=None, help="Bitrate of output file in Megabits.") p.set_defaults(func=process_videoed_cut_video) - - def process_videoed_denoise_image_sequence(arguments): + + def process_videoed_denoise_image_sequence(arguments): from mainscripts import VideoEd VideoEd.denoise_image_sequence (arguments.input_dir, arguments.ext, arguments.factor) p = videoed_parser.add_parser( "denoise-image-sequence", help="Denoise sequence of images, keeping sharp edges. This allows you to make the final fake more believable, since the neural network is not able to make a detailed skin texture, but it makes the edges quite clear. Therefore, if the whole frame is more `blurred`, then a fake will seem more believable. Especially true for scenes of the film, which are usually very clear.") @@ -163,65 +163,65 @@ if __name__ == "__main__": p.add_argument('--ext', dest="ext", default='png', help="Image format (extension) of input files.") p.add_argument('--factor', type=int, dest="factor", default=None, help="Denoise factor (1-20).") p.set_defaults(func=process_videoed_denoise_image_sequence) - - def process_videoed_video_from_sequence(arguments): + + def process_videoed_video_from_sequence(arguments): from mainscripts import VideoEd VideoEd.video_from_sequence (arguments.input_dir, - arguments.output_file, + arguments.output_file, arguments.reference_file, arguments.ext, arguments.fps, arguments.bitrate, arguments.lossless) - + p = videoed_parser.add_parser( "video-from-sequence", help="Make video from image sequence.") p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input file to be processed. Specify .*-extension to find first file.") p.add_argument('--output-file', required=True, action=fixPathAction, dest="output_file", help="Input file to be processed. Specify .*-extension to find first file.") p.add_argument('--reference-file', action=fixPathAction, dest="reference_file", help="Reference file used to determine proper FPS and transfer audio from it. Specify .*-extension to find first file.") p.add_argument('--ext', dest="ext", default='png', help="Image format (extension) of input files.") p.add_argument('--fps', type=int, dest="fps", default=None, help="FPS of output file. Overwritten by reference-file.") - p.add_argument('--bitrate', type=int, dest="bitrate", default=None, help="Bitrate of output file in Megabits.") + p.add_argument('--bitrate', type=int, dest="bitrate", default=None, help="Bitrate of output file in Megabits.") 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): + + def process_labelingtool(arguments): from mainscripts import LabelingTool LabelingTool.main (arguments.input_dir, arguments.output_dir) - + p = subparsers.add_parser( "labelingtool", help="Labeling tool.") 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) - + 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.") """ Suppressing error with keras 2.2.4+ on python exit: - + Exception ignored in: > Traceback (most recent call last): File "D:\DeepFaceLab\_internal\bin\lib\site-packages\tensorflow\python\client\session.py", line 1413, in __del__ AttributeError: 'NoneType' object has no attribute 'raise_exception_on_not_ok_status' - + reproduce: https://github.com/keras-team/keras/issues/11751 ( still no solution ) """ outnull_file = open(os.devnull, 'w') os.dup2 ( outnull_file.fileno(), sys.stderr.fileno() ) sys.stderr = outnull_file - - + + ''' import code code.interact(local=dict(globals(), **locals())) -''' \ No newline at end of file +''' diff --git a/mainscripts/Converter.py b/mainscripts/Converter.py index 68d57b1..1f6ef87 100644 --- a/mainscripts/Converter.py +++ b/mainscripts/Converter.py @@ -18,39 +18,39 @@ from interact import interact as io class ConvertSubprocessor(Subprocessor): class Cli(Subprocessor.Cli): - + #override def on_initialize(self, client_dict): io.log_info ('Running on %s.' % (client_dict['device_name']) ) self.device_idx = client_dict['device_idx'] self.device_name = client_dict['device_name'] self.converter = client_dict['converter'] - self.output_path = Path(client_dict['output_dir']) if 'output_dir' in client_dict.keys() else None + self.output_path = Path(client_dict['output_dir']) if 'output_dir' in client_dict.keys() else None self.alignments = client_dict['alignments'] self.debug = client_dict['debug'] - + #transfer and set stdin in order to work code.interact in debug subprocess stdin_fd = client_dict['stdin_fd'] if stdin_fd is not None: sys.stdin = os.fdopen(stdin_fd) - - from nnlib import nnlib + + from nnlib import nnlib #model process ate all GPU mem, #so we cannot use GPU for any TF operations in converter processes #therefore forcing active_DeviceConfig to CPU only nnlib.active_DeviceConfig = nnlib.DeviceConfig (cpu_only=True) - + return None - + #override def process_data(self, data): filename_path = Path(data) files_processed = 1 faces_processed = 0 - + output_filename_path = self.output_path / (filename_path.stem + '.png') - if self.converter.type == Converter.TYPE_FACE and filename_path.stem not in self.alignments.keys(): + if self.converter.type == Converter.TYPE_FACE and filename_path.stem not in self.alignments.keys(): if not self.debug: self.log_info ( 'no faces found for %s, copying without faces' % (filename_path.name) ) shutil.copy ( str(filename_path), str(output_filename_path) ) @@ -72,12 +72,12 @@ class ConvertSubprocessor(Subprocessor): dflimg = DFLJPG.load ( str(filename_path) ) else: dflimg = None - + if dflimg is not None: image_landmarks = dflimg.get_landmarks() - + image = self.converter.convert_image(image, image_landmarks, self.debug) - + if self.debug: raise NotImplementedError #for img in image: @@ -85,14 +85,14 @@ class ConvertSubprocessor(Subprocessor): # cv2.waitKey(0) faces_processed = 1 else: - self.log_err ("%s is not a dfl image file" % (filename_path.name) ) - + self.log_err ("%s is not a dfl image file" % (filename_path.name) ) + elif self.converter.type == Converter.TYPE_FACE: faces = self.alignments[filename_path.stem] - + if self.debug: debug_images = [] - + for face_num, image_landmarks in enumerate(faces): try: if self.debug: @@ -101,56 +101,56 @@ class ConvertSubprocessor(Subprocessor): if self.debug: debug_images += self.converter.convert_face(image, image_landmarks, self.debug) else: - image = self.converter.convert_face(image, image_landmarks, self.debug) - + image = self.converter.convert_face(image, image_landmarks, self.debug) + except Exception as e: e_str = traceback.format_exc() if 'MemoryError' in e_str: raise Subprocessor.SilenceException else: raise Exception( 'Error while converting face_num [%d] in file [%s]: %s' % (face_num, filename_path, e_str) ) - + if self.debug: return (1, debug_images) - + faces_processed = len(faces) - + if not self.debug: cv2_imwrite (str(output_filename_path), (image*255).astype(np.uint8) ) - - + + return (0, files_processed, faces_processed) - + #overridable def get_data_name (self, data): #return string identificator of your data return data - - #override - def __init__(self, converter, input_path_image_paths, output_path, alignments, debug = False): - super().__init__('Converter', ConvertSubprocessor.Cli, 86400 if debug == True else 60) - self.converter = converter + #override + def __init__(self, converter, input_path_image_paths, output_path, alignments, debug = False): + super().__init__('Converter', ConvertSubprocessor.Cli, 86400 if debug == True else 60) + + self.converter = converter self.host_processor, self.cli_func = SubprocessFunctionCaller.make_pair ( self.converter.predictor_func ) self.process_converter = self.converter.copy_and_set_predictor(self.cli_func) - + self.input_data = self.input_path_image_paths = input_path_image_paths self.output_path = output_path self.alignments = alignments self.debug = debug - + self.files_processed = 0 self.faces_processed = 0 - + #override def process_info_generator(self): r = [0] if self.debug else range(multiprocessing.cpu_count()) for i in r: yield 'CPU%d' % (i), {}, {'device_idx': i, - 'device_name': 'CPU%d' % (i), - 'converter' : self.process_converter, - 'output_dir' : str(self.output_path), + 'device_name': 'CPU%d' % (i), + 'converter' : self.process_converter, + 'output_dir' : str(self.output_path), 'alignments' : self.alignments, 'debug': self.debug, 'stdin_fd': sys.stdin.fileno() if self.debug else None @@ -160,25 +160,25 @@ class ConvertSubprocessor(Subprocessor): def on_clients_initialized(self): if self.debug: io.named_window ("Debug convert") - + io.progress_bar ("Converting", len (self.input_data) ) - + #overridable optional def on_clients_finalized(self): io.progress_bar_close() - + if self.debug: io.destroy_all_windows() - + #override def get_data(self, host_dict): if len (self.input_data) > 0: - return self.input_data.pop(0) + return self.input_data.pop(0) return None - + #override def on_data_return (self, host_dict, data): - self.input_data.insert(0, data) + self.input_data.insert(0, data) #override def on_result (self, host_dict, data, result): @@ -190,25 +190,25 @@ class ConvertSubprocessor(Subprocessor): io.show_image ('Debug convert', (img*255).astype(np.uint8) ) io.wait_any_key() io.progress_bar_inc(1) - + #override def on_tick(self): self.host_processor.process_messages() - + #override def get_result(self): return self.files_processed, self.faces_processed - + def main (args, device_args): io.log_info ("Running converter.\r\n") - + aligned_dir = args.get('aligned_dir', None) - + try: input_path = Path(args['input_dir']) output_path = Path(args['output_dir']) model_path = Path(args['model_dir']) - + if not input_path.exists(): io.log_err('Input directory not found. Please ensure it exists.') return @@ -218,69 +218,69 @@ def main (args, device_args): Path(filename).unlink() else: output_path.mkdir(parents=True, exist_ok=True) - + if not model_path.exists(): io.log_err('Model directory not found. Please ensure it exists.') return - - import models + + import models model = models.import_model( args['model_name'] )(model_path, device_args=device_args) converter = model.get_converter() converter.dummy_predict() alignments = None - + if converter.type == Converter.TYPE_FACE: if aligned_dir is None: io.log_err('Aligned directory not found. Please ensure it exists.') - return - + return + aligned_path = Path(aligned_dir) if not aligned_path.exists(): io.log_err('Aligned directory not found. Please ensure it exists.') - return - + return + alignments = {} - + aligned_path_image_paths = Path_utils.get_image_paths(aligned_path) for filepath in io.progress_bar_generator(aligned_path_image_paths, "Collecting alignments"): filepath = Path(filepath) - + 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) ) + io.log_err ("%s is not a dfl image file" % (filepath.name) ) continue - + source_filename_stem = Path( dflimg.get_source_filename() ).stem if source_filename_stem not in alignments.keys(): alignments[ source_filename_stem ] = [] alignments[ source_filename_stem ].append (dflimg.get_source_landmarks()) - files_processed, faces_processed = ConvertSubprocessor ( + files_processed, faces_processed = ConvertSubprocessor ( converter = converter, - input_path_image_paths = Path_utils.get_image_paths(input_path), + input_path_image_paths = Path_utils.get_image_paths(input_path), output_path = output_path, alignments = alignments, debug = args.get('debug',False) ).run() model.finalize() - + except Exception as e: print ( 'Error: %s' % (str(e))) traceback.print_exc() - -''' + +''' if model_name == 'AVATAR': output_path_image_paths = Path_utils.get_image_paths(output_path) - + last_ok_frame = -1 for filename in output_path_image_paths: filename_path = Path(filename) @@ -289,15 +289,15 @@ if model_name == 'AVATAR': frame = int(stem) except: raise Exception ('Aligned avatars must be created from indexed sequence files.') - + if frame-last_ok_frame > 1: start = last_ok_frame + 1 end = frame - 1 - + print ("Filling gaps: [%d...%d]" % (start, end) ) - for i in range (start, end+1): + for i in range (start, end+1): shutil.copy ( str(filename), str( output_path / ('%.5d%s' % (i, filename_path.suffix )) ) ) - + last_ok_frame = frame ''' #interpolate landmarks @@ -306,28 +306,28 @@ if model_name == 'AVATAR': #a = sorted(alignments.keys()) #a_len = len(a) # -#box_pts = 3 +#box_pts = 3 #box = np.ones(box_pts)/box_pts #for i in range( a_len ): # if i >= box_pts and i <= a_len-box_pts-1: # af0 = alignments[ a[i] ][0] ##first face -# m0 = LandmarksProcessor.get_transform_mat (af0, 256, face_type=FaceType.FULL) -# +# m0 = LandmarksProcessor.get_transform_mat (af0, 256, face_type=FaceType.FULL) +# # points = [] -# +# # for j in range(-box_pts, box_pts+1): # af = alignments[ a[i+j] ][0] ##first face -# m = LandmarksProcessor.get_transform_mat (af, 256, face_type=FaceType.FULL) +# m = LandmarksProcessor.get_transform_mat (af, 256, face_type=FaceType.FULL) # p = LandmarksProcessor.transform_points (af, m) # points.append (p) -# +# # points = np.array(points) # points_len = len(points) # t_points = np.transpose(points, [1,0,2]) -# +# # p1 = np.array ( [ int(np.convolve(x[:,0], box, mode='same')[points_len//2]) for x in t_points ] ) # p2 = np.array ( [ int(np.convolve(x[:,1], box, mode='same')[points_len//2]) for x in t_points ] ) -# +# # new_points = np.concatenate( [np.expand_dims(p1,-1),np.expand_dims(p2,-1)], -1 ) -# +# # alignments[ a[i] ][0] = LandmarksProcessor.transform_points (new_points, m0, True).astype(np.int32) diff --git a/mainscripts/Extractor.py b/mainscripts/Extractor.py index 22cbf35..22a5547 100644 --- a/mainscripts/Extractor.py +++ b/mainscripts/Extractor.py @@ -18,9 +18,9 @@ from facelib import LandmarksProcessor from nnlib import nnlib from joblib import Subprocessor from interact import interact as io - + class ExtractSubprocessor(Subprocessor): - + class Cli(Subprocessor.Cli): #override @@ -32,19 +32,19 @@ class ExtractSubprocessor(Subprocessor): self.face_type = client_dict['face_type'] self.device_idx = client_dict['device_idx'] self.cpu_only = client_dict['device_type'] == 'CPU' - self.output_path = Path(client_dict['output_dir']) if 'output_dir' in client_dict.keys() else None + self.output_path = Path(client_dict['output_dir']) if 'output_dir' in client_dict.keys() else None self.debug_dir = client_dict['debug_dir'] self.detector = client_dict['detector'] - + self.cached_image = (None, None) - + self.e = None device_config = nnlib.DeviceConfig ( cpu_only=self.cpu_only, force_gpu_idx=self.device_idx, allow_growth=True) if self.type == 'rects': if self.detector is not None: if self.detector == 'mt': nnlib.import_all (device_config) - self.e = facelib.MTCExtractor() + self.e = facelib.MTCExtractor() elif self.detector == 'dlib': nnlib.import_dlib (device_config) self.e = facelib.DLIBExtractor(nnlib.dlib) @@ -53,10 +53,10 @@ class ExtractSubprocessor(Subprocessor): self.e = facelib.S3FDExtractor() else: raise ValueError ("Wrong detector type.") - + if self.e is not None: self.e.__enter__() - + elif self.type == 'landmarks': nnlib.import_all (device_config) self.e = facelib.LandmarksExtractor(nnlib.keras) @@ -66,15 +66,15 @@ class ExtractSubprocessor(Subprocessor): self.second_pass_e.__enter__() else: self.second_pass_e = None - + elif self.type == 'final': pass - + #override def on_finalize(self): if self.e is not None: self.e.__exit__() - + #override def process_data(self, data): filename_path = Path( data[0] ) @@ -84,64 +84,64 @@ class ExtractSubprocessor(Subprocessor): image = self.cached_image[1] #cached image for manual extractor else: image = cv2_imread( filename_path_str ) - + if image is None: self.log_err ( 'Failed to extract %s, reason: cv2_imread() fail.' % ( str(filename_path) ) ) return None - + image_shape = image.shape if len(image_shape) == 2: h, w = image.shape - ch = 1 + ch = 1 else: h, w, ch = image.shape - + if ch == 1: image = np.repeat ( image [:,:,np.newaxis], 3, -1 ) elif ch == 4: image = image[:,:,0:3] - + wm = w % 2 hm = h % 2 if wm + hm != 0: #fix odd image image = image[0:h-hm,0:w-wm,:] self.cached_image = ( filename_path_str, image ) - + src_dflimg = None - h, w, ch = image.shape + h, w, ch = image.shape if h == w: #extracting from already extracted jpg image? if filename_path.suffix == '.jpg': src_dflimg = DFLJPG.load ( str(filename_path) ) - + if self.type == 'rects': if min(w,h) < 128: self.log_err ( 'Image is too small %s : [%d, %d]' % ( str(filename_path), w, h ) ) rects = [] - else: + else: rects = self.e.extract_from_bgr (image) - + return [str(filename_path), rects] elif self.type == 'landmarks': rects = data[1] if rects is None: landmarks = None - else: - landmarks = self.e.extract_from_bgr (image, rects, self.second_pass_e if src_dflimg is None else None) - + else: + landmarks = self.e.extract_from_bgr (image, rects, self.second_pass_e if src_dflimg is None else None) + return [str(filename_path), landmarks] elif self.type == 'final': - - + + result = [] faces = data[1] - + if self.debug_dir is not None: debug_output_file = str( Path(self.debug_dir) / (filename_path.stem+'.jpg') ) debug_image = image.copy() - + if src_dflimg is not None and len(faces) != 1: #if re-extracting from dflimg and more than 1 or zero faces detected - dont process and just copy it print("src_dflimg is not None and len(faces) != 1", str(filename_path) ) @@ -151,26 +151,26 @@ class ExtractSubprocessor(Subprocessor): result.append (output_file) else: face_idx = 0 - for face in faces: + for face in faces: rect = np.array(face[0]) image_landmarks = face[1] if image_landmarks is None: continue image_landmarks = np.array(image_landmarks) - if self.face_type == FaceType.MARK_ONLY: + if self.face_type == FaceType.MARK_ONLY: face_image = image face_image_landmarks = image_landmarks else: - image_to_face_mat = LandmarksProcessor.get_transform_mat (image_landmarks, self.image_size, self.face_type) + image_to_face_mat = LandmarksProcessor.get_transform_mat (image_landmarks, self.image_size, self.face_type) face_image = cv2.warpAffine(image, image_to_face_mat, (self.image_size, self.image_size), cv2.INTER_LANCZOS4) face_image_landmarks = LandmarksProcessor.transform_points (image_landmarks, image_to_face_mat) - + landmarks_bbox = LandmarksProcessor.transform_points ( [ (0,0), (0,self.image_size-1), (self.image_size-1, self.image_size-1), (self.image_size-1,0) ], image_to_face_mat, True) - + rect_area = mathlib.polygon_area(np.array(rect[[0,2,2,0]]), np.array(rect[[1,1,3,3]])) landmarks_area = mathlib.polygon_area(landmarks_bbox[:,0], landmarks_bbox[:,1] ) - + if landmarks_area > 4*rect_area: #get rid of faces which umeyama-landmark-area > 4*detector-rect-area continue @@ -192,24 +192,24 @@ class ExtractSubprocessor(Subprocessor): source_rect=rect, source_landmarks=image_landmarks.tolist(), image_to_face_mat=image_to_face_mat - ) - + ) + result.append (output_file) face_idx += 1 - + if self.debug_dir is not None: cv2_imwrite(debug_output_file, debug_image, [int(cv2.IMWRITE_JPEG_QUALITY), 50] ) - - return result - - + + return result + + #overridable def get_data_name (self, data): #return string identificator of your data return data[0] - + #override - def __init__(self, input_data, type, image_size, face_type, debug_dir, multi_gpu=False, cpu_only=False, manual=False, manual_window_size=0, detector=None, output_path=None): + def __init__(self, input_data, type, image_size, face_type, debug_dir, multi_gpu=False, cpu_only=False, manual=False, manual_window_size=0, detector=None, output_path=None): self.input_data = input_data self.type = type self.image_size = image_size @@ -218,8 +218,8 @@ class ExtractSubprocessor(Subprocessor): self.multi_gpu = multi_gpu self.cpu_only = cpu_only self.detector = detector - self.output_path = output_path - self.manual = manual + self.output_path = output_path + self.manual = manual self.manual_window_size = manual_window_size self.result = [] @@ -233,32 +233,32 @@ class ExtractSubprocessor(Subprocessor): io.named_window(self.wnd_name) io.capture_mouse(self.wnd_name) io.capture_keys(self.wnd_name) - + self.cache_original_image = (None, None) self.cache_image = (None, None) self.cache_text_lines_img = (None, None) self.hide_help = False - + self.landmarks = None self.x = 0 self.y = 0 self.rect_size = 100 self.rect_locked = False self.extract_needed = True - + io.progress_bar (None, len (self.input_data)) - + #override def on_clients_finalized(self): if self.manual == True: io.destroy_all_windows() - + io.progress_bar_close() - + def get_devices_for_type (self, type, multi_gpu, cpu_only): if 'cpu' in nnlib.device.backend: cpu_only = True - + if not cpu_only and (type == 'rects' or type == 'landmarks'): if type == 'rects' and (self.detector == 'mt') and nnlib.device.backend == "plaidML": cpu_only = True @@ -269,11 +269,11 @@ class ExtractSubprocessor(Subprocessor): devices = [nnlib.device.getBestValidDeviceIdx()] if len(devices) == 0: devices = [0] - + for idx in devices: dev_name = nnlib.device.getDeviceName(idx) dev_vram = nnlib.device.getDeviceVRAMTotalGb(idx) - + if not self.manual and ( self.type == 'rects' and self.detector != 's3fd' ): for i in range ( int (max (1, dev_vram / 2) ) ): yield (idx, 'GPU', '%s #%d' % (dev_name,i) , dev_vram) @@ -286,21 +286,21 @@ class ExtractSubprocessor(Subprocessor): else: for i in range( min(8, multiprocessing.cpu_count() // 2) ): yield (i, 'CPU', 'CPU%d' % (i), 0 ) - + if type == 'final': for i in range( min(8, multiprocessing.cpu_count()) ): - yield (i, 'CPU', 'CPU%d' % (i), 0 ) - + yield (i, 'CPU', 'CPU%d' % (i), 0 ) + #override def process_info_generator(self): - base_dict = {'type' : self.type, - 'image_size': self.image_size, - 'face_type': self.face_type, - 'debug_dir': self.debug_dir, - 'output_dir': str(self.output_path), + base_dict = {'type' : self.type, + 'image_size': self.image_size, + 'face_type': self.face_type, + 'debug_dir': self.debug_dir, + 'output_dir': str(self.output_path), 'detector': self.detector} - - for (device_idx, device_type, device_name, device_total_vram_gb) in self.get_devices_for_type(self.type, self.multi_gpu, self.cpu_only): + + for (device_idx, device_type, device_name, device_total_vram_gb) in self.get_devices_for_type(self.type, self.multi_gpu, self.cpu_only): client_dict = base_dict.copy() client_dict['device_idx'] = device_idx client_dict['device_name'] = device_name @@ -311,7 +311,7 @@ class ExtractSubprocessor(Subprocessor): def get_data(self, host_dict): if not self.manual: if len (self.input_data) > 0: - return self.input_data.pop(0) + return self.input_data.pop(0) else: need_remark_face = False @@ -327,7 +327,7 @@ class ExtractSubprocessor(Subprocessor): self.rect, self.landmarks = faces.pop() faces.clear() redraw_needed = True - self.rect_locked = True + self.rect_locked = True self.rect_size = ( self.rect[2] - self.rect[0] ) / 2 self.x = ( self.rect[0] + self.rect[2] ) / 2 self.y = ( self.rect[1] + self.rect[3] ) / 2 @@ -338,19 +338,19 @@ class ExtractSubprocessor(Subprocessor): else: self.original_image = cv2_imread( filename ) self.cache_original_image = (filename, self.original_image ) - + (h,w,c) = self.original_image.shape self.view_scale = 1.0 if self.manual_window_size == 0 else self.manual_window_size / ( h * (16.0/9.0) ) - + if self.cache_image[0] == (h,w,c) + (self.view_scale,filename): self.image = self.cache_image[1] - else: - self.image = cv2.resize (self.original_image, ( int(w*self.view_scale), int(h*self.view_scale) ), interpolation=cv2.INTER_LINEAR) + else: + self.image = cv2.resize (self.original_image, ( int(w*self.view_scale), int(h*self.view_scale) ), interpolation=cv2.INTER_LINEAR) self.cache_image = ( (h,w,c) + (self.view_scale,filename), self.image ) - + (h,w,c) = self.image.shape - - sh = (0,0, w, min(100, h) ) + + sh = (0,0, w, min(100, h) ) if self.cache_text_lines_img[0] == sh: self.text_lines_img = self.cache_text_lines_img[1] else: @@ -362,30 +362,30 @@ class ExtractSubprocessor(Subprocessor): '[,] [.]- prev frame, next frame. [Q] - skip remaining frames', '[h] - hide this help' ], (1, 1, 1) )*255).astype(np.uint8) - + self.cache_text_lines_img = (sh, self.text_lines_img) while True: io.process_messages(0.0001) - + new_x = self.x new_y = self.y new_rect_size = self.rect_size - + mouse_events = io.get_mouse_events(self.wnd_name) for ev in mouse_events: (x, y, ev, flags) = ev if ev == io.EVENT_MOUSEWHEEL and not self.rect_locked: - mod = 1 if flags > 0 else -1 + mod = 1 if flags > 0 else -1 diff = 1 if new_rect_size <= 40 else np.clip(new_rect_size / 10, 1, 10) - new_rect_size = max (5, new_rect_size + diff*mod) + new_rect_size = max (5, new_rect_size + diff*mod) elif ev == io.EVENT_LBUTTONDOWN: self.rect_locked = not self.rect_locked self.extract_needed = True elif not self.rect_locked: new_x = np.clip (x, 0, w-1) / self.view_scale new_y = np.clip (y, 0, h-1) / self.view_scale - + key_events = io.get_key_events(self.wnd_name) key, = key_events[-1] if len(key_events) > 0 else (0,) @@ -393,48 +393,48 @@ class ExtractSubprocessor(Subprocessor): #confirm frame is_frame_done = True faces.append ( [(self.rect), self.landmarks] ) - + break elif key == ord(' '): #confirm skip frame is_frame_done = True break - elif key == ord(',') and len(self.result) > 0: - #go prev frame - + elif key == ord(',') and len(self.result) > 0: + #go prev frame + if self.rect_locked: # Only save the face if the rect is still locked faces.append ( [(self.rect), self.landmarks] ) - + self.input_data.insert(0, self.result.pop() ) io.progress_bar_inc(-1) need_remark_face = True - + break - elif key == ord('.'): - #go next frame - + elif key == ord('.'): + #go next frame + if self.rect_locked: # Only save the face if the rect is still locked faces.append ( [(self.rect), self.landmarks] ) need_remark_face = True is_frame_done = True - break + break elif key == ord('q'): #skip remaining - + if self.rect_locked: faces.append ( [(self.rect), self.landmarks] ) while len(self.input_data) > 0: self.result.append( self.input_data.pop(0) ) io.progress_bar_inc(1) - + break - + elif key == ord('h'): self.hide_help = not self.hide_help break - + if self.x != new_x or \ self.y != new_y or \ self.rect_size != new_rect_size or \ @@ -443,33 +443,33 @@ class ExtractSubprocessor(Subprocessor): self.x = new_x self.y = new_y self.rect_size = new_rect_size - self.rect = ( int(self.x-self.rect_size), - int(self.y-self.rect_size), - int(self.x+self.rect_size), + self.rect = ( int(self.x-self.rect_size), + int(self.y-self.rect_size), + int(self.x+self.rect_size), int(self.y+self.rect_size) ) - + if redraw_needed: redraw_needed = False return [filename, None] else: return [filename, [self.rect]] - + else: is_frame_done = True - + if is_frame_done: self.result.append ( data ) self.input_data.pop(0) io.progress_bar_inc(1) self.extract_needed = True - self.rect_locked = False + self.rect_locked = False return None - + #override def on_data_return (self, host_dict, data): if not self.manual: - self.input_data.insert(0, data) + self.input_data.insert(0, data) #override def on_result (self, host_dict, data, result): @@ -477,33 +477,33 @@ class ExtractSubprocessor(Subprocessor): filename, landmarks = result if landmarks is not None: self.landmarks = landmarks[0][1] - + (h,w,c) = self.image.shape - + if not self.hide_help: image = cv2.addWeighted (self.image,1.0,self.text_lines_img,1.0,0) else: image = self.image.copy() - + view_rect = (np.array(self.rect) * self.view_scale).astype(np.int).tolist() view_landmarks = (np.array(self.landmarks) * self.view_scale).astype(np.int).tolist() - + if self.rect_size <= 40: scaled_rect_size = h // 3 if w > h else w // 3 p1 = (self.x - self.rect_size, self.y - self.rect_size) p2 = (self.x + self.rect_size, self.y - self.rect_size) p3 = (self.x - self.rect_size, self.y + self.rect_size) - - wh = h if h < w else w + + wh = h if h < w else w np1 = (w / 2 - wh / 4, h / 2 - wh / 4) np2 = (w / 2 + wh / 4, h / 2 - wh / 4) np3 = (w / 2 - wh / 4, h / 2 + wh / 4) - + mat = cv2.getAffineTransform( np.float32([p1,p2,p3])*self.view_scale, np.float32([np1,np2,np3]) ) - image = cv2.warpAffine(image, mat,(w,h) ) + image = cv2.warpAffine(image, mat,(w,h) ) view_landmarks = LandmarksProcessor.transform_points (view_landmarks, mat) - + landmarks_color = (255,255,0) if self.rect_locked else (0,255,0) LandmarksProcessor.draw_rect_landmarks (image, view_rect, view_landmarks, self.image_size, self.face_type, landmarks_color=landmarks_color) self.extract_needed = False @@ -513,10 +513,10 @@ class ExtractSubprocessor(Subprocessor): if self.type == 'rects': self.result.append ( result ) elif self.type == 'landmarks': - self.result.append ( result ) + self.result.append ( result ) elif self.type == 'final': self.result += result - + io.progress_bar_inc(1) #override @@ -530,47 +530,47 @@ class DeletedFilesSearcherSubprocessor(Subprocessor): def on_initialize(self, client_dict): self.debug_paths_stems = client_dict['debug_paths_stems'] return None - + #override - def process_data(self, data): - input_path_stem = Path(data[0]).stem + def process_data(self, data): + input_path_stem = Path(data[0]).stem return any ( [ input_path_stem == d_stem for d_stem in self.debug_paths_stems] ) - + #override def get_data_name (self, data): #return string identificator of your data return data[0] - + #override - def __init__(self, input_paths, debug_paths ): + def __init__(self, input_paths, debug_paths ): self.input_paths = input_paths - self.debug_paths_stems = [ Path(d).stem for d in debug_paths] + self.debug_paths_stems = [ Path(d).stem for d in debug_paths] self.result = [] - super().__init__('DeletedFilesSearcherSubprocessor', DeletedFilesSearcherSubprocessor.Cli, 60) - + super().__init__('DeletedFilesSearcherSubprocessor', DeletedFilesSearcherSubprocessor.Cli, 60) + #override - def process_info_generator(self): + def process_info_generator(self): for i in range(min(multiprocessing.cpu_count(), 8)): yield 'CPU%d' % (i), {}, {'debug_paths_stems' : self.debug_paths_stems} #override def on_clients_initialized(self): io.progress_bar ("Searching deleted files", len (self.input_paths)) - + #override def on_clients_finalized(self): io.progress_bar_close() #override def get_data(self, host_dict): - if len (self.input_paths) > 0: - return [self.input_paths.pop(0)] + if len (self.input_paths) > 0: + return [self.input_paths.pop(0)] return None - + #override def on_data_return (self, host_dict, data): - self.input_paths.insert(0, data[0]) - + self.input_paths.insert(0, data[0]) + #override def on_result (self, host_dict, data, result): if result == False: @@ -591,40 +591,40 @@ def main(input_dir, image_size=256, face_type='full_face', device_args={}): - + input_path = Path(input_dir) output_path = Path(output_dir) face_type = FaceType.fromString(face_type) - + multi_gpu = device_args.get('multi_gpu', False) cpu_only = device_args.get('cpu_only', False) - + if not input_path.exists(): raise ValueError('Input directory not found. Please ensure it exists.') - + if output_path.exists(): if not manual_output_debug_fix and input_path != output_path: for filename in Path_utils.get_image_paths(output_path): Path(filename).unlink() else: output_path.mkdir(parents=True, exist_ok=True) - + if manual_output_debug_fix: if debug_dir is None: raise ValueError('debug-dir must be specified') detector = 'manual' io.log_info('Performing re-extract frames which were deleted from _debug directory.') - + input_path_image_paths = Path_utils.get_image_unique_filestem_paths(input_path, verbose_print_func=io.log_info) if debug_dir is not None: debug_output_path = Path(debug_dir) - + if manual_output_debug_fix: if not debug_output_path.exists(): raise ValueError("%s not found " % ( str(debug_output_path) )) input_path_image_paths = DeletedFilesSearcherSubprocessor (input_path_image_paths, Path_utils.get_image_paths(debug_output_path) ).run() - input_path_image_paths = sorted (input_path_image_paths) + input_path_image_paths = sorted (input_path_image_paths) else: if debug_output_path.exists(): for filename in Path_utils.get_image_paths(debug_output_path): @@ -634,20 +634,20 @@ def main(input_dir, images_found = len(input_path_image_paths) faces_detected = 0 - if images_found != 0: + if images_found != 0: if detector == 'manual': io.log_info ('Performing manual extract...') extracted_faces = ExtractSubprocessor ([ (filename,[]) for filename in input_path_image_paths ], 'landmarks', image_size, face_type, debug_dir, cpu_only=cpu_only, manual=True, manual_window_size=manual_window_size).run() else: io.log_info ('Performing 1st pass...') extracted_rects = ExtractSubprocessor ([ (x,) for x in input_path_image_paths ], 'rects', image_size, face_type, debug_dir, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False, detector=detector).run() - + io.log_info ('Performing 2nd pass...') extracted_faces = ExtractSubprocessor (extracted_rects, 'landmarks', image_size, face_type, debug_dir, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False).run() - + if manual_fix: io.log_info ('Performing manual fix...') - + if all ( np.array ( [ len(data[1]) > 0 for data in extracted_faces] ) == True ): io.log_info ('All faces are detected, manual fix not needed.') else: @@ -657,8 +657,8 @@ def main(input_dir, io.log_info ('Performing 3rd pass...') final_imgs_paths = ExtractSubprocessor (extracted_faces, 'final', image_size, face_type, debug_dir, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False, output_path=output_path).run() faces_detected = len(final_imgs_paths) - + io.log_info ('-------------------------') io.log_info ('Images found: %d' % (images_found) ) io.log_info ('Faces detected: %d' % (faces_detected) ) - io.log_info ('-------------------------') \ No newline at end of file + io.log_info ('-------------------------') diff --git a/mainscripts/LabelingTool_unfinished.py b/mainscripts/LabelingTool_unfinished.py index 7d43453..60067c2 100644 --- a/mainscripts/LabelingTool_unfinished.py +++ b/mainscripts/LabelingTool_unfinished.py @@ -17,18 +17,18 @@ 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) @@ -39,165 +39,165 @@ def main(input_dir, output_dir): dflimg = DFLJPG.load ( str(filepath) ) else: dflimg = None - + if dflimg is None: - io.log_err ("%s is not a dfl image file" % (filepath.name) ) + 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] + 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 = 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] ) - + #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] ) - + #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] ) - + draw_fill_convex ( mask_fgp, lmrks[27:36] ) + #convex inner of nose half - draw_fill_convex ( mask_fg, lmrks[27:36], scale=0.5 ) - - + 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] ) + #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) - + 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_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 @@ -211,34 +211,34 @@ def main(input_dir, output_dir): # # 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] - + 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 = + + ##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-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) @@ -246,25 +246,25 @@ def main(input_dir, output_dir): ##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, + 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'): @@ -273,15 +273,15 @@ def main(input_dir, output_dir): pass if key == ord('3'): pass - - if ord(' ') in key_events: + + 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) \ No newline at end of file +#cv2.drawContours(original_mask, points, -1, (1,), 1) diff --git a/mainscripts/Sorter.py b/mainscripts/Sorter.py index 8254072..a5c2a78 100644 --- a/mainscripts/Sorter.py +++ b/mainscripts/Sorter.py @@ -15,10 +15,10 @@ from joblib import Subprocessor import multiprocessing from interact import interact as io from imagelib import estimate_sharpness - + class BlurEstimatorSubprocessor(Subprocessor): class Cli(Subprocessor.Cli): - + #override def on_initialize(self, client_dict): self.log_info('Running on %s.' % (client_dict['device_name']) ) @@ -26,58 +26,58 @@ class BlurEstimatorSubprocessor(Subprocessor): #override def process_data(self, data): filepath = Path( data[0] ) - + if filepath.suffix == '.png': dflimg = DFLPNG.load( str(filepath) ) elif filepath.suffix == '.jpg': dflimg = DFLJPG.load ( str(filepath) ) else: dflimg = None - + if dflimg is not None: image = cv2_imread( str(filepath) ) return [ str(filepath), estimate_sharpness(image) ] else: - self.log_err ("%s is not a dfl image file" % (filepath.name) ) + self.log_err ("%s is not a dfl image file" % (filepath.name) ) return [ str(filepath), 0 ] #override def get_data_name (self, data): #return string identificator of your data return data[0] - + #override - def __init__(self, input_data ): + def __init__(self, input_data ): self.input_data = input_data self.img_list = [] self.trash_img_list = [] - super().__init__('BlurEstimator', BlurEstimatorSubprocessor.Cli, 60) + super().__init__('BlurEstimator', BlurEstimatorSubprocessor.Cli, 60) #override def on_clients_initialized(self): io.progress_bar ("", len (self.input_data)) - + #override def on_clients_finalized(self): io.progress_bar_close () - + #override - def process_info_generator(self): + def process_info_generator(self): for i in range(0, multiprocessing.cpu_count() ): yield 'CPU%d' % (i), {}, {'device_idx': i, - 'device_name': 'CPU%d' % (i), + 'device_name': 'CPU%d' % (i), } #override def get_data(self, host_dict): if len (self.input_data) > 0: - return self.input_data.pop(0) - + return self.input_data.pop(0) + return None - + #override def on_data_return (self, host_dict, data): - self.input_data.insert(0, data) + self.input_data.insert(0, data) #override def on_result (self, host_dict, data, result): @@ -85,20 +85,20 @@ class BlurEstimatorSubprocessor(Subprocessor): self.trash_img_list.append ( result ) else: self.img_list.append ( result ) - + io.progress_bar_inc(1) - + #override def get_result(self): return self.img_list, self.trash_img_list - + def sort_by_blur(input_path): - io.log_info ("Sorting by blur...") - - img_list = [ (filename,[]) for filename in Path_utils.get_image_paths(input_path) ] + io.log_info ("Sorting by blur...") + + img_list = [ (filename,[]) for filename in Path_utils.get_image_paths(input_path) ] img_list, trash_img_list = BlurEstimatorSubprocessor (img_list).run() - + io.log_info ("Sorting...") img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True) @@ -111,21 +111,21 @@ def sort_by_face(input_path): trash_img_list = [] for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading"): filepath = Path(filepath) - + 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) ) trash_img_list.append ( [str(filepath)] ) continue img_list.append( [str(filepath), dflimg.get_landmarks()] ) - + img_list_len = len(img_list) for i in io.progress_bar_generator ( range(0, img_list_len-1), "Sorting"): @@ -152,21 +152,21 @@ def sort_by_face_dissim(input_path): trash_img_list = [] for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading"): filepath = Path(filepath) - + 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) ) trash_img_list.append ( [str(filepath)] ) - continue - + continue + img_list.append( [str(filepath), dflimg.get_landmarks(), 0 ] ) - + img_list_len = len(img_list) for i in io.progress_bar_generator( range(img_list_len-1), "Sorting"): score_total = 0 @@ -183,79 +183,79 @@ def sort_by_face_dissim(input_path): img_list = sorted(img_list, key=operator.itemgetter(2), reverse=True) return img_list, trash_img_list - + def sort_by_face_yaw(input_path): io.log_info ("Sorting by face yaw...") img_list = [] trash_img_list = [] for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading"): filepath = Path(filepath) - + 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) ) trash_img_list.append ( [str(filepath)] ) continue - + pitch, yaw = LandmarksProcessor.estimate_pitch_yaw ( dflimg.get_landmarks() ) - + img_list.append( [str(filepath), yaw ] ) io.log_info ("Sorting...") img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True) - + return img_list, trash_img_list - + def sort_by_face_pitch(input_path): io.log_info ("Sorting by face pitch...") img_list = [] trash_img_list = [] for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading"): filepath = Path(filepath) - + 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) ) trash_img_list.append ( [str(filepath)] ) continue - + pitch, yaw = LandmarksProcessor.estimate_pitch_yaw ( dflimg.get_landmarks() ) - + img_list.append( [str(filepath), pitch ] ) io.log_info ("Sorting...") img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True) - + return img_list, trash_img_list class HistSsimSubprocessor(Subprocessor): class Cli(Subprocessor.Cli): #override - def on_initialize(self, client_dict): + def on_initialize(self, client_dict): self.log_info ('Running on %s.' % (client_dict['device_name']) ) #override def process_data(self, data): img_list = [] for x in data: - img = cv2_imread(x) + img = cv2_imread(x) img_list.append ([x, cv2.calcHist([img], [0], None, [256], [0, 256]), cv2.calcHist([img], [1], None, [256], [0, 256]), cv2.calcHist([img], [2], None, [256], [0, 256]) ]) - + img_list_len = len(img_list) for i in range(img_list_len-1): min_score = float("inf") @@ -268,23 +268,23 @@ class HistSsimSubprocessor(Subprocessor): min_score = score j_min_score = j img_list[i+1], img_list[j_min_score] = img_list[j_min_score], img_list[i+1] - + self.progress_bar_inc(1) - - return img_list + + return img_list #override def get_data_name (self, data): return "Bunch of images" - + #override - def __init__(self, img_list ): + def __init__(self, img_list ): self.img_list = img_list self.img_list_len = len(img_list) - + slice_count = 20000 sliced_count = self.img_list_len // slice_count - + if sliced_count > 12: sliced_count = 11.9 slice_count = int(self.img_list_len / sliced_count) @@ -294,10 +294,10 @@ class HistSsimSubprocessor(Subprocessor): [ self.img_list[sliced_count*slice_count:] ] self.result = [] - super().__init__('HistSsim', HistSsimSubprocessor.Cli, 0) + super().__init__('HistSsim', HistSsimSubprocessor.Cli, 0) #override - def process_info_generator(self): + def process_info_generator(self): for i in range( len(self.img_chunks_list) ): yield 'CPU%d' % (i), {'i':i}, {'device_idx': i, 'device_name': 'CPU%d' % (i) @@ -306,21 +306,21 @@ class HistSsimSubprocessor(Subprocessor): def on_clients_initialized(self): io.progress_bar ("Sorting", len(self.img_list)) io.progress_bar_inc(len(self.img_chunks_list)) - + #override def on_clients_finalized(self): io.progress_bar_close() - + #override - def get_data(self, host_dict): - if len (self.img_chunks_list) > 0: - return self.img_chunks_list.pop(0) + def get_data(self, host_dict): + if len (self.img_chunks_list) > 0: + return self.img_chunks_list.pop(0) return None - + #override def on_data_return (self, host_dict, data): raise Exception("Fail to process data. Decrease number of images and try again.") - + #override def on_result (self, host_dict, data, result): self.result += result @@ -329,10 +329,10 @@ class HistSsimSubprocessor(Subprocessor): #override def get_result(self): return self.result - + def sort_by_hist(input_path): io.log_info ("Sorting by histogram similarity...") - img_list = HistSsimSubprocessor(Path_utils.get_image_paths(input_path)).run() + img_list = HistSsimSubprocessor(Path_utils.get_image_paths(input_path)).run() return img_list class HistDissimSubprocessor(Subprocessor): @@ -344,7 +344,7 @@ class HistDissimSubprocessor(Subprocessor): self.img_list_len = len(self.img_list) #override - def process_data(self, data): + def process_data(self, data): i = data[0] score_total = 0 for j in range( 0, self.img_list_len): @@ -358,40 +358,40 @@ class HistDissimSubprocessor(Subprocessor): def get_data_name (self, data): #return string identificator of your data return self.img_list[data[0]][0] - + #override - def __init__(self, img_list ): + def __init__(self, img_list ): self.img_list = img_list self.img_list_range = [i for i in range(0, len(img_list) )] self.result = [] - super().__init__('HistDissim', HistDissimSubprocessor.Cli, 60) + super().__init__('HistDissim', HistDissimSubprocessor.Cli, 60) #override def on_clients_initialized(self): io.progress_bar ("Sorting", len (self.img_list) ) - + #override def on_clients_finalized(self): io.progress_bar_close() - + #override - def process_info_generator(self): + def process_info_generator(self): for i in range(0, min(multiprocessing.cpu_count(), 8) ): yield 'CPU%d' % (i), {}, {'device_idx': i, - 'device_name': 'CPU%d' % (i), + 'device_name': 'CPU%d' % (i), 'img_list' : self.img_list } #override def get_data(self, host_dict): - if len (self.img_list_range) > 0: + if len (self.img_list_range) > 0: return [self.img_list_range.pop(0)] - + return None - + #override def on_data_return (self, host_dict, data): - self.img_list_range.insert(0, data[0]) - + self.img_list_range.insert(0, data[0]) + #override def on_result (self, host_dict, data, result): self.img_list[data[0]][2] = result @@ -400,7 +400,7 @@ class HistDissimSubprocessor(Subprocessor): #override def get_result(self): return self.img_list - + def sort_by_hist_dissim(input_path): io.log_info ("Sorting by histogram dissimilarity...") @@ -408,19 +408,19 @@ def sort_by_hist_dissim(input_path): trash_img_list = [] for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading"): filepath = Path(filepath) - + 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) ) trash_img_list.append ([str(filepath)]) continue - + image = cv2_imread(str(filepath)) face_mask = LandmarksProcessor.get_image_hull_mask (image.shape, dflimg.get_landmarks()) image = (image*face_mask).astype(np.uint8) @@ -428,26 +428,26 @@ def sort_by_hist_dissim(input_path): img_list.append ([str(filepath), cv2.calcHist([cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)], [0], None, [256], [0, 256]), 0 ]) img_list = HistDissimSubprocessor(img_list).run() - + io.log_info ("Sorting...") img_list = sorted(img_list, key=operator.itemgetter(2), reverse=True) return img_list, trash_img_list - + def sort_by_brightness(input_path): io.log_info ("Sorting by brightness...") img_list = [ [x, np.mean ( cv2.cvtColor(cv2_imread(x), cv2.COLOR_BGR2HSV)[...,2].flatten() )] for x in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading") ] io.log_info ("Sorting...") - img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True) + img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True) return img_list - + def sort_by_hue(input_path): io.log_info ("Sorting by hue...") img_list = [ [x, np.mean ( cv2.cvtColor(cv2_imread(x), cv2.COLOR_BGR2HSV)[...,0].flatten() )] for x in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading") ] io.log_info ("Sorting...") - img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True) + img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True) return img_list - + def sort_by_black(input_path): io.log_info ("Sorting by amount of black pixels...") @@ -460,22 +460,22 @@ def sort_by_black(input_path): img_list = sorted(img_list, key=operator.itemgetter(1), reverse=False) return img_list - + def sort_by_origname(input_path): io.log_info ("Sort by original filename...") - + img_list = [] trash_img_list = [] for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading"): filepath = Path(filepath) - + 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) ) trash_img_list.append( [str(filepath)] ) @@ -486,7 +486,7 @@ def sort_by_origname(input_path): io.log_info ("Sorting...") img_list = sorted(img_list, key=operator.itemgetter(1)) return img_list, trash_img_list - + def sort_by_oneface_in_image(input_path): io.log_info ("Sort by one face in images...") image_paths = Path_utils.get_image_paths(input_path) @@ -503,17 +503,17 @@ def sort_by_oneface_in_image(input_path): trash_img_list = [ (image_paths[x],) for x in idxs ] return img_list, trash_img_list return [], [] - + class FinalLoaderSubprocessor(Subprocessor): class Cli(Subprocessor.Cli): #override - def on_initialize(self, client_dict): + def on_initialize(self, client_dict): self.log_info ('Running on %s.' % (client_dict['device_name']) ) self.include_by_blur = client_dict['include_by_blur'] - + #override - def process_data(self, data): - filepath = Path(data[0]) + def process_data(self, data): + filepath = Path(data[0]) try: if filepath.suffix == '.png': @@ -522,40 +522,40 @@ class FinalLoaderSubprocessor(Subprocessor): dflimg = DFLJPG.load( str(filepath) ) else: dflimg = None - + if dflimg is None: self.log_err("%s is not a dfl image file" % (filepath.name)) return [ 1, [str(filepath)] ] - + bgr = cv2_imread(str(filepath)) if bgr is None: - raise Exception ("Unable to load %s" % (filepath.name) ) - + raise Exception ("Unable to load %s" % (filepath.name) ) + gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY) sharpness = estimate_sharpness(gray) if self.include_by_blur else 0 pitch, yaw = LandmarksProcessor.estimate_pitch_yaw ( dflimg.get_landmarks() ) - + hist = cv2.calcHist([gray], [0], None, [256], [0, 256]) except Exception as e: self.log_err (e) return [ 1, [str(filepath)] ] - + return [ 0, [str(filepath), sharpness, hist, yaw ] ] #override def get_data_name (self, data): #return string identificator of your data return data[0] - + #override - def __init__(self, img_list, include_by_blur ): + def __init__(self, img_list, include_by_blur ): self.img_list = img_list self.include_by_blur = include_by_blur self.result = [] self.result_trash = [] - super().__init__('FinalLoader', FinalLoaderSubprocessor.Cli, 60) + super().__init__('FinalLoader', FinalLoaderSubprocessor.Cli, 60) #override def on_clients_initialized(self): @@ -564,9 +564,9 @@ class FinalLoaderSubprocessor(Subprocessor): #override def on_clients_finalized(self): io.progress_bar_close() - + #override - def process_info_generator(self): + def process_info_generator(self): for i in range(0, min(multiprocessing.cpu_count(), 8) ): yield 'CPU%d' % (i), {}, {'device_idx': i, 'device_name': 'CPU%d' % (i), @@ -575,15 +575,15 @@ class FinalLoaderSubprocessor(Subprocessor): #override def get_data(self, host_dict): - if len (self.img_list) > 0: + if len (self.img_list) > 0: return [self.img_list.pop(0)] - + return None - + #override def on_data_return (self, host_dict, data): - self.img_list.insert(0, data[0]) - + self.img_list.insert(0, data[0]) + #override def on_result (self, host_dict, data, result): if result[0] == 0: @@ -599,7 +599,7 @@ class FinalLoaderSubprocessor(Subprocessor): class FinalHistDissimSubprocessor(Subprocessor): class Cli(Subprocessor.Cli): #override - def on_initialize(self, client_dict): + def on_initialize(self, client_dict): self.log_info ('Running on %s.' % (client_dict['device_name']) ) #override @@ -611,25 +611,25 @@ class FinalHistDissimSubprocessor(Subprocessor): if i == j: continue score_total += cv2.compareHist(img_list[i][2], img_list[j][2], cv2.HISTCMP_BHATTACHARYYA) - img_list[i][3] = score_total + img_list[i][3] = score_total img_list = sorted(img_list, key=operator.itemgetter(3), reverse=True) - return idx, img_list + return idx, img_list #override def get_data_name (self, data): return "Bunch of images" - - #override - def __init__(self, yaws_sample_list ): - self.yaws_sample_list = yaws_sample_list - self.yaws_sample_list_len = len(yaws_sample_list) - - self.yaws_sample_list_idxs = [ i for i in range(self.yaws_sample_list_len) if self.yaws_sample_list[i] is not None ] - self.result = [ None for _ in range(self.yaws_sample_list_len) ] - super().__init__('FinalHistDissimSubprocessor', FinalHistDissimSubprocessor.Cli) #override - def process_info_generator(self): + def __init__(self, yaws_sample_list ): + self.yaws_sample_list = yaws_sample_list + self.yaws_sample_list_len = len(yaws_sample_list) + + self.yaws_sample_list_idxs = [ i for i in range(self.yaws_sample_list_len) if self.yaws_sample_list[i] is not None ] + self.result = [ None for _ in range(self.yaws_sample_list_len) ] + super().__init__('FinalHistDissimSubprocessor', FinalHistDissimSubprocessor.Cli) + + #override + def process_info_generator(self): for i in range(min(multiprocessing.cpu_count(), 8) ): yield 'CPU%d' % (i), {'i':i}, {'device_idx': i, 'device_name': 'CPU%d' % (i) @@ -637,38 +637,38 @@ class FinalHistDissimSubprocessor(Subprocessor): #override def on_clients_initialized(self): io.progress_bar ("Sort by hist-dissim", self.yaws_sample_list_len) - + #override def on_clients_finalized(self): io.progress_bar_close() - + #override - def get_data(self, host_dict): + def get_data(self, host_dict): if len (self.yaws_sample_list_idxs) > 0: idx = self.yaws_sample_list_idxs.pop(0) - + return idx, self.yaws_sample_list[idx] return None - + #override def on_data_return (self, host_dict, data): self.yaws_sample_list_idxs.insert(0, data[0]) - + #override def on_result (self, host_dict, data, result): - idx, yaws_sample_list = data + idx, yaws_sample_list = data self.result[idx] = yaws_sample_list io.progress_bar_inc(1) #override def get_result(self): return self.result - + def sort_final(input_path, include_by_blur=True): io.log_info ("Performing final sort.") - + target_count = io.input_int ("Target number of images? (default:2000) : ", 2000) - + img_list, trash_img_list = FinalLoaderSubprocessor( Path_utils.get_image_paths(input_path), include_by_blur ).run() final_img_list = [] @@ -676,12 +676,12 @@ def sort_final(input_path, include_by_blur=True): imgs_per_grad = round (target_count / grads) grads_space = np.linspace (-1.0,1.0,grads) - + yaws_sample_list = [None]*grads - for g in io.progress_bar_generator ( range(grads), "Sort by yaw"): + for g in io.progress_bar_generator ( range(grads), "Sort by yaw"): yaw = grads_space[g] next_yaw = grads_space[g+1] if g < grads-1 else yaw - + yaw_samples = [] for img in img_list: s_yaw = -img[3] @@ -691,17 +691,17 @@ def sort_final(input_path, include_by_blur=True): yaw_samples += [ img ] if len(yaw_samples) > 0: yaws_sample_list[g] = yaw_samples - + total_lack = 0 for g in io.progress_bar_generator ( range(grads), ""): img_list = yaws_sample_list[g] img_list_len = len(img_list) if img_list is not None else 0 - - lack = imgs_per_grad - img_list_len - total_lack += max(lack, 0) - imgs_per_grad += total_lack // grads - + lack = imgs_per_grad - img_list_len + total_lack += max(lack, 0) + + imgs_per_grad += total_lack // grads + if include_by_blur: sharpned_imgs_per_grad = imgs_per_grad*10 for g in io.progress_bar_generator ( range (grads), "Sort by blur"): @@ -709,47 +709,47 @@ def sort_final(input_path, include_by_blur=True): if img_list is None: continue - img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True) - + img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True) + if len(img_list) > sharpned_imgs_per_grad: trash_img_list += img_list[sharpned_imgs_per_grad:] img_list = img_list[0:sharpned_imgs_per_grad] - + yaws_sample_list[g] = img_list - + yaws_sample_list = FinalHistDissimSubprocessor(yaws_sample_list).run() for g in io.progress_bar_generator ( range (grads), "Fetching best"): img_list = yaws_sample_list[g] if img_list is None: continue - + final_img_list += img_list[0:imgs_per_grad] trash_img_list += img_list[imgs_per_grad:] return final_img_list, trash_img_list - + def final_process(input_path, img_list, trash_img_list): if len(trash_img_list) != 0: parent_input_path = input_path.parent trash_path = parent_input_path / (input_path.stem + '_trash') trash_path.mkdir (exist_ok=True) - - io.log_info ("Trashing %d items to %s" % ( len(trash_img_list), str(trash_path) ) ) - + + io.log_info ("Trashing %d items to %s" % ( len(trash_img_list), str(trash_path) ) ) + for filename in Path_utils.get_image_paths(trash_path): Path(filename).unlink() for i in io.progress_bar_generator( range(len(trash_img_list)), "Moving trash", leave=False): - src = Path (trash_img_list[i][0]) + src = Path (trash_img_list[i][0]) dst = trash_path / src.name try: src.rename (dst) except: io.log_info ('fail to trashing %s' % (src.name) ) - + io.log_info ("") - + if len(img_list) != 0: for i in io.progress_bar_generator( [*range(len(img_list))], "Renaming", leave=False): src = Path (img_list[i][0]) @@ -758,24 +758,24 @@ def final_process(input_path, img_list, trash_img_list): src.rename (dst) except: io.log_info ('fail to rename %s' % (src.name) ) - + for i in io.progress_bar_generator( [*range(len(img_list))], "Renaming"): - src = Path (img_list[i][0]) + src = Path (img_list[i][0]) src = input_path / ('%.5d_%s' % (i, src.name)) dst = input_path / ('%.5d%s' % (i, src.suffix)) try: src.rename (dst) except: - io.log_info ('fail to rename %s' % (src.name) ) + io.log_info ('fail to rename %s' % (src.name) ) + - def main (input_path, sort_by_method): input_path = Path(input_path) sort_by_method = sort_by_method.lower() io.log_info ("Running sort tool.\r\n") - + img_list = [] trash_img_list = [] if sort_by_method == 'blur': img_list, trash_img_list = sort_by_blur (input_path) @@ -787,10 +787,10 @@ def main (input_path, sort_by_method): elif sort_by_method == 'hist-dissim': img_list, trash_img_list = sort_by_hist_dissim (input_path) elif sort_by_method == 'brightness': img_list = sort_by_brightness (input_path) elif sort_by_method == 'hue': img_list = sort_by_hue (input_path) - elif sort_by_method == 'black': img_list = sort_by_black (input_path) + elif sort_by_method == 'black': img_list = sort_by_black (input_path) elif sort_by_method == 'origname': img_list, trash_img_list = sort_by_origname (input_path) - elif sort_by_method == 'oneface': img_list, trash_img_list = sort_by_oneface_in_image (input_path) - elif sort_by_method == 'final': img_list, trash_img_list = sort_final (input_path) + elif sort_by_method == 'oneface': img_list, trash_img_list = sort_by_oneface_in_image (input_path) + elif sort_by_method == 'final': img_list, trash_img_list = sort_final (input_path) elif sort_by_method == 'final-no-blur': img_list, trash_img_list = sort_final (input_path, include_by_blur=False) - + final_process (input_path, img_list, trash_img_list) diff --git a/mainscripts/Trainer.py b/mainscripts/Trainer.py index fc8283b..2c2e563 100644 --- a/mainscripts/Trainer.py +++ b/mainscripts/Trainer.py @@ -7,39 +7,39 @@ import numpy as np import itertools from pathlib import Path from utils import Path_utils -from utils import image_utils +from utils import image_utils import cv2 import models from interact import interact as io def trainerThread (s2c, c2s, args, device_args): while True: - try: + try: training_data_src_path = Path( args.get('training_data_src_dir', '') ) training_data_dst_path = Path( args.get('training_data_dst_dir', '') ) model_path = Path( args.get('model_path', '') ) model_name = args.get('model_name', '') - save_interval_min = 15 + save_interval_min = 15 debug = args.get('debug', '') - + if not training_data_src_path.exists(): io.log_err('Training data src directory does not exist.') break - + if not training_data_dst_path.exists(): io.log_err('Training data dst directory does not exist.') break - + if not model_path.exists(): model_path.mkdir(exist_ok=True) - + model = models.import_model(model_name)( - model_path, - training_data_src_path=training_data_src_path, - training_data_dst_path=training_data_dst_path, + model_path, + training_data_src_path=training_data_src_path, + training_data_dst_path=training_data_dst_path, debug=debug, device_args=device_args) - + is_reached_goal = model.is_reached_iter_goal() is_upd_save_time_after_train = False loss_string = "" @@ -49,37 +49,37 @@ def trainerThread (s2c, c2s, args, device_args): model.save() io.log_info(loss_string) is_upd_save_time_after_train = True - + def send_preview(): - if not debug: - previews = model.get_previews() + if not debug: + previews = model.get_previews() c2s.put ( {'op':'show', 'previews': previews, 'iter':model.get_iter(), 'loss_history': model.get_loss_history().copy() } ) else: previews = [( 'debug, press update for new', model.debug_one_iter())] c2s.put ( {'op':'show', 'previews': previews} ) - - + + if model.is_first_run(): model_save() - + if model.get_target_iter() != 0: if is_reached_goal: io.log_info('Model already trained to target iteration. You can use preview.') else: io.log_info('Starting. Target iteration: %d. Press "Enter" to stop training and save model.' % ( model.get_target_iter() ) ) - else: + else: io.log_info('Starting. Press "Enter" to stop training and save model.') - + last_save_time = time.time() - + for i in itertools.count(0,1): if not debug: if not is_reached_goal: - loss_string = model.train_one_iter() + loss_string = model.train_one_iter() if is_upd_save_time_after_train: #save resets plaidML programs, so upd last_save_time only after plaidML rebuild them last_save_time = time.time() - + io.log_info (loss_string, end='\r') if model.get_target_iter() != 0 and model.is_reached_iter_goal(): io.log_info ('Reached target iteration.') @@ -91,77 +91,77 @@ def trainerThread (s2c, c2s, args, device_args): last_save_time = time.time() model_save() send_preview() - + if i==0: if is_reached_goal: - model.pass_one_iter() + model.pass_one_iter() send_preview() - + if debug: time.sleep(0.005) - + while not s2c.empty(): input = s2c.get() op = input['op'] if op == 'save': model_save() - elif op == 'preview': + elif op == 'preview': if is_reached_goal: - model.pass_one_iter() + model.pass_one_iter() send_preview() elif op == 'close': model_save() i = -1 break - + if i == -1: break - - + + model.finalize() - + except Exception as e: print ('Error: %s' % (str(e))) traceback.print_exc() break c2s.put ( {'op':'close'} ) - - + + def main(args, device_args): io.log_info ("Running trainer.\r\n") - + no_preview = args.get('no_preview', False) - + s2c = queue.Queue() c2s = queue.Queue() - + thread = threading.Thread(target=trainerThread, args=(s2c, c2s, args, device_args) ) thread.start() if no_preview: - while True: + while True: if not c2s.empty(): input = c2s.get() op = input.get('op','') if op == 'close': break io.process_messages(0.1) - else: + else: wnd_name = "Training preview" io.named_window(wnd_name) io.capture_keys(wnd_name) - + previews = None loss_history = None selected_preview = 0 update_preview = False is_showing = False is_waiting_preview = False - show_last_history_iters_count = 0 + show_last_history_iters_count = 0 iter = 0 - while True: + while True: if not c2s.empty(): input = c2s.get() op = input['op'] @@ -177,7 +177,7 @@ def main(args, device_args): (h, w, c) = preview_rgb.shape max_h = max (max_h, h) max_w = max (max_w, w) - + max_size = 800 if max_h > max_size: max_w = int( max_w / (max_h / max_size) ) @@ -194,49 +194,49 @@ def main(args, device_args): update_preview = True elif op == 'close': break - + if update_preview: update_preview = False selected_preview_name = previews[selected_preview][0] selected_preview_rgb = previews[selected_preview][1] (h,w,c) = selected_preview_rgb.shape - + # HEAD head_lines = [ '[s]:save [enter]:exit', '[p]:update [space]:next preview [l]:change history range', 'Preview: "%s" [%d/%d]' % (selected_preview_name,selected_preview+1, len(previews) ) - ] + ] head_line_height = 15 head_height = len(head_lines) * head_line_height head = np.ones ( (head_height,w,c) ) * 0.1 - + for i in range(0, len(head_lines)): t = i*head_line_height b = (i+1)*head_line_height head[t:b, 0:w] += image_utils.get_text_image ( (w,head_line_height,c) , head_lines[i], color=[0.8]*c ) - + final = head - - if loss_history is not None: + + if loss_history is not None: if show_last_history_iters_count == 0: loss_history_to_show = loss_history else: loss_history_to_show = loss_history[-show_last_history_iters_count:] - + lh_img = models.ModelBase.get_loss_history_preview(loss_history_to_show, iter, w, c) final = np.concatenate ( [final, lh_img], axis=0 ) final = np.concatenate ( [final, selected_preview_rgb], axis=0 ) final = np.clip(final, 0, 1) - + io.show_image( wnd_name, (final*255).astype(np.uint8) ) is_showing = True - + key_events = io.get_key_events(wnd_name) key, = key_events[-1] if len(key_events) > 0 else (0,) - + if key == ord('\n') or key == ord('\r'): s2c.put ( {'op': 'close'} ) elif key == ord('s'): @@ -253,14 +253,14 @@ def main(args, device_args): elif show_last_history_iters_count == 10000: show_last_history_iters_count = 50000 elif show_last_history_iters_count == 50000: - show_last_history_iters_count = 100000 + show_last_history_iters_count = 100000 elif show_last_history_iters_count == 100000: - show_last_history_iters_count = 0 + show_last_history_iters_count = 0 update_preview = True elif key == ord(' '): selected_preview = (selected_preview + 1) % len(previews) update_preview = True - + io.process_messages(0.1) - - io.destroy_all_windows() \ No newline at end of file + + io.destroy_all_windows() diff --git a/mainscripts/Util.py b/mainscripts/Util.py index 09540ff..93f3a4e 100644 --- a/mainscripts/Util.py +++ b/mainscripts/Util.py @@ -9,30 +9,30 @@ from interact import interact as io def convert_png_to_jpg_file (filepath): filepath = Path(filepath) - - if filepath.suffix != '.png': + + if filepath.suffix != '.png': return - + dflpng = DFLPNG.load (str(filepath) ) if dflpng is None: - io.log_err ("%s is not a dfl image file" % (filepath.name) ) + io.log_err ("%s is not a dfl image file" % (filepath.name) ) return - + dfl_dict = dflpng.getDFLDictData() - + img = cv2_imread (str(filepath)) new_filepath = str(filepath.parent / (filepath.stem + '.jpg')) cv2_imwrite ( new_filepath, img, [int(cv2.IMWRITE_JPEG_QUALITY), 85]) - DFLJPG.embed_data( new_filepath, + DFLJPG.embed_data( new_filepath, face_type=dfl_dict.get('face_type', None), landmarks=dfl_dict.get('landmarks', 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) ) - + filepath.unlink() - + def convert_png_to_jpg_folder (input_path): input_path = Path(input_path) @@ -41,73 +41,73 @@ def convert_png_to_jpg_folder (input_path): for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Converting"): filepath = Path(filepath) convert_png_to_jpg_file(filepath) - + def add_landmarks_debug_images(input_path): io.log_info ("Adding landmarks debug images...") for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Processing"): filepath = Path(filepath) - + img = cv2_imread(str(filepath)) - + 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) ) + io.log_err ("%s is not a dfl image file" % (filepath.name) ) continue if img is not None: face_landmarks = dflimg.get_landmarks() LandmarksProcessor.draw_landmarks(img, face_landmarks, transparent_mask=True) - + output_file = '{}{}'.format( str(Path(str(input_path)) / filepath.stem), '_debug.jpg') cv2_imwrite(output_file, img, [int(cv2.IMWRITE_JPEG_QUALITY), 50] ) - + def recover_original_aligned_filename(input_path): io.log_info ("Recovering original aligned filename...") - + files = [] for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Processing"): filepath = Path(filepath) - + 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) ) + io.log_err ("%s is not a dfl image file" % (filepath.name) ) continue - + files += [ [filepath, None, dflimg.get_source_filename(), False] ] - + files_len = len(files) for i in io.progress_bar_generator( range(files_len), "Sorting" ): fp, _, sf, converted = files[i] - + if converted: continue - + sf_stem = Path(sf).stem - + files[i][1] = fp.parent / ( sf_stem + '_0' + fp.suffix ) files[i][3] = True c = 1 - + for j in range(i+1, files_len): fp_j, _, sf_j, converted_j = files[j] if converted_j: continue - + if sf_j == sf: - files[j][1] = fp_j.parent / ( sf_stem + ('_%d' % (c)) + fp_j.suffix ) + files[j][1] = fp_j.parent / ( sf_stem + ('_%d' % (c)) + fp_j.suffix ) files[j][3] = True c += 1 @@ -118,11 +118,11 @@ def recover_original_aligned_filename(input_path): fs.rename (dst) except: io.log_err ('fail to rename %s' % (fs.name) ) - + for file in io.progress_bar_generator( files, "Renaming" ): fs, fd, _, _ = file fs = fs.parent / ( fs.stem + '_tmp' + fs.suffix ) try: fs.rename (fd) except: - io.log_err ('fail to rename %s' % (fs.name) ) \ No newline at end of file + io.log_err ('fail to rename %s' % (fs.name) ) diff --git a/mainscripts/VideoEd.py b/mainscripts/VideoEd.py index fd5cdca..945109c 100644 --- a/mainscripts/VideoEd.py +++ b/mainscripts/VideoEd.py @@ -8,38 +8,38 @@ from interact import interact as io def extract_video(input_file, output_dir, output_ext=None, fps=None): input_file_path = Path(input_file) output_path = Path(output_dir) - + if not output_path.exists(): output_path.mkdir(exist_ok=True) - - + + if input_file_path.suffix == '.*': input_file_path = Path_utils.get_first_file_by_stem (input_file_path.parent, input_file_path.stem) else: if not input_file_path.exists(): input_file_path = None - + if input_file_path is None: io.log_err("input_file not found.") return - + if output_ext is None: output_ext = io.input_str ("Output image format (extension)? ( default:png ) : ", "png") - + if fps is None: fps = io.input_int ("Enter FPS ( ?:help skip:fullfps ) : ", 0, help_message="How many frames of every second of the video will be extracted.") - + for filename in Path_utils.get_image_paths (output_path, ['.'+output_ext]): Path(filename).unlink() job = ffmpeg.input(str(input_file_path)) - - kwargs = {} + + kwargs = {} if fps != 0: kwargs.update ({'r':str(fps)}) job = job.output( str (output_path / ('%5d.'+output_ext)), **kwargs ) - + try: job = job.run() except: @@ -50,18 +50,18 @@ def cut_video ( input_file, from_time=None, to_time=None, audio_track_id=None, b if input_file_path is None: io.log_err("input_file not found.") return - + output_file_path = input_file_path.parent / (input_file_path.stem + "_cut" + input_file_path.suffix) - + if from_time is None: from_time = io.input_str ("From time (skip: 00:00:00.000) : ", "00:00:00.000") - + if to_time is None: to_time = io.input_str ("To time (skip: 00:00:00.000) : ", "00:00:00.000") - + if audio_track_id is None: audio_track_id = io.input_int ("Specify audio track id. ( skip:0 ) : ", 0) - + if bitrate is None: bitrate = max (1, io.input_int ("Bitrate of output file in MB/s ? (default:25) : ", 25) ) @@ -69,64 +69,64 @@ def cut_video ( input_file, from_time=None, to_time=None, audio_track_id=None, b "b:v": "%dM" %(bitrate), "pix_fmt": "yuv420p", } - + job = ffmpeg.input(str(input_file_path), ss=from_time, to=to_time) - + job_v = job['v:0'] job_a = job['a:' + str(audio_track_id) + '?' ] - + job = ffmpeg.output(job_v, job_a, str(output_file_path), **kwargs).overwrite_output() - + try: job = job.run() except: io.log_err ("ffmpeg fail, job commandline:" + str(job.compile()) ) - + def denoise_image_sequence( input_dir, ext=None, factor=None ): input_path = Path(input_dir) - + if not input_path.exists(): io.log_err("input_dir not found.") return - + if ext is None: ext = io.input_str ("Input image format (extension)? ( default:png ) : ", "png") - + if factor is None: factor = np.clip ( io.input_int ("Denoise factor? (1-20 default:5) : ", 5), 1, 20 ) - + job = ( ffmpeg .input(str ( input_path / ('%5d.'+ext) ) ) .filter("hqdn3d", factor, factor, 5,5) .output(str ( input_path / ('%5d.'+ext) ) ) - ) - + ) + try: job = job.run() except: io.log_err ("ffmpeg fail, job commandline:" + str(job.compile()) ) - + def video_from_sequence( input_dir, output_file, reference_file=None, ext=None, fps=None, bitrate=None, lossless=None ): - input_path = Path(input_dir) + input_path = Path(input_dir) output_file_path = Path(output_file) reference_file_path = Path(reference_file) if reference_file is not None else None - + if not input_path.exists(): io.log_err("input_dir not found.") return - + if not output_file_path.parent.exists(): output_file_path.parent.mkdir(parents=True, exist_ok=True) return - + out_ext = output_file_path.suffix - + if ext is None: ext = io.input_str ("Input image format (extension)? ( default:png ) : ", "png") - + if lossless is None: lossless = io.input_bool ("Use lossless codec ? ( default:no ) : ", False) - + video_id = None audio_id = None ref_in_a = None @@ -136,7 +136,7 @@ def video_from_sequence( input_dir, output_file, reference_file=None, ext=None, else: if not reference_file_path.exists(): reference_file_path = None - + if reference_file_path is None: io.log_err("reference_file not found.") return @@ -149,32 +149,32 @@ def video_from_sequence( input_dir, output_file, reference_file=None, ext=None, if video_id is None and stream['codec_type'] == 'video': video_id = stream['index'] fps = stream['r_frame_rate'] - + if audio_id is None and stream['codec_type'] == 'audio': audio_id = stream['index'] if audio_id is not None: #has audio track ref_in_a = ffmpeg.input (str(reference_file_path))[str(audio_id)] - + if fps is None: #if fps not specified and not overwritten by reference-file fps = max (1, io.input_int ("FPS ? (default:25) : ", 25) ) - + if not lossless and bitrate is None: bitrate = max (1, io.input_int ("Bitrate of output file in MB/s ? (default:16) : ", 16) ) - + i_in = ffmpeg.input(str (input_path / ('%5d.'+ext)), r=fps) - + output_args = [i_in] - + if ref_in_a is not None: output_args += [ref_in_a] - + output_args += [str (output_file_path)] - + output_kwargs = {} - + if lossless: output_kwargs.update ({"c:v": "png" }) @@ -183,15 +183,14 @@ def video_from_sequence( input_dir, output_file, reference_file=None, ext=None, "b:v": "%dM" %(bitrate), "pix_fmt": "yuv420p", }) - + output_kwargs.update ({"c:a": "aac", "b:a": "192k", "ar" : "48000" }) - + job = ( ffmpeg.output(*output_args, **output_kwargs).overwrite_output() ) try: job = job.run() except: io.log_err ("ffmpeg fail, job commandline:" + str(job.compile()) ) - diff --git a/mathlib/__init__.py b/mathlib/__init__.py index 3b038f3..a11e725 100644 --- a/mathlib/__init__.py +++ b/mathlib/__init__.py @@ -7,10 +7,10 @@ def get_power_of_two(x): while (1 << i) < x: i += 1 return i - + def rotationMatrixToEulerAngles(R) : - sy = math.sqrt(R[0,0] * R[0,0] + R[1,0] * R[1,0]) - singular = sy < 1e-6 + sy = math.sqrt(R[0,0] * R[0,0] + R[1,0] * R[1,0]) + singular = sy < 1e-6 if not singular : x = math.atan2(R[2,1] , R[2,2]) y = math.atan2(-R[2,0], sy) @@ -18,8 +18,8 @@ def rotationMatrixToEulerAngles(R) : else : x = math.atan2(-R[1,2], R[1,1]) y = math.atan2(-R[2,0], sy) - z = 0 + z = 0 return np.array([x, y, z]) - + def polygon_area(x,y): - return 0.5*np.abs(np.dot(x,np.roll(y,1))-np.dot(y,np.roll(x,1))) \ No newline at end of file + return 0.5*np.abs(np.dot(x,np.roll(y,1))-np.dot(y,np.roll(x,1))) diff --git a/mathlib/umeyama.py b/mathlib/umeyama.py index aad6235..7c6b2d0 100644 --- a/mathlib/umeyama.py +++ b/mathlib/umeyama.py @@ -68,4 +68,4 @@ def umeyama(src, dst, estimate_scale): T[:dim, dim] = dst_mean - scale * np.dot(T[:dim, :dim], src_mean.T) T[:dim, :dim] *= scale - return T \ No newline at end of file + return T diff --git a/models/ModelBase.py b/models/ModelBase.py index d7b9909..6b4f5ad 100644 --- a/models/ModelBase.py +++ b/models/ModelBase.py @@ -23,11 +23,11 @@ class ModelBase(object): def __init__(self, model_path, training_data_src_path=None, training_data_dst_path=None, debug = False, device_args = None, ask_write_preview_history=True, ask_target_iter=True, ask_batch_size=True, ask_sort_by_yaw=True, ask_random_flip=True, ask_src_scale_mod=True): - + device_args['force_gpu_idx'] = device_args.get('force_gpu_idx',-1) device_args['cpu_only'] = device_args.get('cpu_only',False) - - if device_args['force_gpu_idx'] == -1 and not device_args['cpu_only']: + + if device_args['force_gpu_idx'] == -1 and not device_args['cpu_only']: idxs_names_list = nnlib.device.getValidDevicesIdxsWithNamesList() if len(idxs_names_list) > 1: io.log_info ("You have multi GPUs in a system: ") @@ -36,17 +36,17 @@ class ModelBase(object): device_args['force_gpu_idx'] = io.input_int("Which GPU idx to choose? ( skip: best GPU ) : ", -1, [ x[0] for x in idxs_names_list] ) self.device_args = device_args - + self.device_config = nnlib.DeviceConfig(allow_growth=False, **self.device_args) - + io.log_info ("Loading model...") - + self.model_path = model_path self.model_data_path = Path( self.get_strpath_storage_for_file('data.dat') ) - + self.training_data_src_path = training_data_src_path self.training_data_dst_path = training_data_dst_path - + self.src_images_paths = None self.dst_images_paths = None self.src_yaw_images_paths = None @@ -60,10 +60,10 @@ class ModelBase(object): self.options = {} self.loss_history = [] self.sample_for_preview = None - + model_data = {} - if self.model_data_path.exists(): - model_data = pickle.loads ( self.model_data_path.read_bytes() ) + if self.model_data_path.exists(): + model_data = pickle.loads ( self.model_data_path.read_bytes() ) self.iter = max( model_data.get('iter',0), model_data.get('epoch',0) ) if 'epoch' in self.options: self.options.pop('epoch') @@ -73,101 +73,101 @@ class ModelBase(object): self.sample_for_preview = model_data['sample_for_preview'] if 'sample_for_preview' in model_data.keys() else None ask_override = self.is_training_mode and self.iter != 0 and io.input_in_time ("Press enter in 2 seconds to override model settings.", 2) - + yn_str = {True:'y',False:'n'} - - if self.iter == 0: + + if self.iter == 0: io.log_info ("\nModel first run. Enter model options as default for each run.") - + if ask_write_preview_history and (self.iter == 0 or ask_override): default_write_preview_history = False if self.iter == 0 else self.options.get('write_preview_history',False) self.options['write_preview_history'] = io.input_bool("Write preview history? (y/n ?:help skip:%s) : " % (yn_str[default_write_preview_history]) , default_write_preview_history, help_message="Preview history will be writed to _history folder.") else: self.options['write_preview_history'] = self.options.get('write_preview_history', False) - + if ask_target_iter and (self.iter == 0 or ask_override): self.options['target_iter'] = max(0, io.input_int("Target iteration (skip:unlimited/default) : ", 0)) else: self.options['target_iter'] = max(model_data.get('target_iter',0), self.options.get('target_epoch',0)) if 'target_epoch' in self.options: self.options.pop('target_epoch') - + if ask_batch_size and (self.iter == 0 or ask_override): default_batch_size = 0 if self.iter == 0 else self.options.get('batch_size',0) self.options['batch_size'] = max(0, io.input_int("Batch_size (?:help skip:%d) : " % (default_batch_size), default_batch_size, help_message="Larger batch size is always better for NN's generalization, but it can cause Out of Memory error. Tune this value for your videocard manually.")) else: self.options['batch_size'] = self.options.get('batch_size', 0) - + if ask_sort_by_yaw and (self.iter == 0): self.options['sort_by_yaw'] = io.input_bool("Feed faces to network sorted by yaw? (y/n ?:help skip:n) : ", False, help_message="NN will not learn src face directions that don't match dst face directions." ) else: self.options['sort_by_yaw'] = self.options.get('sort_by_yaw', False) - + if ask_random_flip and (self.iter == 0): self.options['random_flip'] = io.input_bool("Flip faces randomly? (y/n ?:help skip:y) : ", True, help_message="Predicted face will look more naturally without this option, but src faceset should cover all face directions as dst faceset.") else: self.options['random_flip'] = self.options.get('random_flip', True) - + if ask_src_scale_mod and (self.iter == 0): self.options['src_scale_mod'] = np.clip( io.input_int("Src face scale modifier % ( -30...30, ?:help skip:0) : ", 0, help_message="If src face shape is wider than dst, try to decrease this value to get a better result."), -30, 30) - else: + else: self.options['src_scale_mod'] = self.options.get('src_scale_mod', 0) - + self.write_preview_history = self.options['write_preview_history'] if not self.options['write_preview_history']: - self.options.pop('write_preview_history') - + self.options.pop('write_preview_history') + self.target_iter = self.options['target_iter'] if self.options['target_iter'] == 0: - self.options.pop('target_iter') - + self.options.pop('target_iter') + self.batch_size = self.options['batch_size'] - self.sort_by_yaw = self.options['sort_by_yaw'] + self.sort_by_yaw = self.options['sort_by_yaw'] self.random_flip = self.options['random_flip'] - + self.src_scale_mod = self.options['src_scale_mod'] if self.src_scale_mod == 0: - self.options.pop('src_scale_mod') + self.options.pop('src_scale_mod') self.onInitializeOptions(self.iter == 0, ask_override) nnlib.import_all(self.device_config) self.keras = nnlib.keras self.K = nnlib.keras.backend - + self.onInitialize() - + self.options['batch_size'] = self.batch_size - + if self.debug or self.batch_size == 0: - self.batch_size = 1 - + self.batch_size = 1 + if self.is_training_mode: if self.write_preview_history: if self.device_args['force_gpu_idx'] == -1: self.preview_history_path = self.model_path / ( '%s_history' % (self.get_model_name()) ) else: self.preview_history_path = self.model_path / ( '%d_%s_history' % (self.device_args['force_gpu_idx'], self.get_model_name()) ) - + if not self.preview_history_path.exists(): self.preview_history_path.mkdir(exist_ok=True) else: if self.iter == 0: for filename in Path_utils.get_image_paths(self.preview_history_path): Path(filename).unlink() - + if self.generator_list is None: raise ValueError( 'You didnt set_training_data_generators()') else: for i, generator in enumerate(self.generator_list): if not isinstance(generator, SampleGeneratorBase): raise ValueError('training data generator is not subclass of SampleGeneratorBase') - + if (self.sample_for_preview is None) or (self.iter == 0): self.sample_for_preview = self.generate_next_sample() model_summary_text = [] - + model_summary_text += ["===== Model summary ====="] model_summary_text += ["== Model name: " + self.get_model_name()] model_summary_text += ["=="] @@ -179,41 +179,41 @@ class ModelBase(object): if self.device_config.multi_gpu: model_summary_text += ["== |== multi_gpu : True "] - + model_summary_text += ["== Running on:"] if self.device_config.cpu_only: model_summary_text += ["== |== [CPU]"] else: for idx in self.device_config.gpu_idxs: model_summary_text += ["== |== [%d : %s]" % (idx, nnlib.device.getDeviceName(idx))] - + if not self.device_config.cpu_only and self.device_config.gpu_vram_gb[0] == 2: model_summary_text += ["=="] model_summary_text += ["== WARNING: You are using 2GB GPU. Result quality may be significantly decreased."] model_summary_text += ["== If training does not start, close all programs and try again."] model_summary_text += ["== Also you can disable Windows Aero Desktop to get extra free VRAM."] model_summary_text += ["=="] - - model_summary_text += ["========================="] - model_summary_text = "\r\n".join (model_summary_text) - self.model_summary_text = model_summary_text + + model_summary_text += ["========================="] + model_summary_text = "\r\n".join (model_summary_text) + self.model_summary_text = model_summary_text io.log_info(model_summary_text) - + #overridable def onInitializeOptions(self, is_first_run, ask_override): pass - + #overridable def onInitialize(self): ''' initialize your keras models - + store and retrieve your model options in self.options[''] - + check example ''' pass - + #overridable def onSave(self): #save your keras models here @@ -229,59 +229,59 @@ class ModelBase(object): #overridable def onGetPreview(self, sample): #you can return multiple previews - #return [ ('preview_name',preview_rgb), ... ] + #return [ ('preview_name',preview_rgb), ... ] return [] #overridable if you want model name differs from folder name def get_model_name(self): return Path(inspect.getmodule(self).__file__).parent.name.rsplit("_", 1)[1] - + #overridable def get_converter(self): raise NotImplementeError #return existing or your own converter which derived from base - + def get_target_iter(self): return self.target_iter - + def is_reached_iter_goal(self): - return self.target_iter != 0 and self.iter >= self.target_iter - + return self.target_iter != 0 and self.iter >= self.target_iter + #multi gpu in keras actually is fake and doesn't work for training https://github.com/keras-team/keras/issues/11976 #def to_multi_gpu_model_if_possible (self, models_list): # if len(self.device_config.gpu_idxs) > 1: # #make batch_size to divide on GPU count without remainder # self.batch_size = int( self.batch_size / len(self.device_config.gpu_idxs) ) # if self.batch_size == 0: - # self.batch_size = 1 + # self.batch_size = 1 # self.batch_size *= len(self.device_config.gpu_idxs) - # + # # result = [] # for model in models_list: # for i in range( len(model.output_names) ): - # model.output_names = 'output_%d' % (i) - # result += [ nnlib.keras.utils.multi_gpu_model( model, self.device_config.gpu_idxs ) ] - # - # return result + # model.output_names = 'output_%d' % (i) + # result += [ nnlib.keras.utils.multi_gpu_model( model, self.device_config.gpu_idxs ) ] + # + # return result # else: # return models_list - - def get_previews(self): + + def get_previews(self): return self.onGetPreview ( self.last_sample ) - - def get_static_preview(self): + + def get_static_preview(self): return self.onGetPreview (self.sample_for_preview)[0][1] #first preview, and bgr - - def save(self): - Path( self.get_strpath_storage_for_file('summary.txt') ).write_text(self.model_summary_text) + + def save(self): + Path( self.get_strpath_storage_for_file('summary.txt') ).write_text(self.model_summary_text) self.onSave() - + model_data = { 'iter': self.iter, 'options': self.options, 'loss_history': self.loss_history, 'sample_for_preview' : self.sample_for_preview - } + } self.model_data_path.write_bytes( pickle.dumps(model_data) ) def load_weights_safe(self, model_filename_list, optimizer_filename_list=[]): @@ -289,17 +289,17 @@ class ModelBase(object): filename = self.get_strpath_storage_for_file(filename) if Path(filename).exists(): model.load_weights(filename) - + if len(optimizer_filename_list) != 0: opt_filename = self.get_strpath_storage_for_file('opt.h5') if Path(opt_filename).exists(): try: with open(opt_filename, "rb") as f: d = pickle.loads(f.read()) - + for x in optimizer_filename_list: opt, filename = x - if filename in d: + if filename in d: weights = d[filename].get('weights', None) if weights: opt.set_weights(weights) @@ -307,16 +307,16 @@ class ModelBase(object): except Exception as e: print ("Unable to load ", opt_filename) - + def save_weights_safe(self, model_filename_list, optimizer_filename_list=[]): for model, filename in model_filename_list: filename = self.get_strpath_storage_for_file(filename) model.save_weights( filename + '.tmp' ) rename_list = model_filename_list - if len(optimizer_filename_list) != 0: + if len(optimizer_filename_list) != 0: opt_filename = self.get_strpath_storage_for_file('opt.h5') - + try: d = {} for opt, filename in optimizer_filename_list: @@ -324,54 +324,54 @@ class ModelBase(object): symbolic_weights = getattr(opt, 'weights') if symbolic_weights: fd['weights'] = self.K.batch_get_value(symbolic_weights) - + d[filename] = fd - + with open(opt_filename+'.tmp', 'wb') as f: f.write( pickle.dumps(d) ) - + rename_list += [('', 'opt.h5')] except Exception as e: print ("Unable to save ", opt_filename) for _, filename in rename_list: - filename = self.get_strpath_storage_for_file(filename) + filename = self.get_strpath_storage_for_file(filename) source_filename = Path(filename+'.tmp') if source_filename.exists(): target_filename = Path(filename) if target_filename.exists(): - target_filename.unlink() + target_filename.unlink() source_filename.rename ( str(target_filename) ) - + def debug_one_iter(self): images = [] - for generator in self.generator_list: + for generator in self.generator_list: for i,batch in enumerate(next(generator)): if len(batch.shape) == 4: images.append( batch[0] ) - + return image_utils.equalize_and_stack_square (images) - + def generate_next_sample(self): return [next(generator) for generator in self.generator_list] def train_one_iter(self): - sample = self.generate_next_sample() - iter_time = time.time() - losses = self.onTrainOneIter(sample, self.generator_list) + sample = self.generate_next_sample() + iter_time = time.time() + losses = self.onTrainOneIter(sample, self.generator_list) iter_time = time.time() - iter_time self.last_sample = sample - + self.loss_history.append ( [float(loss[1]) for loss in losses] ) if self.write_preview_history: - if self.iter % 10 == 0: + if self.iter % 10 == 0: preview = self.get_static_preview() preview_lh = ModelBase.get_loss_history_preview(self.loss_history, self.iter, preview.shape[1], preview.shape[2]) img = (np.concatenate ( [preview_lh, preview], axis=0 ) * 255).astype(np.uint8) - cv2_imwrite ( str (self.preview_history_path / ('%.6d.jpg' %( self.iter) )), img ) - + cv2_imwrite ( str (self.preview_history_path / ('%.6d.jpg' %( self.iter) )), img ) + self.iter += 1 time_str = time.strftime("[%H:%M:%S]") @@ -383,40 +383,40 @@ class ModelBase(object): loss_string += " %s:%.3f" % (loss_name, loss_value) return loss_string - + def pass_one_iter(self): - self.last_sample = self.generate_next_sample() - + self.last_sample = self.generate_next_sample() + def finalize(self): nnlib.finalize_all() - + def is_first_run(self): return self.iter == 0 - + def is_debug(self): return self.debug - + def set_batch_size(self, batch_size): self.batch_size = batch_size - + def get_batch_size(self): return self.batch_size - + def get_iter(self): return self.iter - + def get_loss_history(self): return self.loss_history - + def set_training_data_generators (self, generator_list): self.generator_list = generator_list - + def get_training_data_generators (self): return self.generator_list - + def get_model_root_path(self): return self.model_path - + def get_strpath_storage_for_file(self, filename): if self.device_args['force_gpu_idx'] == -1: return str( self.model_path / ( self.get_model_name() + '_' + filename) ) @@ -424,65 +424,65 @@ class ModelBase(object): return str( self.model_path / ( str(self.device_args['force_gpu_idx']) + '_' + self.get_model_name() + '_' + filename) ) def set_vram_batch_requirements (self, d): - #example d = {2:2,3:4,4:8,5:16,6:32,7:32,8:32,9:48} + #example d = {2:2,3:4,4:8,5:16,6:32,7:32,8:32,9:48} keys = [x for x in d.keys()] - + if self.device_config.cpu_only: if self.batch_size == 0: self.batch_size = 2 else: - if self.batch_size == 0: + if self.batch_size == 0: for x in keys: if self.device_config.gpu_vram_gb[0] <= x: self.batch_size = d[x] break - + if self.batch_size == 0: self.batch_size = d[ keys[-1] ] - + @staticmethod def get_loss_history_preview(loss_history, iter, w, c): loss_history = np.array (loss_history.copy()) - + lh_height = 100 lh_img = np.ones ( (lh_height,w,c) ) * 0.1 loss_count = len(loss_history[0]) lh_len = len(loss_history) - - l_per_col = lh_len / w + + l_per_col = lh_len / w plist_max = [ [ max (0.0, loss_history[int(col*l_per_col)][p], - *[ loss_history[i_ab][p] - for i_ab in range( int(col*l_per_col), int((col+1)*l_per_col) ) + *[ loss_history[i_ab][p] + for i_ab in range( int(col*l_per_col), int((col+1)*l_per_col) ) ] - ) + ) for p in range(loss_count) - ] + ] for col in range(w) ] plist_min = [ [ min (plist_max[col][p], loss_history[int(col*l_per_col)][p], - *[ loss_history[i_ab][p] - for i_ab in range( int(col*l_per_col), int((col+1)*l_per_col) ) + *[ loss_history[i_ab][p] + for i_ab in range( int(col*l_per_col), int((col+1)*l_per_col) ) ] - ) - for p in range(loss_count) - ] - for col in range(w) + ) + for p in range(loss_count) + ] + for col in range(w) ] plist_abs_max = np.mean(loss_history[ len(loss_history) // 5 : ]) * 2 for col in range(0, w): - for p in range(0,loss_count): + for p in range(0,loss_count): point_color = [1.0]*c point_color[0:3] = colorsys.hsv_to_rgb ( p * (1.0/loss_count), 1.0, 1.0 ) - + ph_max = int ( (plist_max[col][p] / plist_abs_max) * (lh_height-1) ) ph_max = np.clip( ph_max, 0, lh_height-1 ) - + ph_min = int ( (plist_min[col][p] / plist_abs_max) * (lh_height-1) ) ph_min = np.clip( ph_min, 0, lh_height-1 ) - + for ph in range(ph_min, ph_max+1): lh_img[ (lh_height-ph-1), col ] = point_color @@ -490,11 +490,11 @@ class ModelBase(object): lh_line_height = (lh_height-1)/lh_lines for i in range(0,lh_lines+1): lh_img[ int(i*lh_line_height), : ] = (0.8,)*c - + last_line_t = int((lh_lines-1)*lh_line_height) last_line_b = int(lh_lines*lh_line_height) - + lh_text = 'Iter: %d' % (iter) if iter != 0 else '' - + lh_img[last_line_t:last_line_b, 0:w] += image_utils.get_text_image ( (w,last_line_b-last_line_t,c), lh_text, color=[0.8]*c ) - return lh_img \ No newline at end of file + return lh_img diff --git a/models/Model_DF/Model.py b/models/Model_DF/Model.py index 90e91f7..197787b 100644 --- a/models/Model_DF/Model.py +++ b/models/Model_DF/Model.py @@ -9,22 +9,22 @@ from interact import interact as io class Model(ModelBase): #override - def onInitializeOptions(self, is_first_run, ask_override): + 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.") else: self.options['pixel_loss'] = self.options.get('pixel_loss', False) - + #override def onInitialize(self): exec(nnlib.import_all(), locals(), globals()) self.set_vram_batch_requirements( {4.5:4} ) - + ae_input_layer = Input(shape=(128, 128, 3)) mask_layer = Input(shape=(128, 128, 1)) #same as output - - self.encoder, self.decoder_src, self.decoder_dst = self.Build(ae_input_layer) + + self.encoder, self.decoder_src, self.decoder_dst = self.Build(ae_input_layer) if not self.is_first_run(): weights_to_load = [ [self.encoder , 'encoder.h5'], @@ -38,39 +38,39 @@ class Model(ModelBase): self.autoencoder_src.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMSEMaskLoss(mask_layer, is_mse=self.options['pixel_loss']), 'mse'] ) self.autoencoder_dst.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMSEMaskLoss(mask_layer, is_mse=self.options['pixel_loss']), 'mse'] ) - + if self.is_training_mode: f = SampleProcessor.TypeFlags - self.set_training_data_generators ([ - SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None, - debug=self.is_debug(), batch_size=self.batch_size, + self.set_training_data_generators ([ + SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None, + debug=self.is_debug(), batch_size=self.batch_size, sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05])+self.src_scale_mod / 100.0 ), - output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128], - [f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128], + output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128], + [f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128], [f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_M | f.FACE_MASK_FULL, 128] ] ), - - SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, + + SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, sample_process_options=SampleProcessor.Options(random_flip=self.random_flip), - output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128], - [f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128], + output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128], + [f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128], [f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_M | f.FACE_MASK_FULL, 128] ] ) ]) #override - def onSave(self): + def onSave(self): self.save_weights_safe( [[self.encoder, 'encoder.h5'], [self.decoder_src, 'decoder_src.h5'], [self.decoder_dst, 'decoder_dst.h5']] ) - + #override def onTrainOneIter(self, sample, generators_list): warped_src, target_src, target_src_mask = sample[0] - warped_dst, target_dst, target_dst_mask = sample[1] - + warped_dst, target_dst, target_dst_mask = sample[1] + loss_src = self.autoencoder_src.train_on_batch( [warped_src, target_src_mask], [target_src, target_src_mask] ) loss_dst = self.autoencoder_dst.train_on_batch( [warped_dst, target_dst_mask], [target_dst, target_dst_mask] ) - + return ( ('loss_src', loss_src[0]), ('loss_dst', loss_dst[0]) ) - + #override def onGetPreview(self, sample): @@ -78,64 +78,64 @@ class Model(ModelBase): test_A_m = sample[0][2][0:4] #first 4 samples test_B = sample[1][1][0:4] test_B_m = sample[1][2][0:4] - - AA, mAA = self.autoencoder_src.predict([test_A, test_A_m]) + + AA, mAA = self.autoencoder_src.predict([test_A, test_A_m]) AB, mAB = self.autoencoder_src.predict([test_B, test_B_m]) BB, mBB = self.autoencoder_dst.predict([test_B, test_B_m]) - + mAA = np.repeat ( mAA, (3,), -1) mAB = np.repeat ( mAB, (3,), -1) mBB = np.repeat ( mBB, (3,), -1) - + st = [] for i in range(0, len(test_A)): st.append ( np.concatenate ( ( test_A[i,:,:,0:3], AA[i], #mAA[i], - test_B[i,:,:,0:3], - BB[i], - #mBB[i], + test_B[i,:,:,0:3], + BB[i], + #mBB[i], AB[i], #mAB[i] ), axis=1) ) - + return [ ('DF', np.concatenate ( st, axis=0 ) ) ] - + def predictor_func (self, face): - + face_128_bgr = face[...,0:3] face_128_mask = np.expand_dims(face[...,3],-1) - + x, mx = self.autoencoder_src.predict ( [ np.expand_dims(face_128_bgr,0), np.expand_dims(face_128_mask,0) ] ) x, mx = x[0], mx[0] - + return np.concatenate ( (x,mx), -1 ) - + #override def get_converter(self): - from converters import ConverterMasked - return ConverterMasked(self.predictor_func, - predictor_input_size=128, - output_size=128, - face_type=FaceType.FULL, + from converters import ConverterMasked + return ConverterMasked(self.predictor_func, + predictor_input_size=128, + output_size=128, + face_type=FaceType.FULL, base_erode_mask_modifier=30, base_blur_mask_modifier=0) - + def Build(self, input_layer): exec(nnlib.code_import_all, locals(), globals()) - + def downscale (dim): def func(x): return LeakyReLU(0.1)(Conv2D(dim, 5, strides=2, padding='same')(x)) - return func - + return func + def upscale (dim): def func(x): return PixelShuffler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x))) - return func - - def Encoder(input_layer): + return func + + def Encoder(input_layer): x = input_layer x = downscale(128)(x) x = downscale(256)(x) @@ -146,7 +146,7 @@ class Model(ModelBase): x = Dense(8 * 8 * 512)(x) x = Reshape((8, 8, 512))(x) x = upscale(512)(x) - + return Model(input_layer, x) def Decoder(): @@ -155,15 +155,15 @@ class Model(ModelBase): x = upscale(512)(x) x = upscale(256)(x) x = upscale(128)(x) - + y = input_ #mask decoder y = upscale(512)(y) y = upscale(256)(y) y = upscale(128)(y) - + x = Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x) y = Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(y) - + return Model(input_, [x,y]) - - return Encoder(input_layer), Decoder(), Decoder() \ No newline at end of file + + return Encoder(input_layer), Decoder(), Decoder() diff --git a/models/Model_DF/__init__.py b/models/Model_DF/__init__.py index cdb3fe7..0188f11 100644 --- a/models/Model_DF/__init__.py +++ b/models/Model_DF/__init__.py @@ -1 +1 @@ -from .Model import Model \ No newline at end of file +from .Model import Model diff --git a/models/Model_FANSegmentator/Model.py b/models/Model_FANSegmentator/Model.py index a1baff1..7b1a81a 100644 --- a/models/Model_FANSegmentator/Model.py +++ b/models/Model_FANSegmentator/Model.py @@ -10,13 +10,13 @@ from interact import interact as io class Model(ModelBase): def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs, - ask_write_preview_history=False, + super().__init__(*args, **kwargs, + ask_write_preview_history=False, ask_target_iter=False, ask_sort_by_yaw=False, ask_random_flip=False, ask_src_scale_mod=False) - + #override def onInitialize(self): exec(nnlib.import_all(), locals(), globals()) @@ -24,33 +24,33 @@ class Model(ModelBase): self.resolution = 256 self.face_type = FaceType.FULL - - self.fan_seg = FANSegmentator(self.resolution, - FaceType.toString(self.face_type), + + self.fan_seg = FANSegmentator(self.resolution, + FaceType.toString(self.face_type), load_weights=not self.is_first_run(), weights_file_root=self.get_model_root_path() ) if self.is_training_mode: f = SampleProcessor.TypeFlags f_type = f.FACE_ALIGN_FULL - - self.set_training_data_generators ([ - SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, - sample_process_options=SampleProcessor.Options(random_flip=True, normalize_tanh = True ), + + self.set_training_data_generators ([ + SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, + sample_process_options=SampleProcessor.Options(random_flip=True, normalize_tanh = True ), output_sample_types=[ [f.TRANSFORMED | f_type | f.MODE_BGR_SHUFFLE, self.resolution], [f.TRANSFORMED | f_type | f.MODE_M | f.FACE_MASK_FULL, self.resolution] ]), - - SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, - sample_process_options=SampleProcessor.Options(random_flip=True, normalize_tanh = True ), + + SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, + sample_process_options=SampleProcessor.Options(random_flip=True, normalize_tanh = True ), output_sample_types=[ [f.TRANSFORMED | f_type | f.MODE_BGR_SHUFFLE, self.resolution] ]) ]) - + #override - def onSave(self): + def onSave(self): self.fan_seg.save_weights() - + #override def onTrainOneIter(self, generators_samples, generators_list): target_src, target_src_mask = generators_samples[0] @@ -58,20 +58,20 @@ class Model(ModelBase): loss = self.fan_seg.train_on_batch( [target_src], [target_src_mask] ) return ( ('loss', loss), ) - + #override def onGetPreview(self, sample): test_A = sample[0][0][0:4] #first 4 samples test_B = sample[1][0][0:4] #first 4 samples - + mAA = self.fan_seg.extract_from_bgr([test_A]) mBB = self.fan_seg.extract_from_bgr([test_B]) - + test_A, test_B, = [ np.clip( (x + 1.0)/2.0, 0.0, 1.0) for x in [test_A, test_B] ] - + mAA = np.repeat ( mAA, (3,), -1) mBB = np.repeat ( mBB, (3,), -1) - + st = [] for i in range(0, len(test_A)): st.append ( np.concatenate ( ( @@ -79,7 +79,7 @@ class Model(ModelBase): mAA[i], test_A[i,:,:,0:3]*mAA[i], ), axis=1) ) - + st2 = [] for i in range(0, len(test_B)): st2.append ( np.concatenate ( ( @@ -87,7 +87,7 @@ class Model(ModelBase): mBB[i], test_B[i,:,:,0:3]*mBB[i], ), axis=1) ) - + return [ ('FANSegmentator', np.concatenate ( st, axis=0 ) ), ('never seen', np.concatenate ( st2, axis=0 ) ), ] diff --git a/models/Model_FANSegmentator/__init__.py b/models/Model_FANSegmentator/__init__.py index cdb3fe7..0188f11 100644 --- a/models/Model_FANSegmentator/__init__.py +++ b/models/Model_FANSegmentator/__init__.py @@ -1 +1 @@ -from .Model import Model \ No newline at end of file +from .Model import Model diff --git a/models/Model_H128/Model.py b/models/Model_H128/Model.py index 404bdab..f3d5629 100644 --- a/models/Model_H128/Model.py +++ b/models/Model_H128/Model.py @@ -9,7 +9,7 @@ from interact import interact as io class Model(ModelBase): #override - def onInitializeOptions(self, is_first_run, ask_override): + def onInitializeOptions(self, is_first_run, ask_override): if is_first_run: self.options['lighter_ae'] = io.input_bool ("Use lightweight autoencoder? (y/n, ?:help skip:n) : ", False, help_message="Lightweight autoencoder is faster, requires less VRAM, sacrificing overall quality. If your GPU VRAM <= 4, you should to choose this option.") else: @@ -17,18 +17,18 @@ class Model(ModelBase): if 'created_vram_gb' in self.options.keys(): self.options.pop ('created_vram_gb') self.options['lighter_ae'] = self.options.get('lighter_ae', default_lighter_ae) - + 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.") else: self.options['pixel_loss'] = self.options.get('pixel_loss', False) - + #override def onInitialize(self): - exec(nnlib.import_all(), locals(), globals()) + exec(nnlib.import_all(), locals(), globals()) self.set_vram_batch_requirements( {2.5:4} ) - + bgr_shape, mask_shape, self.encoder, self.decoder_src, self.decoder_dst = self.Build( self.options['lighter_ae'] ) if not self.is_first_run(): weights_to_load = [ [self.encoder , 'encoder.h5'], @@ -36,120 +36,120 @@ class Model(ModelBase): [self.decoder_dst, 'decoder_dst.h5'] ] self.load_weights_safe(weights_to_load) - + input_src_bgr = Input(bgr_shape) input_src_mask = Input(mask_shape) input_dst_bgr = Input(bgr_shape) input_dst_mask = Input(mask_shape) - rec_src_bgr, rec_src_mask = self.decoder_src( self.encoder(input_src_bgr) ) + rec_src_bgr, rec_src_mask = self.decoder_src( self.encoder(input_src_bgr) ) rec_dst_bgr, rec_dst_mask = self.decoder_dst( self.encoder(input_dst_bgr) ) self.ae = Model([input_src_bgr,input_src_mask,input_dst_bgr,input_dst_mask], [rec_src_bgr, rec_src_mask, rec_dst_bgr, rec_dst_mask] ) - + self.ae.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[ DSSIMMSEMaskLoss(input_src_mask, is_mse=self.options['pixel_loss']), 'mae', DSSIMMSEMaskLoss(input_dst_mask, is_mse=self.options['pixel_loss']), 'mae' ] ) - + self.src_view = K.function([input_src_bgr],[rec_src_bgr, rec_src_mask]) self.dst_view = K.function([input_dst_bgr],[rec_dst_bgr, rec_dst_mask]) - + if self.is_training_mode: f = SampleProcessor.TypeFlags - self.set_training_data_generators ([ - SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None, - debug=self.is_debug(), batch_size=self.batch_size, - sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05])+self.src_scale_mod / 100.0 ), - output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 128], - [f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 128], + self.set_training_data_generators ([ + SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None, + debug=self.is_debug(), batch_size=self.batch_size, + sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05])+self.src_scale_mod / 100.0 ), + output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 128], + [f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 128], [f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_M | f.FACE_MASK_FULL, 128] ] ), - - SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, - sample_process_options=SampleProcessor.Options(random_flip=self.random_flip), - output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 128], - [f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 128], + + SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, + sample_process_options=SampleProcessor.Options(random_flip=self.random_flip), + output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 128], + [f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 128], [f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_M | f.FACE_MASK_FULL, 128] ] ) ]) - + #override - def onSave(self): + def onSave(self): self.save_weights_safe( [[self.encoder, 'encoder.h5'], [self.decoder_src, 'decoder_src.h5'], [self.decoder_dst, 'decoder_dst.h5']] ) - + #override def onTrainOneIter(self, sample, generators_list): warped_src, target_src, target_src_mask = sample[0] - warped_dst, target_dst, target_dst_mask = sample[1] + warped_dst, target_dst, target_dst_mask = sample[1] total, loss_src_bgr, loss_src_mask, loss_dst_bgr, loss_dst_mask = self.ae.train_on_batch( [warped_src, target_src_mask, warped_dst, target_dst_mask], [target_src, target_src_mask, target_dst, target_dst_mask] ) return ( ('loss_src', loss_src_bgr), ('loss_dst', loss_dst_bgr) ) - + #override def onGetPreview(self, sample): test_A = sample[0][1][0:4] #first 4 samples test_A_m = sample[0][2][0:4] #first 4 samples test_B = sample[1][1][0:4] test_B_m = sample[1][2][0:4] - - AA, mAA = self.src_view([test_A]) + + AA, mAA = self.src_view([test_A]) AB, mAB = self.src_view([test_B]) BB, mBB = self.dst_view([test_B]) - + mAA = np.repeat ( mAA, (3,), -1) mAB = np.repeat ( mAB, (3,), -1) mBB = np.repeat ( mBB, (3,), -1) - + st = [] for i in range(0, len(test_A)): st.append ( np.concatenate ( ( test_A[i,:,:,0:3], AA[i], #mAA[i], - test_B[i,:,:,0:3], - BB[i], - #mBB[i], + test_B[i,:,:,0:3], + BB[i], + #mBB[i], AB[i], #mAB[i] ), axis=1) ) - + return [ ('H128', np.concatenate ( st, axis=0 ) ) ] - def predictor_func (self, face): + def predictor_func (self, face): face_128_bgr = face[...,0:3] face_128_mask = np.expand_dims(face[...,3],-1) - + x, mx = self.src_view ( [ np.expand_dims(face_128_bgr,0) ] ) - x, mx = x[0], mx[0] - + x, mx = x[0], mx[0] + return np.concatenate ( (x,mx), -1 ) #override def get_converter(self): from converters import ConverterMasked - return ConverterMasked(self.predictor_func, - predictor_input_size=128, - output_size=128, + return ConverterMasked(self.predictor_func, + predictor_input_size=128, + output_size=128, face_type=FaceType.HALF, base_erode_mask_modifier=100, base_blur_mask_modifier=100) - + def Build(self, lighter_ae): exec(nnlib.code_import_all, locals(), globals()) - + bgr_shape = (128, 128, 3) mask_shape = (128, 128, 1) - + def downscale (dim): def func(x): return LeakyReLU(0.1)(Conv2D(dim, 5, strides=2, padding='same')(x)) - return func - + return func + def upscale (dim): def func(x): return PixelShuffler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x))) - return func - + return func + def Encoder(input_shape): input_layer = Input(input_shape) x = input_layer @@ -171,7 +171,7 @@ class Model(ModelBase): x = Dense(8 * 8 * 256)(x) x = Reshape((8, 8, 256))(x) x = upscale(256)(x) - + return Model(input_layer, x) def Decoder(): @@ -181,7 +181,7 @@ class Model(ModelBase): x = upscale(512)(x) x = upscale(256)(x) x = upscale(128)(x) - + y = input_ #mask decoder y = upscale(512)(y) y = upscale(256)(y) @@ -192,16 +192,16 @@ class Model(ModelBase): x = upscale(256)(x) x = upscale(128)(x) x = upscale(64)(x) - + y = input_ #mask decoder y = upscale(256)(y) y = upscale(128)(y) y = upscale(64)(y) - + x = Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x) y = Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(y) - - + + return Model(input_, [x,y]) - + return bgr_shape, mask_shape, Encoder(bgr_shape), Decoder(), Decoder() diff --git a/models/Model_H128/__init__.py b/models/Model_H128/__init__.py index cdb3fe7..0188f11 100644 --- a/models/Model_H128/__init__.py +++ b/models/Model_H128/__init__.py @@ -1 +1 @@ -from .Model import Model \ No newline at end of file +from .Model import Model diff --git a/models/Model_H64/Model.py b/models/Model_H64/Model.py index 5bb6075..3822cbc 100644 --- a/models/Model_H64/Model.py +++ b/models/Model_H64/Model.py @@ -9,7 +9,7 @@ from interact import interact as io class Model(ModelBase): #override - def onInitializeOptions(self, is_first_run, ask_override): + def onInitializeOptions(self, is_first_run, ask_override): if is_first_run: self.options['lighter_ae'] = io.input_bool ("Use lightweight autoencoder? (y/n, ?:help skip:n) : ", False, help_message="Lightweight autoencoder is faster, requires less VRAM, sacrificing overall quality. If your GPU VRAM <= 4, you should to choose this option.") else: @@ -17,141 +17,141 @@ class Model(ModelBase): if 'created_vram_gb' in self.options.keys(): self.options.pop ('created_vram_gb') self.options['lighter_ae'] = self.options.get('lighter_ae', default_lighter_ae) - + 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.") else: self.options['pixel_loss'] = self.options.get('pixel_loss', False) - + #override def onInitialize(self): exec(nnlib.import_all(), locals(), globals()) self.set_vram_batch_requirements( {1.5:4} ) - - + + bgr_shape, mask_shape, self.encoder, self.decoder_src, self.decoder_dst = self.Build(self.options['lighter_ae']) - + if not self.is_first_run(): weights_to_load = [ [self.encoder , 'encoder.h5'], [self.decoder_src, 'decoder_src.h5'], [self.decoder_dst, 'decoder_dst.h5'] ] self.load_weights_safe(weights_to_load) - + input_src_bgr = Input(bgr_shape) input_src_mask = Input(mask_shape) input_dst_bgr = Input(bgr_shape) input_dst_mask = Input(mask_shape) - - rec_src_bgr, rec_src_mask = self.decoder_src( self.encoder(input_src_bgr) ) + + rec_src_bgr, rec_src_mask = self.decoder_src( self.encoder(input_src_bgr) ) rec_dst_bgr, rec_dst_mask = self.decoder_dst( self.encoder(input_dst_bgr) ) self.ae = Model([input_src_bgr,input_src_mask,input_dst_bgr,input_dst_mask], [rec_src_bgr, rec_src_mask, rec_dst_bgr, rec_dst_mask] ) - + self.ae.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[ DSSIMMSEMaskLoss(input_src_mask, is_mse=self.options['pixel_loss']), 'mae', DSSIMMSEMaskLoss(input_dst_mask, is_mse=self.options['pixel_loss']), 'mae' ] ) - + self.src_view = K.function([input_src_bgr],[rec_src_bgr, rec_src_mask]) self.dst_view = K.function([input_dst_bgr],[rec_dst_bgr, rec_dst_mask]) - + if self.is_training_mode: f = SampleProcessor.TypeFlags - self.set_training_data_generators ([ - SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None, - debug=self.is_debug(), batch_size=self.batch_size, - sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05])+self.src_scale_mod / 100.0 ), - output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64], - [f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64], + self.set_training_data_generators ([ + SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None, + debug=self.is_debug(), batch_size=self.batch_size, + sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05])+self.src_scale_mod / 100.0 ), + output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64], + [f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64], [f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_M | f.FACE_MASK_FULL, 64] ] ), - + SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, - sample_process_options=SampleProcessor.Options(random_flip=self.random_flip), - output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64], - [f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64], + sample_process_options=SampleProcessor.Options(random_flip=self.random_flip), + output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64], + [f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64], [f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_M | f.FACE_MASK_FULL, 64] ] ) ]) - + #override - def onSave(self): + def onSave(self): self.save_weights_safe( [[self.encoder, 'encoder.h5'], [self.decoder_src, 'decoder_src.h5'], [self.decoder_dst, 'decoder_dst.h5']] ) - + #override def onTrainOneIter(self, sample, generators_list): warped_src, target_src, target_src_full_mask = sample[0] - warped_dst, target_dst, target_dst_full_mask = sample[1] + warped_dst, target_dst, target_dst_full_mask = sample[1] total, loss_src_bgr, loss_src_mask, loss_dst_bgr, loss_dst_mask = self.ae.train_on_batch( [warped_src, target_src_full_mask, warped_dst, target_dst_full_mask], [target_src, target_src_full_mask, target_dst, target_dst_full_mask] ) return ( ('loss_src', loss_src_bgr), ('loss_dst', loss_dst_bgr) ) - + #override def onGetPreview(self, sample): test_A = sample[0][1][0:4] #first 4 samples test_A_m = sample[0][2][0:4] test_B = sample[1][1][0:4] test_B_m = sample[1][2][0:4] - - AA, mAA = self.src_view([test_A]) + + AA, mAA = self.src_view([test_A]) AB, mAB = self.src_view([test_B]) BB, mBB = self.dst_view([test_B]) - + mAA = np.repeat ( mAA, (3,), -1) mAB = np.repeat ( mAB, (3,), -1) mBB = np.repeat ( mBB, (3,), -1) - + st = [] for i in range(0, len(test_A)): st.append ( np.concatenate ( ( test_A[i,:,:,0:3], AA[i], #mAA[i], - test_B[i,:,:,0:3], - BB[i], - #mBB[i], + test_B[i,:,:,0:3], + BB[i], + #mBB[i], AB[i], #mAB[i] ), axis=1) ) - + return [ ('H64', np.concatenate ( st, axis=0 ) ) ] def predictor_func (self, face): - + face_64_bgr = face[...,0:3] face_64_mask = np.expand_dims(face[...,3],-1) - + x, mx = self.src_view ( [ np.expand_dims(face_64_bgr,0) ] ) - x, mx = x[0], mx[0] - + x, mx = x[0], mx[0] + return np.concatenate ( (x,mx), -1 ) #override def get_converter(self): from converters import ConverterMasked return ConverterMasked(self.predictor_func, - predictor_input_size=64, - output_size=64, - face_type=FaceType.HALF, + predictor_input_size=64, + output_size=64, + face_type=FaceType.HALF, base_erode_mask_modifier=100, base_blur_mask_modifier=100) - + def Build(self, lighter_ae): exec(nnlib.code_import_all, locals(), globals()) - + bgr_shape = (64, 64, 3) mask_shape = (64, 64, 1) - + def downscale (dim): def func(x): return LeakyReLU(0.1)(Conv2D(dim, 5, strides=2, padding='same')(x)) - return func - + return func + def upscale (dim): def func(x): return PixelShuffler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x))) - return func - + return func + def Encoder(input_shape): input_layer = Input(input_shape) x = input_layer @@ -183,23 +183,23 @@ class Model(ModelBase): x = upscale(512)(x) x = upscale(256)(x) x = upscale(128)(x) - + else: input_ = Input(shape=(8, 8, 256)) - - x = input_ + + x = input_ x = upscale(256)(x) x = upscale(128)(x) x = upscale(64)(x) - + y = input_ #mask decoder y = upscale(256)(y) y = upscale(128)(y) y = upscale(64)(y) - + x = Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x) y = Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(y) - + return Model(input_, [x,y]) - - return bgr_shape, mask_shape, Encoder(bgr_shape), Decoder(), Decoder() \ No newline at end of file + + return bgr_shape, mask_shape, Encoder(bgr_shape), Decoder(), Decoder() diff --git a/models/Model_H64/__init__.py b/models/Model_H64/__init__.py index cdb3fe7..0188f11 100644 --- a/models/Model_H64/__init__.py +++ b/models/Model_H64/__init__.py @@ -1 +1 @@ -from .Model import Model \ No newline at end of file +from .Model import Model diff --git a/models/Model_LIAEF128/Model.py b/models/Model_LIAEF128/Model.py index 91472ef..6886310 100644 --- a/models/Model_LIAEF128/Model.py +++ b/models/Model_LIAEF128/Model.py @@ -9,13 +9,13 @@ from interact import interact as io class Model(ModelBase): #override - def onInitializeOptions(self, is_first_run, ask_override): + 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.") else: self.options['pixel_loss'] = self.options.get('pixel_loss', False) - + #override def onInitialize(self): exec(nnlib.import_all(), locals(), globals()) @@ -25,7 +25,7 @@ class Model(ModelBase): mask_layer = Input(shape=(128, 128, 1)) #same as output self.encoder, self.decoder, self.inter_B, self.inter_AB = self.Build(ae_input_layer) - + if not self.is_first_run(): weights_to_load = [ [self.encoder, 'encoder.h5'], [self.decoder, 'decoder.h5'], @@ -39,46 +39,46 @@ class Model(ModelBase): B = self.inter_B(code) self.autoencoder_src = Model([ae_input_layer,mask_layer], self.decoder(Concatenate()([AB, AB])) ) self.autoencoder_dst = Model([ae_input_layer,mask_layer], self.decoder(Concatenate()([B, AB])) ) - + self.autoencoder_src.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMSEMaskLoss(mask_layer, is_mse=self.options['pixel_loss']), 'mse'] ) self.autoencoder_dst.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMSEMaskLoss(mask_layer, is_mse=self.options['pixel_loss']), 'mse'] ) - + if self.is_training_mode: f = SampleProcessor.TypeFlags - self.set_training_data_generators ([ - - - SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None, + self.set_training_data_generators ([ + + + SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None, debug=self.is_debug(), batch_size=self.batch_size, - sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05])+self.src_scale_mod / 100.0 ), - output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128], - [f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128], + sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05])+self.src_scale_mod / 100.0 ), + output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128], + [f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128], [f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_M | f.FACE_MASK_FULL, 128] ] ), - - SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, - sample_process_options=SampleProcessor.Options(random_flip=self.random_flip), - output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128], - [f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128], + + SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, + sample_process_options=SampleProcessor.Options(random_flip=self.random_flip), + output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128], + [f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128], [f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_M | f.FACE_MASK_FULL, 128] ] ) ]) - + #override def onSave(self): self.save_weights_safe( [[self.encoder, 'encoder.h5'], [self.decoder, 'decoder.h5'], [self.inter_B, 'inter_B.h5'], [self.inter_AB, 'inter_AB.h5']] ) - + #override def onTrainOneIter(self, sample, generators_list): warped_src, target_src, target_src_mask = sample[0] - warped_dst, target_dst, target_dst_mask = sample[1] + warped_dst, target_dst, target_dst_mask = sample[1] loss_src = self.autoencoder_src.train_on_batch( [warped_src, target_src_mask], [target_src, target_src_mask] ) loss_dst = self.autoencoder_dst.train_on_batch( [warped_dst, target_dst_mask], [target_dst, target_dst_mask] ) - + return ( ('loss_src', loss_src[0]), ('loss_dst', loss_dst[0]) ) - + #override def onGetPreview(self, sample): @@ -86,63 +86,63 @@ class Model(ModelBase): test_A_m = sample[0][2][0:4] #first 4 samples test_B = sample[1][1][0:4] test_B_m = sample[1][2][0:4] - - AA, mAA = self.autoencoder_src.predict([test_A, test_A_m]) + + AA, mAA = self.autoencoder_src.predict([test_A, test_A_m]) AB, mAB = self.autoencoder_src.predict([test_B, test_B_m]) BB, mBB = self.autoencoder_dst.predict([test_B, test_B_m]) - + mAA = np.repeat ( mAA, (3,), -1) mAB = np.repeat ( mAB, (3,), -1) mBB = np.repeat ( mBB, (3,), -1) - + st = [] for i in range(0, len(test_A)): st.append ( np.concatenate ( ( test_A[i,:,:,0:3], AA[i], #mAA[i], - test_B[i,:,:,0:3], - BB[i], - #mBB[i], + test_B[i,:,:,0:3], + BB[i], + #mBB[i], AB[i], #mAB[i] ), axis=1) ) - + return [ ('LIAEF128', np.concatenate ( st, axis=0 ) ) ] - + def predictor_func (self, face): - + face_128_bgr = face[...,0:3] face_128_mask = np.expand_dims(face[...,3],-1) - + x, mx = self.autoencoder_src.predict ( [ np.expand_dims(face_128_bgr,0), np.expand_dims(face_128_mask,0) ] ) x, mx = x[0], mx[0] - + return np.concatenate ( (x,mx), -1 ) - + #override def get_converter(self): from converters import ConverterMasked - return ConverterMasked(self.predictor_func, - predictor_input_size=128, - output_size=128, - face_type=FaceType.FULL, + return ConverterMasked(self.predictor_func, + predictor_input_size=128, + output_size=128, + face_type=FaceType.FULL, base_erode_mask_modifier=30, base_blur_mask_modifier=0) - + def Build(self, input_layer): exec(nnlib.code_import_all, locals(), globals()) - + def downscale (dim): def func(x): return LeakyReLU(0.1)(Conv2D(dim, 5, strides=2, padding='same')(x)) - return func - + return func + def upscale (dim): def func(x): return PixelShuffler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x))) - return func - + return func + def Encoder(): x = input_layer x = downscale(128)(x) @@ -161,20 +161,20 @@ class Model(ModelBase): x = upscale(512)(x) return Model(input_layer, x) - def Decoder(): + def Decoder(): input_ = Input(shape=(16, 16, 1024)) x = input_ x = upscale(512)(x) x = upscale(256)(x) x = upscale(128)(x) x = Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x) - + y = input_ #mask decoder y = upscale(512)(y) y = upscale(256)(y) y = upscale(128)(y) y = Conv2D(1, kernel_size=5, padding='same', activation='sigmoid' )(y) - + return Model(input_, [x,y]) - return Encoder(), Decoder(), Intermediate(), Intermediate() \ No newline at end of file + return Encoder(), Decoder(), Intermediate(), Intermediate() diff --git a/models/Model_LIAEF128/__init__.py b/models/Model_LIAEF128/__init__.py index cdb3fe7..0188f11 100644 --- a/models/Model_LIAEF128/__init__.py +++ b/models/Model_LIAEF128/__init__.py @@ -1 +1 @@ -from .Model import Model \ No newline at end of file +from .Model import Model diff --git a/models/Model_SAE/Model.py b/models/Model_SAE/Model.py index 7807e4b..0429cda 100644 --- a/models/Model_SAE/Model.py +++ b/models/Model_SAE/Model.py @@ -9,51 +9,51 @@ from interact import interact as io #SAE - Styled AutoEncoder class SAEModel(ModelBase): - encoderH5 = 'encoder.h5' + encoderH5 = 'encoder.h5' inter_BH5 = 'inter_B.h5' inter_ABH5 = 'inter_AB.h5' decoderH5 = 'decoder.h5' decodermH5 = 'decoderm.h5' - + decoder_srcH5 = 'decoder_src.h5' decoder_srcmH5 = 'decoder_srcm.h5' decoder_dstH5 = 'decoder_dst.h5' decoder_dstmH5 = 'decoder_dstm.h5' - + #override def onInitializeOptions(self, is_first_run, ask_override): yn_str = {True:'y',False:'n'} - + default_resolution = 128 default_archi = 'df' default_face_type = 'f' - + if is_first_run: resolution = io.input_int("Resolution ( 64-256 ?:help skip:128) : ", default_resolution, help_message="More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 16.") - resolution = np.clip (resolution, 64, 256) + resolution = np.clip (resolution, 64, 256) while np.modf(resolution / 16)[0] != 0.0: resolution -= 1 self.options['resolution'] = resolution - - self.options['face_type'] = io.input_str ("Half or Full face? (h/f, ?:help skip:f) : ", default_face_type, ['h','f'], help_message="Half face has better resolution, but covers less area of cheeks.").lower() + + self.options['face_type'] = io.input_str ("Half or Full face? (h/f, ?:help skip:f) : ", default_face_type, ['h','f'], help_message="Half face has better resolution, but covers less area of cheeks.").lower() self.options['learn_mask'] = io.input_bool ("Learn mask? (y/n, ?:help skip:y) : ", True, help_message="Learning mask can help model to recognize face directions. Learn without mask can reduce model size, in this case converter forced to use 'not predicted mask' that is not smooth as predicted. Model with style values can be learned without mask and produce same quality result.") else: self.options['resolution'] = self.options.get('resolution', default_resolution) self.options['face_type'] = self.options.get('face_type', default_face_type) self.options['learn_mask'] = self.options.get('learn_mask', True) - - + + if is_first_run and 'tensorflow' in self.device_config.backend: def_optimizer_mode = self.options.get('optimizer_mode', 1) self.options['optimizer_mode'] = io.input_int ("Optimizer mode? ( 1,2,3 ?:help skip:%d) : " % (def_optimizer_mode), def_optimizer_mode, help_message="1 - no changes. 2 - allows you to train x2 bigger network consuming RAM. 3 - allows you to train x3 bigger network consuming huge amount of RAM and slower, depends on CPU power.") else: self.options['optimizer_mode'] = self.options.get('optimizer_mode', 1) - + if is_first_run: self.options['archi'] = io.input_str ("AE architecture (df, liae, vg ?:help skip:%s) : " % (default_archi) , default_archi, ['df','liae','vg'], help_message="'df' keeps faces more natural. 'liae' can fix overly different face shapes. 'vg' - currently testing.").lower() else: self.options['archi'] = self.options.get('archi', default_archi) - + default_ae_dims = 256 if self.options['archi'] == 'liae' else 512 default_ed_ch_dims = 42 def_ca_weights = False @@ -65,31 +65,31 @@ class SAEModel(ModelBase): self.options['ae_dims'] = self.options.get('ae_dims', default_ae_dims) self.options['ed_ch_dims'] = self.options.get('ed_ch_dims', default_ed_ch_dims) self.options['ca_weights'] = self.options.get('ca_weights', def_ca_weights) - + if is_first_run: self.options['lighter_encoder'] = io.input_bool ("Use lightweight encoder? (y/n, ?:help skip:n) : ", False, help_message="Lightweight encoder is 35% faster, requires less VRAM, but sacrificing overall quality.") - + if self.options['archi'] != 'vg': self.options['multiscale_decoder'] = io.input_bool ("Use multiscale decoder? (y/n, ?:help skip:n) : ", False, help_message="Multiscale decoder helps to get better details.") else: self.options['lighter_encoder'] = self.options.get('lighter_encoder', False) - + if self.options['archi'] != 'vg': self.options['multiscale_decoder'] = self.options.get('multiscale_decoder', False) - - default_face_style_power = 0.0 - default_bg_style_power = 0.0 + + default_face_style_power = 0.0 + 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 15-25k iters to enhance fine details and decrease face jitter.") - + 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, - help_message="Learn to transfer face style details such as light and color conditions. Warning: Enable it only after 10k iters, when predicted face is clear enough to start learn style. Start from 0.1 value and check history changes."), 0.0, 100.0 ) - + 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, + help_message="Learn to transfer face style details such as light and color conditions. Warning: Enable it only after 10k iters, when predicted face is clear enough to start learn style. Start from 0.1 value and check history changes."), 0.0, 100.0 ) + default_bg_style_power = default_bg_style_power if is_first_run else self.options.get('bg_style_power', default_bg_style_power) - self.options['bg_style_power'] = np.clip ( io.input_number("Background style power ( 0.0 .. 100.0 ?:help skip:%.2f) : " % (default_bg_style_power), default_bg_style_power, - help_message="Learn to transfer image around face. This can make face more like dst."), 0.0, 100.0 ) + self.options['bg_style_power'] = np.clip ( io.input_number("Background style power ( 0.0 .. 100.0 ?:help skip:%.2f) : " % (default_bg_style_power), default_bg_style_power, + help_message="Learn to transfer image around face. This can make face more like dst."), 0.0, 100.0 ) else: self.options['pixel_loss'] = self.options.get('pixel_loss', False) self.options['face_style_power'] = self.options.get('face_style_power', default_face_style_power) @@ -100,7 +100,7 @@ class SAEModel(ModelBase): exec(nnlib.import_all(), locals(), globals()) SAEModel.initialize_nn_functions() self.set_vram_batch_requirements({1.5:4}) - + resolution = self.options['resolution'] ae_dims = self.options['ae_dims'] ed_ch_dims = self.options['ed_ch_dims'] @@ -108,13 +108,13 @@ class SAEModel(ModelBase): mask_shape = (resolution, resolution, 1) self.ms_count = ms_count = 3 if (self.options['archi'] != 'vg' and self.options['multiscale_decoder']) else 1 - + masked_training = True - + warped_src = Input(bgr_shape) target_src = Input(bgr_shape) target_srcm = Input(mask_shape) - + warped_dst = Input(bgr_shape) target_dst = Input(bgr_shape) target_dstm = Input(mask_shape) @@ -124,27 +124,28 @@ class SAEModel(ModelBase): target_dst_ar = [ Input ( ( bgr_shape[0] // (2**i) ,)*2 + (bgr_shape[-1],) ) for i in range(ms_count-1, -1, -1)] target_dstm_ar = [ Input ( ( mask_shape[0] // (2**i) ,)*2 + (mask_shape[-1],) ) for i in range(ms_count-1, -1, -1)] - + use_bn = True + models_list = [] weights_to_load = [] if self.options['archi'] == 'liae': - self.encoder = modelify(SAEModel.LIAEEncFlow(resolution, self.options['lighter_encoder'], ed_ch_dims=ed_ch_dims) ) (Input(bgr_shape)) - - enc_output_Inputs = [ Input(K.int_shape(x)[1:]) for x in self.encoder.outputs ] - - self.inter_B = modelify(SAEModel.LIAEInterFlow(resolution, ae_dims=ae_dims)) (enc_output_Inputs) - self.inter_AB = modelify(SAEModel.LIAEInterFlow(resolution, ae_dims=ae_dims)) (enc_output_Inputs) - - inter_output_Inputs = [ Input( np.array(K.int_shape(x)[1:])*(1,1,2) ) for x in self.inter_B.outputs ] + self.encoder = modelify(SAEModel.LIAEEncFlow(resolution, self.options['lighter_encoder'], ed_ch_dims=ed_ch_dims, use_bn=use_bn) ) (Input(bgr_shape)) + + enc_output_Inputs = [ Input(K.int_shape(x)[1:]) for x in self.encoder.outputs ] + + self.inter_B = modelify(SAEModel.LIAEInterFlow(resolution, ae_dims=ae_dims, use_bn=use_bn)) (enc_output_Inputs) + self.inter_AB = modelify(SAEModel.LIAEInterFlow(resolution, ae_dims=ae_dims, use_bn=use_bn)) (enc_output_Inputs) + + inter_output_Inputs = [ Input( np.array(K.int_shape(x)[1:])*(1,1,2) ) for x in self.inter_B.outputs ] + + self.decoder = modelify(SAEModel.LIAEDecFlow (bgr_shape[2],ed_ch_dims=ed_ch_dims//2, multiscale_count=self.ms_count, use_bn=use_bn )) (inter_output_Inputs) - self.decoder = modelify(SAEModel.LIAEDecFlow (bgr_shape[2],ed_ch_dims=ed_ch_dims//2, multiscale_count=self.ms_count )) (inter_output_Inputs) - models_list += [self.encoder, self.inter_B, self.inter_AB, self.decoder] - + if self.options['learn_mask']: - self.decoderm = modelify(SAEModel.LIAEDecFlow (mask_shape[2],ed_ch_dims=int(ed_ch_dims/1.5) )) (inter_output_Inputs) + self.decoderm = modelify(SAEModel.LIAEDecFlow (mask_shape[2],ed_ch_dims=int(ed_ch_dims/1.5), use_bn=use_bn )) (inter_output_Inputs) models_list += [self.decoderm] - + if not self.is_first_run(): weights_to_load += [ [self.encoder , 'encoder.h5'], [self.inter_B , 'inter_B.h5'], @@ -153,22 +154,22 @@ class SAEModel(ModelBase): ] if self.options['learn_mask']: weights_to_load += [ [self.decoderm, 'decoderm.h5'] ] - - warped_src_code = self.encoder (warped_src) + + warped_src_code = self.encoder (warped_src) warped_src_inter_AB_code = self.inter_AB (warped_src_code) - warped_src_inter_code = Concatenate()([warped_src_inter_AB_code,warped_src_inter_AB_code]) - + warped_src_inter_code = Concatenate()([warped_src_inter_AB_code,warped_src_inter_AB_code]) + warped_dst_code = self.encoder (warped_dst) warped_dst_inter_B_code = self.inter_B (warped_dst_code) warped_dst_inter_AB_code = self.inter_AB (warped_dst_code) warped_dst_inter_code = Concatenate()([warped_dst_inter_B_code,warped_dst_inter_AB_code]) warped_src_dst_inter_code = Concatenate()([warped_dst_inter_AB_code,warped_dst_inter_AB_code]) - + pred_src_src = self.decoder(warped_src_inter_code) - pred_dst_dst = self.decoder(warped_dst_inter_code) + pred_dst_dst = self.decoder(warped_dst_inter_code) pred_src_dst = self.decoder(warped_src_dst_inter_code) - + if self.options['learn_mask']: pred_src_srcm = self.decoderm(warped_src_inter_code) pred_dst_dstm = self.decoderm(warped_dst_inter_code) @@ -177,18 +178,18 @@ class SAEModel(ModelBase): elif self.options['archi'] == 'df': self.encoder = modelify(SAEModel.DFEncFlow(resolution, self.options['lighter_encoder'], ae_dims=ae_dims, ed_ch_dims=ed_ch_dims) ) (Input(bgr_shape)) - dec_Inputs = [ Input(K.int_shape(x)[1:]) for x in self.encoder.outputs ] - + dec_Inputs = [ Input(K.int_shape(x)[1:]) for x in self.encoder.outputs ] + self.decoder_src = modelify(SAEModel.DFDecFlow (bgr_shape[2],ed_ch_dims=ed_ch_dims//2, multiscale_count=self.ms_count )) (dec_Inputs) self.decoder_dst = modelify(SAEModel.DFDecFlow (bgr_shape[2],ed_ch_dims=ed_ch_dims//2, multiscale_count=self.ms_count )) (dec_Inputs) - + models_list += [self.encoder, self.decoder_src, self.decoder_dst] - + if self.options['learn_mask']: self.decoder_srcm = modelify(SAEModel.DFDecFlow (mask_shape[2],ed_ch_dims=int(ed_ch_dims/1.5) )) (dec_Inputs) self.decoder_dstm = modelify(SAEModel.DFDecFlow (mask_shape[2],ed_ch_dims=int(ed_ch_dims/1.5) )) (dec_Inputs) models_list += [self.decoder_srcm, self.decoder_dstm] - + if not self.is_first_run(): weights_to_load += [ [self.encoder , 'encoder.h5'], [self.decoder_src, 'decoder_src.h5'], @@ -198,37 +199,37 @@ class SAEModel(ModelBase): weights_to_load += [ [self.decoder_srcm, 'decoder_srcm.h5'], [self.decoder_dstm, 'decoder_dstm.h5'], ] - - - - - + + + + + warped_src_code = self.encoder (warped_src) warped_dst_code = self.encoder (warped_dst) pred_src_src = self.decoder_src(warped_src_code) pred_dst_dst = self.decoder_dst(warped_dst_code) pred_src_dst = self.decoder_src(warped_dst_code) - + if self.options['learn_mask']: pred_src_srcm = self.decoder_srcm(warped_src_code) pred_dst_dstm = self.decoder_dstm(warped_dst_code) pred_src_dstm = self.decoder_srcm(warped_dst_code) - + elif self.options['archi'] == 'vg': self.encoder = modelify(SAEModel.VGEncFlow(resolution, self.options['lighter_encoder'], ae_dims=ae_dims, ed_ch_dims=ed_ch_dims) ) (Input(bgr_shape)) - dec_Inputs = [ Input(K.int_shape(x)[1:]) for x in self.encoder.outputs ] - + dec_Inputs = [ Input(K.int_shape(x)[1:]) for x in self.encoder.outputs ] + self.decoder_src = modelify(SAEModel.VGDecFlow (bgr_shape[2],ed_ch_dims=ed_ch_dims//2 )) (dec_Inputs) self.decoder_dst = modelify(SAEModel.VGDecFlow (bgr_shape[2],ed_ch_dims=ed_ch_dims//2 )) (dec_Inputs) - + models_list += [self.encoder, self.decoder_src, self.decoder_dst] - + if self.options['learn_mask']: self.decoder_srcm = modelify(SAEModel.VGDecFlow (mask_shape[2],ed_ch_dims=int(ed_ch_dims/1.5) )) (dec_Inputs) self.decoder_dstm = modelify(SAEModel.VGDecFlow (mask_shape[2],ed_ch_dims=int(ed_ch_dims/1.5) )) (dec_Inputs) models_list += [self.decoder_srcm, self.decoder_dstm] - + if not self.is_first_run(): weights_to_load += [ [self.encoder , 'encoder.h5'], [self.decoder_src, 'decoder_src.h5'], @@ -238,7 +239,7 @@ class SAEModel(ModelBase): weights_to_load += [ [self.decoder_srcm, 'decoder_srcm.h5'], [self.decoder_dstm, 'decoder_dstm.h5'], ] - + warped_src_code = self.encoder (warped_src) warped_dst_code = self.encoder (warped_dst) pred_src_src = self.decoder_src(warped_src_code) @@ -250,30 +251,30 @@ class SAEModel(ModelBase): pred_src_srcm = self.decoder_srcm(warped_src_code) pred_dst_dstm = self.decoder_dstm(warped_dst_code) pred_src_dstm = self.decoder_srcm(warped_dst_code) - + if self.is_first_run() and self.options['ca_weights']: io.log_info ("Initializing CA weights...") conv_weights_list = [] for model in models_list: for layer in model.layers: if type(layer) == Conv2D: - conv_weights_list += [layer.weights[0]] #Conv2D kernel_weights + conv_weights_list += [layer.weights[0]] #Conv2D kernel_weights CAInitializerMP ( conv_weights_list ) - - + + pred_src_src, pred_dst_dst, pred_src_dst, = [ [x] if type(x) != list else x for x in [pred_src_src, pred_dst_dst, pred_src_dst, ] ] - + if self.options['learn_mask']: pred_src_srcm, pred_dst_dstm, pred_src_dstm = [ [x] if type(x) != list else x for x in [pred_src_srcm, pred_dst_dstm, pred_src_dstm] ] target_srcm_blurred_ar = [ gaussian_blur( max(1, K.int_shape(x)[1] // 32) )(x) for x in target_srcm_ar] - target_srcm_sigm_ar = target_srcm_blurred_ar #[ x / 2.0 + 0.5 for x in target_srcm_blurred_ar] - target_srcm_anti_sigm_ar = [ 1.0 - x for x in target_srcm_sigm_ar] - + target_srcm_sigm_ar = target_srcm_blurred_ar #[ x / 2.0 + 0.5 for x in target_srcm_blurred_ar] + target_srcm_anti_sigm_ar = [ 1.0 - x for x in target_srcm_sigm_ar] + target_dstm_blurred_ar = [ gaussian_blur( max(1, K.int_shape(x)[1] // 32) )(x) for x in target_dstm_ar] - target_dstm_sigm_ar = target_dstm_blurred_ar#[ x / 2.0 + 0.5 for x in target_dstm_blurred_ar] - target_dstm_anti_sigm_ar = [ 1.0 - x for x in target_dstm_sigm_ar] - + target_dstm_sigm_ar = target_dstm_blurred_ar#[ x / 2.0 + 0.5 for x in target_dstm_blurred_ar] + target_dstm_anti_sigm_ar = [ 1.0 - x for x in target_dstm_sigm_ar] + target_src_sigm_ar = target_src_ar#[ x + 1 for x in target_src_ar] target_dst_sigm_ar = target_dst_ar#[ x + 1 for x in target_dst_ar] @@ -284,32 +285,32 @@ class SAEModel(ModelBase): target_src_masked_ar = [ target_src_sigm_ar[i]*target_srcm_sigm_ar[i] for i in range(len(target_src_sigm_ar))] target_dst_masked_ar = [ target_dst_sigm_ar[i]*target_dstm_sigm_ar[i] for i in range(len(target_dst_sigm_ar))] target_dst_anti_masked_ar = [ target_dst_sigm_ar[i]*target_dstm_anti_sigm_ar[i] for i in range(len(target_dst_sigm_ar))] - + pred_src_src_masked_ar = [ pred_src_src_sigm_ar[i] * target_srcm_sigm_ar[i] for i in range(len(pred_src_src_sigm_ar))] pred_dst_dst_masked_ar = [ pred_dst_dst_sigm_ar[i] * target_dstm_sigm_ar[i] for i in range(len(pred_dst_dst_sigm_ar))] - + target_src_masked_ar_opt = target_src_masked_ar if masked_training else target_src_sigm_ar target_dst_masked_ar_opt = target_dst_masked_ar if masked_training else target_dst_sigm_ar - + pred_src_src_masked_ar_opt = pred_src_src_masked_ar if masked_training else pred_src_src_sigm_ar pred_dst_dst_masked_ar_opt = pred_dst_dst_masked_ar if masked_training else pred_dst_dst_sigm_ar - + psd_target_dst_masked_ar = [ pred_src_dst_sigm_ar[i]*target_dstm_sigm_ar[i] for i in range(len(pred_src_dst_sigm_ar))] psd_target_dst_anti_masked_ar = [ pred_src_dst_sigm_ar[i]*target_dstm_anti_sigm_ar[i] for i in range(len(pred_src_dst_sigm_ar))] - - if self.is_training_mode: + + if self.is_training_mode: self.src_dst_opt = Adam(lr=5e-5, beta_1=0.5, beta_2=0.999, tf_cpu_mode=self.options['optimizer_mode']-1) self.src_dst_mask_opt = Adam(lr=5e-5, beta_1=0.5, beta_2=0.999, tf_cpu_mode=self.options['optimizer_mode']-1) - if self.options['archi'] == 'liae': + if self.options['archi'] == 'liae': src_dst_loss_train_weights = self.encoder.trainable_weights + self.inter_B.trainable_weights + self.inter_AB.trainable_weights + self.decoder.trainable_weights if self.options['learn_mask']: src_dst_mask_loss_train_weights = self.encoder.trainable_weights + self.inter_B.trainable_weights + self.inter_AB.trainable_weights + self.decoderm.trainable_weights - else: + else: src_dst_loss_train_weights = self.encoder.trainable_weights + self.decoder_src.trainable_weights + self.decoder_dst.trainable_weights if self.options['learn_mask']: src_dst_mask_loss_train_weights = self.encoder.trainable_weights + self.decoder_srcm.trainable_weights + self.decoder_dstm.trainable_weights - + if not self.options['pixel_loss']: src_loss_batch = sum([ ( 100*K.square( dssim(kernel_size=int(resolution/11.6),max_value=1.0)( target_src_masked_ar_opt[i], pred_src_src_masked_ar_opt[i] ) )) for i in range(len(target_src_masked_ar_opt)) ]) else: @@ -318,9 +319,9 @@ class SAEModel(ModelBase): src_loss = K.mean(src_loss_batch) face_style_power = self.options['face_style_power'] / 100.0 - - if face_style_power != 0: - src_loss += style_loss(gaussian_blur_radius=resolution//16, loss_weight=face_style_power, wnd_size=0)( psd_target_dst_masked_ar[-1], target_dst_masked_ar[-1] ) + + if face_style_power != 0: + src_loss += style_loss(gaussian_blur_radius=resolution//16, loss_weight=face_style_power, wnd_size=0)( psd_target_dst_masked_ar[-1], target_dst_masked_ar[-1] ) bg_style_power = self.options['bg_style_power'] / 100.0 if bg_style_power != 0: @@ -334,32 +335,32 @@ class SAEModel(ModelBase): dst_loss_batch = sum([ ( 100*K.square(dssim(kernel_size=int(resolution/11.6),max_value=1.0)( target_dst_masked_ar_opt[i], pred_dst_dst_masked_ar_opt[i] ) )) for i in range(len(target_dst_masked_ar_opt)) ]) else: dst_loss_batch = sum([ K.mean ( 100*K.square( target_dst_masked_ar_opt[i] - pred_dst_dst_masked_ar_opt[i] ), axis=[1,2,3]) for i in range(len(target_dst_masked_ar_opt)) ]) - + dst_loss = K.mean(dst_loss_batch) - feed = [warped_src, warped_dst] + feed = [warped_src, warped_dst] feed += target_src_ar[::-1] feed += target_srcm_ar[::-1] feed += target_dst_ar[::-1] feed += target_dstm_ar[::-1] - + self.src_dst_train = K.function (feed,[src_loss,dst_loss], self.src_dst_opt.get_updates(src_loss+dst_loss, src_dst_loss_train_weights) ) if self.options['learn_mask']: src_mask_loss = sum([ K.mean(K.square(target_srcm_ar[-1]-pred_src_srcm[-1])) for i in range(len(target_srcm_ar)) ]) dst_mask_loss = sum([ K.mean(K.square(target_dstm_ar[-1]-pred_dst_dstm[-1])) for i in range(len(target_dstm_ar)) ]) - - feed = [ warped_src, warped_dst] + + feed = [ warped_src, warped_dst] feed += target_srcm_ar[::-1] - feed += target_dstm_ar[::-1] - + feed += target_dstm_ar[::-1] + 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]]) else: - self.AE_view = K.function ([warped_src, warped_dst], [pred_src_src[-1], pred_dst_dst[-1], pred_src_dst[-1] ] ) - + self.AE_view = K.function ([warped_src, warped_dst], [pred_src_src[-1], pred_dst_dst[-1], pred_src_dst[-1] ] ) + self.load_weights_safe(weights_to_load)#, [ [self.src_dst_opt, 'src_dst_opt'], [self.src_dst_mask_opt, 'src_dst_mask_opt']]) else: self.load_weights_safe(weights_to_load) @@ -367,30 +368,30 @@ class SAEModel(ModelBase): self.AE_convert = K.function ([warped_dst],[ pred_src_dst[-1], pred_src_dstm[-1] ]) else: self.AE_convert = K.function ([warped_dst],[ pred_src_dst[-1] ]) - - - if self.is_training_mode: + + + if self.is_training_mode: self.src_sample_losses = [] self.dst_sample_losses = [] - - f = SampleProcessor.TypeFlags + + f = SampleProcessor.TypeFlags face_type = f.FACE_ALIGN_FULL if self.options['face_type'] == 'f' else f.FACE_ALIGN_HALF - - output_sample_types=[ [f.WARPED_TRANSFORMED | face_type | f.MODE_BGR, resolution] ] + + output_sample_types=[ [f.WARPED_TRANSFORMED | face_type | f.MODE_BGR, resolution] ] output_sample_types += [ [f.TRANSFORMED | face_type | f.MODE_BGR, resolution // (2**i) ] for i in range(ms_count)] output_sample_types += [ [f.TRANSFORMED | face_type | f.MODE_M | f.FACE_MASK_FULL, resolution // (2**i) ] for i in range(ms_count)] - - self.set_training_data_generators ([ - SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None, - debug=self.is_debug(), batch_size=self.batch_size, - sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05])+self.src_scale_mod / 100.0 ), + + self.set_training_data_generators ([ + SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None, + debug=self.is_debug(), batch_size=self.batch_size, + sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05])+self.src_scale_mod / 100.0 ), output_sample_types=output_sample_types ), - + SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, - sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, ), + sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, ), output_sample_types=output_sample_types ) ]) - + #override def onSave(self): opt_ar = [ [self.src_dst_opt, 'src_dst_opt'], @@ -413,10 +414,10 @@ class SAEModel(ModelBase): if self.options['learn_mask']: ar += [ [self.decoder_srcm, 'decoder_srcm.h5'], [self.decoder_dstm, 'decoder_dstm.h5'] ] - + self.save_weights_safe(ar) - - + + #override def onTrainOneIter(self, generators_samples, generators_list): src_samples = generators_samples[0] @@ -425,17 +426,17 @@ class SAEModel(ModelBase): feed = [src_samples[0], dst_samples[0] ] + \ src_samples[1:1+self.ms_count*2] + \ dst_samples[1:1+self.ms_count*2] - + src_loss, dst_loss, = self.src_dst_train (feed) - + if self.options['learn_mask']: feed = [ src_samples[0], dst_samples[0] ] + \ src_samples[1+self.ms_count:1+self.ms_count*2] + \ dst_samples[1+self.ms_count:1+self.ms_count*2] src_mask_loss, dst_mask_loss, = self.src_dst_mask_train (feed) - + return ( ('src_loss', src_loss), ('dst_loss', dst_loss) ) - + #override def onGetPreview(self, sample): @@ -454,33 +455,33 @@ class SAEModel(ModelBase): 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],) + # ar += (SDM[i],) st.append ( np.concatenate ( ar, axis=1) ) - + return [ ('SAE', np.concatenate (st, axis=0 )), ] - + def predictor_func (self, face): - + prd = [ x[0] for x in self.AE_convert ( [ face[np.newaxis,:,:,0:3] ] ) ] if not self.options['learn_mask']: - prd += [ face[...,3:4] ] - + prd += [ face[...,3:4] ] + return np.concatenate ( prd, -1 ) - + #override def get_converter(self): base_erode_mask_modifier = 30 if self.options['face_type'] == 'f' else 100 base_blur_mask_modifier = 0 if self.options['face_type'] == 'f' else 100 - + default_erode_mask_modifier = 0 default_blur_mask_modifier = 100 if (self.options['face_style_power'] or self.options['bg_style_power']) and \ self.options['face_type'] == 'f' else 0 - + face_type = FaceType.FULL if self.options['face_type'] == 'f' else FaceType.HALF - + from converters import ConverterMasked - return ConverterMasked(self.predictor_func, + return ConverterMasked(self.predictor_func, predictor_input_size=self.options['resolution'], output_size=self.options['resolution'], face_type=face_type, @@ -490,30 +491,34 @@ class SAEModel(ModelBase): default_erode_mask_modifier=default_erode_mask_modifier, default_blur_mask_modifier=default_blur_mask_modifier, clip_hborder_mask_per=0.0625 if self.options['face_type'] == 'f' else 0) - + @staticmethod def initialize_nn_functions(): exec (nnlib.import_all(), locals(), globals()) - + + def BatchNorm(): + return BatchNormalization(axis=-1, gamma_initializer=RandomNormal(1., 0.02) ) + + class ResidualBlock(object): def __init__(self, filters, kernel_size=3, padding='same', use_reflection_padding=False): self.filters = filters self.kernel_size = kernel_size self.padding = padding #if not use_reflection_padding else 'valid' self.use_reflection_padding = use_reflection_padding - + def __call__(self, inp): var_x = LeakyReLU(alpha=0.2)(inp) - + #if self.use_reflection_padding: # #var_x = ReflectionPadding2D(stride=1, kernel_size=kernel_size)(var_x) - var_x = Conv2D(self.filters, kernel_size=self.kernel_size, padding=self.padding, kernel_initializer=RandomNormal(0, 0.02) )(var_x) + var_x = Conv2D(self.filters, kernel_size=self.kernel_size, padding=self.padding, kernel_initializer=RandomNormal(0, 0.02) )(var_x) var_x = LeakyReLU(alpha=0.2)(var_x) - + #if self.use_reflection_padding: # #var_x = ReflectionPadding2D(stride=1, kernel_size=kernel_size)(var_x) - + var_x = Conv2D(self.filters, kernel_size=self.kernel_size, padding=self.padding, kernel_initializer=RandomNormal(0, 0.02) )(var_x) var_x = Scale(gamma_init=keras.initializers.Constant(value=0.1))(var_x) var_x = Add()([var_x, inp]) @@ -521,99 +526,108 @@ class SAEModel(ModelBase): return var_x SAEModel.ResidualBlock = ResidualBlock - def downscale (dim): + def downscale (dim, use_bn=False): def func(x): - return LeakyReLU(0.1)(Conv2D(dim, kernel_size=5, strides=2, padding='same', kernel_initializer=RandomNormal(0, 0.02))(x)) - return func + if use_bn: + return LeakyReLU(0.1)(BatchNorm()(Conv2D(dim, kernel_size=5, strides=2, padding='same', kernel_initializer=RandomNormal(0, 0.02), use_bias=False)(x))) + else: + return LeakyReLU(0.1)(Conv2D(dim, kernel_size=5, strides=2, padding='same', kernel_initializer=RandomNormal(0, 0.02))(x)) + return func SAEModel.downscale = downscale - - def downscale_sep (dim): + + def downscale_sep (dim, use_bn=False): def func(x): - return LeakyReLU(0.1)(SeparableConv2D(dim, kernel_size=5, strides=2, padding='same', depthwise_initializer=RandomNormal(0, 0.02), pointwise_initializer=RandomNormal(0, 0.02) )(x)) - return func + if use_bn: + return LeakyReLU(0.1)(BatchNorm()(SeparableConv2D(dim, kernel_size=5, strides=2, padding='same', depthwise_initializer=RandomNormal(0, 0.02), pointwise_initializer=RandomNormal(0, 0.02), use_bias=False )(x))) + else: + return LeakyReLU(0.1)(SeparableConv2D(dim, kernel_size=5, strides=2, padding='same', depthwise_initializer=RandomNormal(0, 0.02), pointwise_initializer=RandomNormal(0, 0.02) )(x)) + return func SAEModel.downscale_sep = downscale_sep - - def upscale (dim): + + def upscale (dim, use_bn=False): def func(x): - return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, kernel_size=3, strides=1, padding='same', kernel_initializer=RandomNormal(0, 0.02) )(x))) - return func + if use_bn: + return SubpixelUpscaler()(LeakyReLU(0.1)(BatchNorm()(Conv2D(dim * 4, kernel_size=3, strides=1, padding='same', kernel_initializer=RandomNormal(0,0.02), use_bias=False )(x)))) + else: + return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, kernel_size=3, strides=1, padding='same', kernel_initializer=RandomNormal(0, 0.02) )(x))) + return func SAEModel.upscale = upscale - + def to_bgr (output_nc): def func(x): return Conv2D(output_nc, kernel_size=5, padding='same', activation='sigmoid', kernel_initializer=RandomNormal(0, 0.02) )(x) return func SAEModel.to_bgr = to_bgr - - + + @staticmethod - def LIAEEncFlow(resolution, light_enc, ed_ch_dims=42): + def LIAEEncFlow(resolution, light_enc, ed_ch_dims=42, use_bn=False): exec (nnlib.import_all(), locals(), globals()) upscale = SAEModel.upscale downscale = SAEModel.downscale downscale_sep = SAEModel.downscale_sep - + def func(input): ed_dims = K.int_shape(input)[-1]*ed_ch_dims - - x = input + + x = input x = downscale(ed_dims)(x) - if not light_enc: - x = downscale(ed_dims*2)(x) - x = downscale(ed_dims*4)(x) - x = downscale(ed_dims*8)(x) + if not light_enc: + x = downscale(ed_dims*2, use_bn=use_bn)(x) + x = downscale(ed_dims*4, use_bn=use_bn)(x) + x = downscale(ed_dims*8, use_bn=use_bn)(x) else: - x = downscale_sep(ed_dims*2)(x) - x = downscale(ed_dims*4)(x) - x = downscale_sep(ed_dims*8)(x) - - x = Flatten()(x) + x = downscale_sep(ed_dims*2, use_bn=use_bn)(x) + x = downscale(ed_dims*4, use_bn=use_bn)(x) + x = downscale_sep(ed_dims*8, use_bn=use_bn)(x) + + x = Flatten()(x) return x return func - + @staticmethod - def LIAEInterFlow(resolution, ae_dims=256): + def LIAEInterFlow(resolution, ae_dims=256, use_bn=False): exec (nnlib.import_all(), locals(), globals()) upscale = SAEModel.upscale lowest_dense_res=resolution // 16 - - def func(input): + + def func(input): x = input[0] x = Dense(ae_dims)(x) x = Dense(lowest_dense_res * lowest_dense_res * ae_dims*2)(x) x = Reshape((lowest_dense_res, lowest_dense_res, ae_dims*2))(x) - x = upscale(ae_dims*2)(x) + x = upscale(ae_dims*2, use_bn=use_bn)(x) return x return func - + @staticmethod - def LIAEDecFlow(output_nc,ed_ch_dims=21, multiscale_count=1): + def LIAEDecFlow(output_nc,ed_ch_dims=21, multiscale_count=1, use_bn=False): exec (nnlib.import_all(), locals(), globals()) upscale = SAEModel.upscale to_bgr = SAEModel.to_bgr ed_dims = output_nc * ed_ch_dims - - def func(input): + + def func(input): x = input[0] - + outputs = [] - x1 = upscale(ed_dims*8)( x ) - + x1 = upscale(ed_dims*8, use_bn=use_bn)( x ) + if multiscale_count >= 3: - outputs += [ to_bgr(output_nc) ( x1 ) ] - - x2 = upscale(ed_dims*4)( x1 ) - + outputs += [ to_bgr(output_nc) ( x1 ) ] + + x2 = upscale(ed_dims*4, use_bn=use_bn)( x1 ) + if multiscale_count >= 2: outputs += [ to_bgr(output_nc) ( x2 ) ] - - x3 = upscale(ed_dims*2)( x2 ) - + + x3 = upscale(ed_dims*2, use_bn=use_bn)( x2 ) + outputs += [ to_bgr(output_nc) ( x3 ) ] - + return outputs return func - + @staticmethod def DFEncFlow(resolution, light_enc, ae_dims=512, ed_ch_dims=42): exec (nnlib.import_all(), locals(), globals()) @@ -622,11 +636,11 @@ class SAEModel(ModelBase): downscale_sep = SAEModel.downscale_sep lowest_dense_res = resolution // 16 - def func(input): + def func(input): x = input - + ed_dims = K.int_shape(input)[-1]*ed_ch_dims - + x = downscale(ed_dims)(x) if not light_enc: x = downscale(ed_dims*2)(x) @@ -636,15 +650,15 @@ class SAEModel(ModelBase): x = downscale_sep(ed_dims*2)(x) x = downscale_sep(ed_dims*4)(x) x = downscale_sep(ed_dims*8)(x) - + x = Dense(ae_dims)(Flatten()(x)) x = Dense(lowest_dense_res * lowest_dense_res * ae_dims)(x) x = Reshape((lowest_dense_res, lowest_dense_res, ae_dims))(x) x = upscale(ae_dims)(x) - + return x return func - + @staticmethod def DFDecFlow(output_nc, ed_ch_dims=21, multiscale_count=1): exec (nnlib.import_all(), locals(), globals()) @@ -652,29 +666,29 @@ class SAEModel(ModelBase): to_bgr = SAEModel.to_bgr ed_dims = output_nc * ed_ch_dims - def func(input): + def func(input): x = input[0] - + outputs = [] - x1 = upscale(ed_dims*8)( x ) - + x1 = upscale(ed_dims*8)( x ) + if multiscale_count >= 3: - outputs += [ to_bgr(output_nc) ( x1 ) ] - - x2 = upscale(ed_dims*4)( x1 ) - + outputs += [ to_bgr(output_nc) ( x1 ) ] + + x2 = upscale(ed_dims*4)( x1 ) + if multiscale_count >= 2: outputs += [ to_bgr(output_nc) ( x2 ) ] - + x3 = upscale(ed_dims*2)( x2 ) - + outputs += [ to_bgr(output_nc) ( x3 ) ] - - return outputs + + return outputs return func - - - + + + @staticmethod def VGEncFlow(resolution, light_enc, ae_dims=512, ed_ch_dims=42): exec (nnlib.import_all(), locals(), globals()) @@ -683,78 +697,78 @@ class SAEModel(ModelBase): downscale_sep = SAEModel.downscale_sep ResidualBlock = SAEModel.ResidualBlock lowest_dense_res = resolution // 16 - + def func(input): x = input ed_dims = K.int_shape(input)[-1]*ed_ch_dims while np.modf(ed_dims / 4)[0] != 0.0: ed_dims -= 1 - + in_conv_filters = ed_dims# if resolution <= 128 else ed_dims + (resolution//128)*ed_ch_dims - + x = tmp_x = Conv2D (in_conv_filters, kernel_size=5, strides=2, padding='same') (x) for _ in range ( 8 if light_enc else 16 ): x = ResidualBlock(ed_dims)(x) - + x = Add()([x, tmp_x]) x = downscale(ed_dims)(x) x = SubpixelUpscaler()(x) - + x = downscale(ed_dims)(x) x = SubpixelUpscaler()(x) - - x = downscale(ed_dims)(x) + + x = downscale(ed_dims)(x) if light_enc: x = downscale_sep (ed_dims*2)(x) else: x = downscale (ed_dims*2)(x) - + x = downscale(ed_dims*4)(x) - + if light_enc: x = downscale_sep (ed_dims*8)(x) else: x = downscale (ed_dims*8)(x) - + x = Dense(ae_dims)(Flatten()(x)) x = Dense(lowest_dense_res * lowest_dense_res * ae_dims)(x) x = Reshape((lowest_dense_res, lowest_dense_res, ae_dims))(x) x = upscale(ae_dims)(x) return x - + return func - + @staticmethod def VGDecFlow(output_nc, ed_ch_dims=21, multiscale_count=1): exec (nnlib.import_all(), locals(), globals()) - upscale = SAEModel.upscale + upscale = SAEModel.upscale to_bgr = SAEModel.to_bgr ResidualBlock = SAEModel.ResidualBlock ed_dims = output_nc * ed_ch_dims - + def func(input): x = input[0] - + x = upscale( ed_dims*8 )(x) x = ResidualBlock( ed_dims*8 )(x) - + x = upscale( ed_dims*4 )(x) x = ResidualBlock( ed_dims*4 )(x) - + x = upscale( ed_dims*2 )(x) x = ResidualBlock( ed_dims*2 )(x) - - x = to_bgr(output_nc) (x) + + x = to_bgr(output_nc) (x) return x - + return func - + Model = SAEModel - + # 'worst' sample booster gives no good result, or I dont know how to filter worst samples properly. # ##gathering array of sample_losses @@ -769,7 +783,7 @@ Model = SAEModel # idxs = (x[:,0][ np.argwhere ( b [ b > (np.mean(b)+np.std(b)) ] )[:,0] ]).astype(np.uint) # generators_list[0].repeat_sample_idxs(idxs) #ask generator to repeat these sample idxs # print ("src repeated %d" % (len(idxs)) ) -# +# #if len(self.dst_sample_losses) >= 128: #array is big enough # #fetching idxs which losses are bigger than average # x = np.array (self.dst_sample_losses) @@ -777,4 +791,4 @@ Model = SAEModel # b = x[:,1] # idxs = (x[:,0][ np.argwhere ( b [ b > (np.mean(b)+np.std(b)) ] )[:,0] ]).astype(np.uint) # generators_list[1].repeat_sample_idxs(idxs) #ask generator to repeat these sample idxs -# print ("dst repeated %d" % (len(idxs)) ) \ No newline at end of file +# print ("dst repeated %d" % (len(idxs)) ) diff --git a/models/Model_SAE/__init__.py b/models/Model_SAE/__init__.py index cdb3fe7..0188f11 100644 --- a/models/Model_SAE/__init__.py +++ b/models/Model_SAE/__init__.py @@ -1 +1 @@ -from .Model import Model \ No newline at end of file +from .Model import Model diff --git a/models/__init__.py b/models/__init__.py index 22a2f3d..971091d 100644 --- a/models/__init__.py +++ b/models/__init__.py @@ -2,4 +2,4 @@ from .ModelBase import ModelBase def import_model(name): module = __import__('Model_'+name, globals(), locals(), [], 1) - return getattr(module, 'Model') \ No newline at end of file + return getattr(module, 'Model') diff --git a/nnlib/CAInitializer.py b/nnlib/CAInitializer.py index 20c03c1..f81dd06 100644 --- a/nnlib/CAInitializer.py +++ b/nnlib/CAInitializer.py @@ -61,7 +61,7 @@ def _scale_filters(filters, variance): def CAGenerateWeights ( shape, floatx, data_format, eps_std=0.05, seed=None ): if seed is not None: np.random.seed(seed) - + fan_in, fan_out = _compute_fans(shape, data_format) variance = 2 / fan_in @@ -109,4 +109,4 @@ def CAGenerateWeights ( shape, floatx, data_format, eps_std=0.05, seed=None ): # Format of array is now: filters, stack, row, column init = np.array(init) init = _scale_filters(init, variance) - return init.transpose(transpose_dimensions) \ No newline at end of file + return init.transpose(transpose_dimensions) diff --git a/nnlib/__init__.py b/nnlib/__init__.py index 98e1677..14793f7 100644 --- a/nnlib/__init__.py +++ b/nnlib/__init__.py @@ -1 +1 @@ -from .nnlib import nnlib \ No newline at end of file +from .nnlib import nnlib diff --git a/nnlib/device.py b/nnlib/device.py index a0d9e4f..672a63c 100644 --- a/nnlib/device.py +++ b/nnlib/device.py @@ -5,11 +5,11 @@ from .pynvml import * #you can set DFL_TF_MIN_REQ_CAP manually for your build #the reason why we cannot check tensorflow.version is it requires import tensorflow -tf_min_req_cap = int(os.environ.get("DFL_TF_MIN_REQ_CAP", 35)) - +tf_min_req_cap = int(os.environ.get("DFL_TF_MIN_REQ_CAP", 35)) + class device: backend = None - class Config(): + class Config(): force_gpu_idx = -1 multi_gpu = False force_gpu_idxs = None @@ -22,36 +22,36 @@ class device: use_fp16 = False cpu_only = False backend = None - def __init__ (self, force_gpu_idx = -1, - multi_gpu = False, - force_gpu_idxs = None, + def __init__ (self, force_gpu_idx = -1, + multi_gpu = False, + force_gpu_idxs = None, choose_worst_gpu = False, allow_growth = True, use_fp16 = False, cpu_only = False, **in_options): - + self.backend = device.backend self.use_fp16 = use_fp16 self.cpu_only = cpu_only - + if not self.cpu_only: self.cpu_only = (self.backend == "tensorflow-cpu") - + if not self.cpu_only: self.force_gpu_idx = force_gpu_idx self.multi_gpu = multi_gpu self.force_gpu_idxs = force_gpu_idxs - self.choose_worst_gpu = choose_worst_gpu + self.choose_worst_gpu = choose_worst_gpu self.allow_growth = allow_growth - + self.gpu_idxs = [] if force_gpu_idxs is not None: for idx in force_gpu_idxs.split(','): idx = int(idx) if device.isValidDeviceIdx(idx): - self.gpu_idxs.append(idx) + self.gpu_idxs.append(idx) else: gpu_idx = force_gpu_idx if (force_gpu_idx >= 0 and device.isValidDeviceIdx(force_gpu_idx)) else device.getBestValidDeviceIdx() if not choose_worst_gpu else device.getWorstValidDeviceIdx() if gpu_idx != -1: @@ -61,10 +61,10 @@ class device: self.multi_gpu = False else: self.gpu_idxs = [gpu_idx] - + self.cpu_only = (len(self.gpu_idxs) == 0) - - + + if not self.cpu_only: self.gpu_names = [] self.gpu_compute_caps = [] @@ -78,10 +78,10 @@ class device: self.gpu_names = ['CPU'] self.gpu_compute_caps = [99] self.gpu_vram_gb = [0] - + if self.cpu_only: self.backend = "tensorflow-cpu" - + @staticmethod def getValidDeviceIdxsEnumerator(): if device.backend == "plaidML": @@ -94,8 +94,8 @@ class device: yield gpu_idx elif device.backend == "tensorflow-generic": yield 0 - - + + @staticmethod def getValidDevicesWithAtLeastTotalMemoryGB(totalmemsize_gb): result = [] @@ -111,9 +111,9 @@ class device: result.append (i) elif device.backend == "tensorflow-generic": return [0] - + return result - + @staticmethod def getAllDevicesIdxsList(): if device.backend == "plaidML": @@ -121,8 +121,8 @@ class device: elif device.backend == "tensorflow": return [ *range(nvmlDeviceGetCount() ) ] elif device.backend == "tensorflow-generic": - return [0] - + return [0] + @staticmethod def getValidDevicesIdxsWithNamesList(): if device.backend == "plaidML": @@ -137,17 +137,17 @@ class device: @staticmethod def getDeviceVRAMTotalGb (idx): if device.backend == "plaidML": - if idx < plaidML_devices_count: + if idx < plaidML_devices_count: return plaidML_devices[idx]['globalMemSize'] / (1024*1024*1024) elif device.backend == "tensorflow": - if idx < nvmlDeviceGetCount(): + if idx < nvmlDeviceGetCount(): memInfo = nvmlDeviceGetMemoryInfo( nvmlDeviceGetHandleByIndex(idx) ) return round ( memInfo.total / (1024*1024*1024) ) return 0 elif device.backend == "tensorflow-generic": return 2 - + @staticmethod def getBestValidDeviceIdx(): if device.backend == "plaidML": @@ -172,7 +172,7 @@ class device: return idx elif device.backend == "tensorflow-generic": return 0 - + @staticmethod def getWorstValidDeviceIdx(): if device.backend == "plaidML": @@ -197,7 +197,7 @@ class device: return idx elif device.backend == "tensorflow-generic": return 0 - + @staticmethod def isValidDeviceIdx(idx): if device.backend == "plaidML": @@ -206,11 +206,11 @@ class device: return idx in [*device.getValidDeviceIdxsEnumerator()] elif device.backend == "tensorflow-generic": return (idx == 0) - + @staticmethod def getDeviceIdxsEqualModel(idx): if device.backend == "plaidML": - result = [] + result = [] idx_name = plaidML_devices[idx]['description'] for i in device.getValidDeviceIdxsEnumerator(): if plaidML_devices[i]['description'] == idx_name: @@ -218,7 +218,7 @@ class device: return result elif device.backend == "tensorflow": - result = [] + result = [] idx_name = nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(idx)).decode() for i in device.getValidDeviceIdxsEnumerator(): if nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(i)).decode() == idx_name: @@ -226,60 +226,60 @@ class device: return result elif device.backend == "tensorflow-generic": - return [0] if idx == 0 else [] - + return [0] if idx == 0 else [] + @staticmethod def getDeviceName (idx): if device.backend == "plaidML": - if idx < plaidML_devices_count: + if idx < plaidML_devices_count: return plaidML_devices[idx]['description'] elif device.backend == "tensorflow": - if idx < nvmlDeviceGetCount(): + if idx < nvmlDeviceGetCount(): return nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(idx)).decode() elif device.backend == "tensorflow-generic": if idx == 0: return "Generic GeForce GPU" - + return None - + @staticmethod def getDeviceID (idx): if device.backend == "plaidML": - if idx < plaidML_devices_count: + if idx < plaidML_devices_count: return plaidML_devices[idx]['id'].decode() - return None - + return None + @staticmethod def getDeviceComputeCapability(idx): result = 0 if device.backend == "plaidML": return 99 elif device.backend == "tensorflow": - if idx < nvmlDeviceGetCount(): + if idx < nvmlDeviceGetCount(): result = nvmlDeviceGetCudaComputeCapability(nvmlDeviceGetHandleByIndex(idx)) elif device.backend == "tensorflow-generic": - return 99 if idx == 0 else 0 - + return 99 if idx == 0 else 0 + return result[0] * 10 + result[1] - + force_plaidML = os.environ.get("DFL_FORCE_PLAIDML", "0") == "1" #for OpenCL build , forcing using plaidML even if NVIDIA found force_tf_cpu = os.environ.get("DFL_FORCE_TF_CPU", "0") == "1" #for OpenCL build , forcing using tf-cpu if plaidML failed has_nvml = False has_nvml_cap = False -#use DFL_FORCE_HAS_NVIDIA_DEVICE=1 if +#use DFL_FORCE_HAS_NVIDIA_DEVICE=1 if #- your NVIDIA cannot be seen by OpenCL #- CUDA build of DFL -has_nvidia_device = os.environ.get("DFL_FORCE_HAS_NVIDIA_DEVICE", "0") == "1" +has_nvidia_device = os.environ.get("DFL_FORCE_HAS_NVIDIA_DEVICE", "0") == "1" plaidML_devices = [] # Using plaidML OpenCL backend to determine system devices and has_nvidia_device -try: +try: os.environ['PLAIDML_EXPERIMENTAL'] = 'false' #this enables work plaidML without run 'plaidml-setup' - import plaidml + import plaidml ctx = plaidml.Context() for d in plaidml.devices(ctx, return_all=True)[0]: details = json.loads(d.details) @@ -288,13 +288,13 @@ try: if 'nvidia' in details['vendor'].lower(): has_nvidia_device = True plaidML_devices += [ {'id':d.id, - 'globalMemSize' : int(details['globalMemSize']), + 'globalMemSize' : int(details['globalMemSize']), 'description' : d.description.decode() }] ctx.shutdown() except: pass - + plaidML_devices_count = len(plaidML_devices) #choosing backend @@ -306,11 +306,11 @@ if device.backend is None and not force_tf_cpu: nvmlInit() has_nvml = True device.backend = "tensorflow" #set tensorflow backend in order to use device.*device() functions - + gpu_idxs = device.getAllDevicesIdxsList() gpu_caps = np.array ( [ device.getDeviceComputeCapability(gpu_idx) for gpu_idx in gpu_idxs ] ) - - if len ( np.ndarray.flatten ( np.argwhere (gpu_caps >= tf_min_req_cap) ) ) == 0: + + if len ( np.ndarray.flatten ( np.argwhere (gpu_caps >= tf_min_req_cap) ) ) == 0: if not force_plaidML: print ("No CUDA devices found with minimum required compute capability: %d.%d. Falling back to OpenCL mode." % (tf_min_req_cap // 10, tf_min_req_cap % 10) ) device.backend = None @@ -320,7 +320,7 @@ if device.backend is None and not force_tf_cpu: except: #if no NVSMI installed exception will occur device.backend = None - has_nvml = False + has_nvml = False if force_plaidML or (device.backend is None and not has_nvidia_device): #tensorflow backend was failed without has_nvidia_device , or forcing plaidML, trying to use plaidML backend @@ -333,7 +333,7 @@ if force_plaidML or (device.backend is None and not has_nvidia_device): if device.backend is None: if force_tf_cpu: device.backend = "tensorflow-cpu" - elif not has_nvml: + elif not has_nvml: if has_nvidia_device: #some notebook systems have NVIDIA card without NVSMI in official drivers #in that case considering we have system with one capable GPU and let tensorflow to choose best GPU @@ -348,4 +348,3 @@ if device.backend is None: else: #has NVSMI, no capable CUDA-devices, also plaidML was failed, then CPU only device.backend = "tensorflow-cpu" - diff --git a/nnlib/nnlib.py b/nnlib/nnlib.py index 138eb38..338f0fc 100644 --- a/nnlib/nnlib.py +++ b/nnlib/nnlib.py @@ -541,7 +541,6 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator result = CAInitializerMPSubprocessor ( [ (i, K.int_shape(conv_weights)) for i, conv_weights in enumerate(conv_weights_list) ], K.floatx(), K.image_data_format() ).run() for idx, weights in result: K.set_value ( conv_weights_list[idx], weights ) - nnlib.CAInitializerMP = CAInitializerMP diff --git a/nnlib/pynvml.py b/nnlib/pynvml.py index 95b8875..b5cf0ec 100644 --- a/nnlib/pynvml.py +++ b/nnlib/pynvml.py @@ -3,7 +3,7 @@ # # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions are met: -# +# # * Redistributions of source code must retain the above copyright notice, # this list of conditions and the following disclaimer. # * Redistributions in binary form must reproduce the above copyright @@ -18,11 +18,11 @@ # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE # ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE # LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR -# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF -# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS # INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN -# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) -# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF +# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF # THE POSSIBILITY OF SUCH DAMAGE. ##### @@ -35,7 +35,7 @@ import sys import os import threading import string - + ## C Type mappings ## ## Enums _nvmlEnableState_t = c_uint @@ -155,9 +155,9 @@ NVML_FAN_FAILED = 1 _nvmlLedColor_t = c_uint NVML_LED_COLOR_GREEN = 0 NVML_LED_COLOR_AMBER = 1 - + _nvmlGpuOperationMode_t = c_uint -NVML_GOM_ALL_ON = 0 +NVML_GOM_ALL_ON = 0 NVML_GOM_COMPUTE = 1 NVML_GOM_LOW_DP = 2 @@ -173,7 +173,7 @@ NVML_RESTRICTED_API_COUNT = 2 _nvmlBridgeChipType_t = c_uint NVML_BRIDGE_CHIP_PLX = 0 -NVML_BRIDGE_CHIP_BRO4 = 1 +NVML_BRIDGE_CHIP_BRO4 = 1 NVML_MAX_PHYSICAL_BRIDGE = 128 _nvmlValueType_t = c_uint @@ -317,7 +317,7 @@ def _nvmlGetFunctionPointer(name): if name in _nvmlGetFunctionPointer_cache: return _nvmlGetFunctionPointer_cache[name] - + libLoadLock.acquire() try: # ensure library was loaded @@ -364,7 +364,7 @@ def nvmlFriendlyObjectToStruct(obj, model): class struct_c_nvmlUnit_t(Structure): pass # opaque handle c_nvmlUnit_t = POINTER(struct_c_nvmlUnit_t) - + class _PrintableStructure(Structure): """ Abstract class that produces nicer __str__ output than ctypes.Structure. @@ -373,7 +373,7 @@ class _PrintableStructure(Structure): this class will print class_name(field_name: formatted_value, field_name: formatted_value) - + _fmt_ dictionary of -> e.g. class that has _field_ 'hex_value', c_uint could be formatted with _fmt_ = {"hex_value" : "%08X"} @@ -397,7 +397,7 @@ class _PrintableStructure(Structure): fmt = self._fmt_[""] result.append(("%s: " + fmt) % (key, value)) return self.__class__.__name__ + "(" + string.join(result, ", ") + ")" - + class c_nvmlUnitInfo_t(_PrintableStructure): _fields_ = [ ('name', c_char * 96), @@ -444,7 +444,7 @@ class nvmlPciInfo_t(_PrintableStructure): ('bus', c_uint), ('device', c_uint), ('pciDeviceId', c_uint), - + # Added in 2.285 ('pciSubSystemId', c_uint), ('reserved0', c_uint), @@ -503,7 +503,7 @@ class c_nvmlBridgeChipHierarchy_t(_PrintableStructure): _fields_ = [ ('bridgeCount', c_uint), ('bridgeChipInfo', c_nvmlBridgeChipInfo_t * 128), - ] + ] class c_nvmlEccErrorCounts_t(_PrintableStructure): _fields_ = [ @@ -582,7 +582,7 @@ nvmlClocksThrottleReasonAll = ( nvmlClocksThrottleReasonSwPowerCap | nvmlClocksThrottleReasonHwSlowdown | nvmlClocksThrottleReasonUnknown - ) + ) class c_nvmlEventData_t(_PrintableStructure): _fields_ = [ @@ -606,31 +606,31 @@ class c_nvmlAccountingStats_t(_PrintableStructure): ## C function wrappers ## def nvmlInit(): _LoadNvmlLibrary() - + # # Initialize the library # fn = _nvmlGetFunctionPointer("nvmlInit_v2") ret = fn() _nvmlCheckReturn(ret) - + # Atomically update refcount global _nvmlLib_refcount libLoadLock.acquire() _nvmlLib_refcount += 1 libLoadLock.release() return None - + def _LoadNvmlLibrary(): ''' Load the library if it isn't loaded already ''' global nvmlLib - + if (nvmlLib == None): # lock to ensure only one caller loads the library libLoadLock.acquire() - + try: # ensure the library still isn't loaded if (nvmlLib == None): @@ -649,7 +649,7 @@ def _LoadNvmlLibrary(): finally: # lock is always freed libLoadLock.release() - + def nvmlShutdown(): # # Leave the library loaded, but shutdown the interface @@ -657,7 +657,7 @@ def nvmlShutdown(): fn = _nvmlGetFunctionPointer("nvmlShutdown") ret = fn() _nvmlCheckReturn(ret) - + # Atomically update refcount global _nvmlLib_refcount libLoadLock.acquire() @@ -701,19 +701,19 @@ def nvmlSystemGetHicVersion(): c_count = c_uint(0) hics = None fn = _nvmlGetFunctionPointer("nvmlSystemGetHicVersion") - + # get the count ret = fn(byref(c_count), None) - + # this should only fail with insufficient size if ((ret != NVML_SUCCESS) and (ret != NVML_ERROR_INSUFFICIENT_SIZE)): raise NVMLError(ret) - + # if there are no hics if (c_count.value == 0): return [] - + hic_array = c_nvmlHwbcEntry_t * c_count.value hics = hic_array() ret = fn(byref(c_count), hics) @@ -770,7 +770,7 @@ def nvmlUnitGetFanSpeedInfo(unit): ret = fn(unit, byref(c_speeds)) _nvmlCheckReturn(ret) return c_speeds - + # added to API def nvmlUnitGetDeviceCount(unit): c_count = c_uint(0) @@ -822,7 +822,7 @@ def nvmlDeviceGetHandleByUUID(uuid): ret = fn(c_uuid, byref(device)) _nvmlCheckReturn(ret) return device - + def nvmlDeviceGetHandleByPciBusId(pciBusId): c_busId = c_char_p(pciBusId) device = c_nvmlDevice_t() @@ -858,7 +858,7 @@ def nvmlDeviceGetBrand(handle): ret = fn(handle, byref(c_type)) _nvmlCheckReturn(ret) return c_type.value - + def nvmlDeviceGetSerial(handle): c_serial = create_string_buffer(NVML_DEVICE_SERIAL_BUFFER_SIZE) fn = _nvmlGetFunctionPointer("nvmlDeviceGetSerial") @@ -892,14 +892,14 @@ def nvmlDeviceGetMinorNumber(handle): ret = fn(handle, byref(c_minor_number)) _nvmlCheckReturn(ret) return c_minor_number.value - + def nvmlDeviceGetUUID(handle): c_uuid = create_string_buffer(NVML_DEVICE_UUID_BUFFER_SIZE) fn = _nvmlGetFunctionPointer("nvmlDeviceGetUUID") ret = fn(handle, c_uuid, c_uint(NVML_DEVICE_UUID_BUFFER_SIZE)) _nvmlCheckReturn(ret) return c_uuid.value - + def nvmlDeviceGetInforomVersion(handle, infoRomObject): c_version = create_string_buffer(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE) fn = _nvmlGetFunctionPointer("nvmlDeviceGetInforomVersion") @@ -929,7 +929,7 @@ def nvmlDeviceValidateInforom(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceValidateInforom") ret = fn(handle) _nvmlCheckReturn(ret) - return None + return None def nvmlDeviceGetDisplayMode(handle): c_mode = _nvmlEnableState_t() @@ -937,29 +937,29 @@ def nvmlDeviceGetDisplayMode(handle): ret = fn(handle, byref(c_mode)) _nvmlCheckReturn(ret) return c_mode.value - + def nvmlDeviceGetDisplayActive(handle): c_mode = _nvmlEnableState_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetDisplayActive") ret = fn(handle, byref(c_mode)) _nvmlCheckReturn(ret) return c_mode.value - - + + def nvmlDeviceGetPersistenceMode(handle): c_state = _nvmlEnableState_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPersistenceMode") ret = fn(handle, byref(c_state)) _nvmlCheckReturn(ret) return c_state.value - + def nvmlDeviceGetPciInfo(handle): c_info = nvmlPciInfo_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPciInfo_v2") ret = fn(handle, byref(c_info)) _nvmlCheckReturn(ret) return c_info - + def nvmlDeviceGetClockInfo(handle, type): c_clock = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetClockInfo") @@ -997,7 +997,7 @@ def nvmlDeviceGetSupportedMemoryClocks(handle): c_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedMemoryClocks") ret = fn(handle, byref(c_count), None) - + if (ret == NVML_SUCCESS): # special case, no clocks return [] @@ -1005,11 +1005,11 @@ def nvmlDeviceGetSupportedMemoryClocks(handle): # typical case clocks_array = c_uint * c_count.value c_clocks = clocks_array() - + # make the call again ret = fn(handle, byref(c_count), c_clocks) _nvmlCheckReturn(ret) - + procs = [] for i in range(c_count.value): procs.append(c_clocks[i]) @@ -1025,7 +1025,7 @@ def nvmlDeviceGetSupportedGraphicsClocks(handle, memoryClockMHz): c_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedGraphicsClocks") ret = fn(handle, c_uint(memoryClockMHz), byref(c_count), None) - + if (ret == NVML_SUCCESS): # special case, no clocks return [] @@ -1033,11 +1033,11 @@ def nvmlDeviceGetSupportedGraphicsClocks(handle, memoryClockMHz): # typical case clocks_array = c_uint * c_count.value c_clocks = clocks_array() - + # make the call again ret = fn(handle, c_uint(memoryClockMHz), byref(c_count), c_clocks) _nvmlCheckReturn(ret) - + procs = [] for i in range(c_count.value): procs.append(c_clocks[i]) @@ -1053,7 +1053,7 @@ def nvmlDeviceGetFanSpeed(handle): ret = fn(handle, byref(c_speed)) _nvmlCheckReturn(ret) return c_speed.value - + def nvmlDeviceGetTemperature(handle, sensor): c_temp = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetTemperature") @@ -1075,7 +1075,7 @@ def nvmlDeviceGetPowerState(handle): ret = fn(handle, byref(c_pstate)) _nvmlCheckReturn(ret) return c_pstate.value - + def nvmlDeviceGetPerformanceState(handle): c_pstate = _nvmlPstates_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPerformanceState") @@ -1089,7 +1089,7 @@ def nvmlDeviceGetPowerManagementMode(handle): ret = fn(handle, byref(c_pcapMode)) _nvmlCheckReturn(ret) return c_pcapMode.value - + def nvmlDeviceGetPowerManagementLimit(handle): c_limit = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementLimit") @@ -1113,7 +1113,7 @@ def nvmlDeviceGetPowerManagementDefaultLimit(handle): ret = fn(handle, byref(c_limit)) _nvmlCheckReturn(ret) return c_limit.value - + # Added in 331 def nvmlDeviceGetEnforcedPowerLimit(handle): @@ -1146,7 +1146,7 @@ def nvmlDeviceGetCurrentGpuOperationMode(handle): # Added in 4.304 def nvmlDeviceGetPendingGpuOperationMode(handle): return nvmlDeviceGetGpuOperationMode(handle)[1] - + def nvmlDeviceGetMemoryInfo(handle): c_memory = c_nvmlMemory_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryInfo") @@ -1160,14 +1160,14 @@ def nvmlDeviceGetBAR1MemoryInfo(handle): ret = fn(handle, byref(c_bar1_memory)) _nvmlCheckReturn(ret) return c_bar1_memory - + def nvmlDeviceGetComputeMode(handle): c_mode = _nvmlComputeMode_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeMode") ret = fn(handle, byref(c_mode)) _nvmlCheckReturn(ret) return c_mode.value - + def nvmlDeviceGetEccMode(handle): c_currState = _nvmlEnableState_t() c_pendingState = _nvmlEnableState_t() @@ -1200,7 +1200,7 @@ def nvmlDeviceGetDetailedEccErrors(handle, errorType, counterType): _nvmlEccCounterType_t(counterType), byref(c_counts)) _nvmlCheckReturn(ret) return c_counts - + # Added in 4.304 def nvmlDeviceGetMemoryErrorCounter(handle, errorType, counterType, locationType): c_count = c_ulonglong() @@ -1212,7 +1212,7 @@ def nvmlDeviceGetMemoryErrorCounter(handle, errorType, counterType, locationType byref(c_count)) _nvmlCheckReturn(ret) return c_count.value - + def nvmlDeviceGetUtilizationRates(handle): c_util = c_nvmlUtilization_t() fn = _nvmlGetFunctionPointer("nvmlDeviceGetUtilizationRates") @@ -1273,7 +1273,7 @@ def nvmlDeviceGetComputeRunningProcesses(handle): c_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeRunningProcesses") ret = fn(handle, byref(c_count), None) - + if (ret == NVML_SUCCESS): # special case, no running processes return [] @@ -1283,11 +1283,11 @@ def nvmlDeviceGetComputeRunningProcesses(handle): c_count.value = c_count.value * 2 + 5 proc_array = c_nvmlProcessInfo_t * c_count.value c_procs = proc_array() - + # make the call again ret = fn(handle, byref(c_count), c_procs) _nvmlCheckReturn(ret) - + procs = [] for i in range(c_count.value): # use an alternative struct for this object @@ -1317,11 +1317,11 @@ def nvmlDeviceGetGraphicsRunningProcesses(handle): c_count.value = c_count.value * 2 + 5 proc_array = c_nvmlProcessInfo_t * c_count.value c_procs = proc_array() - + # make the call again ret = fn(handle, byref(c_count), c_procs) _nvmlCheckReturn(ret) - + procs = [] for i in range(c_count.value): # use an alternative struct for this object @@ -1351,19 +1351,19 @@ def nvmlUnitSetLedState(unit, color): ret = fn(unit, _nvmlLedColor_t(color)) _nvmlCheckReturn(ret) return None - + def nvmlDeviceSetPersistenceMode(handle, mode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetPersistenceMode") ret = fn(handle, _nvmlEnableState_t(mode)) _nvmlCheckReturn(ret) return None - + def nvmlDeviceSetComputeMode(handle, mode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetComputeMode") ret = fn(handle, _nvmlComputeMode_t(mode)) _nvmlCheckReturn(ret) return None - + def nvmlDeviceSetEccMode(handle, mode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetEccMode") ret = fn(handle, _nvmlEnableState_t(mode)) @@ -1381,15 +1381,15 @@ def nvmlDeviceSetDriverModel(handle, model): ret = fn(handle, _nvmlDriverModel_t(model)) _nvmlCheckReturn(ret) return None - -def nvmlDeviceSetAutoBoostedClocksEnabled(handle, enabled): + +def nvmlDeviceSetAutoBoostedClocksEnabled(handle, enabled): fn = _nvmlGetFunctionPointer("nvmlDeviceSetAutoBoostedClocksEnabled") ret = fn(handle, _nvmlEnableState_t(enabled)) _nvmlCheckReturn(ret) return None #Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks -def nvmlDeviceSetDefaultAutoBoostedClocksEnabled(handle, enabled, flags): +def nvmlDeviceSetDefaultAutoBoostedClocksEnabled(handle, enabled, flags): fn = _nvmlGetFunctionPointer("nvmlDeviceSetDefaultAutoBoostedClocksEnabled") ret = fn(handle, _nvmlEnableState_t(enabled), c_uint(flags)) _nvmlCheckReturn(ret) @@ -1402,7 +1402,7 @@ def nvmlDeviceSetApplicationsClocks(handle, maxMemClockMHz, maxGraphicsClockMHz) ret = fn(handle, c_uint(maxMemClockMHz), c_uint(maxGraphicsClockMHz)) _nvmlCheckReturn(ret) return None - + # Added in 4.304 def nvmlDeviceResetApplicationsClocks(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceResetApplicationsClocks") @@ -1416,7 +1416,7 @@ def nvmlDeviceSetPowerManagementLimit(handle, limit): ret = fn(handle, c_uint(limit)) _nvmlCheckReturn(ret) return None - + # Added in 4.304 def nvmlDeviceSetGpuOperationMode(handle, mode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetGpuOperationMode") @@ -1534,7 +1534,7 @@ def nvmlDeviceGetAccountingMode(handle): ret = fn(handle, byref(c_mode)) _nvmlCheckReturn(ret) return c_mode.value - + def nvmlDeviceSetAccountingMode(handle, mode): fn = _nvmlGetFunctionPointer("nvmlDeviceSetAccountingMode") ret = fn(handle, _nvmlEnableState_t(mode)) @@ -1563,7 +1563,7 @@ def nvmlDeviceGetAccountingPids(handle): fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingPids") ret = fn(handle, byref(count), pids) _nvmlCheckReturn(ret) - return map(int, pids[0:count.value]) + return map(int, pids[0:count.value]) def nvmlDeviceGetAccountingBufferSize(handle): bufferSize = c_uint() @@ -1576,10 +1576,10 @@ def nvmlDeviceGetRetiredPages(device, sourceFilter): c_source = _nvmlPageRetirementCause_t(sourceFilter) c_count = c_uint(0) fn = _nvmlGetFunctionPointer("nvmlDeviceGetRetiredPages") - + # First call will get the size ret = fn(device, c_source, byref(c_count), None) - + # this should only fail with insufficient size if ((ret != NVML_SUCCESS) and (ret != NVML_ERROR_INSUFFICIENT_SIZE)): @@ -1651,7 +1651,7 @@ def nvmlDeviceGetViolationStatus(device, perfPolicyType): ret = fn(device, c_perfPolicy_type, byref(c_violTime)) _nvmlCheckReturn(ret) return c_violTime - + def nvmlDeviceGetPcieThroughput(device, counter): c_util = c_uint() fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieThroughput") @@ -1704,17 +1704,17 @@ def nvmlDeviceGetTopologyCommonAncestor(device1, device2): def nvmlDeviceGetCudaComputeCapability(device): c_major = c_int() c_minor = c_int() - + try: fn = _nvmlGetFunctionPointer("nvmlDeviceGetCudaComputeCapability") except: return 9, 9 - + # get the count ret = fn(device, byref(c_major), byref(c_minor)) - + # this should only fail with insufficient size if (ret != NVML_SUCCESS): raise NVMLError(ret) - return c_major.value, c_minor.value \ No newline at end of file + return c_major.value, c_minor.value diff --git a/samples/Sample.py b/samples/Sample.py index ae18400..f9953ed 100644 --- a/samples/Sample.py +++ b/samples/Sample.py @@ -5,17 +5,17 @@ from utils.cv2_utils import * class SampleType(IntEnum): IMAGE = 0 #raw image - + FACE_BEGIN = 1 FACE = 1 #aligned face unsorted FACE_YAW_SORTED = 2 #sorted by yaw FACE_YAW_SORTED_AS_TARGET = 3 #sorted by yaw and included only yaws which exist in TARGET also automatic mirrored FACE_WITH_CLOSE_TO_SELF = 4 FACE_END = 4 - + QTY = 5 - -class Sample(object): + +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): self.sample_type = sample_type if sample_type is not None else SampleType.IMAGE self.filename = filename @@ -26,19 +26,19 @@ class Sample(object): 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): - 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(), - 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, + 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(), + 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, close_target_list=close_target_list if close_target_list is not None else self.close_target_list) - + def load_bgr(self): img = cv2_imread (self.filename).astype(np.float32) / 255.0 if self.mirror: @@ -48,4 +48,4 @@ class Sample(object): def get_random_close_target_sample(self): if self.close_target_list is None: return None - return self.close_target_list[randint (0, len(self.close_target_list)-1)] \ No newline at end of file + return self.close_target_list[randint (0, len(self.close_target_list)-1)] diff --git a/samples/SampleGeneratorBase.py b/samples/SampleGeneratorBase.py index 505fdb7..dec741e 100644 --- a/samples/SampleGeneratorBase.py +++ b/samples/SampleGeneratorBase.py @@ -4,22 +4,21 @@ from pathlib import Path You can implement your own SampleGenerator ''' class SampleGeneratorBase(object): - - + + def __init__ (self, samples_path, debug, batch_size): if samples_path is None: raise Exception('samples_path is None') - + self.samples_path = Path(samples_path) self.debug = debug - self.batch_size = 1 if self.debug else batch_size - + self.batch_size = 1 if self.debug else batch_size + #overridable def __iter__(self): #implement your own iterator return self - + def __next__(self): #implement your own iterator return None - diff --git a/samples/SampleGeneratorFace.py b/samples/SampleGeneratorFace.py index 228421e..ed549cb 100644 --- a/samples/SampleGeneratorFace.py +++ b/samples/SampleGeneratorFace.py @@ -12,9 +12,9 @@ from samples import SampleLoader from samples import SampleGeneratorBase ''' -arg +arg output_sample_types = [ - [SampleProcessor.TypeFlags, size, (optional)random_sub_size] , + [SampleProcessor.TypeFlags, size, (optional)random_sub_size] , ... ] ''' @@ -26,7 +26,7 @@ class SampleGeneratorFace(SampleGeneratorBase): self.add_sample_idx = add_sample_idx self.add_pitch = add_pitch self.add_yaw = add_yaw - + if sort_by_yaw_target_samples_path is not None: self.sample_type = SampleType.FACE_YAW_SORTED_AS_TARGET elif sort_by_yaw: @@ -34,9 +34,9 @@ class SampleGeneratorFace(SampleGeneratorBase): elif with_close_to_self: self.sample_type = SampleType.FACE_WITH_CLOSE_TO_SELF else: - self.sample_type = SampleType.FACE - - self.samples = SampleLoader.load (self.sample_type, self.samples_path, sort_by_yaw_target_samples_path) + self.sample_type = SampleType.FACE + + self.samples = SampleLoader.load (self.sample_type, self.samples_path, sort_by_yaw_target_samples_path) if self.debug: self.generators_count = 1 @@ -46,24 +46,24 @@ class SampleGeneratorFace(SampleGeneratorBase): self.generators = [iter_utils.SubprocessGenerator ( self.batch_func, i ) for i in range(self.generators_count) ] self.generators_sq = [ multiprocessing.Queue() for _ in range(self.generators_count) ] - + self.generator_counter = -1 - + def __iter__(self): return self - + def __next__(self): self.generator_counter += 1 generator = self.generators[self.generator_counter % len(self.generators) ] return next(generator) - + #forces to repeat these sample idxs as fast as possible #currently unused def repeat_sample_idxs(self, idxs): # [ idx, ... ] #send idxs list to all sub generators. for gen_sq in self.generators_sq: - gen_sq.put (idxs) - + gen_sq.put (idxs) + def batch_func(self, generator_id): gen_sq = self.generators_sq[generator_id] samples = self.samples @@ -73,11 +73,11 @@ class SampleGeneratorFace(SampleGeneratorBase): if len(samples_idxs) == 0: raise ValueError('No training data provided.') - + if self.sample_type == SampleType.FACE_YAW_SORTED or self.sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET: if all ( [ samples[idx] == None for idx in samples_idxs] ): raise ValueError('Not enough training data. Gather more faces!') - + if self.sample_type == SampleType.FACE or self.sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF: shuffle_idxs = [] elif self.sample_type == SampleType.FACE_YAW_SORTED or self.sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET: @@ -89,25 +89,25 @@ class SampleGeneratorFace(SampleGeneratorBase): idxs = gen_sq.get() for idx in idxs: if idx in samples_idxs: - repeat_samples_idxs.append(idx) + repeat_samples_idxs.append(idx) batches = None for n_batch in range(self.batch_size): while True: sample = None - + if len(repeat_samples_idxs) > 0: - idx = repeat_samples_idxs.pop() - if self.sample_type == SampleType.FACE or self.sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF: + idx = repeat_samples_idxs.pop() + if self.sample_type == SampleType.FACE or self.sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF: sample = samples[idx] elif self.sample_type == SampleType.FACE_YAW_SORTED or self.sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET: sample = samples[(idx >> 16) & 0xFFFF][idx & 0xFFFF] - else: + else: if self.sample_type == SampleType.FACE or self.sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF: if len(shuffle_idxs) == 0: shuffle_idxs = samples_idxs.copy() np.random.shuffle(shuffle_idxs) - + idx = shuffle_idxs.pop() sample = samples[ idx ] @@ -120,18 +120,18 @@ class SampleGeneratorFace(SampleGeneratorBase): if samples[idx] != None: if len(shuffle_idxs_2D[idx]) == 0: shuffle_idxs_2D[idx] = random.sample( range(len(samples[idx])), len(samples[idx]) ) - - idx2 = shuffle_idxs_2D[idx].pop() + + idx2 = shuffle_idxs_2D[idx].pop() sample = samples[idx][idx2] - + idx = (idx << 16) | (idx2 & 0xFFFF) - - if sample is not None: + + if sample is not None: try: x = SampleProcessor.process (sample, self.sample_process_options, self.output_sample_types, self.debug) except: raise Exception ("Exception occured in sample %s. Error: %s" % (sample.filename, traceback.format_exc() ) ) - + if type(x) != tuple and type(x) != list: raise Exception('SampleProcessor.process returns NOT tuple/list') @@ -144,23 +144,23 @@ class SampleGeneratorFace(SampleGeneratorBase): batches += [ [] ] i_pitch = len(batches)-1 if self.add_yaw: - batches += [ [] ] + batches += [ [] ] i_yaw = len(batches)-1 - + for i in range(len(x)): batches[i].append ( x[i] ) if self.add_sample_idx: batches[i_sample_idx].append (idx) - + if self.add_pitch or self.add_yaw: pitch, yaw = LandmarksProcessor.estimate_pitch_yaw (sample.landmarks) - + if self.add_pitch: batches[i_pitch].append ([pitch]) - + if self.add_yaw: batches[i_yaw].append ([yaw]) - + break yield [ np.array(batch) for batch in batches] diff --git a/samples/SampleGeneratorImageTemporal.py b/samples/SampleGeneratorImageTemporal.py index cd8a92f..8e647d9 100644 --- a/samples/SampleGeneratorImageTemporal.py +++ b/samples/SampleGeneratorImageTemporal.py @@ -11,36 +11,36 @@ from samples import SampleLoader from samples import SampleGeneratorBase ''' -output_sample_types = [ - [SampleProcessor.TypeFlags, size, (optional)random_sub_size] , +output_sample_types = [ + [SampleProcessor.TypeFlags, size, (optional)random_sub_size] , ... ] ''' class SampleGeneratorImageTemporal(SampleGeneratorBase): def __init__ (self, samples_path, debug, batch_size, temporal_image_count, sample_process_options=SampleProcessor.Options(), output_sample_types=[], **kwargs): super().__init__(samples_path, debug, batch_size) - + self.temporal_image_count = temporal_image_count self.sample_process_options = sample_process_options self.output_sample_types = output_sample_types - self.samples = SampleLoader.load (SampleType.IMAGE, self.samples_path) - + self.samples = SampleLoader.load (SampleType.IMAGE, self.samples_path) + self.generator_samples = [ self.samples ] self.generators = [iter_utils.ThisThreadGenerator ( self.batch_func, 0 )] if self.debug else \ [iter_utils.SubprocessGenerator ( self.batch_func, 0 )] - + self.generator_counter = -1 - + def __iter__(self): return self - + def __next__(self): self.generator_counter += 1 generator = self.generators[self.generator_counter % len(self.generators) ] return next(generator) - - def batch_func(self, generator_id): + + def batch_func(self, generator_id): samples = self.generator_samples[generator_id] samples_len = len(samples) if samples_len == 0: @@ -48,20 +48,20 @@ class SampleGeneratorImageTemporal(SampleGeneratorBase): if samples_len - self.temporal_image_count < 0: raise ValueError('Not enough samples to fit temporal line.') - + shuffle_idxs = [] samples_sub_len = samples_len - self.temporal_image_count + 1 - - while True: - + + while True: + batches = None for n_batch in range(self.batch_size): if len(shuffle_idxs) == 0: shuffle_idxs = random.sample( range(samples_sub_len), samples_sub_len ) - + idx = shuffle_idxs.pop() - + temporal_samples = [] for i in range( self.temporal_image_count ): @@ -70,11 +70,11 @@ class SampleGeneratorImageTemporal(SampleGeneratorBase): temporal_samples += SampleProcessor.process (sample, self.sample_process_options, self.output_sample_types, self.debug) except: raise Exception ("Exception occured in sample %s. Error: %s" % (sample.filename, traceback.format_exc() ) ) - + if batches is None: batches = [ [] for _ in range(len(temporal_samples)) ] - + for i in range(len(temporal_samples)): batches[i].append ( temporal_samples[i] ) - + yield [ np.array(batch) for batch in batches] diff --git a/samples/SampleLoader.py b/samples/SampleLoader.py index 9c2966a..083d81c 100644 --- a/samples/SampleLoader.py +++ b/samples/SampleLoader.py @@ -17,44 +17,44 @@ from interact import interact as io class SampleLoader: cache = dict() - + @staticmethod def load(sample_type, samples_path, target_samples_path=None): cache = SampleLoader.cache - + if str(samples_path) not in cache.keys(): cache[str(samples_path)] = [None]*SampleType.QTY - + datas = cache[str(samples_path)] if sample_type == SampleType.IMAGE: - if datas[sample_type] is None: + if datas[sample_type] is None: datas[sample_type] = [ Sample(filename=filename) for filename in io.progress_bar_generator( Path_utils.get_image_paths(samples_path), "Loading") ] elif sample_type == SampleType.FACE: - if datas[sample_type] is None: + if datas[sample_type] is None: datas[sample_type] = SampleLoader.upgradeToFaceSamples( [ Sample(filename=filename) for filename in Path_utils.get_image_paths(samples_path) ] ) - + elif sample_type == SampleType.FACE_YAW_SORTED: if datas[sample_type] is None: datas[sample_type] = SampleLoader.upgradeToFaceYawSortedSamples( SampleLoader.load(SampleType.FACE, samples_path) ) - - elif sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET: + + elif sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET: if datas[sample_type] is None: if target_samples_path is None: raise Exception('target_samples_path is None for FACE_YAW_SORTED_AS_TARGET') datas[sample_type] = SampleLoader.upgradeToFaceYawSortedAsTargetSamples( SampleLoader.load(SampleType.FACE_YAW_SORTED, samples_path), SampleLoader.load(SampleType.FACE_YAW_SORTED, target_samples_path) ) - elif sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF: + elif sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF: if datas[sample_type] is None: datas[sample_type] = SampleLoader.upgradeToFaceCloseToSelfSamples( SampleLoader.load(SampleType.FACE, samples_path) ) - - + + return datas[sample_type] - + @staticmethod def upgradeToFaceSamples ( samples ): sample_list = [] - + for s in io.progress_bar_generator(samples, "Loading"): s_filename_path = Path(s.filename) try: @@ -64,57 +64,57 @@ class SampleLoader: dflimg = DFLJPG.load ( str(s_filename_path) ) else: dflimg = None - + if dflimg is None: - print ("%s is not a dfl image file required for training" % (s_filename_path.name) ) + print ("%s is not a dfl image file required for training" % (s_filename_path.name) ) continue - + pitch, yaw = LandmarksProcessor.estimate_pitch_yaw ( dflimg.get_landmarks() ) sample_list.append( s.copy_and_set(sample_type=SampleType.FACE, face_type=FaceType.fromString (dflimg.get_face_type()), - shape=dflimg.get_shape(), + shape=dflimg.get_shape(), landmarks=dflimg.get_landmarks(), pitch=pitch, yaw=yaw) ) except: print ("Unable to load %s , error: %s" % (str(s_filename_path), traceback.format_exc() ) ) - - return sample_list - + + return sample_list + @staticmethod def upgradeToFaceCloseToSelfSamples (samples): yaw_samples = SampleLoader.upgradeToFaceYawSortedSamples(samples) yaw_samples_len = len(yaw_samples) - + sample_list = [] for i in io.progress_bar_generator( range(yaw_samples_len), "Sorting"): if yaw_samples[i] is not None: for s in yaw_samples[i]: s_t = [] - - for n in range(2000): + + for n in range(2000): yaw_idx = np.clip ( i-10 +np.random.randint(20), 0, yaw_samples_len-1 ) if yaw_samples[yaw_idx] is None: continue - + yaw_idx_samples_len = len(yaw_samples[yaw_idx]) - + yaw_idx_sample = yaw_samples[yaw_idx][ np.random.randint(yaw_idx_samples_len) ] if s.filename == yaw_idx_sample.filename: continue - + s_t.append ( yaw_idx_sample ) if len(s_t) >= 50: break - + if len(s_t) == 0: s_t = [s] - + sample_list.append( s.copy_and_set(close_target_list = s_t) ) return sample_list - + @staticmethod def upgradeToFaceYawSortedSamples( samples ): @@ -123,50 +123,50 @@ class SampleLoader: diff_rot_per_grad = abs(highest_yaw-lowest_yaw) / gradations yaws_sample_list = [None]*gradations - + for i in io.progress_bar_generator(range(gradations), "Sorting"): yaw = lowest_yaw + i*diff_rot_per_grad next_yaw = lowest_yaw + (i+1)*diff_rot_per_grad yaw_samples = [] - for s in samples: + for s in samples: s_yaw = s.yaw if (i == 0 and s_yaw < next_yaw) or \ (i < gradations-1 and s_yaw >= yaw and s_yaw < next_yaw) or \ (i == gradations-1 and s_yaw >= yaw): yaw_samples.append ( s.copy_and_set(sample_type=SampleType.FACE_YAW_SORTED) ) - + if len(yaw_samples) > 0: yaws_sample_list[i] = yaw_samples - + return yaws_sample_list - + @staticmethod def upgradeToFaceYawSortedAsTargetSamples (s, t): l = len(s) if l != len(t): raise Exception('upgradeToFaceYawSortedAsTargetSamples() s_len != t_len') b = l // 2 - + s_idxs = np.argwhere ( np.array ( [ 1 if x != None else 0 for x in s] ) == 1 )[:,0] t_idxs = np.argwhere ( np.array ( [ 1 if x != None else 0 for x in t] ) == 1 )[:,0] - - new_s = [None]*l - + + new_s = [None]*l + for t_idx in t_idxs: - search_idxs = [] + search_idxs = [] for i in range(0,l): search_idxs += [t_idx - i, (l-t_idx-1) - i, t_idx + i, (l-t_idx-1) + i] - for search_idx in search_idxs: + for search_idx in search_idxs: if search_idx in s_idxs: mirrored = ( t_idx != search_idx and ((t_idx < b and search_idx >= b) or (search_idx < b and t_idx >= b)) ) new_s[t_idx] = [ sample.copy_and_set(sample_type=SampleType.FACE_YAW_SORTED_AS_TARGET, - mirror=True, - yaw=-sample.yaw, + mirror=True, + yaw=-sample.yaw, landmarks=LandmarksProcessor.mirror_landmarks (sample.landmarks, sample.shape[1] )) - for sample in s[search_idx] - ] if mirrored else s[search_idx] + for sample in s[search_idx] + ] if mirrored else s[search_idx] break - - return new_s \ No newline at end of file + + return new_s diff --git a/samples/SampleProcessor.py b/samples/SampleProcessor.py index 93fe1c6..f10f5e9 100644 --- a/samples/SampleProcessor.py +++ b/samples/SampleProcessor.py @@ -13,61 +13,61 @@ class SampleProcessor(object): WARPED_TRANSFORMED = 0x00000004, TRANSFORMED = 0x00000008, LANDMARKS_ARRAY = 0x00000010, #currently unused - + RANDOM_CLOSE = 0x00000020, MORPH_TO_RANDOM_CLOSE = 0x00000040, - + FACE_ALIGN_HALF = 0x00000100, FACE_ALIGN_FULL = 0x00000200, FACE_ALIGN_HEAD = 0x00000400, - FACE_ALIGN_AVATAR = 0x00000800, - + FACE_ALIGN_AVATAR = 0x00000800, + FACE_MASK_FULL = 0x00001000, FACE_MASK_EYES = 0x00002000, - + MODE_BGR = 0x01000000, #BGR MODE_G = 0x02000000, #Grayscale - MODE_GGG = 0x04000000, #3xGrayscale + MODE_GGG = 0x04000000, #3xGrayscale MODE_M = 0x08000000, #mask only MODE_BGR_SHUFFLE = 0x10000000, #BGR shuffle - - class Options(object): + + class Options(object): def __init__(self, random_flip = True, normalize_tanh = False, rotation_range=[-10,10], scale_range=[-0.05, 0.05], tx_range=[-0.05, 0.05], ty_range=[-0.05, 0.05]): - self.random_flip = random_flip + self.random_flip = random_flip self.normalize_tanh = normalize_tanh self.rotation_range = rotation_range self.scale_range = scale_range self.tx_range = tx_range - self.ty_range = ty_range - + self.ty_range = ty_range + @staticmethod def process (sample, sample_process_options, output_sample_types, debug): sample_bgr = sample.load_bgr() h,w,c = sample_bgr.shape - is_face_sample = sample.landmarks is not None - + is_face_sample = sample.landmarks is not None + if debug and is_face_sample: LandmarksProcessor.draw_landmarks (sample_bgr, sample.landmarks, (0, 1, 0)) - + close_sample = sample.close_target_list[ np.random.randint(0, len(sample.close_target_list)) ] if sample.close_target_list is not None else None close_sample_bgr = close_sample.load_bgr() if close_sample is not None else None - + if debug and close_sample_bgr is not None: - LandmarksProcessor.draw_landmarks (close_sample_bgr, close_sample.landmarks, (0, 1, 0)) - + LandmarksProcessor.draw_landmarks (close_sample_bgr, close_sample.landmarks, (0, 1, 0)) + params = image_utils.gen_warp_params(sample_bgr, sample_process_options.random_flip, rotation_range=sample_process_options.rotation_range, scale_range=sample_process_options.scale_range, tx_range=sample_process_options.tx_range, ty_range=sample_process_options.ty_range ) images = [[None]*3 for _ in range(30)] - + sample_rnd_seed = np.random.randint(0x80000000) - - outputs = [] + + outputs = [] for sample_type in output_sample_types: f = sample_type[0] size = sample_type[1] random_sub_size = 0 if len (sample_type) < 3 else min( sample_type[2] , size) - + if f & SampleProcessor.TypeFlags.SOURCE != 0: img_type = 0 elif f & SampleProcessor.TypeFlags.WARPED != 0: @@ -77,53 +77,53 @@ class SampleProcessor(object): elif f & SampleProcessor.TypeFlags.TRANSFORMED != 0: img_type = 3 elif f & SampleProcessor.TypeFlags.LANDMARKS_ARRAY != 0: - img_type = 4 + img_type = 4 else: raise ValueError ('expected SampleTypeFlags type') - + if f & SampleProcessor.TypeFlags.RANDOM_CLOSE != 0: img_type += 10 elif f & SampleProcessor.TypeFlags.MORPH_TO_RANDOM_CLOSE != 0: img_type += 20 - + face_mask_type = 0 if f & SampleProcessor.TypeFlags.FACE_MASK_FULL != 0: - face_mask_type = 1 + face_mask_type = 1 elif f & SampleProcessor.TypeFlags.FACE_MASK_EYES != 0: face_mask_type = 2 - + target_face_type = -1 if f & SampleProcessor.TypeFlags.FACE_ALIGN_HALF != 0: - target_face_type = FaceType.HALF + target_face_type = FaceType.HALF elif f & SampleProcessor.TypeFlags.FACE_ALIGN_FULL != 0: target_face_type = FaceType.FULL elif f & SampleProcessor.TypeFlags.FACE_ALIGN_HEAD != 0: target_face_type = FaceType.HEAD elif f & SampleProcessor.TypeFlags.FACE_ALIGN_AVATAR != 0: target_face_type = FaceType.AVATAR - + if img_type == 4: - l = sample.landmarks + l = sample.landmarks l = np.concatenate ( [ np.expand_dims(l[:,0] / w,-1), np.expand_dims(l[:,1] / h,-1) ], -1 ) l = np.clip(l, 0.0, 1.0) img = l - else: + else: if images[img_type][face_mask_type] is None: if img_type >= 10 and img_type <= 19: #RANDOM_CLOSE img_type -= 10 img = close_sample_bgr cur_sample = close_sample - + elif img_type >= 20 and img_type <= 29: #MORPH_TO_RANDOM_CLOSE img_type -= 20 res = sample.shape[0] - - s_landmarks = sample.landmarks.copy() - d_landmarks = close_sample.landmarks.copy() - idxs = list(range(len(s_landmarks))) + + s_landmarks = sample.landmarks.copy() + d_landmarks = close_sample.landmarks.copy() + idxs = list(range(len(s_landmarks))) #remove landmarks near boundaries for i in idxs[:]: - s_l = s_landmarks[i] + s_l = s_landmarks[i] d_l = d_landmarks[i] if s_l[0] < 5 or s_l[1] < 5 or s_l[0] >= res-5 or s_l[1] >= res-5 or \ d_l[0] < 5 or d_l[1] < 5 or d_l[0] >= res-5 or d_l[1] >= res-5: @@ -139,39 +139,39 @@ class SampleProcessor(object): diff_l = np.abs(s_l - s_l_2) if np.sqrt(diff_l.dot(diff_l)) < 5: idxs.remove(i) - break + break s_landmarks = s_landmarks[idxs] d_landmarks = d_landmarks[idxs] - s_landmarks = np.concatenate ( [s_landmarks, [ [0,0], [ res // 2, 0], [ res-1, 0], [0, res//2], [res-1, res//2] ,[0,res-1] ,[res//2, res-1] ,[res-1,res-1] ] ] ) + s_landmarks = np.concatenate ( [s_landmarks, [ [0,0], [ res // 2, 0], [ res-1, 0], [0, res//2], [res-1, res//2] ,[0,res-1] ,[res//2, res-1] ,[res-1,res-1] ] ] ) d_landmarks = np.concatenate ( [d_landmarks, [ [0,0], [ res // 2, 0], [ res-1, 0], [0, res//2], [res-1, res//2] ,[0,res-1] ,[res//2, res-1] ,[res-1,res-1] ] ] ) img = image_utils.morph_by_points (sample_bgr, s_landmarks, d_landmarks) cur_sample = close_sample else: img = sample_bgr cur_sample = sample - + 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) ), -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) mask[mask > 0.0] = 1.0 - img = np.concatenate( (img, mask ), -1 ) + img = np.concatenate( (img, mask ), -1 ) images[img_type][face_mask_type] = image_utils.warp_by_params (params, img, (img_type==1 or img_type==2), (img_type==2 or img_type==3), img_type != 0, face_mask_type == 0) - + img = images[img_type][face_mask_type] - + if is_face_sample and target_face_type != -1: if target_face_type > sample.face_type: raise Exception ('sample %s type %s does not match model requirement %s. Consider extract necessary type of faces.' % (sample.filename, sample.face_type, target_face_type) ) img = cv2.warpAffine( img, LandmarksProcessor.get_transform_mat (sample.landmarks, size, target_face_type), (size,size), flags=cv2.INTER_CUBIC ) else: img = cv2.resize( img, (size,size), cv2.INTER_CUBIC ) - + if random_sub_size != 0: - sub_size = size - random_sub_size + sub_size = size - random_sub_size rnd_state = np.random.RandomState (sample_rnd_seed+random_sub_size) start_x = rnd_state.randint(sub_size+1) start_y = rnd_state.randint(sub_size+1) @@ -195,7 +195,7 @@ class SampleProcessor(object): img = img_mask else: raise ValueError ('expected SampleTypeFlags mode') - + if not debug: if sample_process_options.normalize_tanh: img = np.clip (img * 2.0 - 1.0, -1.0, 1.0) @@ -213,6 +213,6 @@ class SampleProcessor(object): elif output.shape[2] == 4: result += [output[...,0:3]*output[...,3:4],] - return result + return result else: - return outputs \ No newline at end of file + return outputs diff --git a/samples/__init__.py b/samples/__init__.py index ed99faa..98d3b05 100644 --- a/samples/__init__.py +++ b/samples/__init__.py @@ -4,4 +4,4 @@ from .SampleLoader import SampleLoader from .SampleProcessor import SampleProcessor from .SampleGeneratorBase import SampleGeneratorBase from .SampleGeneratorFace import SampleGeneratorFace -from .SampleGeneratorImageTemporal import SampleGeneratorImageTemporal \ No newline at end of file +from .SampleGeneratorImageTemporal import SampleGeneratorImageTemporal diff --git a/utils/DFLJPG.py b/utils/DFLJPG.py index efd4370..362cf5d 100644 --- a/utils/DFLJPG.py +++ b/utils/DFLJPG.py @@ -11,7 +11,7 @@ class DFLJPG(object): self.chunks = [] self.dfl_dict = None self.shape = (0,0,0) - + @staticmethod def load_raw(filename): try: @@ -19,7 +19,7 @@ class DFLJPG(object): data = f.read() except: raise FileNotFoundError(data) - + try: inst = DFLJPG() inst.data = data @@ -30,23 +30,23 @@ class DFLJPG(object): while data_counter < inst_length: chunk_m_l, chunk_m_h = struct.unpack ("BB", data[data_counter:data_counter+2]) data_counter += 2 - + if chunk_m_l != 0xFF: raise ValueError("No Valid JPG info") - + chunk_name = None chunk_size = None chunk_data = None chunk_ex_data = None is_unk_chunk = False - - if chunk_m_h & 0xF0 == 0xD0: + + if chunk_m_h & 0xF0 == 0xD0: n = chunk_m_h & 0x0F - - if n >= 0 and n <= 7: + + if n >= 0 and n <= 7: chunk_name = "RST%d" % (n) chunk_size = 0 - elif n == 0x8: + elif n == 0x8: chunk_name = "SOI" chunk_size = 0 if len(chunks) != 0: @@ -54,73 +54,73 @@ class DFLJPG(object): elif n == 0x9: chunk_name = "EOI" chunk_size = 0 - elif n == 0xA: - chunk_name = "SOS" - elif n == 0xB: + elif n == 0xA: + chunk_name = "SOS" + elif n == 0xB: chunk_name = "DQT" elif n == 0xD: chunk_name = "DRI" chunk_size = 2 else: is_unk_chunk = True - elif chunk_m_h & 0xF0 == 0xC0: - n = chunk_m_h & 0x0F - if n == 0: + elif chunk_m_h & 0xF0 == 0xC0: + n = chunk_m_h & 0x0F + if n == 0: chunk_name = "SOF0" - elif n == 2: + elif n == 2: chunk_name = "SOF2" - elif n == 4: + elif n == 4: chunk_name = "DHT" else: is_unk_chunk = True - elif chunk_m_h & 0xF0 == 0xE0: + elif chunk_m_h & 0xF0 == 0xE0: n = chunk_m_h & 0x0F chunk_name = "APP%d" % (n) else: is_unk_chunk = True - + if is_unk_chunk: - raise ValueError("Unknown chunk %X" % (chunk_m_h) ) - + raise ValueError("Unknown chunk %X" % (chunk_m_h) ) + if chunk_size == None: #variable size chunk_size, = struct.unpack (">H", data[data_counter:data_counter+2]) chunk_size -= 2 data_counter += 2 - + if chunk_size > 0: chunk_data = data[data_counter:data_counter+chunk_size] data_counter += chunk_size - + if chunk_name == "SOS": - c = data_counter + c = data_counter while c < inst_length and (data[c] != 0xFF or data[c+1] != 0xD9): c += 1 - + chunk_ex_data = data[data_counter:c] data_counter = c - + chunks.append ({'name' : chunk_name, 'm_h' : chunk_m_h, 'data' : chunk_data, 'ex_data' : chunk_ex_data, - }) + }) inst.chunks = chunks - + return inst except Exception as e: raise Exception ("Corrupted JPG file: %s" % (str(e))) - + @staticmethod def load(filename): try: inst = DFLJPG.load_raw (filename) inst.dfl_dict = None - + for chunk in inst.chunks: if chunk['name'] == 'APP0': d, c = chunk['data'], 0 c, id, _ = struct_unpack (d, c, "=4sB") - + if id == b"JFIF": c, ver_major, ver_minor, units, Xdensity, Ydensity, Xthumbnail, Ythumbnail = struct_unpack (d, c, "=BBBHHBB") #if units == 0: @@ -131,22 +131,22 @@ class DFLJPG(object): d, c = chunk['data'], 0 c, precision, height, width = struct_unpack (d, c, ">BHH") inst.shape = (height, width, 3) - + elif chunk['name'] == 'APP15': if type(chunk['data']) == bytes: inst.dfl_dict = pickle.loads(chunk['data']) if (inst.dfl_dict is not None) and ('face_type' not in inst.dfl_dict.keys()): inst.dfl_dict['face_type'] = FaceType.toString (FaceType.FULL) - + if inst.dfl_dict == None: return None - + return inst except Exception as e: print (e) return None - + @staticmethod def embed_data(filename, face_type=None, landmarks=None, @@ -155,7 +155,7 @@ class DFLJPG(object): source_landmarks=None, image_to_face_mat=None ): - + inst = DFLJPG.load_raw (filename) inst.setDFLDictData ({ 'face_type': face_type, @@ -165,41 +165,41 @@ class DFLJPG(object): 'source_landmarks': source_landmarks, 'image_to_face_mat': image_to_face_mat }) - + try: with open(filename, "wb") as f: f.write ( inst.dump() ) except: raise Exception( 'cannot save %s' % (filename) ) - + def dump(self): data = b"" - + for chunk in self.chunks: data += struct.pack ("BB", 0xFF, chunk['m_h'] ) chunk_data = chunk['data'] if chunk_data is not None: data += struct.pack (">H", len(chunk_data)+2 ) data += chunk_data - + chunk_ex_data = chunk['ex_data'] - if chunk_ex_data is not None: + if chunk_ex_data is not None: data += chunk_ex_data return data - - def get_shape(self): + + def get_shape(self): return self.shape - + def get_height(self): for chunk in self.chunks: if type(chunk) == IHDR: return chunk.height return 0 - + def getDFLDictData(self): return self.dfl_dict - + def setDFLDictData (self, dict_data=None): self.dfl_dict = dict_data @@ -211,17 +211,17 @@ class DFLJPG(object): last_app_chunk = 0 for i, chunk in enumerate (self.chunks): if chunk['m_h'] & 0xF0 == 0xE0: - last_app_chunk = i - + last_app_chunk = i + dflchunk = {'name' : 'APP15', 'm_h' : 0xEF, 'data' : pickle.dumps(dict_data), 'ex_data' : None, } self.chunks.insert (last_app_chunk+1, dflchunk) - + def get_face_type(self): return self.dfl_dict['face_type'] - def get_landmarks(self): return np.array ( self.dfl_dict['landmarks'] ) - def get_source_filename(self): return self.dfl_dict['source_filename'] - def get_source_rect(self): return self.dfl_dict['source_rect'] + def get_landmarks(self): return np.array ( self.dfl_dict['landmarks'] ) + 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'] ) diff --git a/utils/DFLPNG.py b/utils/DFLPNG.py index 604e44b..5cee5ce 100644 --- a/utils/DFLPNG.py +++ b/utils/DFLPNG.py @@ -110,7 +110,7 @@ class Chunk(object): def __str__(self): return "".format(**self.__dict__) - + class IHDR(Chunk): """IHDR Chunk width, height, bit_depth, color_type, compression_method, @@ -189,24 +189,24 @@ class IEND(Chunk): class DFLChunk(Chunk): def __init__(self, dict_data=None): super().__init__("fcWp") - self.dict_data = dict_data + self.dict_data = dict_data def setDictData(self, dict_data): self.dict_data = dict_data - + def getDictData(self): return self.dict_data - + @classmethod def load(cls, data): inst = super().load(data) - inst.dict_data = pickle.loads( inst.data ) + inst.dict_data = pickle.loads( inst.data ) return inst - + def dump(self): self.data = pickle.dumps (self.dict_data) return super().dump() - + chunk_map = { b"IHDR": IHDR, b"fcWp": DFLChunk, @@ -219,7 +219,7 @@ class DFLPNG(object): self.length = 0 self.chunks = [] self.fcwp_dict = None - + @staticmethod def load_raw(filename): try: @@ -227,11 +227,11 @@ class DFLPNG(object): data = f.read() except: raise FileNotFoundError(data) - + inst = DFLPNG() inst.data = data inst.length = len(data) - + if data[0:8] != PNG_HEADER: msg = "No Valid PNG header" raise ValueError(msg) @@ -244,26 +244,26 @@ class DFLPNG(object): chunk = chunk_map.get(chunk_name, Chunk).load(data[chunk_start:chunk_end]) inst.chunks.append(chunk) chunk_start = chunk_end - + return inst - + @staticmethod def load(filename): try: inst = DFLPNG.load_raw (filename) inst.fcwp_dict = inst.getDFLDictData() - + if (inst.fcwp_dict is not None) and ('face_type' not in inst.fcwp_dict.keys()): inst.fcwp_dict['face_type'] = FaceType.toString (FaceType.FULL) - + if inst.fcwp_dict == None: return None - + return inst except Exception as e: print(e) return None - + @staticmethod def embed_data(filename, face_type=None, landmarks=None, @@ -271,7 +271,7 @@ class DFLPNG(object): source_rect=None, source_landmarks=None ): - + inst = DFLPNG.load_raw (filename) inst.setDFLDictData ({ 'face_type': face_type, @@ -280,7 +280,7 @@ class DFLPNG(object): 'source_rect': source_rect, 'source_landmarks': source_landmarks }) - + try: with open(filename, "wb") as f: f.write ( inst.dump() ) @@ -292,7 +292,7 @@ class DFLPNG(object): for chunk in self.chunks: data += chunk.dump() return data - + def get_shape(self): for chunk in self.chunks: if type(chunk) == IHDR: @@ -301,34 +301,34 @@ class DFLPNG(object): h = chunk.height return (h,w,c) return (0,0,0) - + def get_height(self): for chunk in self.chunks: if type(chunk) == IHDR: return chunk.height return 0 - - def getDFLDictData(self): + + def getDFLDictData(self): for chunk in self.chunks: if type(chunk) == DFLChunk: return chunk.getDictData() return None - + def setDFLDictData (self, dict_data=None): for chunk in self.chunks: if type(chunk) == DFLChunk: self.chunks.remove(chunk) break - + if not dict_data is None: chunk = DFLChunk(dict_data) self.chunks.insert(-1, chunk) - - def get_face_type(self): return self.fcwp_dict['face_type'] + + def get_face_type(self): return self.fcwp_dict['face_type'] def get_landmarks(self): return np.array ( self.fcwp_dict['landmarks'] ) 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 __str__(self): return "".format(len(self.chunks), **self.__dict__) diff --git a/utils/Path_utils.py b/utils/Path_utils.py index 14a181d..70abc9a 100644 --- a/utils/Path_utils.py +++ b/utils/Path_utils.py @@ -5,8 +5,8 @@ image_extensions = [".jpg", ".jpeg", ".png", ".tif", ".tiff"] def get_image_paths(dir_path, image_extensions=image_extensions): dir_path = Path (dir_path) - - result = [] + + result = [] if dir_path.exists(): for x in list(scandir(str(dir_path))): if any([x.name.lower().endswith(ext) for ext in image_extensions]): @@ -14,25 +14,25 @@ def get_image_paths(dir_path, image_extensions=image_extensions): return result def get_image_unique_filestem_paths(dir_path, verbose_print_func=None): - result = get_image_paths(dir_path) - result_dup = set() - + result = get_image_paths(dir_path) + result_dup = set() + for f in result[:]: f_stem = Path(f).stem - if f_stem in result_dup: + if f_stem in result_dup: result.remove(f) if verbose_print_func is not None: - verbose_print_func ("Duplicate filenames are not allowed, skipping: %s" % Path(f).name ) - continue + verbose_print_func ("Duplicate filenames are not allowed, skipping: %s" % Path(f).name ) + continue result_dup.add(f_stem) - + return result - + def get_all_dir_names_startswith (dir_path, startswith): dir_path = Path (dir_path) startswith = startswith.lower() - - result = [] + + result = [] if dir_path.exists(): for x in list(scandir(str(dir_path))): if x.name.lower().startswith(startswith): @@ -42,7 +42,7 @@ def get_all_dir_names_startswith (dir_path, startswith): def get_first_file_by_stem (dir_path, stem, exts=None): dir_path = Path (dir_path) stem = stem.lower() - + if dir_path.exists(): for x in list(scandir(str(dir_path))): if not x.is_file(): @@ -50,5 +50,5 @@ def get_first_file_by_stem (dir_path, stem, exts=None): xp = Path(x.path) if xp.stem.lower() == stem and (exts is None or xp.suffix.lower() in exts): return xp - - return None \ No newline at end of file + + return None diff --git a/utils/cv2_utils.py b/utils/cv2_utils.py index f4412ee..ff8d82a 100644 --- a/utils/cv2_utils.py +++ b/utils/cv2_utils.py @@ -11,7 +11,7 @@ def cv2_imread(filename, flags=cv2.IMREAD_UNCHANGED): return cv2.imdecode(numpyarray, flags) except: return None - + def cv2_imwrite(filename, img, *args): ret, buf = cv2.imencode( Path(filename).suffix, img, *args) if ret == True: @@ -19,4 +19,4 @@ def cv2_imwrite(filename, img, *args): with open(filename, "wb") as stream: stream.write( buf ) except: - pass \ No newline at end of file + pass diff --git a/utils/image_utils.py b/utils/image_utils.py index 98e6e4a..1afffde 100644 --- a/utils/image_utils.py +++ b/utils/image_utils.py @@ -21,7 +21,7 @@ def reinhard_color_transfer(target, source, clip=False, preserve_paper=False, so OpenCV image in BGR color space (the source image) target: NumPy array OpenCV image in BGR color space (the target image) - clip: Should components of L*a*b* image be scaled by np.clip before + clip: Should components of L*a*b* image be scaled by np.clip before converting back to BGR color space? If False then components will be min-max scaled appropriately. Clipping will keep target image brightness truer to the input. @@ -32,7 +32,7 @@ def reinhard_color_transfer(target, source, clip=False, preserve_paper=False, so aesthetically pleasing results. If False then L*a*b* components will scaled using the reciprocal of the scaling factor proposed in the paper. This method seems to produce - more consistently aesthetically pleasing results + more consistently aesthetically pleasing results Returns: ------- @@ -40,13 +40,13 @@ def reinhard_color_transfer(target, source, clip=False, preserve_paper=False, so OpenCV image (w, h, 3) NumPy array (uint8) """ - + # convert the images from the RGB to L*ab* color space, being # sure to utilizing the floating point data type (note: OpenCV # expects floats to be 32-bit, so use that instead of 64-bit) source = cv2.cvtColor(source, cv2.COLOR_BGR2LAB).astype(np.float32) - target = cv2.cvtColor(target, cv2.COLOR_BGR2LAB).astype(np.float32) - + target = cv2.cvtColor(target, cv2.COLOR_BGR2LAB).astype(np.float32) + # compute color statistics for the source and target images src_input = source if source_mask is None else source*source_mask tgt_input = target if target_mask is None else target*target_mask @@ -86,7 +86,7 @@ def reinhard_color_transfer(target, source, clip=False, preserve_paper=False, so # type transfer = cv2.merge([l, a, b]) transfer = cv2.cvtColor(transfer.astype(np.uint8), cv2.COLOR_LAB2BGR) - + # return the color transferred image return transfer @@ -127,7 +127,7 @@ def linear_color_transfer(target_img, source_img, mode='pca', eps=1e-5): matched_img[matched_img>1] = 1 matched_img[matched_img<0] = 0 return matched_img - + def lab_image_stats(image): # compute the mean and standard deviation of each channel (l, a, b) = cv2.split(image) @@ -137,7 +137,7 @@ def lab_image_stats(image): # return the color statistics return (lMean, lStd, aMean, aStd, bMean, bStd) - + def _scale_array(arr, clip=True): if clip: return np.clip(arr, 0, 255) @@ -145,12 +145,12 @@ def _scale_array(arr, clip=True): mn = arr.min() mx = arr.max() scale_range = (max([mn, 0]), min([mx, 255])) - + if mn < scale_range[0] or mx > scale_range[1]: return (scale_range[1] - scale_range[0]) * (arr - mn) / (mx - mn) + scale_range[0] return arr - + def channel_hist_match(source, template, hist_match_threshold=255, mask=None): # Code borrowed from: # https://stackoverflow.com/questions/32655686/histogram-matching-of-two-images-in-python-2-x @@ -179,22 +179,22 @@ def channel_hist_match(source, template, hist_match_threshold=255, mask=None): t_quantiles = 255 * t_quantiles / t_quantiles[-1] interp_t_values = np.interp(s_quantiles, t_quantiles, t_values) - return interp_t_values[bin_idx].reshape(oldshape) + return interp_t_values[bin_idx].reshape(oldshape) def color_hist_match(src_im, tar_im, hist_match_threshold=255): h,w,c = src_im.shape matched_R = channel_hist_match(src_im[:,:,0], tar_im[:,:,0], hist_match_threshold, None) matched_G = channel_hist_match(src_im[:,:,1], tar_im[:,:,1], hist_match_threshold, None) matched_B = channel_hist_match(src_im[:,:,2], tar_im[:,:,2], hist_match_threshold, None) - + to_stack = (matched_R, matched_G, matched_B) for i in range(3, c): to_stack += ( src_im[:,:,i],) - - + + matched = np.stack(to_stack, axis=-1).astype(src_im.dtype) return matched - + pil_fonts = {} def _get_pil_font (font, size): @@ -204,65 +204,65 @@ def _get_pil_font (font, size): if font_str_id not in pil_fonts.keys(): pil_fonts[font_str_id] = ImageFont.truetype(font + ".ttf", size=size, encoding="unic") pil_font = pil_fonts[font_str_id] - return pil_font + return pil_font except: return ImageFont.load_default() - + def get_text_image( shape, text, color=(1,1,1), border=0.2, font=None): - try: + try: size = shape[1] pil_font = _get_pil_font( localization.get_default_ttf_font_name() , size) text_width, text_height = pil_font.getsize(text) - + canvas = Image.new('RGB', shape[0:2], (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: + if shape[2] != 3: result = np.concatenate ( (result, np.ones ( (shape[1],) + (shape[0],) + (shape[2]-3,)) ), axis=2 ) return result - except: + except: return np.zeros ( (shape[1], shape[0], shape[2]), dtype=np.float32 ) - + def draw_text( image, rect, text, color=(1,1,1), border=0.2, font=None): h,w,c = image.shape - + l,t,r,b = rect l = np.clip (l, 0, w-1) r = np.clip (r, 0, w-1) t = np.clip (t, 0, h-1) b = np.clip (b, 0, h-1) - + image[t:b, l:r] += get_text_image ( (r-l,b-t,c) , text, color, border, font ) - + def draw_text_lines (image, rect, text_lines, color=(1,1,1), border=0.2, font=None): text_lines_len = len(text_lines) if text_lines_len == 0: return - + l,t,r,b = rect h = b-t h_per_line = h // text_lines_len - + for i in range(0, text_lines_len): draw_text (image, (l, i*h_per_line, r, (i+1)*h_per_line), text_lines[i], color, border, font) - + def get_draw_text_lines ( image, rect, text_lines, color=(1,1,1), border=0.2, font=None): image = np.zeros ( image.shape, dtype=np.float ) draw_text_lines ( image, rect, text_lines, color, border, font) return image - - + + def draw_polygon (image, points, color, thickness = 1): points_len = len(points) for i in range (0, points_len): p0 = tuple( points[i] ) p1 = tuple( points[ (i+1) % points_len] ) cv2.line (image, p0, p1, color, thickness=thickness) - + def draw_rect(image, rect, color, thickness=1): l,t,r,b = rect draw_polygon (image, [ (l,t), (r,t), (r,b), (l,b ) ], color, thickness) @@ -272,40 +272,40 @@ def rectContains(rect, point) : def applyAffineTransform(src, srcTri, dstTri, size) : warpMat = cv2.getAffineTransform( np.float32(srcTri), np.float32(dstTri) ) - return cv2.warpAffine( src, warpMat, (size[0], size[1]), None, flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT_101 ) - -def morphTriangle(dst_img, src_img, st, dt) : + return cv2.warpAffine( src, warpMat, (size[0], size[1]), None, flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT_101 ) + +def morphTriangle(dst_img, src_img, st, dt) : (h,w,c) = dst_img.shape sr = np.array( cv2.boundingRect(np.float32(st)) ) dr = np.array( cv2.boundingRect(np.float32(dt)) ) sRect = st - sr[0:2] dRect = dt - dr[0:2] d_mask = np.zeros((dr[3], dr[2], c), dtype = np.float32) - cv2.fillConvexPoly(d_mask, np.int32(dRect), (1.0,)*c, 8, 0); - imgRect = src_img[sr[1]:sr[1] + sr[3], sr[0]:sr[0] + sr[2]] - size = (dr[2], dr[3]) - warpImage1 = applyAffineTransform(imgRect, sRect, dRect, size) + cv2.fillConvexPoly(d_mask, np.int32(dRect), (1.0,)*c, 8, 0); + imgRect = src_img[sr[1]:sr[1] + sr[3], sr[0]:sr[0] + sr[2]] + size = (dr[2], dr[3]) + warpImage1 = applyAffineTransform(imgRect, sRect, dRect, size) if c == 1: warpImage1 = np.expand_dims( warpImage1, -1 ) dst_img[dr[1]:dr[1]+dr[3], dr[0]:dr[0]+dr[2]] = dst_img[dr[1]:dr[1]+dr[3], dr[0]:dr[0]+dr[2]]*(1-d_mask) + warpImage1 * d_mask - + def morph_by_points (image, sp, dp): if sp.shape != dp.shape: raise ValueError ('morph_by_points() sp.shape != dp.shape') - (h,w,c) = image.shape + (h,w,c) = image.shape result_image = np.zeros(image.shape, dtype = image.dtype) - for tri in Delaunay(dp).simplices: + for tri in Delaunay(dp).simplices: morphTriangle(result_image, image, sp[tri], dp[tri]) - + return result_image - + def equalize_and_stack_square (images, axis=1): max_c = max ([ 1 if len(image.shape) == 2 else image.shape[2] for image in images ] ) - + target_wh = 99999 for i,image in enumerate(images): if len(image.shape) == 2: @@ -313,113 +313,112 @@ def equalize_and_stack_square (images, axis=1): c = 1 else: h,w,c = image.shape - + if h < target_wh: target_wh = h - + if w < target_wh: target_wh = w - + for i,image in enumerate(images): if len(image.shape) == 2: h,w = image.shape c = 1 else: h,w,c = image.shape - + if c < max_c: if c == 1: if len(image.shape) == 2: - image = np.expand_dims ( image, -1 ) + image = np.expand_dims ( image, -1 ) image = np.concatenate ( (image,)*max_c, -1 ) elif c == 2: #GA image = np.expand_dims ( image[...,0], -1 ) - image = np.concatenate ( (image,)*max_c, -1 ) + image = np.concatenate ( (image,)*max_c, -1 ) else: image = np.concatenate ( (image, np.ones((h,w,max_c - c))), -1 ) if h != target_wh or w != target_wh: image = cv2.resize ( image, (target_wh, target_wh) ) h,w,c = image.shape - + images[i] = image - + return np.concatenate ( images, axis = 1 ) - -def bgr2hsv (img): + +def bgr2hsv (img): return cv2.cvtColor(img, cv2.COLOR_BGR2HSV) - -def hsv2bgr (img): + +def hsv2bgr (img): return cv2.cvtColor(img, cv2.COLOR_HSV2BGR) - -def bgra2hsva (img): + +def bgra2hsva (img): return np.concatenate ( (cv2.cvtColor(img[...,0:3], cv2.COLOR_BGR2HSV ), np.expand_dims (img[...,3], -1)), -1 ) def bgra2hsva_list (imgs): return [ bgra2hsva(img) for img in imgs ] - + def hsva2bgra (img): return np.concatenate ( (cv2.cvtColor(img[...,0:3], cv2.COLOR_HSV2BGR ), np.expand_dims (img[...,3], -1)), -1 ) def hsva2bgra_list (imgs): return [ hsva2bgra(img) for img in imgs ] - + def gen_warp_params (source, flip, rotation_range=[-10,10], scale_range=[-0.5, 0.5], tx_range=[-0.05, 0.05], ty_range=[-0.05, 0.05] ): h,w,c = source.shape if (h != w) or (w != 64 and w != 128 and w != 256 and w != 512 and w != 1024): raise ValueError ('TrainingDataGenerator accepts only square power of 2 images.') - + rotation = np.random.uniform( rotation_range[0], rotation_range[1] ) scale = np.random.uniform(1 +scale_range[0], 1 +scale_range[1]) tx = np.random.uniform( tx_range[0], tx_range[1] ) - ty = np.random.uniform( ty_range[0], ty_range[1] ) - + ty = np.random.uniform( ty_range[0], ty_range[1] ) + #random warp by grid cell_size = [ w // (2**i) for i in range(1,4) ] [ np.random.randint(3) ] cell_count = w // cell_size + 1 - + grid_points = np.linspace( 0, w, cell_count) mapx = np.broadcast_to(grid_points, (cell_count, cell_count)).copy() mapy = mapx.T - + mapx[1:-1,1:-1] = mapx[1:-1,1:-1] + random_utils.random_normal( size=(cell_count-2, cell_count-2) )*(cell_size*0.24) mapy[1:-1,1:-1] = mapy[1:-1,1:-1] + random_utils.random_normal( size=(cell_count-2, cell_count-2) )*(cell_size*0.24) half_cell_size = cell_size // 2 - + mapx = cv2.resize(mapx, (w+cell_size,)*2 )[half_cell_size:-half_cell_size-1,half_cell_size:-half_cell_size-1].astype(np.float32) mapy = cv2.resize(mapy, (w+cell_size,)*2 )[half_cell_size:-half_cell_size-1,half_cell_size:-half_cell_size-1].astype(np.float32) - + #random transform random_transform_mat = cv2.getRotationMatrix2D((w // 2, w // 2), rotation, scale) random_transform_mat[:, 2] += (tx*w, ty*w) - + params = dict() params['mapx'] = mapx params['mapy'] = mapy params['rmat'] = random_transform_mat - params['w'] = w + params['w'] = w params['flip'] = flip and np.random.randint(10) < 4 - + return params - + def warp_by_params (params, img, warp, transform, flip, is_border_replicate): if warp: img = cv2.remap(img, params['mapx'], params['mapy'], cv2.INTER_CUBIC ) if transform: - img = cv2.warpAffine( img, params['rmat'], (params['w'], params['w']), borderMode=(cv2.BORDER_REPLICATE if is_border_replicate else cv2.BORDER_CONSTANT), flags=cv2.INTER_CUBIC ) + img = cv2.warpAffine( img, params['rmat'], (params['w'], params['w']), borderMode=(cv2.BORDER_REPLICATE if is_border_replicate else cv2.BORDER_CONSTANT), flags=cv2.INTER_CUBIC ) if flip and params['flip']: img = img[:,::-1,:] return img - + #n_colors = [0..256] def reduce_colors (img_bgr, n_colors): img_rgb = (cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) * 255.0).astype(np.uint8) img_rgb_pil = Image.fromarray(img_rgb) img_rgb_pil_p = img_rgb_pil.convert('P', palette=Image.ADAPTIVE, colors=n_colors) - + img_rgb_p = img_rgb_pil_p.convert('RGB') img_bgr = cv2.cvtColor( np.array(img_rgb_p, dtype=np.float32) / 255.0, cv2.COLOR_RGB2BGR ) - + return img_bgr - \ No newline at end of file diff --git a/utils/iter_utils.py b/utils/iter_utils.py index a1a55e6..b197fc0 100644 --- a/utils/iter_utils.py +++ b/utils/iter_utils.py @@ -5,7 +5,7 @@ import time class ThisThreadGenerator(object): - def __init__(self, generator_func, user_param=None): + def __init__(self, generator_func, user_param=None): super().__init__() self.generator_func = generator_func self.user_param = user_param @@ -13,30 +13,30 @@ class ThisThreadGenerator(object): def __iter__(self): return self - + def __next__(self): - if not self.initialized: + if not self.initialized: self.initialized = True self.generator_func = self.generator_func(self.user_param) return next(self.generator_func) class SubprocessGenerator(object): - def __init__(self, generator_func, user_param=None, prefetch=2): - super().__init__() + def __init__(self, generator_func, user_param=None, prefetch=2): + super().__init__() self.prefetch = prefetch self.generator_func = generator_func self.user_param = user_param self.sc_queue = multiprocessing.Queue() self.cs_queue = multiprocessing.Queue() self.p = None - + def process_func(self): self.generator_func = self.generator_func(self.user_param) - while True: + while True: while self.prefetch > -1: try: - gen_data = next (self.generator_func) + gen_data = next (self.generator_func) except StopIteration: self.cs_queue.put (None) return @@ -47,17 +47,17 @@ class SubprocessGenerator(object): def __iter__(self): return self - + def __next__(self): if self.p == None: self.p = multiprocessing.Process(target=self.process_func, args=()) self.p.daemon = True self.p.start() - + gen_data = self.cs_queue.get() if gen_data is None: self.p.terminate() self.p.join() raise StopIteration() - self.sc_queue.put (1) - return gen_data \ No newline at end of file + self.sc_queue.put (1) + return gen_data diff --git a/utils/os_utils.py b/utils/os_utils.py index 9f77121..ff1bd98 100644 --- a/utils/os_utils.py +++ b/utils/os_utils.py @@ -4,12 +4,12 @@ import sys if sys.platform[0:3] == 'win': from ctypes import windll from ctypes import wintypes - + def set_process_lowest_prio(): try: if sys.platform[0:3] == 'win': GetCurrentProcess = windll.kernel32.GetCurrentProcess - GetCurrentProcess.restype = wintypes.HANDLE + GetCurrentProcess.restype = wintypes.HANDLE SetPriorityClass = windll.kernel32.SetPriorityClass SetPriorityClass.argtypes = (wintypes.HANDLE, wintypes.DWORD) SetPriorityClass ( GetCurrentProcess(), 0x00000040 ) @@ -19,7 +19,7 @@ def set_process_lowest_prio(): os.nice(20) except: print("Unable to set lowest process priority") - + def set_process_dpi_aware(): if sys.platform[0:3] == 'win': - windll.user32.SetProcessDPIAware(True) \ No newline at end of file + windll.user32.SetProcessDPIAware(True) diff --git a/utils/random_utils.py b/utils/random_utils.py index 0bbc6f7..7b3af6e 100644 --- a/utils/random_utils.py +++ b/utils/random_utils.py @@ -3,12 +3,12 @@ import numpy as np def random_normal( size=(1,), trunc_val = 2.5 ): len = np.array(size).prod() result = np.empty ( (len,) , dtype=np.float32) - + for i in range (len): while True: x = np.random.normal() if x >= -trunc_val and x <= trunc_val: break result[i] = (x / trunc_val) - - return result.reshape ( size ) \ No newline at end of file + + return result.reshape ( size ) diff --git a/utils/std_utils.py b/utils/std_utils.py index 81af906..2f23be9 100644 --- a/utils/std_utils.py +++ b/utils/std_utils.py @@ -11,26 +11,26 @@ class suppress_stdout_stderr(object): self.old_stdout_fileno = os.dup ( sys.stdout.fileno() ) self.old_stderr_fileno = os.dup ( sys.stderr.fileno() ) - + self.old_stdout = sys.stdout self.old_stderr = sys.stderr - + os.dup2 ( self.outnull_file.fileno(), self.old_stdout_fileno_undup ) os.dup2 ( self.errnull_file.fileno(), self.old_stderr_fileno_undup ) - - sys.stdout = self.outnull_file + + sys.stdout = self.outnull_file sys.stderr = self.errnull_file return self - - def __exit__(self, *_): + + def __exit__(self, *_): sys.stdout = self.old_stdout sys.stderr = self.old_stderr - + os.dup2 ( self.old_stdout_fileno, self.old_stdout_fileno_undup ) os.dup2 ( self.old_stderr_fileno, self.old_stderr_fileno_undup ) os.close ( self.old_stdout_fileno ) os.close ( self.old_stderr_fileno ) - + self.outnull_file.close() - self.errnull_file.close() \ No newline at end of file + self.errnull_file.close() diff --git a/utils/struct_utils.py b/utils/struct_utils.py index b12bacc..cc63559 100644 --- a/utils/struct_utils.py +++ b/utils/struct_utils.py @@ -1,6 +1,5 @@ import struct -def struct_unpack(data, counter, fmt): +def struct_unpack(data, counter, fmt): fmt_size = struct.calcsize(fmt) return (counter+fmt_size,) + struct.unpack (fmt, data[counter:counter+fmt_size]) - \ No newline at end of file