mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-07-06 13:02:15 -07:00
5ac7e5d7f1
Pixel loss may help to enhance fine details and stabilize face color. Use it only if quality does not improve over time. SAE: previous SAE model will not work with this update. Greatly decreased chance of model collapse. Increased model accuracy. Residual blocks now default and this option has been removed. Improved 'learn mask'. Added masked preview (switch by space key) Converter: fixed rct/lct in seamless mode added mask mode (6) learned*FAN-prd*FAN-dst added mask editor, its created for refining dataset for FANSeg model, and not for production, but you can spend your time and test it in regular fakes with face obstructions
347 lines
11 KiB
Python
347 lines
11 KiB
Python
import os
|
|
import sys
|
|
import time
|
|
import types
|
|
import multiprocessing
|
|
import cv2
|
|
from tqdm import tqdm
|
|
|
|
try:
|
|
import IPython #if success we are in colab
|
|
from IPython.display import display, clear_output
|
|
import PIL
|
|
import matplotlib.pyplot as plt
|
|
is_colab = True
|
|
except:
|
|
is_colab = False
|
|
|
|
class InteractBase(object):
|
|
EVENT_LBUTTONDOWN = 1
|
|
EVENT_LBUTTONUP = 2
|
|
EVENT_MBUTTONDOWN = 3
|
|
EVENT_MBUTTONUP = 4
|
|
EVENT_RBUTTONDOWN = 5
|
|
EVENT_RBUTTONUP = 6
|
|
EVENT_MOUSEWHEEL = 10
|
|
|
|
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 is_colab(self):
|
|
return False
|
|
|
|
def on_destroy_all_windows(self):
|
|
raise NotImplemented
|
|
|
|
def on_create_window (self, wnd_name):
|
|
raise NotImplemented
|
|
|
|
def on_show_image (self, wnd_name, img):
|
|
raise NotImplemented
|
|
|
|
def on_capture_mouse (self, wnd_name):
|
|
raise NotImplemented
|
|
|
|
def on_capture_keys (self, wnd_name):
|
|
raise NotImplemented
|
|
|
|
def on_process_messages(self, sleep_time=0):
|
|
raise NotImplemented
|
|
|
|
def on_wait_any_key(self):
|
|
raise NotImplemented
|
|
|
|
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:
|
|
self.on_destroy_all_windows()
|
|
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
|
|
self.on_create_window(wnd_name)
|
|
if wnd_name in self.capture_mouse_windows:
|
|
self.capture_mouse(wnd_name)
|
|
self.on_show_image(wnd_name,img)
|
|
else: print("show_image: named_window ", wnd_name, " not found.")
|
|
|
|
def capture_mouse(self, wnd_name):
|
|
if wnd_name in self.named_windows:
|
|
self.capture_mouse_windows[wnd_name] = True
|
|
if self.named_windows[wnd_name] == 1:
|
|
self.on_capture_mouse(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
|
|
self.on_capture_keys(wnd_name)
|
|
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):
|
|
self.on_process_messages(sleep_time)
|
|
|
|
def wait_any_key(self):
|
|
self.on_wait_any_key()
|
|
|
|
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
|
|
|
|
|
|
|
|
class InteractDesktop(InteractBase):
|
|
|
|
def on_destroy_all_windows(self):
|
|
cv2.destroyAllWindows()
|
|
|
|
def on_create_window (self, wnd_name):
|
|
cv2.namedWindow(wnd_name)
|
|
|
|
def on_show_image (self, wnd_name, img):
|
|
cv2.imshow (wnd_name, img)
|
|
|
|
def on_capture_mouse (self, wnd_name):
|
|
def onMouse(event, x, y, flags, param):
|
|
(inst, wnd_name) = param
|
|
if event == cv2.EVENT_LBUTTONDOWN: ev = InteractBase.EVENT_LBUTTONDOWN
|
|
elif event == cv2.EVENT_LBUTTONUP: ev = InteractBase.EVENT_LBUTTONUP
|
|
elif event == cv2.EVENT_RBUTTONDOWN: ev = InteractBase.EVENT_RBUTTONDOWN
|
|
elif event == cv2.EVENT_RBUTTONUP: ev = InteractBase.EVENT_RBUTTONUP
|
|
elif event == cv2.EVENT_MBUTTONDOWN: ev = InteractBase.EVENT_MBUTTONDOWN
|
|
elif event == cv2.EVENT_MBUTTONUP: ev = InteractBase.EVENT_MBUTTONUP
|
|
elif event == cv2.EVENT_MOUSEWHEEL: ev = InteractBase.EVENT_MOUSEWHEEL
|
|
|
|
else: ev = 0
|
|
inst.add_mouse_event (wnd_name, x, y, ev, flags)
|
|
cv2.setMouseCallback(wnd_name, onMouse, (self,wnd_name) )
|
|
|
|
def on_capture_keys (self, wnd_name):
|
|
pass
|
|
|
|
def on_process_messages(self, sleep_time=0):
|
|
|
|
has_windows = False
|
|
has_capture_keys = False
|
|
|
|
if len(self.named_windows) != 0:
|
|
has_windows = True
|
|
|
|
if len(self.capture_keys_windows) != 0:
|
|
has_capture_keys = True
|
|
|
|
if has_windows or has_capture_keys:
|
|
wait_key_time = max(1, int(sleep_time*1000) )
|
|
key = cv2.waitKey(wait_key_time) & 0xFF
|
|
else:
|
|
if sleep_time != 0:
|
|
time.sleep(sleep_time)
|
|
|
|
if has_capture_keys and key != 255:
|
|
for wnd_name in self.capture_keys_windows:
|
|
self.add_key_event (wnd_name, key)
|
|
|
|
def on_wait_any_key(self):
|
|
cv2.waitKey(0)
|
|
|
|
class InteractColab(InteractBase):
|
|
|
|
def is_colab(self):
|
|
return True
|
|
|
|
def on_destroy_all_windows(self):
|
|
pass
|
|
#clear_output()
|
|
|
|
def on_create_window (self, wnd_name):
|
|
pass
|
|
#clear_output()
|
|
|
|
def on_show_image (self, wnd_name, img):
|
|
pass
|
|
# # cv2 stores colors as BGR; convert to RGB
|
|
# if img.ndim == 3:
|
|
# if img.shape[2] == 4:
|
|
# img = cv2.cvtColor(img, cv2.COLOR_BGRA2RGBA)
|
|
# else:
|
|
# img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
|
# img = PIL.Image.fromarray(img)
|
|
# plt.imshow(img)
|
|
# plt.show()
|
|
|
|
def on_capture_mouse (self, wnd_name):
|
|
pass
|
|
#print("on_capture_mouse(): Colab does not support")
|
|
|
|
def on_capture_keys (self, wnd_name):
|
|
pass
|
|
#print("on_capture_keys(): Colab does not support")
|
|
|
|
def on_process_messages(self, sleep_time=0):
|
|
time.sleep(sleep_time)
|
|
|
|
def on_wait_any_key(self):
|
|
pass
|
|
#print("on_wait_any_key(): Colab does not support")
|
|
|
|
if is_colab:
|
|
interact = InteractColab()
|
|
else:
|
|
interact = InteractDesktop()
|