manual extractor: increased FPS,

sort by final : now you can specify target number of images,
converter: fix seamless mask and exception,
huge refactoring
This commit is contained in:
iperov 2019-02-28 11:56:31 +04:00
parent 7db469a1da
commit 438213e97c
30 changed files with 1834 additions and 1718 deletions

View file

@ -3,21 +3,15 @@ import copy
You can implement your own Converter, check example ConverterMasked.py You can implement your own Converter, check example ConverterMasked.py
''' '''
class ConverterBase(object): class Converter(object):
MODE_FACE = 0 TYPE_FACE = 0 #calls convert_face
MODE_IMAGE = 1 TYPE_IMAGE = 1 #calls convert_image without landmarks
MODE_IMAGE_WITH_LANDMARKS = 2 TYPE_IMAGE_WITH_LANDMARKS = 2 #calls convert_image with landmarks
#overridable #overridable
def __init__(self, predictor): def __init__(self, predictor_func, type):
self.predictor = predictor self.predictor_func = predictor_func
self.type = type
#overridable
def get_mode(self):
#MODE_FACE calls convert_face
#MODE_IMAGE calls convert_image without landmarks
#MODE_IMAGE_WITH_LANDMARKS calls convert_image with landmarks
return ConverterBase.MODE_FACE
#overridable #overridable
def convert_face (self, img_bgr, img_face_landmarks, debug): def convert_face (self, img_bgr, img_face_landmarks, debug):
@ -40,7 +34,7 @@ class ConverterBase(object):
def copy(self): def copy(self):
return copy.copy(self) return copy.copy(self)
def copy_and_set_predictor(self, predictor): def copy_and_set_predictor(self, predictor_func):
result = self.copy() result = self.copy()
result.predictor = predictor result.predictor_func = predictor_func
return result return result

View file

@ -1,4 +1,4 @@
from models import ConverterBase from .Converter import Converter
from facelib import LandmarksProcessor from facelib import LandmarksProcessor
from facelib import FaceType from facelib import FaceType
@ -7,38 +7,33 @@ import numpy as np
from utils import image_utils from utils import image_utils
''' '''
predictor: predictor_func:
input: [predictor_input_size, predictor_input_size, BGR] input: [predictor_input_size, predictor_input_size, BGR]
output: [predictor_input_size, predictor_input_size, BGR] output: [predictor_input_size, predictor_input_size, BGR]
''' '''
class ConverterImage(ConverterBase): class ConverterImage(Converter):
#override #override
def __init__(self, predictor, def __init__(self, predictor_func,
predictor_input_size=0, predictor_input_size=0,
output_size=0, output_size=0):
**in_options):
super().__init__(predictor) super().__init__(predictor_func, Converter.TYPE_IMAGE)
self.predictor_input_size = predictor_input_size self.predictor_input_size = predictor_input_size
self.output_size = output_size self.output_size = output_size
#override
def get_mode(self):
return ConverterBase.MODE_IMAGE
#override #override
def dummy_predict(self): def dummy_predict(self):
self.predictor ( np.zeros ( (self.predictor_input_size, self.predictor_input_size,3), dtype=np.float32) ) self.predictor_func ( np.zeros ( (self.predictor_input_size, self.predictor_input_size,3), dtype=np.float32) )
#override #override
def convert_image (self, img_bgr, img_landmarks, debug): def convert_image (self, img_bgr, img_landmarks, debug):
img_size = img_bgr.shape[1], img_bgr.shape[0] img_size = img_bgr.shape[1], img_bgr.shape[0]
predictor_input_bgr = cv2.resize ( img_bgr, (self.predictor_input_size, self.predictor_input_size), cv2.INTER_LANCZOS4 ) predictor_input_bgr = cv2.resize ( img_bgr, (self.predictor_input_size, self.predictor_input_size), cv2.INTER_LANCZOS4 )
predicted_bgr = self.predictor ( predictor_input_bgr ) predicted_bgr = self.predictor_func ( predictor_input_bgr )
output = cv2.resize ( predicted_bgr, (self.output_size, self.output_size), cv2.INTER_LANCZOS4 ) output = cv2.resize ( predicted_bgr, (self.output_size, self.output_size), cv2.INTER_LANCZOS4 )
if debug: if debug:

View file

@ -1,10 +1,10 @@
from models import ConverterBase from .Converter import Converter
from facelib import LandmarksProcessor from facelib import LandmarksProcessor
from facelib import FaceType from facelib import FaceType
import cv2 import cv2
import numpy as np import numpy as np
from utils import image_utils from utils import image_utils
from utils.console_utils import * from interact import interact as io
''' '''
default_mode = {1:'overlay', default_mode = {1:'overlay',
@ -14,10 +14,10 @@ default_mode = {1:'overlay',
5:'seamless-hist-match', 5:'seamless-hist-match',
6:'raw'} 6:'raw'}
''' '''
class ConverterMasked(ConverterBase): class ConverterMasked(Converter):
#override #override
def __init__(self, predictor, def __init__(self, predictor_func,
predictor_input_size=0, predictor_input_size=0,
output_size=0, output_size=0,
face_type=FaceType.FULL, face_type=FaceType.FULL,
@ -26,16 +26,15 @@ class ConverterMasked(ConverterBase):
base_blur_mask_modifier = 0, base_blur_mask_modifier = 0,
default_erode_mask_modifier = 0, default_erode_mask_modifier = 0,
default_blur_mask_modifier = 0, default_blur_mask_modifier = 0,
clip_hborder_mask_per = 0, clip_hborder_mask_per = 0):
**in_options):
super().__init__(predictor) super().__init__(predictor_func, Converter.TYPE_FACE)
self.predictor_input_size = predictor_input_size self.predictor_input_size = predictor_input_size
self.output_size = output_size self.output_size = output_size
self.face_type = face_type self.face_type = face_type
self.clip_hborder_mask_per = clip_hborder_mask_per self.clip_hborder_mask_per = clip_hborder_mask_per
mode = input_int ("Choose mode: (1) overlay, (2) hist match, (3) hist match bw, (4) seamless, (5) seamless hist match, (6) raw. Default - %d : " % (default_mode) , default_mode) mode = io.input_int ("Choose mode: (1) overlay, (2) hist match, (3) hist match bw, (4) seamless, (5) seamless hist match, (6) raw. Default - %d : " % (default_mode) , default_mode)
mode_dict = {1:'overlay', mode_dict = {1:'overlay',
2:'hist-match', 2:'hist-match',
@ -47,7 +46,7 @@ class ConverterMasked(ConverterBase):
self.mode = mode_dict.get (mode, mode_dict[default_mode] ) self.mode = mode_dict.get (mode, mode_dict[default_mode] )
if self.mode == 'raw': if self.mode == 'raw':
mode = input_int ("Choose raw mode: (1) rgb, (2) rgb+mask (default), (3) mask only, (4) predicted only : ", 2) mode = io.input_int ("Choose raw mode: (1) rgb, (2) rgb+mask (default), (3) mask only, (4) predicted only : ", 2)
self.raw_mode = {1:'rgb', self.raw_mode = {1:'rgb',
2:'rgb-mask', 2:'rgb-mask',
3:'mask-only', 3:'mask-only',
@ -55,37 +54,33 @@ class ConverterMasked(ConverterBase):
if self.mode != 'raw': if self.mode != 'raw':
if self.mode == 'hist-match' or self.mode == 'hist-match-bw': if self.mode == 'hist-match' or self.mode == 'hist-match-bw':
self.masked_hist_match = input_bool("Masked hist match? (y/n skip:y) : ", True) self.masked_hist_match = io.input_bool("Masked hist match? (y/n skip:y) : ", True)
if self.mode == 'hist-match' or self.mode == 'hist-match-bw' or self.mode == 'seamless-hist-match': if self.mode == 'hist-match' or self.mode == 'hist-match-bw' or self.mode == 'seamless-hist-match':
self.hist_match_threshold = np.clip ( input_int("Hist match threshold [0..255] (skip:255) : ", 255), 0, 255) self.hist_match_threshold = np.clip ( io.input_int("Hist match threshold [0..255] (skip:255) : ", 255), 0, 255)
self.use_predicted_mask = input_bool("Use predicted mask? (y/n skip:y) : ", True) self.use_predicted_mask = io.input_bool("Use predicted mask? (y/n skip:y) : ", True)
if self.mode != 'raw': if self.mode != 'raw':
self.erode_mask_modifier = base_erode_mask_modifier + np.clip ( input_int ("Choose erode mask modifier [-200..200] (skip:%d) : " % (default_erode_mask_modifier), default_erode_mask_modifier), -200, 200) self.erode_mask_modifier = base_erode_mask_modifier + np.clip ( io.input_int ("Choose erode mask modifier [-200..200] (skip:%d) : " % (default_erode_mask_modifier), default_erode_mask_modifier), -200, 200)
self.blur_mask_modifier = base_blur_mask_modifier + np.clip ( input_int ("Choose blur mask modifier [-200..200] (skip:%d) : " % (default_blur_mask_modifier), default_blur_mask_modifier), -200, 200) self.blur_mask_modifier = base_blur_mask_modifier + np.clip ( io.input_int ("Choose blur mask modifier [-200..200] (skip:%d) : " % (default_blur_mask_modifier), default_blur_mask_modifier), -200, 200)
self.seamless_erode_mask_modifier = 0 self.seamless_erode_mask_modifier = 0
if self.mode == 'seamless' or self.mode == 'seamless-hist-match': if self.mode == 'seamless' or self.mode == 'seamless-hist-match':
self.seamless_erode_mask_modifier = np.clip ( input_int ("Choose seamless erode mask modifier [-100..100] (skip:0) : ", 0), -100, 100) self.seamless_erode_mask_modifier = np.clip ( io.input_int ("Choose seamless erode mask modifier [-100..100] (skip:0) : ", 0), -100, 100)
self.output_face_scale = np.clip ( 1.0 + input_int ("Choose output face scale modifier [-50..50] (skip:0) : ", 0)*0.01, 0.5, 1.5) self.output_face_scale = np.clip ( 1.0 + io.input_int ("Choose output face scale modifier [-50..50] (skip:0) : ", 0)*0.01, 0.5, 1.5)
self.color_transfer_mode = input_str ("Apply color transfer to predicted face? Choose mode ( rct/lct skip:None ) : ", None, ['rct','lct']) self.color_transfer_mode = io.input_str ("Apply color transfer to predicted face? Choose mode ( rct/lct skip:None ) : ", None, ['rct','lct'])
if self.mode != 'raw': if self.mode != 'raw':
self.final_image_color_degrade_power = np.clip ( input_int ("Degrade color power of final image [0..100] (skip:0) : ", 0), 0, 100) self.final_image_color_degrade_power = np.clip ( io.input_int ("Degrade color power of final image [0..100] (skip:0) : ", 0), 0, 100)
self.alpha = input_bool("Export png with alpha channel? (y/n skip:n) : ", False) self.alpha = io.input_bool("Export png with alpha channel? (y/n skip:n) : ", False)
print ("") io.log_info ("")
#override
def get_mode(self):
return ConverterBase.MODE_FACE
#override #override
def dummy_predict(self): def dummy_predict(self):
self.predictor ( np.zeros ( (self.predictor_input_size,self.predictor_input_size,4), dtype=np.float32 ) ) self.predictor_func ( np.zeros ( (self.predictor_input_size,self.predictor_input_size,4), dtype=np.float32 ) )
#override #override
def convert_face (self, img_bgr, img_face_landmarks, debug): def convert_face (self, img_bgr, img_face_landmarks, debug):
@ -106,7 +101,7 @@ class ConverterMasked(ConverterBase):
predictor_input_mask_a_0 = cv2.resize (dst_face_mask_a_0, (self.predictor_input_size,self.predictor_input_size)) predictor_input_mask_a_0 = cv2.resize (dst_face_mask_a_0, (self.predictor_input_size,self.predictor_input_size))
predictor_input_mask_a = np.expand_dims (predictor_input_mask_a_0, -1) predictor_input_mask_a = np.expand_dims (predictor_input_mask_a_0, -1)
predicted_bgra = self.predictor ( np.concatenate( (predictor_input_bgr, predictor_input_mask_a), -1) ) predicted_bgra = self.predictor_func ( np.concatenate( (predictor_input_bgr, predictor_input_mask_a), -1) )
prd_face_bgr = np.clip (predicted_bgra[:,:,0:3], 0, 1.0 ) prd_face_bgr = np.clip (predicted_bgra[:,:,0:3], 0, 1.0 )
prd_face_mask_a_0 = np.clip (predicted_bgra[:,:,3], 0.0, 1.0) prd_face_mask_a_0 = np.clip (predicted_bgra[:,:,3], 0.0, 1.0)
@ -129,10 +124,12 @@ class ConverterMasked(ConverterBase):
img_face_mask_aaa [ img_face_mask_aaa <= 0.1 ] = 0.0 img_face_mask_aaa [ img_face_mask_aaa <= 0.1 ] = 0.0
img_face_mask_flatten_aaa = img_face_mask_aaa.copy() if self.mode == 'seamless' or self.mode == 'seamless-hist-match':
img_face_mask_flatten_aaa[img_face_mask_flatten_aaa > 0.9] = 1.0 img_face_seamless_mask_aaa = img_face_mask_aaa.copy()
img_face_seamless_mask_aaa[img_face_seamless_mask_aaa > 0.9] = 1.0
img_face_seamless_mask_aaa[img_face_seamless_mask_aaa <= 0.9] = 0.0
maxregion = np.argwhere(img_face_mask_flatten_aaa==1.0) maxregion = np.argwhere(img_face_mask_aaa > 0.9)
out_img = img_bgr.copy() out_img = img_bgr.copy()
@ -155,24 +152,24 @@ class ConverterMasked(ConverterBase):
maxy,maxx = maxregion.max(axis=0)[:2] maxy,maxx = maxregion.max(axis=0)[:2]
if debug: if debug:
print ("maxregion.size: %d, minx:%d, maxx:%d miny:%d, maxy:%d" % (maxregion.size, minx, maxx, miny, maxy ) ) io.log_info ("maxregion.size: %d, minx:%d, maxx:%d miny:%d, maxy:%d" % (maxregion.size, minx, maxx, miny, maxy ) )
lenx = maxx - minx lenx = maxx - minx
leny = maxy - miny leny = maxy - miny
if lenx >= 4 and leny >= 4: if lenx >= 4 and leny >= 4:
masky = int(minx+(lenx//2)) maskx = int(minx+(lenx//2))
maskx = int(miny+(leny//2)) masky = int(miny+(leny//2))
lowest_len = min (lenx, leny) lowest_len = min (lenx, leny)
if debug: if debug:
print ("lowest_len = %f" % (lowest_len) ) io.log_info ("lowest_len = %f" % (lowest_len) )
img_mask_blurry_aaa = img_face_mask_aaa img_mask_blurry_aaa = img_face_mask_aaa
if self.erode_mask_modifier != 0: if self.erode_mask_modifier != 0:
ero = int( lowest_len * ( 0.126 - lowest_len * 0.00004551365 ) * 0.01*self.erode_mask_modifier ) ero = int( lowest_len * ( 0.126 - lowest_len * 0.00004551365 ) * 0.01*self.erode_mask_modifier )
if debug: if debug:
print ("erode_size = %d" % (ero) ) io.log_info ("erode_size = %d" % (ero) )
if ero > 0: if ero > 0:
img_mask_blurry_aaa = cv2.erode(img_mask_blurry_aaa, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(ero,ero)), iterations = 1 ) img_mask_blurry_aaa = cv2.erode(img_mask_blurry_aaa, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(ero,ero)), iterations = 1 )
elif ero < 0: elif ero < 0:
@ -181,11 +178,12 @@ class ConverterMasked(ConverterBase):
if self.seamless_erode_mask_modifier != 0: if self.seamless_erode_mask_modifier != 0:
ero = int( lowest_len * ( 0.126 - lowest_len * 0.00004551365 ) * 0.01*self.seamless_erode_mask_modifier ) ero = int( lowest_len * ( 0.126 - lowest_len * 0.00004551365 ) * 0.01*self.seamless_erode_mask_modifier )
if debug: if debug:
print ("seamless_erode_size = %d" % (ero) ) io.log_info ("seamless_erode_size = %d" % (ero) )
if ero > 0: if ero > 0:
img_face_mask_flatten_aaa = cv2.erode(img_face_mask_flatten_aaa, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(ero,ero)), iterations = 1 ) img_face_seamless_mask_aaa = cv2.erode(img_face_seamless_mask_aaa, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(ero,ero)), iterations = 1 )
elif ero < 0: elif ero < 0:
img_face_mask_flatten_aaa = cv2.dilate(img_face_mask_flatten_aaa, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(-ero,-ero)), iterations = 1 ) img_face_seamless_mask_aaa = cv2.dilate(img_face_seamless_mask_aaa, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(-ero,-ero)), iterations = 1 )
img_face_seamless_mask_aaa = np.clip (img_face_seamless_mask_aaa, 0, 1)
if self.clip_hborder_mask_per > 0: #clip hborder before blur if self.clip_hborder_mask_per > 0: #clip hborder before blur
prd_hborder_rect_mask_a = np.ones ( prd_face_mask_a.shape, dtype=np.float32) prd_hborder_rect_mask_a = np.ones ( prd_face_mask_a.shape, dtype=np.float32)
@ -205,7 +203,7 @@ class ConverterMasked(ConverterBase):
if self.blur_mask_modifier > 0: if self.blur_mask_modifier > 0:
blur = int( lowest_len * 0.10 * 0.01*self.blur_mask_modifier ) blur = int( lowest_len * 0.10 * 0.01*self.blur_mask_modifier )
if debug: if debug:
print ("blur_size = %d" % (blur) ) io.log_info ("blur_size = %d" % (blur) )
if blur > 0: if blur > 0:
img_mask_blurry_aaa = cv2.blur(img_mask_blurry_aaa, (blur, blur) ) img_mask_blurry_aaa = cv2.blur(img_mask_blurry_aaa, (blur, blur) )
@ -276,8 +274,12 @@ class ConverterMasked(ConverterBase):
if debug: if debug:
debugs += [out_img.copy()] debugs += [out_img.copy()]
out_img = cv2.seamlessClone( (out_img*255).astype(np.uint8), (img_bgr*255).astype(np.uint8), (img_face_mask_flatten_aaa*255).astype(np.uint8), (masky,maskx) , cv2.NORMAL_CLONE ) try:
out_img = cv2.seamlessClone( (out_img*255).astype(np.uint8), (img_bgr*255).astype(np.uint8), (img_face_seamless_mask_aaa*255).astype(np.uint8), (maskx,masky) , cv2.NORMAL_CLONE )
out_img = out_img.astype(dtype=np.float32) / 255.0 out_img = out_img.astype(dtype=np.float32) / 255.0
except:
#seamlessClone may fail in some cases
pass
if debug: if debug:
debugs += [out_img.copy()] debugs += [out_img.copy()]

3
converters/__init__.py Normal file
View file

@ -0,0 +1,3 @@
from .Converter import Converter
from .ConverterMasked import ConverterMasked
from .ConverterImage import ConverterImage

View file

@ -3,10 +3,6 @@ import os
import cv2 import cv2
from pathlib import Path from pathlib import Path
from utils import std_utils
def transform(point, center, scale, resolution): def transform(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 h = 200.0 * scale
@ -123,7 +119,6 @@ class LandmarksExtractor(object):
image = crop(input_image, center, scale).transpose ( (2,0,1) ).astype(np.float32) / 255.0 image = crop(input_image, center, scale).transpose ( (2,0,1) ).astype(np.float32) / 255.0
image = np.expand_dims(image, 0) image = np.expand_dims(image, 0)
with std_utils.suppress_stdout_stderr():
predicted = self.keras_model.predict (image) predicted = self.keras_model.predict (image)
pts_img = get_pts_from_predict ( predicted[-1][0], center, scale) pts_img = get_pts_from_predict ( predicted[-1][0], center, scale)

View file

@ -134,7 +134,7 @@ def get_transform_mat (image_landmarks, output_size, face_type, scale=1.0):
elif face_type == FaceType.HEAD: elif face_type == FaceType.HEAD:
padding = (output_size / 64) * 24 padding = (output_size / 64) * 24
else: else:
raise ValueError ('wrong face_type') raise ValueError ('wrong face_type: ', face_type)
mat = umeyama(image_landmarks[17:], landmarks_2D, True)[0:2] mat = umeyama(image_landmarks[17:], landmarks_2D, True)[0:2]
mat = mat * (output_size - 2 * padding) mat = mat * (output_size - 2 * padding)

1
interact/__init__.py Normal file
View file

@ -0,0 +1 @@
from .interact import interact

236
interact/interact.py Normal file
View file

@ -0,0 +1,236 @@
import os
import sys
import time
import types
import multiprocessing
import cv2
from tqdm import tqdm
class Interact(object):
EVENT_LBUTTONDOWN = 1
EVENT_MOUSEWHEEL = 2
def __init__(self):
self.named_windows = {}
self.capture_mouse_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)
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()
self.named_windows = {}
self.capture_mouse_windows = {}
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)
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
if event == cv2.EVENT_LBUTTONDOWN: ev = Interact.EVENT_LBUTTONDOWN
elif event == cv2.EVENT_MOUSEWHEEL: ev = Interact.EVENT_MOUSEWHEEL
else: ev = 0
inst.add_mouse_event (wnd_name, x, y, ev, flags)
if wnd_name in self.named_windows:
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.")
def capture_keys(self, wnd_name):
if wnd_name in self.named_windows:
if wnd_name not in self.capture_keys_windows:
self.capture_keys_windows[wnd_name] = True
else: print("capture_keys: already set for window ", wnd_name)
else: print("capture_keys: named_window ", wnd_name, " not found.")
def progress_bar(self, desc, total, leave=True):
if self.pg_bar is None:
self.pg_bar = tqdm( total=total, desc=desc, leave=leave, ascii=True )
else: print("progress_bar: already set.")
def progress_bar_inc(self, c):
if self.pg_bar is not None:
self.pg_bar.n += c
self.pg_bar.refresh()
else: print("progress_bar not set.")
def progress_bar_close(self):
if self.pg_bar is not None:
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 = int(sleep_time / 1000) if sleep_time != 0 else 1
key = cv2.waitKey(1) & 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:
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:
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] = []
return ar
def get_key_events(self, wnd_name):
ar = self.key_events.get(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
return i
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
return i
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:
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:
inp = input (str)
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))
p.start()
t = time.time()
inp = False
while True:
if not sq.empty():
inp = sq.get()
break
if time.time() - t > max_time_sec:
break
p.terminate()
sys.stdin = os.fdopen( sys.stdin.fileno() )
return inp
interact = Interact()

View file

@ -0,0 +1,42 @@
import time
import multiprocessing
class SubprocessFunctionCaller(object):
class CliFunction(object):
def __init__(self, s2c, c2s, lock):
self.s2c = s2c
self.c2s = c2s
self.lock = lock
def __call__(self, value):
self.lock.acquire()
self.c2s.put (value)
while True:
if not self.s2c.empty():
obj = self.s2c.get()
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

255
joblib/SubprocessorBase.py Normal file
View file

@ -0,0 +1,255 @@
import traceback
import multiprocessing
import time
import sys
from interact import interact as io
class Subprocessor(object):
class Cli(object):
def __init__ ( self, client_dict ):
self.s2c = multiprocessing.Queue()
self.c2s = multiprocessing.Queue()
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
try:
self.on_initialize(client_dict)
c2s.put ( {'op': 'init_ok'} )
while True:
msg = s2c.get()
op = msg.get('op','')
if op == 'data':
data = msg['data']
result = self.process_data (data)
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'} )
except Exception as e:
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()) )
c2s.put ( {'op': 'error', 'data' : data} )
#overridable
def __init__(self, name, SubprocessorCli_class, no_response_time_sec = 60):
if not issubclass(SubprocessorCli_class, Subprocessor.Cli):
raise ValueError("SubprocessorCli_class must be subclass of Subprocessor.Cli")
self.name = name
self.SubprocessorCli_class = SubprocessorCli_class
self.no_response_time_sec = no_response_time_sec
#overridable
def process_info_generator(self):
#yield per process (name, host_dict, client_dict)
raise NotImplementedError
#overridable optional
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
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
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():
try:
cli = self.SubprocessorCli_class(client_dict)
cli.state = 1
cli.sent_time = time.time()
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','')
if op == 'init_ok':
cli.state = 0
elif op == 'error':
cli.kill()
self.clis.remove(cli)
break
if all ([cli.state == 0 for cli in self.clis]):
break
io.process_messages(0.005)
if len(self.clis) == 0:
raise Exception ( "Unable to start subprocesses." )
#ok some processes survived, initialize host logic
self.on_clients_initialized()
#main loop of data processing
while True:
for cli in self.clis[:]:
while not cli.c2s.empty():
obj = cli.c2s.get()
op = obj.get('op','')
if op == 'success':
#success processed data, return data and result to on_result
self.on_result (cli.host_dict, obj['data'], obj['result'])
self.sent_data = None
cli.state = 0
elif op == 'error':
#some error occured while process data, returning chunk to on_data_return
if 'data' in obj.keys():
self.on_data_return (cli.host_dict, obj['data'] )
#and killing process
cli.kill()
self.clis.remove(cli)
elif op == 'log_info':
io.log_info(obj['msg'])
elif op == 'log_err':
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
data = self.get_data(cli.host_dict)
if data is not None:
#and send it to subprocess
cli.s2c.put ( {'op': 'data', 'data' : data} )
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
print ( '%s doesnt response, terminating it.' % (cli.name) )
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
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']
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:
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()

2
joblib/__init__.py Normal file
View file

@ -0,0 +1,2 @@
from .SubprocessorBase import Subprocessor
from .SubprocessFunctionCaller import SubprocessFunctionCaller

90
main.py
View file

@ -1,5 +1,6 @@
import os import os
import sys import sys
import time
import argparse import argparse
from utils import Path_utils from utils import Path_utils
from utils import os_utils from utils import os_utils
@ -12,34 +13,25 @@ class fixPathAction(argparse.Action):
def __call__(self, parser, namespace, values, option_string=None): def __call__(self, parser, namespace, values, option_string=None):
setattr(namespace, self.dest, os.path.abspath(os.path.expanduser(values))) setattr(namespace, self.dest, os.path.abspath(os.path.expanduser(values)))
def str2bool(v):
if v.lower() in ('yes', 'true', 't', 'y', '1'):
return True
elif v.lower() in ('no', 'false', 'f', 'n', '0'):
return False
else:
raise argparse.ArgumentTypeError('Boolean value expected.')
if __name__ == "__main__": if __name__ == "__main__":
os_utils.set_process_lowest_prio() os_utils.set_process_lowest_prio()
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('--tf-suppress-std', action="store_true", dest="tf_suppress_std", default=False, help="Suppress tensorflow initialization info. May not works on some python builds such as anaconda python 3.6.4. If you can fix it, you are welcome.")
subparsers = parser.add_subparsers() subparsers = parser.add_subparsers()
def process_extract(arguments): def process_extract(arguments):
from mainscripts import Extractor from mainscripts import Extractor
Extractor.main ( Extractor.main( arguments.input_dir,
input_dir=arguments.input_dir, arguments.output_dir,
output_dir=arguments.output_dir, arguments.debug,
debug=arguments.debug, arguments.detector,
arguments.manual_fix,
arguments.manual_output_debug_fix,
arguments.manual_window_size,
face_type=arguments.face_type, face_type=arguments.face_type,
detector=arguments.detector, device_args={'cpu_only' : arguments.cpu_only,
multi_gpu=arguments.multi_gpu, 'multi_gpu' : arguments.multi_gpu,
cpu_only=arguments.cpu_only, }
manual_fix=arguments.manual_fix,
manual_output_debug_fix=arguments.manual_output_debug_fix,
manual_window_size=arguments.manual_window_size
) )
extract_parser = subparsers.add_parser( "extract", help="Extract the faces from a pictures.") extract_parser = subparsers.add_parser( "extract", help="Extract the faces from a pictures.")
@ -80,23 +72,25 @@ if __name__ == "__main__":
util_parser.set_defaults (func=process_util) util_parser.set_defaults (func=process_util)
def process_train(arguments): def process_train(arguments):
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,
}
device_args = {'cpu_only' : arguments.cpu_only,
'force_gpu_idx' : arguments.force_gpu_idx,
}
from mainscripts import Trainer from mainscripts import Trainer
Trainer.main ( Trainer.main(args, device_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,
debug = arguments.debug,
#**options
force_gpu_idx = arguments.force_gpu_idx,
cpu_only = arguments.cpu_only
)
train_parser = subparsers.add_parser( "train", help="Trainer") train_parser = subparsers.add_parser( "train", help="Trainer")
train_parser.add_argument('--training-data-src-dir', required=True, action=fixPathAction, dest="training_data_src_dir", help="Dir of src-set.") train_parser.add_argument('--training-data-src-dir', required=True, action=fixPathAction, dest="training_data_src_dir", help="Dir of src-set.")
train_parser.add_argument('--training-data-dst-dir', required=True, action=fixPathAction, dest="training_data_dst_dir", help="Dir of dst-set.") train_parser.add_argument('--training-data-dst-dir', required=True, action=fixPathAction, dest="training_data_dst_dir", help="Dir of dst-set.")
train_parser.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Model dir.") train_parser.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Model dir.")
train_parser.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") train_parser.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")
train_parser.add_argument('--no-preview', action="store_true", dest="no_preview", default=False, help="Disable preview window.")
train_parser.add_argument('--debug', action="store_true", dest="debug", default=False, help="Debug samples.") train_parser.add_argument('--debug', action="store_true", dest="debug", default=False, help="Debug samples.")
train_parser.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Train on CPU.") train_parser.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Train on CPU.")
train_parser.add_argument('--force-gpu-idx', type=int, dest="force_gpu_idx", default=-1, help="Force to choose this GPU idx.") train_parser.add_argument('--force-gpu-idx', type=int, dest="force_gpu_idx", default=-1, help="Force to choose this GPU idx.")
@ -104,17 +98,18 @@ if __name__ == "__main__":
train_parser.set_defaults (func=process_train) train_parser.set_defaults (func=process_train)
def process_convert(arguments): def process_convert(arguments):
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,
}
device_args = {'cpu_only' : arguments.cpu_only,
'force_gpu_idx' : arguments.force_gpu_idx,
}
from mainscripts import Converter from mainscripts import Converter
Converter.main ( Converter.main (args, device_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,
force_gpu_idx = arguments.force_gpu_idx,
cpu_only = arguments.cpu_only
)
convert_parser = subparsers.add_parser( "convert", help="Converter") convert_parser = subparsers.add_parser( "convert", help="Converter")
convert_parser.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.") convert_parser.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
@ -134,8 +129,6 @@ if __name__ == "__main__":
parser.set_defaults(func=bad_args) parser.set_defaults(func=bad_args)
arguments = parser.parse_args() arguments = parser.parse_args()
if arguments.tf_suppress_std:
os.environ['TF_SUPPRESS_STD'] = '1'
#os.environ['force_plaidML'] = '1' #os.environ['force_plaidML'] = '1'
@ -143,6 +136,21 @@ if __name__ == "__main__":
print ("Done.") print ("Done.")
"""
Suppressing error with keras 2.2.4+ on python exit:
Exception ignored in: <bound method BaseSession._Callable.__del__ of <tensorflow.python.client.session.BaseSession._Callable object at 0x000000001BDEA9B0>>
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 import code
code.interact(local=dict(globals(), **locals())) code.interact(local=dict(globals(), **locals()))

View file

@ -4,7 +4,6 @@ import traceback
from pathlib import Path from pathlib import Path
from utils import Path_utils from utils import Path_utils
import cv2 import cv2
from tqdm import tqdm
from utils.DFLPNG import DFLPNG from utils.DFLPNG import DFLPNG
from utils.DFLJPG import DFLJPG from utils.DFLJPG import DFLJPG
from utils.cv2_utils import * from utils.cv2_utils import *
@ -13,116 +12,16 @@ import shutil
import numpy as np import numpy as np
import time import time
import multiprocessing import multiprocessing
from models import ConverterBase from converters import Converter
from joblib import Subprocessor, SubprocessFunctionCaller
from interact import interact as io
class model_process_predictor(object): class ConvertSubprocessor(Subprocessor):
def __init__(self, sq, cq, lock): class Cli(Subprocessor.Cli):
self.sq = sq
self.cq = cq
self.lock = lock
def __call__(self, face):
self.lock.acquire()
self.sq.put ( {'op': 'predict', 'face' : face} )
while True:
if not self.cq.empty():
obj = self.cq.get()
obj_op = obj['op']
if obj_op == 'predict_result':
self.lock.release()
return obj['result']
time.sleep(0.005)
def model_process(stdin_fd, model_name, model_dir, in_options, sq, cq):
sys.stdin = os.fdopen(stdin_fd)
try:
model_path = Path(model_dir)
import models
model = models.import_model(model_name)(model_path, **in_options)
converter = model.get_converter(**in_options)
converter.dummy_predict()
cq.put ( {'op':'init', 'converter' : converter.copy_and_set_predictor( None ) } )
while True:
while not sq.empty():
obj = sq.get()
obj_op = obj['op']
if obj_op == 'predict':
result = converter.predictor ( obj['face'] )
cq.put ( {'op':'predict_result', 'result':result} )
time.sleep(0.005)
except Exception as e:
print ( 'Error: %s' % (str(e)))
traceback.print_exc()
from utils.SubprocessorBase import SubprocessorBase
class ConvertSubprocessor(SubprocessorBase):
#override #override
def __init__(self, converter, input_path_image_paths, output_path, alignments, debug = False, **in_options): def on_initialize(self, client_dict):
super().__init__('Converter', 86400 if debug == True else 60) io.log_info ('Running on %s.' % (client_dict['device_name']) )
self.converter = converter
self.input_path_image_paths = input_path_image_paths
self.output_path = output_path
self.alignments = alignments
self.debug = debug
self.in_options = in_options
self.input_data = self.input_path_image_paths
self.files_processed = 0
self.faces_processed = 0
#override
def process_info_generator(self):
r = [0] if self.debug else range( min(multiprocessing.cpu_count(), 6) )
for i in r:
yield 'CPU%d' % (i), {}, {'device_idx': i,
'device_name': 'CPU%d' % (i),
'converter' : self.converter,
'output_dir' : str(self.output_path),
'alignments' : self.alignments,
'debug': self.debug,
'in_options': self.in_options
}
#override
def get_no_process_started_message(self):
return 'Unable to start CPU processes.'
#override
def onHostGetProgressBarDesc(self):
return "Converting"
#override
def onHostGetProgressBarLen(self):
return len (self.input_data)
#override
def onHostGetData(self, host_dict):
if len (self.input_data) > 0:
return self.input_data.pop(0)
return None
#override
def onHostDataReturn (self, host_dict, data):
self.input_data.insert(0, data)
#overridable
def onClientGetDataName (self, data):
#return string identificator of your data
return data
#override
def onClientInitialize(self, client_dict):
print ('Running on %s.' % (client_dict['device_name']) )
self.device_idx = client_dict['device_idx'] self.device_idx = client_dict['device_idx']
self.device_name = client_dict['device_name'] self.device_name = client_dict['device_name']
self.converter = client_dict['converter'] self.converter = client_dict['converter']
@ -132,100 +31,177 @@ class ConvertSubprocessor(SubprocessorBase):
from nnlib import nnlib from nnlib import nnlib
#model process ate all GPU mem, #model process ate all GPU mem,
#so we cannot use GPU for any TF operations in converter processes (for example image_utils.TFLabConverter) #so we cannot use GPU for any TF operations in converter processes
#therefore forcing active_DeviceConfig to CPU only #therefore forcing active_DeviceConfig to CPU only
nnlib.active_DeviceConfig = nnlib.DeviceConfig (cpu_only=True) nnlib.active_DeviceConfig = nnlib.DeviceConfig (cpu_only=True)
return None return None
#override #override
def onClientFinalize(self): def process_data(self, data):
pass
#override
def onClientProcessData(self, data):
filename_path = Path(data) filename_path = Path(data)
files_processed = 1 files_processed = 1
faces_processed = 0 faces_processed = 0
output_filename_path = self.output_path / (filename_path.stem + '.png') output_filename_path = self.output_path / (filename_path.stem + '.png')
if self.converter.get_mode() == ConverterBase.MODE_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: if not self.debug:
print ( 'no faces found for %s, copying without faces' % (filename_path.name) ) self.log_info ( 'no faces found for %s, copying without faces' % (filename_path.name) )
shutil.copy ( str(filename_path), str(output_filename_path) ) shutil.copy ( str(filename_path), str(output_filename_path) )
else: else:
image = (cv2_imread(str(filename_path)) / 255.0).astype(np.float32) image = (cv2_imread(str(filename_path)) / 255.0).astype(np.float32)
if self.converter.get_mode() == ConverterBase.MODE_IMAGE: if self.converter.type == Converter.TYPE_IMAGE:
image = self.converter.convert_image(image, None, self.debug) image = self.converter.convert_image(image, None, self.debug)
if self.debug: if self.debug:
for img in image: raise NotImplementedError
cv2.imshow ('Debug convert', img ) #for img in image:
cv2.waitKey(0) # io.show_image ('Debug convert', img )
# cv2.waitKey(0)
faces_processed = 1 faces_processed = 1
elif self.converter.get_mode() == ConverterBase.MODE_IMAGE_WITH_LANDMARKS: elif self.converter.type == Converter.TYPE_IMAGE_WITH_LANDMARKS:
if filename_path.suffix == '.png': if filename_path.suffix == '.png':
dflimg = DFLPNG.load( str(filename_path), throw_on_no_embedded_data=True ) dflimg = DFLPNG.load( str(filename_path) )
elif filename_path.suffix == '.jpg': elif filename_path.suffix == '.jpg':
dflimg = DFLJPG.load ( str(filename_path), throw_on_no_embedded_data=True ) dflimg = DFLJPG.load ( str(filename_path) )
else: else:
raise Exception ("%s is not a dfl image file" % (filename_path.name) ) dflimg = None
if dflimg is not None:
image_landmarks = dflimg.get_landmarks() image_landmarks = dflimg.get_landmarks()
image = self.converter.convert_image(image, image_landmarks, self.debug) image = self.converter.convert_image(image, image_landmarks, self.debug)
if self.debug: if self.debug:
for img in image: raise NotImplementedError
cv2.imshow ('Debug convert', img ) #for img in image:
cv2.waitKey(0) # io.show_image ('Debug convert', img )
# cv2.waitKey(0)
faces_processed = 1 faces_processed = 1
elif self.converter.get_mode() == ConverterBase.MODE_FACE: else:
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] faces = self.alignments[filename_path.stem]
if self.debug:
debug_images = []
for face_num, image_landmarks in enumerate(faces): for face_num, image_landmarks in enumerate(faces):
try: try:
if self.debug: if self.debug:
print ( '\nConverting face_num [%d] in file [%s]' % (face_num, filename_path) ) self.log_info ( '\nConverting face_num [%d] in file [%s]' % (face_num, filename_path) )
image = self.converter.convert_face(image, image_landmarks, self.debug)
if self.debug: if self.debug:
for img in image: debug_images += self.converter.convert_face(image, image_landmarks, self.debug)
cv2.imshow ('Debug convert', (img*255).astype(np.uint8) ) else:
cv2.waitKey(0) image = self.converter.convert_face(image, image_landmarks, self.debug)
except Exception as e: except Exception as e:
print ( 'Error while converting face_num [%d] in file [%s]: %s' % (face_num, filename_path, str(e)) ) self.log_info ( 'Error while converting face_num [%d] in file [%s]: %s' % (face_num, filename_path, str(e)) )
traceback.print_exc() traceback.print_exc()
if self.debug:
return (1, debug_images)
faces_processed = len(faces) faces_processed = len(faces)
if not self.debug: if not self.debug:
cv2_imwrite (str(output_filename_path), (image*255).astype(np.uint8) ) cv2_imwrite (str(output_filename_path), (image*255).astype(np.uint8) )
return (files_processed, faces_processed) return (0, files_processed, faces_processed)
#overridable
def get_data_name (self, data):
#return string identificator of your data
return data
#override #override
def onHostResult (self, host_dict, data, result): 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( min(multiprocessing.cpu_count(), 6) )
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),
'alignments' : self.alignments,
'debug': self.debug
}
#overridable optional
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 None
#override
def on_data_return (self, host_dict, data):
self.input_data.insert(0, data)
#override
def on_result (self, host_dict, data, result):
if result[0] == 0:
self.files_processed += result[0] self.files_processed += result[0]
self.faces_processed += result[1] self.faces_processed += result[1]
return 1 elif result[0] == 1:
for img in result[1]:
io.show_image ('Debug convert', (img*255).astype(np.uint8) )
io.wait_any_key()
io.progress_bar_inc(1)
#override #override
def onFinalizeAndGetResult(self): def on_tick(self):
self.host_processor.process_messages()
#override
def get_result(self):
return self.files_processed, self.faces_processed return self.files_processed, self.faces_processed
def main (input_dir, output_dir, model_dir, model_name, aligned_dir=None, **in_options): def main (args, device_args):
print ("Running converter.\r\n") io.log_info ("Running converter.\r\n")
debug = in_options['debug'] aligned_dir = args.get('aligned_dir', None)
try: try:
input_path = Path(input_dir) input_path = Path(args['input_dir'])
output_path = Path(output_dir) output_path = Path(args['output_dir'])
model_path = Path(model_dir) model_path = Path(args['model_dir'])
if not input_path.exists(): if not input_path.exists():
print('Input directory not found. Please ensure it exists.') io.log_err('Input directory not found. Please ensure it exists.')
return return
if output_path.exists(): if output_path.exists():
@ -234,50 +210,43 @@ def main (input_dir, output_dir, model_dir, model_name, aligned_dir=None, **in_o
else: else:
output_path.mkdir(parents=True, exist_ok=True) output_path.mkdir(parents=True, exist_ok=True)
if not model_path.exists(): if not model_path.exists():
print('Model directory not found. Please ensure it exists.') io.log_err('Model directory not found. Please ensure it exists.')
return return
model_sq = multiprocessing.Queue() import models
model_cq = multiprocessing.Queue() model = models.import_model( args['model_name'] )(model_path, device_args=device_args)
model_lock = multiprocessing.Lock() converter = model.get_converter()
model_p = multiprocessing.Process(target=model_process, args=( sys.stdin.fileno(), model_name, model_dir, in_options, model_sq, model_cq)) converter.dummy_predict()
model_p.start()
while True:
if not model_cq.empty():
obj = model_cq.get()
obj_op = obj['op']
if obj_op == 'init':
converter = obj['converter']
break
alignments = None alignments = None
if converter.get_mode() == ConverterBase.MODE_FACE: if converter.type == Converter.TYPE_FACE:
if aligned_dir is None: if aligned_dir is None:
print('Aligned directory not found. Please ensure it exists.') io.log_err('Aligned directory not found. Please ensure it exists.')
return return
aligned_path = Path(aligned_dir) aligned_path = Path(aligned_dir)
if not aligned_path.exists(): if not aligned_path.exists():
print('Aligned directory not found. Please ensure it exists.') io.log_err('Aligned directory not found. Please ensure it exists.')
return return
alignments = {} alignments = {}
aligned_path_image_paths = Path_utils.get_image_paths(aligned_path) aligned_path_image_paths = Path_utils.get_image_paths(aligned_path)
for filepath in tqdm(aligned_path_image_paths, desc="Collecting alignments", ascii=True ): for filepath in io.progress_bar_generator(aligned_path_image_paths, "Collecting alignments"):
filepath = Path(filepath) filepath = Path(filepath)
if filepath.suffix == '.png': if filepath.suffix == '.png':
dflimg = DFLPNG.load( str(filepath), print_on_no_embedded_data=True ) dflimg = DFLPNG.load( str(filepath) )
elif filepath.suffix == '.jpg': elif filepath.suffix == '.jpg':
dflimg = DFLJPG.load ( str(filepath), print_on_no_embedded_data=True ) dflimg = DFLJPG.load ( str(filepath) )
else: else:
print ("%s is not a dfl image file" % (filepath.name) ) dflimg = None
if dflimg is None:
io.log_err ("%s is not a dfl image file" % (filepath.name) )
continue
source_filename_stem = Path( dflimg.get_source_filename() ).stem source_filename_stem = Path( dflimg.get_source_filename() ).stem
if source_filename_stem not in alignments.keys(): if source_filename_stem not in alignments.keys():
@ -285,50 +254,22 @@ def main (input_dir, output_dir, model_dir, model_name, aligned_dir=None, **in_o
alignments[ source_filename_stem ].append (dflimg.get_source_landmarks()) alignments[ source_filename_stem ].append (dflimg.get_source_landmarks())
#interpolate landmarks
#from facelib import LandmarksProcessor
#from facelib import FaceType
#a = sorted(alignments.keys())
#a_len = len(a)
#
#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)
#
# 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)
# 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)
files_processed, faces_processed = ConvertSubprocessor ( files_processed, faces_processed = ConvertSubprocessor (
converter = converter.copy_and_set_predictor( model_process_predictor(model_sq,model_cq,model_lock) ), 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, output_path = output_path,
alignments = alignments, alignments = alignments,
**in_options ).process() debug = args.get('debug',False)
).run()
model_p.terminate() model.finalize()
''' except Exception as e:
if model_name == 'AVATAR': print ( 'Error: %s' % (str(e)))
traceback.print_exc()
'''
if model_name == 'AVATAR':
output_path_image_paths = Path_utils.get_image_paths(output_path) output_path_image_paths = Path_utils.get_image_paths(output_path)
last_ok_frame = -1 last_ok_frame = -1
@ -349,10 +290,35 @@ def main (input_dir, output_dir, model_dir, model_name, aligned_dir=None, **in_o
shutil.copy ( str(filename), str( output_path / ('%.5d%s' % (i, filename_path.suffix )) ) ) shutil.copy ( str(filename), str( output_path / ('%.5d%s' % (i, filename_path.suffix )) ) )
last_ok_frame = frame last_ok_frame = frame
''' '''
#interpolate landmarks
except Exception as e: #from facelib import LandmarksProcessor
print ( 'Error: %s' % (str(e))) #from facelib import FaceType
traceback.print_exc() #a = sorted(alignments.keys())
#a_len = len(a)
#
#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)
#
# 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)
# 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)

View file

@ -15,240 +15,16 @@ import facelib
from facelib import FaceType from facelib import FaceType
from facelib import LandmarksProcessor from facelib import LandmarksProcessor
from nnlib import nnlib from nnlib import nnlib
from joblib import Subprocessor
from interact import interact as io
from utils.SubprocessorBase import SubprocessorBase class ExtractSubprocessor(Subprocessor):
class ExtractSubprocessor(SubprocessorBase):
class Cli(Subprocessor.Cli):
#override #override
def __init__(self, input_data, type, image_size, face_type, debug, multi_gpu=False, cpu_only=False, manual=False, manual_window_size=0, detector=None, output_path=None ): def on_initialize(self, client_dict):
self.input_data = input_data self.log_info ('Running on %s.' % (client_dict['device_name']) )
self.type = type
self.image_size = image_size
self.face_type = face_type
self.debug = debug
self.multi_gpu = multi_gpu
self.cpu_only = cpu_only
self.detector = detector
self.output_path = output_path
self.manual = manual
self.manual_window_size = manual_window_size
self.result = []
no_response_time_sec = 60 if not self.manual else 999999
super().__init__('Extractor', no_response_time_sec)
#override
def onHostClientsInitialized(self):
if self.manual == True:
self.wnd_name = 'Manual pass'
cv2.namedWindow(self.wnd_name)
self.landmarks = None
self.param_x = -1
self.param_y = -1
self.param_rect_size = -1
self.param = {'x': 0, 'y': 0, 'rect_size' : 100, 'rect_locked' : False, 'redraw_needed' : False }
def onMouse(event, x, y, flags, param):
if event == cv2.EVENT_MOUSEWHEEL:
mod = 1 if flags > 0 else -1
diff = 1 if param['rect_size'] <= 40 else np.clip(param['rect_size'] / 10, 1, 10)
param['rect_size'] = max (5, param['rect_size'] + diff*mod)
elif event == cv2.EVENT_LBUTTONDOWN:
param['rect_locked'] = not param['rect_locked']
param['redraw_needed'] = True
elif not param['rect_locked']:
param['x'] = x
param['y'] = y
cv2.setMouseCallback(self.wnd_name, onMouse, self.param)
def get_devices_for_type (self, type, multi_gpu, cpu_only):
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
else:
if multi_gpu:
devices = nnlib.device.getValidDevicesWithAtLeastTotalMemoryGB(2)
if not multi_gpu or len(devices) == 0:
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 == 'mt':
for i in range ( int (max (1, dev_vram / 2) ) ):
yield (idx, 'GPU', '%s #%d' % (dev_name,i) , dev_vram)
else:
yield (idx, 'GPU', dev_name, dev_vram)
if cpu_only and (type == 'rects' or type == 'landmarks'):
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 )
#override
def process_info_generator(self):
base_dict = {'type' : self.type,
'image_size': self.image_size,
'face_type': self.face_type,
'debug': self.debug,
'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):
client_dict = base_dict.copy()
client_dict['device_idx'] = device_idx
client_dict['device_name'] = device_name
client_dict['device_type'] = device_type
yield client_dict['device_name'], {}, client_dict
#override
def get_no_process_started_message(self):
if (self.type == 'rects' or self.type == 'landmarks'):
print ( 'You have no capable GPUs. Try to close programs which can consume VRAM, and run again.')
elif self.type == 'final':
print ( 'Unable to start CPU processes.')
#override
def onHostGetProgressBarDesc(self):
return None
#override
def onHostGetProgressBarLen(self):
return len (self.input_data)
#override
def onHostGetData(self, host_dict):
if not self.manual:
if len (self.input_data) > 0:
return self.input_data.pop(0)
else:
skip_remaining = False
allow_remark_faces = False
while len (self.input_data) > 0:
data = self.input_data[0]
filename, faces = data
is_frame_done = False
go_to_prev_frame = False
# Can we mark an image that already has a marked face?
if allow_remark_faces:
allow_remark_faces = False
# If there was already a face then lock the rectangle to it until the mouse is clicked
if len(faces) > 0:
self.rect, self.landmarks = faces.pop()
self.param['rect_locked'] = True
faces.clear()
self.param['rect_size'] = ( self.rect[2] - self.rect[0] ) / 2
self.param['x'] = ( ( self.rect[0] + self.rect[2] ) / 2 ) * self.view_scale
self.param['y'] = ( ( self.rect[1] + self.rect[3] ) / 2 ) * self.view_scale
if len(faces) == 0:
self.original_image = cv2_imread(filename)
(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) )
self.original_image = cv2.resize (self.original_image, ( int(w*self.view_scale), int(h*self.view_scale) ), interpolation=cv2.INTER_LINEAR)
(h,w,c) = self.original_image.shape
self.text_lines_img = (image_utils.get_draw_text_lines ( self.original_image, (0,0, self.original_image.shape[1], min(100, self.original_image.shape[0]) ),
[ 'Match landmarks with face exactly. Click to confirm/unconfirm selection',
'[Enter] - confirm face landmarks and continue',
'[Space] - confirm as unmarked frame and continue',
'[Mouse wheel] - change rect',
'[,] [.]- prev frame, next frame',
'[Q] - skip remaining frames'
], (1, 1, 1) )*255).astype(np.uint8)
while True:
key = cv2.waitKey(1) & 0xFF
if key == ord('\r') or key == ord('\n'):
faces.append ( [(self.rect), self.landmarks] )
is_frame_done = True
break
elif key == ord(' '):
is_frame_done = True
break
elif key == ord('.'):
allow_remark_faces = True
# Only save the face if the rect is still locked
if self.param['rect_locked']:
faces.append ( [(self.rect), self.landmarks] )
is_frame_done = True
break
elif key == ord(',') and len(self.result) > 0:
# Only save the face if the rect is still locked
if self.param['rect_locked']:
faces.append ( [(self.rect), self.landmarks] )
go_to_prev_frame = True
break
elif key == ord('q'):
skip_remaining = True
break
new_param_x = np.clip (self.param['x'], 0, w-1) / self.view_scale
new_param_y = np.clip (self.param['y'], 0, h-1) / self.view_scale
new_param_rect_size = self.param['rect_size']
if self.param_x != new_param_x or \
self.param_y != new_param_y or \
self.param_rect_size != new_param_rect_size or \
self.param['redraw_needed']:
self.param_x = new_param_x
self.param_y = new_param_y
self.param_rect_size = new_param_rect_size
self.rect = ( int(self.param_x-self.param_rect_size),
int(self.param_y-self.param_rect_size),
int(self.param_x+self.param_rect_size),
int(self.param_y+self.param_rect_size) )
return [filename, [self.rect]]
else:
is_frame_done = True
if is_frame_done:
self.result.append ( data )
self.input_data.pop(0)
self.inc_progress_bar(1)
self.param['redraw_needed'] = True
self.param['rect_locked'] = False
elif go_to_prev_frame:
self.input_data.insert(0, self.result.pop() )
self.inc_progress_bar(-1)
allow_remark_faces = True
self.param['redraw_needed'] = True
self.param['rect_locked'] = False
elif skip_remaining:
if self.param['rect_locked']:
faces.append ( [(self.rect), self.landmarks] )
while len(self.input_data) > 0:
self.result.append( self.input_data.pop(0) )
self.inc_progress_bar(1)
return None
#override
def onHostDataReturn (self, host_dict, data):
if not self.manual:
self.input_data.insert(0, data)
#override
def onClientInitialize(self, client_dict):
self.safe_print ('Running on %s.' % (client_dict['device_name']) )
self.type = client_dict['type'] self.type = client_dict['type']
self.image_size = client_dict['image_size'] self.image_size = client_dict['image_size']
self.face_type = client_dict['face_type'] self.face_type = client_dict['face_type']
@ -258,6 +34,8 @@ class ExtractSubprocessor(SubprocessorBase):
self.debug = client_dict['debug'] self.debug = client_dict['debug']
self.detector = client_dict['detector'] self.detector = client_dict['detector']
self.cached_image = (None, None)
self.e = None self.e = None
device_config = nnlib.DeviceConfig ( cpu_only=self.cpu_only, force_gpu_idx=self.device_idx, allow_growth=True) device_config = nnlib.DeviceConfig ( cpu_only=self.cpu_only, force_gpu_idx=self.device_idx, allow_growth=True)
if self.type == 'rects': if self.type == 'rects':
@ -268,6 +46,10 @@ class ExtractSubprocessor(SubprocessorBase):
elif self.detector == 'dlib': elif self.detector == 'dlib':
nnlib.import_dlib (device_config) nnlib.import_dlib (device_config)
self.e = facelib.DLIBExtractor(nnlib.dlib) self.e = facelib.DLIBExtractor(nnlib.dlib)
else:
raise ValueError ("Wrond detector type.")
if self.e is not None:
self.e.__enter__() self.e.__enter__()
elif self.type == 'landmarks': elif self.type == 'landmarks':
@ -278,21 +60,24 @@ class ExtractSubprocessor(SubprocessorBase):
elif self.type == 'final': elif self.type == 'final':
pass pass
return None
#override #override
def onClientFinalize(self): def on_finalize(self):
if self.e is not None: if self.e is not None:
self.e.__exit__() self.e.__exit__()
#override #override
def onClientProcessData(self, data): def process_data(self, data):
filename_path = Path( data[0] ) filename_path = Path( data[0] )
image = cv2_imread( str(filename_path) ) filename_path_str = str(filename_path)
if self.cached_image[0] == filename_path_str:
image = self.cached_image[1]
else:
image = cv2_imread( filename_path_str )
self.cached_image = ( filename_path_str, image )
if image is None: if image is None:
print ( 'Failed to extract %s, reason: cv2_imread() fail.' % ( str(filename_path) ) ) self.log_err ( 'Failed to extract %s, reason: cv2_imread() fail.' % ( str(filename_path) ) )
else: else:
if self.type == 'rects': if self.type == 'rects':
rects = self.e.extract_from_bgr (image) rects = self.e.extract_from_bgr (image)
@ -357,26 +142,273 @@ class ExtractSubprocessor(SubprocessorBase):
return None return None
#overridable #overridable
def onClientGetDataName (self, data): def get_data_name (self, data):
#return string identificator of your data #return string identificator of your data
return data[0] return data[0]
#override #override
def onHostResult (self, host_dict, data, result): def __init__(self, input_data, type, image_size, face_type, debug, 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
self.face_type = face_type
self.debug = debug
self.multi_gpu = multi_gpu
self.cpu_only = cpu_only
self.detector = detector
self.output_path = output_path
self.manual = manual
self.manual_window_size = manual_window_size
self.result = []
no_response_time_sec = 60 if not self.manual else 999999
super().__init__('Extractor', ExtractSubprocessor.Cli, no_response_time_sec)
#override
def on_clients_initialized(self):
if self.manual == True:
self.wnd_name = 'Manual pass'
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.landmarks = None
self.x = 0
self.y = 0
self.rect_size = 100
self.rect_locked = False
self.redraw_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 not cpu_only and (type == 'rects' or type == 'landmarks'):
if type == 'rects' and self.detector == 'mt' and nnlib.device.backend == "plaidML":
cpu_only = True
else:
if multi_gpu:
devices = nnlib.device.getValidDevicesWithAtLeastTotalMemoryGB(2)
if not multi_gpu or len(devices) == 0:
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 == 'mt':
for i in range ( int (max (1, dev_vram / 2) ) ):
yield (idx, 'GPU', '%s #%d' % (dev_name,i) , dev_vram)
else:
yield (idx, 'GPU', dev_name, dev_vram)
if cpu_only and (type == 'rects' or type == 'landmarks'):
if self.manual:
yield (0, 'CPU', 'CPU', 0 )
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 )
#override
def process_info_generator(self):
base_dict = {'type' : self.type,
'image_size': self.image_size,
'face_type': self.face_type,
'debug': self.debug,
'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):
client_dict = base_dict.copy()
client_dict['device_idx'] = device_idx
client_dict['device_name'] = device_name
client_dict['device_type'] = device_type
yield client_dict['device_name'], {}, client_dict
#override
def get_data(self, host_dict):
if not self.manual:
if len (self.input_data) > 0:
return self.input_data.pop(0)
else:
skip_remaining = False
allow_remark_faces = False
while len (self.input_data) > 0:
data = self.input_data[0]
filename, faces = data
is_frame_done = False
go_to_prev_frame = False
# Can we mark an image that already has a marked face?
if allow_remark_faces:
allow_remark_faces = False
# If there was already a face then lock the rectangle to it until the mouse is clicked
if len(faces) > 0:
self.rect, self.landmarks = faces.pop()
self.rect_locked = True
self.redraw_needed = True
faces.clear()
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
if len(faces) == 0:
if self.cache_original_image[0] == filename:
self.original_image = self.cache_original_image[1]
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)
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) )
if self.cache_text_lines_img[0] == sh:
self.text_lines_img = self.cache_text_lines_img[1]
else:
self.text_lines_img = (image_utils.get_draw_text_lines ( self.image, sh,
[ 'Match landmarks with face exactly. Click to confirm/unconfirm selection',
'[Enter] - confirm face landmarks and continue',
'[Space] - confirm as unmarked frame and continue',
'[Mouse wheel] - change rect',
'[,] [.]- prev frame, next frame',
'[Q] - skip remaining frames'
], (1, 1, 1) )*255).astype(np.uint8)
self.cache_text_lines_img = (sh, self.text_lines_img)
while True:
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
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)
elif ev == io.EVENT_LBUTTONDOWN:
self.rect_locked = not self.rect_locked
self.redraw_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,)
if key == ord('\r') or key == ord('\n'):
faces.append ( [(self.rect), self.landmarks] )
is_frame_done = True
break
elif key == ord(' '):
is_frame_done = True
break
elif key == ord('.'):
allow_remark_faces = True
# Only save the face if the rect is still locked
if self.rect_locked:
faces.append ( [(self.rect), self.landmarks] )
is_frame_done = True
break
elif key == ord(',') and len(self.result) > 0:
# Only save the face if the rect is still locked
if self.rect_locked:
faces.append ( [(self.rect), self.landmarks] )
go_to_prev_frame = True
break
elif key == ord('q'):
skip_remaining = True
break
if self.x != new_x or \
self.y != new_y or \
self.rect_size != new_rect_size or \
self.redraw_needed:
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),
int(self.y+self.rect_size) )
return [filename, [self.rect]]
io.process_messages(0.0001)
else:
is_frame_done = True
if is_frame_done:
self.result.append ( data )
self.input_data.pop(0)
io.progress_bar_inc(1)
self.redraw_needed = True
self.rect_locked = False
elif go_to_prev_frame:
self.input_data.insert(0, self.result.pop() )
io.progress_bar_inc(-1)
allow_remark_faces = True
self.redraw_needed = True
self.rect_locked = False
elif 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)
return None
#override
def on_data_return (self, host_dict, data):
if not self.manual:
self.input_data.insert(0, data)
#override
def on_result (self, host_dict, data, result):
if self.manual == True: if self.manual == True:
self.landmarks = result[1][0][1] self.landmarks = result[1][0][1]
(h,w,c) = self.original_image.shape (h,w,c) = self.image.shape
image = cv2.addWeighted (self.original_image,1.0,self.text_lines_img,1.0,0) image = cv2.addWeighted (self.image,1.0,self.text_lines_img,1.0,0)
view_rect = (np.array(self.rect) * self.view_scale).astype(np.int).tolist() 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() view_landmarks = (np.array(self.landmarks) * self.view_scale).astype(np.int).tolist()
if self.param_rect_size <= 40: if self.rect_size <= 40:
scaled_rect_size = h // 3 if w > h else w // 3 scaled_rect_size = h // 3 if w > h else w // 3
p1 = (self.param_x - self.param_rect_size, self.param_y - self.param_rect_size) p1 = (self.x - self.rect_size, self.y - self.rect_size)
p2 = (self.param_x + self.param_rect_size, self.param_y - self.param_rect_size) p2 = (self.x + self.rect_size, self.y - self.rect_size)
p3 = (self.param_x - self.param_rect_size, self.param_y + self.param_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) np1 = (w / 2 - wh / 4, h / 2 - wh / 4)
@ -389,12 +421,11 @@ class ExtractSubprocessor(SubprocessorBase):
LandmarksProcessor.draw_rect_landmarks (image, view_rect, view_landmarks, self.image_size, self.face_type) LandmarksProcessor.draw_rect_landmarks (image, view_rect, view_landmarks, self.image_size, self.face_type)
if self.param['rect_locked']: if self.rect_locked:
LandmarksProcessor.draw_landmarks(image, view_landmarks, (255,255,0) ) LandmarksProcessor.draw_landmarks(image, view_landmarks, (255,255,0) )
self.param['redraw_needed'] = False self.redraw_needed = False
cv2.imshow (self.wnd_name, image) io.show_image (self.wnd_name, image)
return 0
else: else:
if self.type == 'rects': if self.type == 'rects':
self.result.append ( result ) self.result.append ( result )
@ -403,97 +434,90 @@ class ExtractSubprocessor(SubprocessorBase):
elif self.type == 'final': elif self.type == 'final':
self.result += result self.result += result
return 1 io.progress_bar_inc(1)
#override #override
def onFinalizeAndGetResult(self): def get_result(self):
if self.manual == True:
cv2.destroyAllWindows()
return self.result return self.result
class DeletedFilesSearcherSubprocessor(SubprocessorBase):
class DeletedFilesSearcherSubprocessor(Subprocessor):
class Cli(Subprocessor.Cli):
#override
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
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 #override
def __init__(self, input_paths, debug_paths ): def __init__(self, input_paths, debug_paths ):
self.input_paths = input_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 = [] self.result = []
super().__init__('DeletedFilesSearcherSubprocessor', 60) super().__init__('DeletedFilesSearcherSubprocessor', DeletedFilesSearcherSubprocessor.Cli, 60)
#override #override
def process_info_generator(self): def process_info_generator(self):
for i in range(0, min(multiprocessing.cpu_count(), 8) ): for i in range(min(multiprocessing.cpu_count(), 8)):
yield 'CPU%d' % (i), {}, {'device_idx': i, yield 'CPU%d' % (i), {}, {'debug_paths_stems' : self.debug_paths_stems}
'device_name': 'CPU%d' % (i),
'debug_paths_stems' : self.debug_paths_stems
}
#override #override
def get_no_process_started_message(self): def on_clients_initialized(self):
print ( 'Unable to start CPU processes.') io.progress_bar ("Searching deleted files", len (self.input_paths))
#override #override
def onHostGetProgressBarDesc(self): def on_clients_finalized(self):
return "Searching deleted files" io.progress_bar_close()
#override #override
def onHostGetProgressBarLen(self): def get_data(self, host_dict):
return len (self.input_paths)
#override
def onHostGetData(self, host_dict):
if len (self.input_paths) > 0: if len (self.input_paths) > 0:
return [self.input_paths.pop(0)] return [self.input_paths.pop(0)]
return None return None
#override #override
def onHostDataReturn (self, host_dict, data): def on_data_return (self, host_dict, data):
self.input_paths.insert(0, data[0]) self.input_paths.insert(0, data[0])
#override #override
def onClientInitialize(self, client_dict): def on_result (self, host_dict, data, result):
self.debug_paths_stems = client_dict['debug_paths_stems']
return None
#override
def onClientProcessData(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 onClientGetDataName (self, data):
#return string identificator of your data
return data[0]
#override
def onHostResult (self, host_dict, data, result):
if result == False: if result == False:
self.result.append( data[0] ) self.result.append( data[0] )
return 1 io.progress_bar_inc(1)
#override #override
def onFinalizeAndGetResult(self): def get_result(self):
return self.result return self.result
''' def main(input_dir,
detector output_dir,
'dlib' debug=False,
'mt' detector='mt',
'manual' manual_fix=False,
manual_output_debug_fix=False,
face_type manual_window_size=1368,
'full_face' image_size=256,
'avatar' face_type='full_face',
''' device_args={}):
def main (input_dir, output_dir, debug, detector='mt', multi_gpu=True, cpu_only=False, manual_fix=False, manual_output_debug_fix=False, manual_window_size=1368, image_size=256, face_type='full_face'):
print ("Running extractor.\r\n")
input_path = Path(input_dir) input_path = Path(input_dir)
output_path = Path(output_dir) output_path = Path(output_dir)
face_type = FaceType.fromString(face_type) 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(): if not input_path.exists():
print('Input directory not found. Please ensure it exists.') raise ValueError('Input directory not found. Please ensure it exists.')
return
if output_path.exists(): if output_path.exists():
if not manual_output_debug_fix: if not manual_output_debug_fix:
@ -505,19 +529,17 @@ def main (input_dir, output_dir, debug, detector='mt', multi_gpu=True, cpu_only=
if manual_output_debug_fix: if manual_output_debug_fix:
debug = True debug = True
detector = 'manual' detector = 'manual'
print('Performing re-extract frames which were deleted from _debug directory.') 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=True)
input_path_image_paths = Path_utils.get_image_unique_filestem_paths(input_path, verbose_print_func=io.log_info)
if debug: if debug:
debug_output_path = Path(str(output_path) + '_debug') debug_output_path = Path(str(output_path) + '_debug')
if manual_output_debug_fix: if manual_output_debug_fix:
if not debug_output_path.exists(): if not debug_output_path.exists():
print ("%s not found " % ( str(debug_output_path) )) raise ValueError("%s not found " % ( str(debug_output_path) ))
return
input_path_image_paths = DeletedFilesSearcherSubprocessor ( input_path_image_paths, Path_utils.get_image_paths(debug_output_path) ).process() 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: else:
if debug_output_path.exists(): if debug_output_path.exists():
@ -530,29 +552,29 @@ def main (input_dir, output_dir, debug, detector='mt', multi_gpu=True, cpu_only=
faces_detected = 0 faces_detected = 0
if images_found != 0: if images_found != 0:
if detector == 'manual': if detector == 'manual':
print ('Performing manual extract...') io.log_info ('Performing manual extract...')
extracted_faces = ExtractSubprocessor ([ (filename,[]) for filename in input_path_image_paths ], 'landmarks', image_size, face_type, debug, cpu_only=cpu_only, manual=True, manual_window_size=manual_window_size).process() extracted_faces = ExtractSubprocessor ([ (filename,[]) for filename in input_path_image_paths ], 'landmarks', image_size, face_type, debug, cpu_only=cpu_only, manual=True, manual_window_size=manual_window_size).run()
else: else:
print ('Performing 1st pass...') io.log_info ('Performing 1st pass...')
extracted_rects = ExtractSubprocessor ([ (x,) for x in input_path_image_paths ], 'rects', image_size, face_type, debug, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False, detector=detector).process() extracted_rects = ExtractSubprocessor ([ (x,) for x in input_path_image_paths ], 'rects', image_size, face_type, debug, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False, detector=detector).run()
print ('Performing 2nd pass...') io.log_info ('Performing 2nd pass...')
extracted_faces = ExtractSubprocessor (extracted_rects, 'landmarks', image_size, face_type, debug, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False).process() extracted_faces = ExtractSubprocessor (extracted_rects, 'landmarks', image_size, face_type, debug, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False).run()
if manual_fix: if manual_fix:
print ('Performing manual fix...') io.log_info ('Performing manual fix...')
if all ( np.array ( [ len(data[1]) > 0 for data in extracted_faces] ) == True ): if all ( np.array ( [ len(data[1]) > 0 for data in extracted_faces] ) == True ):
print ('All faces are detected, manual fix not needed.') io.log_info ('All faces are detected, manual fix not needed.')
else: else:
extracted_faces = ExtractSubprocessor (extracted_faces, 'landmarks', image_size, face_type, debug, manual=True, manual_window_size=manual_window_size).process() extracted_faces = ExtractSubprocessor (extracted_faces, 'landmarks', image_size, face_type, debug, manual=True, manual_window_size=manual_window_size).run()
if len(extracted_faces) > 0: if len(extracted_faces) > 0:
print ('Performing 3rd pass...') io.log_info ('Performing 3rd pass...')
final_imgs_paths = ExtractSubprocessor (extracted_faces, 'final', image_size, face_type, debug, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False, output_path=output_path).process() final_imgs_paths = ExtractSubprocessor (extracted_faces, 'final', image_size, face_type, debug, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False, output_path=output_path).run()
faces_detected = len(final_imgs_paths) faces_detected = len(final_imgs_paths)
print('-------------------------') io.log_info ('-------------------------')
print('Images found: %d' % (images_found) ) io.log_info ('Images found: %d' % (images_found) )
print('Faces detected: %d' % (faces_detected) ) io.log_info ('Faces detected: %d' % (faces_detected) )
print('-------------------------') io.log_info ('-------------------------')

File diff suppressed because it is too large Load diff

View file

@ -1,32 +1,34 @@
import sys import sys
import traceback import traceback
import queue import queue
import colorsys import threading
import time import time
import numpy as np import numpy as np
import itertools import itertools
from pathlib import Path from pathlib import Path
from utils import Path_utils from utils import Path_utils
from utils import image_utils from utils import image_utils
import cv2 import cv2
import models import models
from interact import interact as io
def trainerThread (input_queue, output_queue, training_data_src_dir, training_data_dst_dir, model_path, model_name, save_interval_min=15, debug=False, **in_options): def trainerThread (s2c, c2s, args, device_args):
while True: while True:
try: try:
training_data_src_path = Path(training_data_src_dir) training_data_src_path = Path( args.get('training_data_src_dir', '') )
training_data_dst_path = Path(training_data_dst_dir) training_data_dst_path = Path( args.get('training_data_dst_dir', '') )
model_path = Path(model_path) model_path = Path( args.get('model_path', '') )
model_name = args.get('model_name', '')
save_interval_min = 15
debug = args.get('debug', '')
if not training_data_src_path.exists(): if not training_data_src_path.exists():
print( 'Training data src directory does not exist.') io.log_err('Training data src directory does not exist.')
return break
if not training_data_dst_path.exists(): if not training_data_dst_path.exists():
print( 'Training data dst directory does not exist.') io.log_err('Training data dst directory does not exist.')
return break
if not model_path.exists(): if not model_path.exists():
model_path.mkdir(exist_ok=True) model_path.mkdir(exist_ok=True)
@ -36,7 +38,7 @@ def trainerThread (input_queue, output_queue, training_data_src_dir, training_da
training_data_src_path=training_data_src_path, training_data_src_path=training_data_src_path,
training_data_dst_path=training_data_dst_path, training_data_dst_path=training_data_dst_path,
debug=debug, debug=debug,
**in_options) device_args=device_args)
is_reached_goal = model.is_reached_epoch_goal() is_reached_goal = model.is_reached_epoch_goal()
is_upd_save_time_after_train = False is_upd_save_time_after_train = False
@ -48,10 +50,10 @@ def trainerThread (input_queue, output_queue, training_data_src_dir, training_da
def send_preview(): def send_preview():
if not debug: if not debug:
previews = model.get_previews() previews = model.get_previews()
output_queue.put ( {'op':'show', 'previews': previews, 'epoch':model.get_epoch(), 'loss_history': model.get_loss_history().copy() } ) c2s.put ( {'op':'show', 'previews': previews, 'epoch':model.get_epoch(), 'loss_history': model.get_loss_history().copy() } )
else: else:
previews = [( 'debug, press update for new', model.debug_one_epoch())] previews = [( 'debug, press update for new', model.debug_one_epoch())]
output_queue.put ( {'op':'show', 'previews': previews} ) c2s.put ( {'op':'show', 'previews': previews} )
if model.is_first_run(): if model.is_first_run():
@ -59,11 +61,11 @@ def trainerThread (input_queue, output_queue, training_data_src_dir, training_da
if model.get_target_epoch() != 0: if model.get_target_epoch() != 0:
if is_reached_goal: if is_reached_goal:
print ('Model already trained to target epoch. You can use preview.') io.log_info('Model already trained to target epoch. You can use preview.')
else: else:
print('Starting. Target epoch: %d. Press "Enter" to stop training and save model.' % ( model.get_target_epoch() ) ) io.log_info('Starting. Target epoch: %d. Press "Enter" to stop training and save model.' % ( model.get_target_epoch() ) )
else: else:
print('Starting. Press "Enter" to stop training and save model.') io.log_info('Starting. Press "Enter" to stop training and save model.')
last_save_time = time.time() last_save_time = time.time()
@ -75,12 +77,12 @@ def trainerThread (input_queue, output_queue, training_data_src_dir, training_da
#save resets plaidML programs, so upd last_save_time only after plaidML rebuild them #save resets plaidML programs, so upd last_save_time only after plaidML rebuild them
last_save_time = time.time() last_save_time = time.time()
print (loss_string, end='\r') io.log_info (loss_string, end='\r')
if model.get_target_epoch() != 0 and model.is_reached_epoch_goal(): if model.get_target_epoch() != 0 and model.is_reached_epoch_goal():
print ('Reached target epoch.') io.log_info ('Reached target epoch.')
model_save() model_save()
is_reached_goal = True is_reached_goal = True
print ('You can use preview now.') io.log_info ('You can use preview now.')
if not is_reached_goal and (time.time() - last_save_time) >= save_interval_min*60: if not is_reached_goal and (time.time() - last_save_time) >= save_interval_min*60:
last_save_time = time.time() last_save_time = time.time()
@ -95,8 +97,8 @@ def trainerThread (input_queue, output_queue, training_data_src_dir, training_da
if debug: if debug:
time.sleep(0.005) time.sleep(0.005)
while not input_queue.empty(): while not s2c.empty():
input = input_queue.get() input = s2c.get()
op = input['op'] op = input['op']
if op == 'save': if op == 'save':
model_save() model_save()
@ -120,11 +122,34 @@ def trainerThread (input_queue, output_queue, training_data_src_dir, training_da
print ('Error: %s' % (str(e))) print ('Error: %s' % (str(e)))
traceback.print_exc() traceback.print_exc()
break break
output_queue.put ( {'op':'close'} ) c2s.put ( {'op':'close'} )
def previewThread (input_queue, output_queue):
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:
if not c2s.empty():
input = c2s.get()
op = input.get('op','')
if op == 'close':
break
io.process_messages(0.1)
else:
wnd_name = "Training preview"
io.named_window(wnd_name)
io.capture_keys(wnd_name)
previews = None previews = None
loss_history = None loss_history = None
selected_preview = 0 selected_preview = 0
@ -134,8 +159,8 @@ def previewThread (input_queue, output_queue):
show_last_history_epochs_count = 0 show_last_history_epochs_count = 0
epoch = 0 epoch = 0
while True: while True:
if not input_queue.empty(): if not c2s.empty():
input = input_queue.get() input = c2s.get()
op = input['op'] op = input['op']
if op == 'show': if op == 'show':
is_waiting_preview = False is_waiting_preview = False
@ -203,23 +228,20 @@ def previewThread (input_queue, output_queue):
final = np.concatenate ( [final, selected_preview_rgb], axis=0 ) final = np.concatenate ( [final, selected_preview_rgb], axis=0 )
final = np.clip(final, 0, 1) final = np.clip(final, 0, 1)
cv2.imshow ( 'Training preview', (final*255).astype(np.uint8) ) io.show_image( wnd_name, (final*255).astype(np.uint8) )
is_showing = True is_showing = True
if is_showing: key_events = io.get_key_events(wnd_name)
key = cv2.waitKey(100) key, = key_events[-1] if len(key_events) > 0 else (0,)
else:
time.sleep(0.1)
key = 0
if key == ord('\n') or key == ord('\r'): if key == ord('\n') or key == ord('\r'):
output_queue.put ( {'op': 'close'} ) s2c.put ( {'op': 'close'} )
elif key == ord('s'): elif key == ord('s'):
output_queue.put ( {'op': 'save'} ) s2c.put ( {'op': 'save'} )
elif key == ord('p'): elif key == ord('p'):
if not is_waiting_preview: if not is_waiting_preview:
is_waiting_preview = True is_waiting_preview = True
output_queue.put ( {'op': 'preview'} ) s2c.put ( {'op': 'preview'} )
elif key == ord('l'): elif key == ord('l'):
if show_last_history_epochs_count == 0: if show_last_history_epochs_count == 0:
show_last_history_epochs_count = 5000 show_last_history_epochs_count = 5000
@ -236,15 +258,6 @@ def previewThread (input_queue, output_queue):
selected_preview = (selected_preview + 1) % len(previews) selected_preview = (selected_preview + 1) % len(previews)
update_preview = True update_preview = True
cv2.destroyAllWindows() io.process_messages(0.1)
def main (training_data_src_dir, training_data_dst_dir, model_path, model_name, **in_options): io.destroy_all_windows()
print ("Running trainer.\r\n")
output_queue = queue.Queue()
input_queue = queue.Queue()
import threading
thread = threading.Thread(target=trainerThread, args=(output_queue, input_queue, training_data_src_dir, training_data_dst_dir, model_path, model_name), kwargs=in_options )
thread.start()
previewThread (input_queue, output_queue)

View file

@ -1,28 +1,21 @@
import os import cv2
import sys
import operator
import numpy as np
import cv2
from tqdm import tqdm
from shutil import copyfile
from pathlib import Path from pathlib import Path
from utils import Path_utils from utils import Path_utils
from utils import image_utils
from utils.DFLPNG import DFLPNG from utils.DFLPNG import DFLPNG
from utils.DFLJPG import DFLJPG from utils.DFLJPG import DFLJPG
from utils.cv2_utils import * from utils.cv2_utils import *
from facelib import LandmarksProcessor from facelib import LandmarksProcessor
from utils.SubprocessorBase import SubprocessorBase from interact import interact as io
import multiprocessing
def convert_png_to_jpg_file (filepath): def convert_png_to_jpg_file (filepath):
filepath = Path(filepath) filepath = Path(filepath)
if filepath.suffix != '.png': if filepath.suffix != '.png':
return return
dflpng = DFLPNG.load (str(filepath), print_on_no_embedded_data=True)
dflpng = DFLPNG.load (str(filepath) )
if dflpng is None: if dflpng is None:
print ("%s is not a dfl image file" % (filepath.name) )
return return
dfl_dict = dflpng.getDFLDictData() dfl_dict = dflpng.getDFLDictData()
@ -45,27 +38,30 @@ def convert_png_to_jpg_folder (input_path):
print ("Converting PNG to JPG...\r\n") print ("Converting PNG to JPG...\r\n")
for filepath in tqdm( Path_utils.get_image_paths(input_path), desc="Converting", ascii=True): for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Converting"):
filepath = Path(filepath) filepath = Path(filepath)
convert_png_to_jpg_file(filepath) convert_png_to_jpg_file(filepath)
def add_landmarks_debug_images(input_path): def add_landmarks_debug_images(input_path):
print ("Adding landmarks debug images...") print ("Adding landmarks debug images...")
for filepath in tqdm( Path_utils.get_image_paths(input_path), desc="Processing", ascii=True): for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Processing"):
filepath = Path(filepath) filepath = Path(filepath)
img = cv2_imread(str(filepath)) img = cv2_imread(str(filepath))
if filepath.suffix == '.png': if filepath.suffix == '.png':
dflimg = DFLPNG.load( str(filepath), print_on_no_embedded_data=True ) dflimg = DFLPNG.load( str(filepath) )
elif filepath.suffix == '.jpg': elif filepath.suffix == '.jpg':
dflimg = DFLJPG.load ( str(filepath), print_on_no_embedded_data=True ) dflimg = DFLJPG.load ( str(filepath) )
else: else:
dflimg = None
if dflimg is None:
print ("%s is not a dfl image file" % (filepath.name) ) print ("%s is not a dfl image file" % (filepath.name) )
continue continue
if not (dflimg is None or img is None): if img is not None:
face_landmarks = dflimg.get_landmarks() face_landmarks = dflimg.get_landmarks()
LandmarksProcessor.draw_landmarks(img, face_landmarks) LandmarksProcessor.draw_landmarks(img, face_landmarks)

View file

@ -7,31 +7,34 @@ from pathlib import Path
from utils import Path_utils from utils import Path_utils
from utils import std_utils from utils import std_utils
from utils import image_utils from utils import image_utils
from utils.console_utils import *
from utils.cv2_utils import * from utils.cv2_utils import *
import numpy as np import numpy as np
import cv2 import cv2
from samples import SampleGeneratorBase from samples import SampleGeneratorBase
from nnlib import nnlib from nnlib import nnlib
from interact import interact as io
''' '''
You can implement your own model. Check examples. You can implement your own model. Check examples.
''' '''
class ModelBase(object): class ModelBase(object):
#DONT OVERRIDE #DONT OVERRIDE
def __init__(self, model_path, training_data_src_path=None, training_data_dst_path=None, debug = False, force_gpu_idx=-1, **in_options): def __init__(self, model_path, training_data_src_path=None, training_data_dst_path=None, debug = False, device_args = None):
if force_gpu_idx == -1: device_args['force_gpu_idx'] = device_args.get('force_gpu_idx',-1)
if device_args['force_gpu_idx'] == -1:
idxs_names_list = nnlib.device.getValidDevicesIdxsWithNamesList() idxs_names_list = nnlib.device.getValidDevicesIdxsWithNamesList()
if len(idxs_names_list) > 1: if len(idxs_names_list) > 1:
print ("You have multi GPUs in a system: ") io.log_info ("You have multi GPUs in a system: ")
for idx, name in idxs_names_list: for idx, name in idxs_names_list:
print ("[%d] : %s" % (idx, name) ) io.log_info ("[%d] : %s" % (idx, name) )
force_gpu_idx = input_int("Which GPU idx to choose? ( skip: best GPU ) : ", -1, [ x[0] for x in idxs_names_list] ) 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.force_gpu_idx = force_gpu_idx self.device_args = device_args
io.log_info ("Loading model...")
print ("Loading model...")
self.model_path = model_path self.model_path = model_path
self.model_data_path = Path( self.get_strpath_storage_for_file('data.dat') ) self.model_data_path = Path( self.get_strpath_storage_for_file('data.dat') )
@ -47,8 +50,6 @@ class ModelBase(object):
self.debug = debug self.debug = debug
self.is_training_mode = (training_data_src_path is not None and training_data_dst_path is not None) self.is_training_mode = (training_data_src_path is not None and training_data_dst_path is not None)
self.supress_std_once = os.environ.get('TF_SUPPRESS_STD', '0') == '1'
self.epoch = 0 self.epoch = 0
self.options = {} self.options = {}
self.loss_history = [] self.loss_history = []
@ -61,40 +62,40 @@ class ModelBase(object):
self.loss_history = model_data['loss_history'] if 'loss_history' in model_data.keys() else [] self.loss_history = model_data['loss_history'] if 'loss_history' in model_data.keys() else []
self.sample_for_preview = model_data['sample_for_preview'] if 'sample_for_preview' in model_data.keys() else None 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.epoch != 0 and input_in_time ("Press enter in 2 seconds to override model settings.", 2) ask_override = self.is_training_mode and self.epoch != 0 and io.input_in_time ("Press enter in 2 seconds to override model settings.", 2)
if self.epoch == 0: if self.epoch == 0:
print ("\nModel first run. Enter model options as default for each run.") io.log_info ("\nModel first run. Enter model options as default for each run.")
if self.epoch == 0 or ask_override: if self.epoch == 0 or ask_override:
default_write_preview_history = False if self.epoch == 0 else self.options.get('write_preview_history',False) default_write_preview_history = False if self.epoch == 0 else self.options.get('write_preview_history',False)
self.options['write_preview_history'] = input_bool("Write preview history? (y/n ?:help skip:n/default) : ", default_write_preview_history, help_message="Preview history will be writed to <ModelName>_history folder.") self.options['write_preview_history'] = io.input_bool("Write preview history? (y/n ?:help skip:n/default) : ", default_write_preview_history, help_message="Preview history will be writed to <ModelName>_history folder.")
else: else:
self.options['write_preview_history'] = self.options.get('write_preview_history', False) self.options['write_preview_history'] = self.options.get('write_preview_history', False)
if self.epoch == 0 or ask_override: if self.epoch == 0 or ask_override:
self.options['target_epoch'] = max(0, input_int("Target epoch (skip:unlimited/default) : ", 0)) self.options['target_epoch'] = max(0, io.input_int("Target epoch (skip:unlimited/default) : ", 0))
else: else:
self.options['target_epoch'] = self.options.get('target_epoch', 0) self.options['target_epoch'] = self.options.get('target_epoch', 0)
if self.epoch == 0 or ask_override: if self.epoch == 0 or ask_override:
default_batch_size = 0 if self.epoch == 0 else self.options.get('batch_size',0) default_batch_size = 0 if self.epoch == 0 else self.options.get('batch_size',0)
self.options['batch_size'] = max(0, input_int("Batch_size (?:help skip:0/default) : ", 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.")) self.options['batch_size'] = max(0, io.input_int("Batch_size (?:help skip:0/default) : ", 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: else:
self.options['batch_size'] = self.options.get('batch_size', 0) self.options['batch_size'] = self.options.get('batch_size', 0)
if self.epoch == 0: if self.epoch == 0:
self.options['sort_by_yaw'] = 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." ) 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: else:
self.options['sort_by_yaw'] = self.options.get('sort_by_yaw', False) self.options['sort_by_yaw'] = self.options.get('sort_by_yaw', False)
if self.epoch == 0: if self.epoch == 0:
self.options['random_flip'] = 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.") 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: else:
self.options['random_flip'] = self.options.get('random_flip', True) self.options['random_flip'] = self.options.get('random_flip', True)
if self.epoch == 0: if self.epoch == 0:
self.options['src_scale_mod'] = np.clip( 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) 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.options['src_scale_mod'] = self.options.get('src_scale_mod', 0)
@ -116,10 +117,10 @@ class ModelBase(object):
self.onInitializeOptions(self.epoch == 0, ask_override) self.onInitializeOptions(self.epoch == 0, ask_override)
nnlib.import_all ( nnlib.DeviceConfig(allow_growth=False, force_gpu_idx=self.force_gpu_idx, **in_options) ) nnlib.import_all ( nnlib.DeviceConfig(allow_growth=False, **self.device_args) )
self.device_config = nnlib.active_DeviceConfig self.device_config = nnlib.active_DeviceConfig
self.onInitialize(**in_options) self.onInitialize()
self.options['batch_size'] = self.batch_size self.options['batch_size'] = self.batch_size
@ -128,10 +129,10 @@ class ModelBase(object):
if self.is_training_mode: if self.is_training_mode:
if self.write_preview_history: if self.write_preview_history:
if self.force_gpu_idx == -1: if self.device_args['force_gpu_idx'] == -1:
self.preview_history_path = self.model_path / ( '%s_history' % (self.get_model_name()) ) self.preview_history_path = self.model_path / ( '%s_history' % (self.get_model_name()) )
else: else:
self.preview_history_path = self.model_path / ( '%d_%s_history' % (self.force_gpu_idx, self.get_model_name()) ) 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(): if not self.preview_history_path.exists():
self.preview_history_path.mkdir(exist_ok=True) self.preview_history_path.mkdir(exist_ok=True)
@ -141,11 +142,11 @@ class ModelBase(object):
Path(filename).unlink() Path(filename).unlink()
if self.generator_list is None: if self.generator_list is None:
raise Exception( 'You didnt set_training_data_generators()') raise ValueError( 'You didnt set_training_data_generators()')
else: else:
for i, generator in enumerate(self.generator_list): for i, generator in enumerate(self.generator_list):
if not isinstance(generator, SampleGeneratorBase): if not isinstance(generator, SampleGeneratorBase):
raise Exception('training data generator is not subclass of SampleGeneratorBase') raise ValueError('training data generator is not subclass of SampleGeneratorBase')
if (self.sample_for_preview is None) or (self.epoch == 0): if (self.sample_for_preview is None) or (self.epoch == 0):
self.sample_for_preview = self.generate_next_sample() self.sample_for_preview = self.generate_next_sample()
@ -181,14 +182,14 @@ class ModelBase(object):
model_summary_text += ["========================="] model_summary_text += ["========================="]
model_summary_text = "\r\n".join (model_summary_text) model_summary_text = "\r\n".join (model_summary_text)
self.model_summary_text = model_summary_text self.model_summary_text = model_summary_text
print(model_summary_text) io.log_info(model_summary_text)
#overridable #overridable
def onInitializeOptions(self, is_first_run, ask_override): def onInitializeOptions(self, is_first_run, ask_override):
pass pass
#overridable #overridable
def onInitialize(self, **in_options): def onInitialize(self):
''' '''
initialize your keras models initialize your keras models
@ -221,10 +222,9 @@ class ModelBase(object):
return Path(inspect.getmodule(self).__file__).parent.name.rsplit("_", 1)[1] return Path(inspect.getmodule(self).__file__).parent.name.rsplit("_", 1)[1]
#overridable #overridable
def get_converter(self, **in_options): def get_converter(self):
raise NotImplementeError
#return existing or your own converter which derived from base #return existing or your own converter which derived from base
from .ConverterBase import ConverterBase
return ConverterBase(self, **in_options)
def get_target_epoch(self): def get_target_epoch(self):
return self.target_epoch return self.target_epoch
@ -258,18 +258,11 @@ class ModelBase(object):
return self.onGetPreview (self.sample_for_preview)[0][1] #first preview, and bgr return self.onGetPreview (self.sample_for_preview)[0][1] #first preview, and bgr
def save(self): def save(self):
print ("Saving...") io.log_info ("Saving...")
if self.supress_std_once:
supressor = std_utils.suppress_stdout_stderr()
supressor.__enter__()
Path( self.get_strpath_storage_for_file('summary.txt') ).write_text(self.model_summary_text) Path( self.get_strpath_storage_for_file('summary.txt') ).write_text(self.model_summary_text)
self.onSave() self.onSave()
if self.supress_std_once:
supressor.__exit__()
model_data = { model_data = {
'epoch': self.epoch, 'epoch': self.epoch,
'options': self.options, 'options': self.options,
@ -311,10 +304,6 @@ class ModelBase(object):
return [next(generator) for generator in self.generator_list] return [next(generator) for generator in self.generator_list]
def train_one_epoch(self): def train_one_epoch(self):
if self.supress_std_once:
supressor = std_utils.suppress_stdout_stderr()
supressor.__enter__()
sample = self.generate_next_sample() sample = self.generate_next_sample()
epoch_time = time.time() epoch_time = time.time()
losses = self.onTrainOneEpoch(sample, self.generator_list) losses = self.onTrainOneEpoch(sample, self.generator_list)
@ -323,10 +312,6 @@ class ModelBase(object):
self.loss_history.append ( [float(loss[1]) for loss in losses] ) self.loss_history.append ( [float(loss[1]) for loss in losses] )
if self.supress_std_once:
supressor.__exit__()
self.supress_std_once = False
if self.write_preview_history: if self.write_preview_history:
if self.epoch % 10 == 0: if self.epoch % 10 == 0:
preview = self.get_static_preview() preview = self.get_static_preview()
@ -377,10 +362,10 @@ class ModelBase(object):
return self.generator_list return self.generator_list
def get_strpath_storage_for_file(self, filename): def get_strpath_storage_for_file(self, filename):
if self.force_gpu_idx == -1: if self.device_args['force_gpu_idx'] == -1:
return str( self.model_path / ( self.get_model_name() + '_' + filename) ) return str( self.model_path / ( self.get_model_name() + '_' + filename) )
else: else:
return str( self.model_path / ( str(self.force_gpu_idx) + '_' + self.get_model_name() + '_' + filename) ) return str( self.model_path / ( str(self.device_args['force_gpu_idx']) + '_' + self.get_model_name() + '_' + filename) )
def set_vram_batch_requirements (self, d): 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}

View file

@ -4,7 +4,7 @@ from nnlib import nnlib
from models import ModelBase from models import ModelBase
from facelib import FaceType from facelib import FaceType
from samples import * from samples import *
from utils.console_utils import * from interact import interact as io
class Model(ModelBase): class Model(ModelBase):
@ -12,12 +12,12 @@ class Model(ModelBase):
def onInitializeOptions(self, is_first_run, ask_override): def onInitializeOptions(self, is_first_run, ask_override):
if is_first_run or ask_override: if is_first_run or ask_override:
def_pixel_loss = self.options.get('pixel_loss', False) def_pixel_loss = self.options.get('pixel_loss', False)
self.options['pixel_loss'] = 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 epochs to enhance fine details and decrease face jitter.") self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 20k epochs to enhance fine details and decrease face jitter.")
else: else:
self.options['pixel_loss'] = self.options.get('pixel_loss', False) self.options['pixel_loss'] = self.options.get('pixel_loss', False)
#override #override
def onInitialize(self, **in_options): def onInitialize(self):
exec(nnlib.import_all(), locals(), globals()) exec(nnlib.import_all(), locals(), globals())
self.set_vram_batch_requirements( {4.5:4} ) self.set_vram_batch_requirements( {4.5:4} )
@ -113,15 +113,14 @@ class Model(ModelBase):
return np.concatenate ( (x,mx), -1 ) return np.concatenate ( (x,mx), -1 )
#override #override
def get_converter(self, **in_options): def get_converter(self):
from models import ConverterMasked from converters import ConverterMasked
return ConverterMasked(self.predictor_func, return ConverterMasked(self.predictor_func,
predictor_input_size=128, predictor_input_size=128,
output_size=128, output_size=128,
face_type=FaceType.FULL, face_type=FaceType.FULL,
base_erode_mask_modifier=30, base_erode_mask_modifier=30,
base_blur_mask_modifier=0, base_blur_mask_modifier=0)
**in_options)
def Build(self, input_layer): def Build(self, input_layer):
exec(nnlib.code_import_all, locals(), globals()) exec(nnlib.code_import_all, locals(), globals())

View file

@ -4,14 +4,14 @@ from nnlib import nnlib
from models import ModelBase from models import ModelBase
from facelib import FaceType from facelib import FaceType
from samples import * from samples import *
from utils.console_utils import * from interact import interact as io
class Model(ModelBase): class Model(ModelBase):
#override #override
def onInitializeOptions(self, is_first_run, ask_override): def onInitializeOptions(self, is_first_run, ask_override):
if is_first_run: if is_first_run:
self.options['lighter_ae'] = 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.") 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: else:
default_lighter_ae = self.options.get('created_vram_gb', 99) <= 4 #temporally support old models, deprecate in future default_lighter_ae = self.options.get('created_vram_gb', 99) <= 4 #temporally support old models, deprecate in future
if 'created_vram_gb' in self.options.keys(): if 'created_vram_gb' in self.options.keys():
@ -20,12 +20,12 @@ class Model(ModelBase):
if is_first_run or ask_override: if is_first_run or ask_override:
def_pixel_loss = self.options.get('pixel_loss', False) def_pixel_loss = self.options.get('pixel_loss', False)
self.options['pixel_loss'] = 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 epochs to enhance fine details and decrease face jitter.") self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 20k epochs to enhance fine details and decrease face jitter.")
else: else:
self.options['pixel_loss'] = self.options.get('pixel_loss', False) self.options['pixel_loss'] = self.options.get('pixel_loss', False)
#override #override
def onInitialize(self, **in_options): def onInitialize(self):
exec(nnlib.import_all(), locals(), globals()) exec(nnlib.import_all(), locals(), globals())
self.set_vram_batch_requirements( {2.5:4} ) self.set_vram_batch_requirements( {2.5:4} )
@ -125,15 +125,14 @@ class Model(ModelBase):
return np.concatenate ( (x,mx), -1 ) return np.concatenate ( (x,mx), -1 )
#override #override
def get_converter(self, **in_options): def get_converter(self):
from models import ConverterMasked from converters import ConverterMasked
return ConverterMasked(self.predictor_func, return ConverterMasked(self.predictor_func,
predictor_input_size=128, predictor_input_size=128,
output_size=128, output_size=128,
face_type=FaceType.HALF, face_type=FaceType.HALF,
base_erode_mask_modifier=100, base_erode_mask_modifier=100,
base_blur_mask_modifier=100, base_blur_mask_modifier=100)
**in_options)
def Build(self, lighter_ae): def Build(self, lighter_ae):
exec(nnlib.code_import_all, locals(), globals()) exec(nnlib.code_import_all, locals(), globals())

View file

@ -4,14 +4,14 @@ from nnlib import nnlib
from models import ModelBase from models import ModelBase
from facelib import FaceType from facelib import FaceType
from samples import * from samples import *
from utils.console_utils import * from interact import interact as io
class Model(ModelBase): class Model(ModelBase):
#override #override
def onInitializeOptions(self, is_first_run, ask_override): def onInitializeOptions(self, is_first_run, ask_override):
if is_first_run: if is_first_run:
self.options['lighter_ae'] = 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.") 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: else:
default_lighter_ae = self.options.get('created_vram_gb', 99) <= 4 #temporally support old models, deprecate in future default_lighter_ae = self.options.get('created_vram_gb', 99) <= 4 #temporally support old models, deprecate in future
if 'created_vram_gb' in self.options.keys(): if 'created_vram_gb' in self.options.keys():
@ -20,12 +20,12 @@ class Model(ModelBase):
if is_first_run or ask_override: if is_first_run or ask_override:
def_pixel_loss = self.options.get('pixel_loss', False) def_pixel_loss = self.options.get('pixel_loss', False)
self.options['pixel_loss'] = 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 epochs to enhance fine details and decrease face jitter.") self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 20k epochs to enhance fine details and decrease face jitter.")
else: else:
self.options['pixel_loss'] = self.options.get('pixel_loss', False) self.options['pixel_loss'] = self.options.get('pixel_loss', False)
#override #override
def onInitialize(self, **in_options): def onInitialize(self):
exec(nnlib.import_all(), locals(), globals()) exec(nnlib.import_all(), locals(), globals())
self.set_vram_batch_requirements( {1.5:4} ) self.set_vram_batch_requirements( {1.5:4} )
@ -127,15 +127,14 @@ class Model(ModelBase):
return np.concatenate ( (x,mx), -1 ) return np.concatenate ( (x,mx), -1 )
#override #override
def get_converter(self, **in_options): def get_converter(self):
from models import ConverterMasked from converters import ConverterMasked
return ConverterMasked(self.predictor_func, return ConverterMasked(self.predictor_func,
predictor_input_size=64, predictor_input_size=64,
output_size=64, output_size=64,
face_type=FaceType.HALF, face_type=FaceType.HALF,
base_erode_mask_modifier=100, base_erode_mask_modifier=100,
base_blur_mask_modifier=100, base_blur_mask_modifier=100)
**in_options)
def Build(self, lighter_ae): def Build(self, lighter_ae):
exec(nnlib.code_import_all, locals(), globals()) exec(nnlib.code_import_all, locals(), globals())

View file

@ -4,7 +4,7 @@ from nnlib import nnlib
from models import ModelBase from models import ModelBase
from facelib import FaceType from facelib import FaceType
from samples import * from samples import *
from utils.console_utils import * from interact import interact as io
class Model(ModelBase): class Model(ModelBase):
@ -12,12 +12,12 @@ class Model(ModelBase):
def onInitializeOptions(self, is_first_run, ask_override): def onInitializeOptions(self, is_first_run, ask_override):
if is_first_run or ask_override: if is_first_run or ask_override:
def_pixel_loss = self.options.get('pixel_loss', False) def_pixel_loss = self.options.get('pixel_loss', False)
self.options['pixel_loss'] = 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 epochs to enhance fine details and decrease face jitter.") self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 20k epochs to enhance fine details and decrease face jitter.")
else: else:
self.options['pixel_loss'] = self.options.get('pixel_loss', False) self.options['pixel_loss'] = self.options.get('pixel_loss', False)
#override #override
def onInitialize(self, **in_options): def onInitialize(self):
exec(nnlib.import_all(), locals(), globals()) exec(nnlib.import_all(), locals(), globals())
self.set_vram_batch_requirements( {4.5:4} ) self.set_vram_batch_requirements( {4.5:4} )
@ -121,15 +121,14 @@ class Model(ModelBase):
return np.concatenate ( (x,mx), -1 ) return np.concatenate ( (x,mx), -1 )
#override #override
def get_converter(self, **in_options): def get_converter(self):
from models import ConverterMasked from converters import ConverterMasked
return ConverterMasked(self.predictor_func, return ConverterMasked(self.predictor_func,
predictor_input_size=128, predictor_input_size=128,
output_size=128, output_size=128,
face_type=FaceType.FULL, face_type=FaceType.FULL,
base_erode_mask_modifier=30, base_erode_mask_modifier=30,
base_blur_mask_modifier=0, base_blur_mask_modifier=0)
**in_options)
def Build(self, input_layer): def Build(self, input_layer):
exec(nnlib.code_import_all, locals(), globals()) exec(nnlib.code_import_all, locals(), globals())

View file

@ -4,7 +4,7 @@ from nnlib import nnlib
from models import ModelBase from models import ModelBase
from facelib import FaceType from facelib import FaceType
from samples import * from samples import *
from utils.console_utils import * from interact import interact as io
#SAE - Styled AutoEncoder #SAE - Styled AutoEncoder
class SAEModel(ModelBase): class SAEModel(ModelBase):
@ -27,15 +27,15 @@ class SAEModel(ModelBase):
default_face_type = 'f' default_face_type = 'f'
if is_first_run: if is_first_run:
resolution = 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 = 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: while np.modf(resolution / 16)[0] != 0.0:
resolution -= 1 resolution -= 1
self.options['resolution'] = resolution self.options['resolution'] = resolution
self.options['face_type'] = 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'] = 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.") 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.")
self.options['archi'] = 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() 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: else:
self.options['resolution'] = self.options.get('resolution', default_resolution) self.options['resolution'] = self.options.get('resolution', default_resolution)
self.options['face_type'] = self.options.get('face_type', default_face_type) self.options['face_type'] = self.options.get('face_type', default_face_type)
@ -45,17 +45,17 @@ class SAEModel(ModelBase):
default_ae_dims = 256 if self.options['archi'] == 'liae' else 512 default_ae_dims = 256 if self.options['archi'] == 'liae' else 512
default_ed_ch_dims = 42 default_ed_ch_dims = 42
if is_first_run: if is_first_run:
self.options['ae_dims'] = np.clip ( input_int("AutoEncoder dims (32-1024 ?:help skip:%d) : " % (default_ae_dims) , default_ae_dims, help_message="More dims are better, but requires more VRAM. You can fine-tune model size to fit your GPU." ), 32, 1024 ) self.options['ae_dims'] = np.clip ( io.input_int("AutoEncoder dims (32-1024 ?:help skip:%d) : " % (default_ae_dims) , default_ae_dims, help_message="More dims are better, but requires more VRAM. You can fine-tune model size to fit your GPU." ), 32, 1024 )
self.options['ed_ch_dims'] = np.clip ( input_int("Encoder/Decoder dims per channel (21-85 ?:help skip:%d) : " % (default_ed_ch_dims) , default_ed_ch_dims, help_message="More dims are better, but requires more VRAM. You can fine-tune model size to fit your GPU." ), 21, 85 ) self.options['ed_ch_dims'] = np.clip ( io.input_int("Encoder/Decoder dims per channel (21-85 ?:help skip:%d) : " % (default_ed_ch_dims) , default_ed_ch_dims, help_message="More dims are better, but requires more VRAM. You can fine-tune model size to fit your GPU." ), 21, 85 )
else: else:
self.options['ae_dims'] = self.options.get('ae_dims', default_ae_dims) 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['ed_ch_dims'] = self.options.get('ed_ch_dims', default_ed_ch_dims)
if is_first_run: if is_first_run:
self.options['lighter_encoder'] = 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.") 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': if self.options['archi'] != 'vg':
self.options['multiscale_decoder'] = input_bool ("Use multiscale decoder? (y/n, ?:help skip:y) : ", True, help_message="Multiscale decoder helps to get better details.") self.options['multiscale_decoder'] = io.input_bool ("Use multiscale decoder? (y/n, ?:help skip:y) : ", True, help_message="Multiscale decoder helps to get better details.")
else: else:
self.options['lighter_encoder'] = self.options.get('lighter_encoder', False) self.options['lighter_encoder'] = self.options.get('lighter_encoder', False)
@ -66,14 +66,14 @@ class SAEModel(ModelBase):
default_bg_style_power = 0.0 default_bg_style_power = 0.0
if is_first_run or ask_override: if is_first_run or ask_override:
def_pixel_loss = self.options.get('pixel_loss', False) def_pixel_loss = self.options.get('pixel_loss', False)
self.options['pixel_loss'] = 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 15-25k epochs to enhance fine details and decrease face jitter.") self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 15-25k epochs 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) 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 ( input_number("Face style power ( 0.0 .. 100.0 ?:help skip:%.2f) : " % (default_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 epochs, when predicted face is clear enough to start learn style. Start from 0.1 value and check history changes."), 0.0, 100.0 ) help_message="Learn to transfer face style details such as light and color conditions. Warning: Enable it only after 10k epochs, 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) 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 ( input_number("Background style power ( 0.0 .. 100.0 ?:help skip:%.2f) : " % (default_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 ) help_message="Learn to transfer image around face. This can make face more like dst."), 0.0, 100.0 )
else: else:
self.options['pixel_loss'] = self.options.get('pixel_loss', False) self.options['pixel_loss'] = self.options.get('pixel_loss', False)
@ -81,7 +81,7 @@ class SAEModel(ModelBase):
self.options['bg_style_power'] = self.options.get('bg_style_power', default_bg_style_power) self.options['bg_style_power'] = self.options.get('bg_style_power', default_bg_style_power)
#override #override
def onInitialize(self, **in_options): def onInitialize(self):
exec(nnlib.import_all(), locals(), globals()) exec(nnlib.import_all(), locals(), globals())
SAEModel.initialize_nn_functions() SAEModel.initialize_nn_functions()
@ -427,9 +427,7 @@ class SAEModel(ModelBase):
return np.concatenate ( [prd[0], prd[1]], -1 ) return np.concatenate ( [prd[0], prd[1]], -1 )
#override #override
def get_converter(self, **in_options): def get_converter(self):
from models import ConverterMasked
base_erode_mask_modifier = 30 if self.options['face_type'] == 'f' else 100 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 base_blur_mask_modifier = 0 if self.options['face_type'] == 'f' else 100
@ -439,6 +437,7 @@ class SAEModel(ModelBase):
face_type = FaceType.FULL if self.options['face_type'] == 'f' else FaceType.HALF 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'], predictor_input_size=self.options['resolution'],
output_size=self.options['resolution'], output_size=self.options['resolution'],
@ -448,8 +447,7 @@ class SAEModel(ModelBase):
base_blur_mask_modifier=base_blur_mask_modifier, base_blur_mask_modifier=base_blur_mask_modifier,
default_erode_mask_modifier=default_erode_mask_modifier, default_erode_mask_modifier=default_erode_mask_modifier,
default_blur_mask_modifier=default_blur_mask_modifier, default_blur_mask_modifier=default_blur_mask_modifier,
clip_hborder_mask_per=0.0625 if self.options['face_type'] == 'f' else 0, clip_hborder_mask_per=0.0625 if self.options['face_type'] == 'f' else 0)
**in_options)
@staticmethod @staticmethod
def initialize_nn_functions(): def initialize_nn_functions():

View file

@ -1,7 +1,4 @@
from .ModelBase import ModelBase from .ModelBase import ModelBase
from .ConverterBase import ConverterBase
from .ConverterMasked import ConverterMasked
from .ConverterImage import ConverterImage
def import_model(name): def import_model(name):
module = __import__('Model_'+name, globals(), locals(), [], 1) module = __import__('Model_'+name, globals(), locals(), [], 1)

View file

@ -2,7 +2,6 @@ import traceback
from enum import IntEnum from enum import IntEnum
import cv2 import cv2
import numpy as np import numpy as np
from tqdm import tqdm
from pathlib import Path from pathlib import Path
from utils import Path_utils from utils import Path_utils
@ -14,6 +13,7 @@ from .Sample import SampleType
from facelib import FaceType from facelib import FaceType
from facelib import LandmarksProcessor from facelib import LandmarksProcessor
from interact import interact as io
class SampleLoader: class SampleLoader:
cache = dict() cache = dict()
@ -29,7 +29,7 @@ class SampleLoader:
if sample_type == SampleType.IMAGE: 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 tqdm( Path_utils.get_image_paths(samples_path), desc="Loading", ascii=True ) ] 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: elif sample_type == SampleType.FACE:
if datas[sample_type] is None: if datas[sample_type] is None:
@ -55,16 +55,17 @@ class SampleLoader:
def upgradeToFaceSamples ( samples ): def upgradeToFaceSamples ( samples ):
sample_list = [] sample_list = []
for s in tqdm( samples, desc="Loading", ascii=True ): for s in io.progress_bar_generator(samples, "Loading"):
s_filename_path = Path(s.filename) s_filename_path = Path(s.filename)
try: try:
if s_filename_path.suffix == '.png': if s_filename_path.suffix == '.png':
dflimg = DFLPNG.load ( str(s_filename_path), print_on_no_embedded_data=True ) dflimg = DFLPNG.load ( str(s_filename_path) )
if dflimg is None: continue
elif s_filename_path.suffix == '.jpg': elif s_filename_path.suffix == '.jpg':
dflimg = DFLJPG.load ( str(s_filename_path), print_on_no_embedded_data=True ) dflimg = DFLJPG.load ( str(s_filename_path) )
if dflimg is None: continue
else: 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 continue
@ -87,7 +88,7 @@ class SampleLoader:
yaw_samples_len = len(yaw_samples) yaw_samples_len = len(yaw_samples)
sample_list = [] sample_list = []
for i in tqdm( range(yaw_samples_len), desc="Sorting", ascii=True ): for i in io.progress_bar_generator( range(yaw_samples_len), "Sorting"):
if yaw_samples[i] is not None: if yaw_samples[i] is not None:
for s in yaw_samples[i]: for s in yaw_samples[i]:
s_t = [] s_t = []
@ -123,7 +124,7 @@ class SampleLoader:
yaws_sample_list = [None]*gradations yaws_sample_list = [None]*gradations
for i in tqdm( range(0, gradations), desc="Sorting", ascii=True ): for i in io.progress_bar_generator(range(gradations), "Sorting"):
yaw = lowest_yaw + i*diff_rot_per_grad yaw = lowest_yaw + i*diff_rot_per_grad
next_yaw = lowest_yaw + (i+1)*diff_rot_per_grad next_yaw = lowest_yaw + (i+1)*diff_rot_per_grad

View file

@ -113,7 +113,7 @@ class DFLJPG(object):
return None return None
@staticmethod @staticmethod
def load(filename, print_on_no_embedded_data=False, throw_on_no_embedded_data=False): def load(filename):
inst = DFLJPG.load_raw (filename) inst = DFLJPG.load_raw (filename)
inst.dfl_dict = None inst.dfl_dict = None
@ -141,10 +141,6 @@ class DFLJPG(object):
inst.dfl_dict['face_type'] = FaceType.toString (FaceType.FULL) inst.dfl_dict['face_type'] = FaceType.toString (FaceType.FULL)
if inst.dfl_dict == None: if inst.dfl_dict == None:
if print_on_no_embedded_data:
print ( "No DFL data found in %s" % (filename) )
if throw_on_no_embedded_data:
raise ValueError("No DFL data found in %s" % (filename) )
return None return None
return inst return inst

View file

@ -248,7 +248,7 @@ class DFLPNG(object):
return inst return inst
@staticmethod @staticmethod
def load(filename, print_on_no_embedded_data=False, throw_on_no_embedded_data=False): def load(filename):
inst = DFLPNG.load_raw (filename) inst = DFLPNG.load_raw (filename)
inst.fcwp_dict = inst.getDFLDictData() inst.fcwp_dict = inst.getDFLDictData()
@ -256,10 +256,6 @@ class DFLPNG(object):
inst.fcwp_dict['face_type'] = FaceType.toString (FaceType.FULL) inst.fcwp_dict['face_type'] = FaceType.toString (FaceType.FULL)
if inst.fcwp_dict == None: if inst.fcwp_dict == None:
if print_on_no_embedded_data:
print ( "No DFL data found in %s" % (filename) )
if throw_on_no_embedded_data:
raise ValueError("No DFL data found in %s" % (filename) )
return None return None
return inst return inst

View file

@ -13,7 +13,7 @@ def get_image_paths(dir_path):
result.append(x.path) result.append(x.path)
return result return result
def get_image_unique_filestem_paths(dir_path, verbose=False): def get_image_unique_filestem_paths(dir_path, verbose_print_func=None):
result = get_image_paths(dir_path) result = get_image_paths(dir_path)
result_dup = set() result_dup = set()
@ -21,8 +21,8 @@ def get_image_unique_filestem_paths(dir_path, verbose=False):
f_stem = Path(f).stem f_stem = Path(f).stem
if f_stem in result_dup: if f_stem in result_dup:
result.remove(f) result.remove(f)
if verbose: if verbose_print_func is not None:
print ("Duplicate filenames are not allowed, skipping: %s" % Path(f).name ) verbose_print_func ("Duplicate filenames are not allowed, skipping: %s" % Path(f).name )
continue continue
result_dup.add(f_stem) result_dup.add(f_stem)

View file

@ -1,257 +0,0 @@
import traceback
from tqdm import tqdm
import multiprocessing
import time
import sys
class SubprocessorBase(object):
#overridable
def __init__(self, name, no_response_time_sec = 60):
self.name = name
self.no_response_time_sec = no_response_time_sec
self.is_host = True
#overridable
def process_info_generator(self):
#yield name, host_dict, client_dict - per process
yield 'first process', {}, {}
#overridable
def get_no_process_started_message(self):
return "No process started."
#overridable
def onHostGetProgressBarDesc(self):
return "Processing"
#overridable
def onHostGetProgressBarLen(self):
return 0
#overridable
def onHostGetData(self, host_dict):
#return data here
return None
#overridable
def onHostDataReturn (self, host_dict, data):
#input_data.insert(0, obj['data'])
pass
#overridable
def onClientInitialize(self, client_dict):
#return fail message or None if ok
return None
#overridable
def onClientFinalize(self):
pass
#overridable
def onClientProcessData(self, data):
#return result object
return None
#overridable
def onClientGetDataName (self, data):
#return string identificator of your data
return "undefined"
#overridable
def onHostClientsInitialized(self):
pass
#overridable
def onHostResult (self, host_dict, data, result):
#return count of progress bar update
return 1
#overridable
def onFinalizeAndGetResult(self):
return None
def inc_progress_bar(self, c):
if self.is_host:
self.progress_bar.n += c
self.progress_bar.refresh()
else:
self.cq.put ( {'op': 'inc_bar', 'c':c} )
def safe_print(self, msg):
self.print_lock.acquire()
print (msg)
self.print_lock.release()
def process(self):
#returns start_return
self.processes = []
self.print_lock = multiprocessing.Lock()
for name, host_dict, client_dict in self.process_info_generator():
sq = multiprocessing.Queue()
cq = multiprocessing.Queue()
client_dict.update ( {'print_lock' : self.print_lock} )
try:
p = multiprocessing.Process(target=self.subprocess, args=(sq,cq,client_dict))
p.daemon = True
p.start()
self.processes.append ( { 'process' : p,
'sq' : sq,
'cq' : cq,
'state' : 'busy',
'sent_time': time.time(),
'name': name,
'host_dict' : host_dict
} )
except:
print ("Unable to start subprocess %s" % (name))
if len(self.processes) == 0:
raise Exception ("Unable to start Subprocessor '%s' " % (self.name))
while True:
for p in self.processes[:]:
while not p['cq'].empty():
obj = p['cq'].get()
obj_op = obj['op']
if obj_op == 'init_ok':
p['state'] = 'free'
elif obj_op == 'error':
if obj['close'] == True:
p['process'].terminate()
p['process'].join()
self.processes.remove(p)
break
if all ([ p['state'] == 'free' for p in self.processes ] ):
break
if len(self.processes) == 0:
print ( self.get_no_process_started_message() )
return None
self.progress_bar = tqdm( total=self.onHostGetProgressBarLen(), desc=self.onHostGetProgressBarDesc(), ascii=True )
self.onHostClientsInitialized()
try:
while True:
for p in self.processes[:]:
while not p['cq'].empty():
obj = p['cq'].get()
obj_op = obj['op']
if obj_op == 'success':
data = obj['data']
result = obj['result']
c = self.onHostResult (p['host_dict'], data, result)
if c > 0:
self.progress_bar.update(c)
p['state'] = 'free'
elif obj_op == 'inc_bar':
self.inc_progress_bar(obj['c'])
elif obj_op == 'error':
if 'data' in obj.keys():
self.onHostDataReturn (p['host_dict'], obj['data'] )
if obj['close'] == True:
p['sq'].put ( {'op': 'close'} )
p['process'].join()
self.processes.remove(p)
break
p['state'] = 'free'
for p in self.processes[:]:
if p['state'] == 'free':
data = self.onHostGetData(p['host_dict'])
if data is not None:
p['sq'].put ( {'op': 'data', 'data' : data} )
p['sent_time'] = time.time()
p['sent_data'] = data
p['state'] = 'busy'
elif p['state'] == 'busy':
if self.no_response_time_sec != 0 and (time.time() - p['sent_time']) > self.no_response_time_sec:
print ( '%s doesnt response, terminating it.' % (p['name']) )
self.onHostDataReturn (p['host_dict'], p['sent_data'] )
p['process'].terminate()
self.processes.remove(p)
if all ([p['state'] == 'free' for p in self.processes]):
break
time.sleep(0.005)
except:
print ("Exception occured in Subprocessor.start(): %s" % (traceback.format_exc()) )
self.progress_bar.close()
for p in self.processes[:]:
p['sq'].put ( {'op': 'close'} )
p['sent_time'] = time.time()
while True:
for p in self.processes[:]:
terminate_it = False
while not p['cq'].empty():
obj = p['cq'].get()
obj_op = obj['op']
if obj_op == 'finalized':
terminate_it = True
break
if self.no_response_time_sec != 0 and (time.time() - p['sent_time']) > self.no_response_time_sec:
terminate_it = True
if terminate_it:
p['state'] = 'finalized'
p['process'].terminate()
if all ([p['state'] == 'finalized' for p in self.processes]):
break
return self.onFinalizeAndGetResult()
def subprocess(self, sq, cq, client_dict):
self.is_host = False
self.print_lock = client_dict['print_lock']
self.cq = cq
try:
fail_message = self.onClientInitialize(client_dict)
except:
fail_message = 'Exception while initialization: %s' % (traceback.format_exc())
if fail_message is None:
cq.put ( {'op': 'init_ok'} )
else:
print (fail_message)
cq.put ( {'op': 'error', 'close': True} )
return
while True:
obj = sq.get()
obj_op = obj['op']
if obj_op == 'data':
data = obj['data']
try:
result = self.onClientProcessData (data)
cq.put ( {'op': 'success', 'data' : data, 'result' : result} )
except:
print ( 'Exception while process data [%s]: %s' % (self.onClientGetDataName(data), traceback.format_exc()) )
cq.put ( {'op': 'error', 'close': True, 'data' : data } )
elif obj_op == 'close':
break
time.sleep(0.005)
self.onClientFinalize()
cq.put ( {'op': 'finalized'} )
while True:
time.sleep(0.1)

View file

@ -1,100 +0,0 @@
import os
import sys
import time
import multiprocessing
def input_number(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
return i
except:
print (default_value)
return default_value
def input_int(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
return i
except:
print (default_value)
return default_value
def input_bool(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(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
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(stdin_fd, sq, str):
sys.stdin = os.fdopen(stdin_fd)
try:
inp = input (str)
sq.put (True)
except:
sq.put (False)
def input_in_time (str, max_time_sec):
sq = multiprocessing.Queue()
p = multiprocessing.Process(target=input_process, args=( sys.stdin.fileno(), sq, str))
p.start()
t = time.time()
inp = False
while True:
if not sq.empty():
inp = sq.get()
break
if time.time() - t > max_time_sec:
break
p.terminate()
sys.stdin = os.fdopen( sys.stdin.fileno() )
return inp