from pathlib import Path

from .DFLJPG import DFLJPG

class DFLIMG():

    @staticmethod
    def load(filepath, loader_func=None):
        if filepath.suffix == '.jpg':
            return DFLJPG.load ( str(filepath), loader_func=loader_func )
        else:
            return None

DFLIMG/DFLIMG.py
oo92/TorchFaceLab
import pickle
import struct
import traceback

import cv2
import numpy as np

from core import imagelib
from core.cv2ex import *
from core.imagelib import SegIEPolys
from core.interact import interact as io
from core.structex import *
from facelib import FaceType


class DFLJPG(object):
    def __init__(self, filename):
        self.filename = filename
        self.data = b""
        self.length = 0
        self.chunks = []
        self.dfl_dict = None
        self.shape = None
        self.img = None

    @staticmethod
    def load_raw(filename, loader_func=None):
        try:
            if loader_func is not None:
                data = loader_func(filename)
            else:
                with open(filename, "rb") as f:
                    data = f.read()
        except:
            raise FileNotFoundError(filename)

        try:
            inst = DFLJPG(filename)
            inst.data = data
            inst.length = len(data)
            inst_length = inst.length
            chunks = []
            data_counter = 0
            while data_counter < inst_length:
                chunk_m_l, chunk_m_h = struct.unpack ("BB", data[data_counter:data_counter+2])
                data_counter += 2

                if chunk_m_l != 0xFF:
                    raise ValueError(f"No Valid JPG info in {filename}")

                chunk_name = None
                chunk_size = None
                chunk_data = None
                chunk_ex_data = None
                is_unk_chunk = False

                if chunk_m_h & 0xF0 == 0xD0:
                    n = chunk_m_h & 0x0F

                    if n >= 0 and n <= 7:
                        chunk_name = "RST%d" % (n)
                        chunk_size = 0
                    elif n == 0x8:
                        chunk_name = "SOI"
                        chunk_size = 0
                        if len(chunks) != 0:
                            raise Exception("")
                    elif n == 0x9:
                        chunk_name = "EOI"
                        chunk_size = 0
                    elif n == 0xA:
                        chunk_name = "SOS"
                    elif n == 0xB:
                        chunk_name = "DQT"
                    elif n == 0xD:
                        chunk_name = "DRI"
                        chunk_size = 2
                    else:
                        is_unk_chunk = True
                elif chunk_m_h & 0xF0 == 0xC0:
                    n = chunk_m_h & 0x0F
                    if n == 0:
                        chunk_name = "SOF0"
                    elif n == 2:
                        chunk_name = "SOF2"
                    elif n == 4:
                        chunk_name = "DHT"
                    else:
                        is_unk_chunk = True
                elif chunk_m_h & 0xF0 == 0xE0:
                    n = chunk_m_h & 0x0F
                    chunk_name = "APP%d" % (n)
                else:
                    is_unk_chunk = True

                #if is_unk_chunk:
                #    #raise ValueError(f"Unknown chunk {chunk_m_h} in {filename}")
                #    io.log_info(f"Unknown chunk {chunk_m_h} in {filename}")

                if chunk_size == None: #variable size
                    chunk_size, = struct.unpack (">H", data[data_counter:data_counter+2])
                    chunk_size -= 2
                    data_counter += 2

                if chunk_size > 0:
                    chunk_data = data[data_counter:data_counter+chunk_size]
                    data_counter += chunk_size

                if chunk_name == "SOS":
                    c = data_counter
                    while c < inst_length and (data[c] != 0xFF or data[c+1] != 0xD9):
                        c += 1

                    chunk_ex_data = data[data_counter:c]
                    data_counter = c

                chunks.append ({'name' : chunk_name,
                                'm_h' : chunk_m_h,
                                'data' : chunk_data,
                                'ex_data' : chunk_ex_data,
                                })
            inst.chunks = chunks

            return inst
        except Exception as e:
            raise Exception (f"Corrupted JPG file {filename} {e}")

    @staticmethod
    def load(filename, loader_func=None):
        try:
            inst = DFLJPG.load_raw (filename, loader_func=loader_func)
            inst.dfl_dict = {}

            for chunk in inst.chunks:
                if chunk['name'] == 'APP0':
                    d, c = chunk['data'], 0
                    c, id, _ = struct_unpack (d, c, "=4sB")

                    if id == b"JFIF":
                        c, ver_major, ver_minor, units, Xdensity, Ydensity, Xthumbnail, Ythumbnail = struct_unpack (d, c, "=BBBHHBB")
                    else:
                        raise Exception("Unknown jpeg ID: %s" % (id) )
                elif chunk['name'] == 'SOF0' or chunk['name'] == 'SOF2':
                    d, c = chunk['data'], 0
                    c, precision, height, width = struct_unpack (d, c, ">BHH")
                    inst.shape = (height, width, 3)

                elif chunk['name'] == 'APP15':
                    if type(chunk['data']) == bytes:
                        inst.dfl_dict = pickle.loads(chunk['data'])

            return inst
        except Exception as e:
            io.log_err (f'Exception occured while DFLJPG.load : {traceback.format_exc()}')
            return None

    def has_data(self):
        return len(self.dfl_dict.keys()) != 0

    def save(self):
        try:
            with open(self.filename, "wb") as f:
                f.write ( self.dump() )
        except:
            raise Exception( f'cannot save {self.filename}' )

    def dump(self):
        data = b""

        dict_data = self.dfl_dict

        # Remove None keys
        for key in list(dict_data.keys()):
            if dict_data[key] is None:
                dict_data.pop(key)

        for chunk in self.chunks:
            if chunk['name'] == 'APP15':
                self.chunks.remove(chunk)
                break

        last_app_chunk = 0
        for i, chunk in enumerate (self.chunks):
            if chunk['m_h'] & 0xF0 == 0xE0:
                last_app_chunk = i

        dflchunk = {'name' : 'APP15',
                    'm_h' : 0xEF,
                    'data' : pickle.dumps(dict_data),
                    'ex_data' : None,
                    }
        self.chunks.insert (last_app_chunk+1, dflchunk)


        for chunk in self.chunks:
            data += struct.pack ("BB", 0xFF, chunk['m_h'] )
            chunk_data = chunk['data']
            if chunk_data is not None:
                data += struct.pack (">H", len(chunk_data)+2 )
                data += chunk_data

            chunk_ex_data = chunk['ex_data']
            if chunk_ex_data is not None:
                data += chunk_ex_data

        return data

    def get_img(self):
        if self.img is None:
            self.img = cv2_imread(self.filename)
        return self.img

    def get_shape(self):
        if self.shape is None:
            img = self.get_img()
            if img is not None:
                self.shape = img.shape
        return self.shape

    def get_height(self):
        for chunk in self.chunks:
            if type(chunk) == IHDR:
                return chunk.height
        return 0

    def get_dict(self):
        return self.dfl_dict

    def set_dict (self, dict_data=None):
        self.dfl_dict = dict_data

    def get_face_type(self):            return self.dfl_dict.get('face_type', FaceType.toString (FaceType.FULL) )
    def set_face_type(self, face_type): self.dfl_dict['face_type'] = face_type

    def get_landmarks(self):            return np.array ( self.dfl_dict['landmarks'] )
    def set_landmarks(self, landmarks): self.dfl_dict['landmarks'] = landmarks

    def get_eyebrows_expand_mod(self):                      return self.dfl_dict.get ('eyebrows_expand_mod', 1.0)
    def set_eyebrows_expand_mod(self, eyebrows_expand_mod): self.dfl_dict['eyebrows_expand_mod'] = eyebrows_expand_mod

    def get_source_filename(self):                  return self.dfl_dict.get ('source_filename', None)
    def set_source_filename(self, source_filename): self.dfl_dict['source_filename'] = source_filename

    def get_source_rect(self):              return self.dfl_dict.get ('source_rect', None)
    def set_source_rect(self, source_rect): self.dfl_dict['source_rect'] = source_rect

    def get_source_landmarks(self):                     return np.array ( self.dfl_dict.get('source_landmarks', None) )
    def set_source_landmarks(self, source_landmarks):   self.dfl_dict['source_landmarks'] = source_landmarks

    def get_image_to_face_mat(self):
        mat = self.dfl_dict.get ('image_to_face_mat', None)
        if mat is not None:
            return np.array (mat)
        return None
    def set_image_to_face_mat(self, image_to_face_mat):   self.dfl_dict['image_to_face_mat'] = image_to_face_mat

    def has_seg_ie_polys(self):
        return self.dfl_dict.get('seg_ie_polys',None) is not None

    def get_seg_ie_polys(self):
        d = self.dfl_dict.get('seg_ie_polys',None)
        if d is not None:
            d = SegIEPolys.load(d)
        else:
            d = SegIEPolys()

        return d

    def set_seg_ie_polys(self, seg_ie_polys):
        if seg_ie_polys is not None:
            if not isinstance(seg_ie_polys, SegIEPolys):
                raise ValueError('seg_ie_polys should be instance of SegIEPolys')

            if seg_ie_polys.has_polys():
                seg_ie_polys = seg_ie_polys.dump()
            else:
                seg_ie_polys = None

        self.dfl_dict['seg_ie_polys'] = seg_ie_polys

    def has_xseg_mask(self):
        return self.dfl_dict.get('xseg_mask',None) is not None

    def get_xseg_mask_compressed(self):
        mask_buf = self.dfl_dict.get('xseg_mask',None)
        if mask_buf is None:
            return None

        return mask_buf
        
    def get_xseg_mask(self):
        mask_buf = self.dfl_dict.get('xseg_mask',None)
        if mask_buf is None:
            return None

        img = cv2.imdecode(mask_buf, cv2.IMREAD_UNCHANGED)
        if len(img.shape) == 2:
            img = img[...,None]

        return img.astype(np.float32) / 255.0


    def set_xseg_mask(self, mask_a):
        if mask_a is None:
            self.dfl_dict['xseg_mask'] = None
            return

        mask_a = imagelib.normalize_channels(mask_a, 1)
        img_data = np.clip( mask_a*255, 0, 255 ).astype(np.uint8)

        data_max_len = 50000

        ret, buf = cv2.imencode('.png', img_data)

        if not ret or len(buf) > data_max_len:
            for jpeg_quality in range(100,-1,-1):
                ret, buf = cv2.imencode( '.jpg', img_data, [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_quality] )
                if ret and len(buf) <= data_max_len:
                    break

        if not ret:
            raise Exception("set_xseg_mask: unable to generate image data for set_xseg_mask")

        self.dfl_dict['xseg_mask'] = buf

DFLIMG/DFLJPG.py
oo92/TorchFaceLab
from .DFLIMG import DFLIMG
from .DFLJPG import DFLJPG
DFLIMG/__init__.py
oo92/TorchFaceLab
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *
                                       
class QCursorDB():
    @staticmethod
    def initialize(cursor_path):
        QCursorDB.cross_red = QCursor ( QPixmap ( str(cursor_path / 'cross_red.png') ) )
        QCursorDB.cross_green = QCursor ( QPixmap ( str(cursor_path / 'cross_green.png') ) )
        QCursorDB.cross_blue = QCursor ( QPixmap ( str(cursor_path / 'cross_blue.png') ) )

XSegEditor/QCursorDB.py
oo92/TorchFaceLab
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *


class QIconDB():
    @staticmethod
    def initialize(icon_path):
        QIconDB.app_icon         = QIcon ( str(icon_path / 'app_icon.png') )
        QIconDB.delete_poly      = QIcon ( str(icon_path / 'delete_poly.png') )
        QIconDB.undo_pt          = QIcon ( str(icon_path / 'undo_pt.png') )
        QIconDB.redo_pt          = QIcon ( str(icon_path / 'redo_pt.png') )
        QIconDB.poly_color_red   = QIcon ( str(icon_path / 'poly_color_red.png') )
        QIconDB.poly_color_green = QIcon ( str(icon_path / 'poly_color_green.png') )
        QIconDB.poly_color_blue  = QIcon ( str(icon_path / 'poly_color_blue.png') )
        QIconDB.poly_type_include = QIcon ( str(icon_path / 'poly_type_include.png') )
        QIconDB.poly_type_exclude = QIcon ( str(icon_path / 'poly_type_exclude.png') )
        QIconDB.left  = QIcon ( str(icon_path / 'left.png') )
        QIconDB.right = QIcon ( str(icon_path / 'right.png') )
        QIconDB.trashcan = QIcon ( str(icon_path / 'trashcan.png') )
        QIconDB.pt_edit_mode = QIcon ( str(icon_path / 'pt_edit_mode.png') )
        QIconDB.view_lock_center = QIcon ( str(icon_path / 'view_lock_center.png') )
        QIconDB.view_baked = QIcon ( str(icon_path / 'view_baked.png') )
        QIconDB.view_xseg = QIcon ( str(icon_path / 'view_xseg.png') )
        QIconDB.view_xseg_overlay = QIcon ( str(icon_path / 'view_xseg_overlay.png') )
        
XSegEditor/QIconDB.py
oo92/TorchFaceLab
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *

class QImageDB():
    @staticmethod
    def initialize(image_path):
        QImageDB.intro = QImage ( str(image_path / 'intro.png') )

XSegEditor/QImageDB.py
oo92/TorchFaceLab
from localization import system_language


class QStringDB():
    
    @staticmethod
    def initialize():
        lang = system_language
        
        if lang not in ['en','ru','zh']:
            lang = 'en'
        
        QStringDB.btn_poly_color_red_tip = {    'en' : 'Poly color scheme red',
                                                'ru' : 'Красная цветовая схема полигонов',
                                                'zh' : '选区配色方案红色',
                                           }[lang]
                                           
        QStringDB.btn_poly_color_green_tip = {  'en' : 'Poly color scheme green',
                                                'ru' : 'Зелёная цветовая схема полигонов',
                                                'zh' : '选区配色方案绿色',
                                           }[lang]
                                           
        QStringDB.btn_poly_color_blue_tip = {   'en' : 'Poly color scheme blue',
                                                'ru' : 'Синяя цветовая схема полигонов',
                                                'zh' : '选区配色方案蓝色',
                                           }[lang]
                                           
        QStringDB.btn_view_baked_mask_tip = {   'en' : 'View baked mask',
                                                'ru' : 'Посмотреть запечёную маску',
                                                'zh' : '查看遮罩通道',
                                           }[lang]
                                           
        QStringDB.btn_view_xseg_mask_tip =  {   'en' : 'View trained XSeg mask',
                                                'ru' : 'Посмотреть тренированную XSeg маску',
                                                'zh' : '查看导入后的XSeg遮罩',
                                            }[lang]
                                            
        QStringDB.btn_view_xseg_overlay_mask_tip =  {   'en' : 'View trained XSeg mask overlay face',
                                                        'ru' : 'Посмотреть тренированную XSeg маску поверх лица',
                                                        'zh' : '查看导入后的XSeg遮罩于脸上方',
                                                    }[lang]
                                           
        QStringDB.btn_poly_type_include_tip = { 'en' : 'Poly include mode',
                                                'ru' : 'Режим полигонов - включение',
                                                'zh' : '包含选区模式',
                                           }[lang]
                                           
        QStringDB.btn_poly_type_exclude_tip = { 'en' : 'Poly exclude mode',
                                                'ru' : 'Режим полигонов - исключение',
                                                'zh' : '排除选区模式',
                                           }[lang]        
                                           
        QStringDB.btn_undo_pt_tip = {   'en' : 'Undo point',
                                        'ru' : 'Отменить точку',
                                        'zh' : '撤消点',
                                    }[lang]      
                                     
        QStringDB.btn_redo_pt_tip = {   'en' : 'Redo point',
                                        'ru' : 'Повторить точку',
                                        'zh' : '重做点',
                                     }[lang]      
                                      
        QStringDB.btn_delete_poly_tip = {   'en' : 'Delete poly',
                                            'ru' : 'Удалить полигон',
                                            'zh' : '删除选区',
                                           }[lang]     
                                              
        QStringDB.btn_pt_edit_mode_tip = {  'en' : 'Add/delete point mode ( HOLD CTRL )',
                                            'ru' : 'Режим добавления/удаления точек ( удерживайте CTRL )',
                                            'zh' : '点加/删除模式 ( 按住CTRL )',
                                           }[lang]    
                                           
        QStringDB.btn_view_lock_center_tip = { 'en' : 'Lock cursor at the center ( HOLD SHIFT )',
                                               'ru' : 'Заблокировать курсор в центре ( удерживайте SHIFT )',
                                               'zh' : '将光标锁定在中心 ( 按住SHIFT )',
                                             }[lang]                    
                                        
                                           
        QStringDB.btn_prev_image_tip = {    'en' : 'Save and Prev image\nHold SHIFT : accelerate\nHold CTRL : skip non masked\n',
                                            'ru' : 'Сохранить и предыдущее изображение\nУдерживать SHIFT : ускорить\nУдерживать CTRL : пропустить неразмеченные\n',
                                            'zh' : '保存并转到上一张图片\n按住SHIFT : 加快\n按住CTRL : 跳过未标记的\n',
                                           }[lang]   
        QStringDB.btn_next_image_tip = {    'en' : 'Save and Next image\nHold SHIFT : accelerate\nHold CTRL : skip non masked\n',
                                            'ru' : 'Сохранить и следующее изображение\nУдерживать SHIFT : ускорить\nУдерживать CTRL : пропустить неразмеченные\n',
                                            'zh' : '保存并转到下一张图片\n按住SHIFT : 加快\n按住CTRL : 跳过未标记的\n',
                                           }[lang]  
        
        QStringDB.btn_delete_image_tip = {  'en' : 'Move to _trash and Next image\n',
                                            'ru' : 'Переместить в _trash и следующее изображение\n',
                                            'zh' : '移至_trash，转到下一张图片 ',
                                           }[lang]  
                                           
        QStringDB.loading_tip = {'en' : 'Loading',
                                 'ru' : 'Загрузка',
                                 'zh' : '正在载入',
                                }[lang]   
                                
        QStringDB.labeled_tip = {'en' : 'labeled',
                                 'ru' : 'размечено',
                                 'zh' : '标记的',
                                }[lang]   
                                           

XSegEditor/QStringDB.py
oo92/TorchFaceLab
import json
import multiprocessing
import os
import pickle
import sys
import tempfile
import time
import traceback
from enum import IntEnum
from types import SimpleNamespace as sn

import cv2
import numpy as np
import numpy.linalg as npla
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *

from core import imagelib, pathex
from core.cv2ex import *
from core.imagelib import SegIEPoly, SegIEPolys, SegIEPolyType, sd
from core.qtex import *
from DFLIMG import *
from localization import StringsDB, system_language
from samplelib import PackedFaceset

from .QCursorDB import QCursorDB
from .QIconDB import QIconDB
from .QStringDB import QStringDB
from .QImageDB import QImageDB

class OpMode(IntEnum):
    NONE = 0
    DRAW_PTS = 1
    EDIT_PTS = 2
    VIEW_BAKED = 3
    VIEW_XSEG_MASK = 4

class PTEditMode(IntEnum):
    MOVE = 0
    ADD_DEL = 1

class DragType(IntEnum):
    NONE = 0
    IMAGE_LOOK = 1
    POLY_PT = 2

class ViewLock(IntEnum):
    NONE = 0
    CENTER = 1

class QUIConfig():
    @staticmethod
    def initialize(icon_size = 48, icon_spacer_size=16, preview_bar_icon_size=64):
        QUIConfig.icon_q_size = QSize(icon_size, icon_size)
        QUIConfig.icon_spacer_q_size = QSize(icon_spacer_size, icon_spacer_size)
        QUIConfig.preview_bar_icon_q_size = QSize(preview_bar_icon_size, preview_bar_icon_size)

class ImagePreviewSequenceBar(QFrame):
    def __init__(self, preview_images_count, icon_size):
        super().__init__()
        self.preview_images_count = preview_images_count = max(1, preview_images_count + (preview_images_count % 2 -1) )

        self.icon_size = icon_size

        black_q_img = QImage(np.zeros( (icon_size,icon_size,3) ).data, icon_size, icon_size, 3*icon_size, QImage.Format_RGB888)
        self.black_q_pixmap = QPixmap.fromImage(black_q_img)

        self.image_containers = [ QLabel() for i in range(preview_images_count)]

        main_frame_l_cont_hl = QGridLayout()
        main_frame_l_cont_hl.setContentsMargins(0,0,0,0)
        #main_frame_l_cont_hl.setSpacing(0)



        for i in range(len(self.image_containers)):
            q_label = self.image_containers[i]
            q_label.setScaledContents(True)
            if i == preview_images_count//2:
                q_label.setMinimumSize(icon_size+16, icon_size+16 )
                q_label.setMaximumSize(icon_size+16, icon_size+16 )
            else:
                q_label.setMinimumSize(icon_size, icon_size )
                q_label.setMaximumSize(icon_size, icon_size )
                opacity_effect = QGraphicsOpacityEffect()
                opacity_effect.setOpacity(0.5)
                q_label.setGraphicsEffect(opacity_effect)

            q_label.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed)

            main_frame_l_cont_hl.addWidget (q_label, 0, i)

        self.setLayout(main_frame_l_cont_hl)

        self.prev_img_conts = self.image_containers[(preview_images_count//2) -1::-1]
        self.next_img_conts = self.image_containers[preview_images_count//2:]

        self.update_images()

    def get_preview_images_count(self):
        return self.preview_images_count

    def update_images(self, prev_imgs=None, next_imgs=None):
        # Fix arrays
        if prev_imgs is None:
            prev_imgs = []
        prev_img_conts_len = len(self.prev_img_conts)
        prev_q_imgs_len = len(prev_imgs)
        if prev_q_imgs_len < prev_img_conts_len:
            for i in range ( prev_img_conts_len - prev_q_imgs_len ):
                prev_imgs.append(None)
        elif prev_q_imgs_len > prev_img_conts_len:
            prev_imgs = prev_imgs[:prev_img_conts_len]

        if next_imgs is None:
            next_imgs = []
        next_img_conts_len = len(self.next_img_conts)
        next_q_imgs_len = len(next_imgs)
        if next_q_imgs_len < next_img_conts_len:
            for i in range ( next_img_conts_len - next_q_imgs_len ):
                next_imgs.append(None)
        elif next_q_imgs_len > next_img_conts_len:
            next_imgs = next_imgs[:next_img_conts_len]

        for i,img in enumerate(prev_imgs):
            self.prev_img_conts[i].setPixmap( QPixmap.fromImage( QImage_from_np(img) ) if img is not None else self.black_q_pixmap )

        for i,img in enumerate(next_imgs):
            self.next_img_conts[i].setPixmap( QPixmap.fromImage( QImage_from_np(img) ) if img is not None else self.black_q_pixmap )

class ColorScheme():
    def __init__(self, unselected_color, selected_color, outline_color, outline_width, pt_outline_color, cross_cursor):
        self.poly_unselected_brush = QBrush(unselected_color)
        self.poly_selected_brush = QBrush(selected_color)

        self.poly_outline_solid_pen = QPen(outline_color, outline_width, Qt.SolidLine, Qt.RoundCap, Qt.RoundJoin)
        self.poly_outline_dot_pen = QPen(outline_color, outline_width, Qt.DotLine, Qt.RoundCap, Qt.RoundJoin)

        self.pt_outline_pen = QPen(pt_outline_color)
        self.cross_cursor = cross_cursor

class CanvasConfig():

    def __init__(self,
                 pt_radius=4,
                 pt_select_radius=8,
                 color_schemes=None,
                 **kwargs):
        self.pt_radius = pt_radius
        self.pt_select_radius = pt_select_radius

        if color_schemes is None:
            color_schemes = [
                    ColorScheme( QColor(192,0,0,alpha=0), QColor(192,0,0,alpha=72), QColor(192,0,0), 2, QColor(255,255,255), QCursorDB.cross_red ),
                    ColorScheme( QColor(0,192,0,alpha=0), QColor(0,192,0,alpha=72), QColor(0,192,0), 2, QColor(255,255,255), QCursorDB.cross_green ),
                    ColorScheme( QColor(0,0,192,alpha=0), QColor(0,0,192,alpha=72), QColor(0,0,192), 2, QColor(255,255,255), QCursorDB.cross_blue ),
                    ]
        self.color_schemes = color_schemes

class QCanvasControlsLeftBar(QFrame):

    def __init__(self):
        super().__init__()
        #==============================================
        btn_poly_type_include = QToolButton()
        self.btn_poly_type_include_act = QActionEx( QIconDB.poly_type_include, QStringDB.btn_poly_type_include_tip, shortcut='Q', shortcut_in_tooltip=True, is_checkable=True)
        btn_poly_type_include.setDefaultAction(self.btn_poly_type_include_act)
        btn_poly_type_include.setIconSize(QUIConfig.icon_q_size)

        btn_poly_type_exclude = QToolButton()
        self.btn_poly_type_exclude_act = QActionEx( QIconDB.poly_type_exclude, QStringDB.btn_poly_type_exclude_tip, shortcut='W', shortcut_in_tooltip=True, is_checkable=True)
        btn_poly_type_exclude.setDefaultAction(self.btn_poly_type_exclude_act)
        btn_poly_type_exclude.setIconSize(QUIConfig.icon_q_size)

        self.btn_poly_type_act_grp = QActionGroup (self)
        self.btn_poly_type_act_grp.addAction(self.btn_poly_type_include_act)
        self.btn_poly_type_act_grp.addAction(self.btn_poly_type_exclude_act)
        self.btn_poly_type_act_grp.setExclusive(True)
        #==============================================
        btn_undo_pt = QToolButton()
        self.btn_undo_pt_act = QActionEx( QIconDB.undo_pt, QStringDB.btn_undo_pt_tip, shortcut='Ctrl+Z', shortcut_in_tooltip=True, is_auto_repeat=True)
        btn_undo_pt.setDefaultAction(self.btn_undo_pt_act)
        btn_undo_pt.setIconSize(QUIConfig.icon_q_size)

        btn_redo_pt = QToolButton()
        self.btn_redo_pt_act = QActionEx( QIconDB.redo_pt, QStringDB.btn_redo_pt_tip, shortcut='Ctrl+Shift+Z',  shortcut_in_tooltip=True, is_auto_repeat=True)
        btn_redo_pt.setDefaultAction(self.btn_redo_pt_act)
        btn_redo_pt.setIconSize(QUIConfig.icon_q_size)

        btn_delete_poly = QToolButton()
        self.btn_delete_poly_act = QActionEx( QIconDB.delete_poly, QStringDB.btn_delete_poly_tip, shortcut='Delete', shortcut_in_tooltip=True)
        btn_delete_poly.setDefaultAction(self.btn_delete_poly_act)
        btn_delete_poly.setIconSize(QUIConfig.icon_q_size)
        #==============================================
        btn_pt_edit_mode = QToolButton()
        self.btn_pt_edit_mode_act = QActionEx( QIconDB.pt_edit_mode, QStringDB.btn_pt_edit_mode_tip, shortcut_in_tooltip=True, is_checkable=True)
        btn_pt_edit_mode.setDefaultAction(self.btn_pt_edit_mode_act)
        btn_pt_edit_mode.setIconSize(QUIConfig.icon_q_size)
        #==============================================

        controls_bar_frame2_l = QVBoxLayout()
        controls_bar_frame2_l.addWidget ( btn_poly_type_include )
        controls_bar_frame2_l.addWidget ( btn_poly_type_exclude )
        controls_bar_frame2 = QFrame()
        controls_bar_frame2.setFrameShape(QFrame.StyledPanel)
        controls_bar_frame2.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed)
        controls_bar_frame2.setLayout(controls_bar_frame2_l)

        controls_bar_frame3_l = QVBoxLayout()
        controls_bar_frame3_l.addWidget ( btn_undo_pt )
        controls_bar_frame3_l.addWidget ( btn_redo_pt )
        controls_bar_frame3_l.addWidget ( btn_delete_poly )
        controls_bar_frame3 = QFrame()
        controls_bar_frame3.setFrameShape(QFrame.StyledPanel)
        controls_bar_frame3.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed)
        controls_bar_frame3.setLayout(controls_bar_frame3_l)

        controls_bar_frame4_l = QVBoxLayout()
        controls_bar_frame4_l.addWidget ( btn_pt_edit_mode )
        controls_bar_frame4 = QFrame()
        controls_bar_frame4.setFrameShape(QFrame.StyledPanel)
        controls_bar_frame4.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed)
        controls_bar_frame4.setLayout(controls_bar_frame4_l)

        controls_bar_l = QVBoxLayout()
        controls_bar_l.setContentsMargins(0,0,0,0)
        controls_bar_l.addWidget(controls_bar_frame2)
        controls_bar_l.addWidget(controls_bar_frame3)
        controls_bar_l.addWidget(controls_bar_frame4)

        self.setSizePolicy ( QSizePolicy.Fixed, QSizePolicy.Expanding )
        self.setLayout(controls_bar_l)

class QCanvasControlsRightBar(QFrame):

    def __init__(self):
        super().__init__()
        #==============================================
        btn_poly_color_red = QToolButton()
        self.btn_poly_color_red_act = QActionEx( QIconDB.poly_color_red, QStringDB.btn_poly_color_red_tip, shortcut='1', shortcut_in_tooltip=True, is_checkable=True)
        btn_poly_color_red.setDefaultAction(self.btn_poly_color_red_act)
        btn_poly_color_red.setIconSize(QUIConfig.icon_q_size)

        btn_poly_color_green = QToolButton()
        self.btn_poly_color_green_act = QActionEx( QIconDB.poly_color_green, QStringDB.btn_poly_color_green_tip, shortcut='2', shortcut_in_tooltip=True, is_checkable=True)
        btn_poly_color_green.setDefaultAction(self.btn_poly_color_green_act)
        btn_poly_color_green.setIconSize(QUIConfig.icon_q_size)

        btn_poly_color_blue = QToolButton()
        self.btn_poly_color_blue_act = QActionEx( QIconDB.poly_color_blue, QStringDB.btn_poly_color_blue_tip, shortcut='3', shortcut_in_tooltip=True, is_checkable=True)
        btn_poly_color_blue.setDefaultAction(self.btn_poly_color_blue_act)
        btn_poly_color_blue.setIconSize(QUIConfig.icon_q_size)

        btn_view_baked_mask = QToolButton()
        self.btn_view_baked_mask_act = QActionEx( QIconDB.view_baked, QStringDB.btn_view_baked_mask_tip, shortcut='4', shortcut_in_tooltip=True, is_checkable=True)
        btn_view_baked_mask.setDefaultAction(self.btn_view_baked_mask_act)
        btn_view_baked_mask.setIconSize(QUIConfig.icon_q_size)

        btn_view_xseg_mask = QToolButton()
        self.btn_view_xseg_mask_act = QActionEx( QIconDB.view_xseg, QStringDB.btn_view_xseg_mask_tip, shortcut='5', shortcut_in_tooltip=True, is_checkable=True)
        btn_view_xseg_mask.setDefaultAction(self.btn_view_xseg_mask_act)
        btn_view_xseg_mask.setIconSize(QUIConfig.icon_q_size)

        btn_view_xseg_overlay_mask = QToolButton()
        self.btn_view_xseg_overlay_mask_act = QActionEx( QIconDB.view_xseg_overlay, QStringDB.btn_view_xseg_overlay_mask_tip, shortcut='`', shortcut_in_tooltip=True, is_checkable=True)
        btn_view_xseg_overlay_mask.setDefaultAction(self.btn_view_xseg_overlay_mask_act)
        btn_view_xseg_overlay_mask.setIconSize(QUIConfig.icon_q_size)

        self.btn_poly_color_act_grp = QActionGroup (self)
        self.btn_poly_color_act_grp.addAction(self.btn_poly_color_red_act)
        self.btn_poly_color_act_grp.addAction(self.btn_poly_color_green_act)
        self.btn_poly_color_act_grp.addAction(self.btn_poly_color_blue_act)
        self.btn_poly_color_act_grp.addAction(self.btn_view_baked_mask_act)
        self.btn_poly_color_act_grp.addAction(self.btn_view_xseg_mask_act)
        self.btn_poly_color_act_grp.setExclusive(True)
        #==============================================
        btn_view_lock_center = QToolButton()
        self.btn_view_lock_center_act = QActionEx( QIconDB.view_lock_center, QStringDB.btn_view_lock_center_tip, shortcut_in_tooltip=True, is_checkable=True)
        btn_view_lock_center.setDefaultAction(self.btn_view_lock_center_act)
        btn_view_lock_center.setIconSize(QUIConfig.icon_q_size)

        controls_bar_frame2_l = QVBoxLayout()
        controls_bar_frame2_l.addWidget ( btn_view_xseg_overlay_mask )
        controls_bar_frame2 = QFrame()
        controls_bar_frame2.setFrameShape(QFrame.StyledPanel)
        controls_bar_frame2.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed)
        controls_bar_frame2.setLayout(controls_bar_frame2_l)
        
        controls_bar_frame1_l = QVBoxLayout()
        controls_bar_frame1_l.addWidget ( btn_poly_color_red )
        controls_bar_frame1_l.addWidget ( btn_poly_color_green )
        controls_bar_frame1_l.addWidget ( btn_poly_color_blue )
        controls_bar_frame1_l.addWidget ( btn_view_baked_mask )
        controls_bar_frame1_l.addWidget ( btn_view_xseg_mask )
        controls_bar_frame1 = QFrame()
        controls_bar_frame1.setFrameShape(QFrame.StyledPanel)
        controls_bar_frame1.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed)
        controls_bar_frame1.setLayout(controls_bar_frame1_l)
        
        controls_bar_frame3_l = QVBoxLayout()
        controls_bar_frame3_l.addWidget ( btn_view_lock_center )
        controls_bar_frame3 = QFrame()
        controls_bar_frame3.setFrameShape(QFrame.StyledPanel)
        controls_bar_frame3.setSizePolicy (QSizePolicy.Fixed, QSizePolicy.Fixed)
        controls_bar_frame3.setLayout(controls_bar_frame3_l)

        controls_bar_l = QVBoxLayout()
        controls_bar_l.setContentsMargins(0,0,0,0)
        controls_bar_l.addWidget(controls_bar_frame2)
        controls_bar_l.addWidget(controls_bar_frame1)
        controls_bar_l.addWidget(controls_bar_frame3)
        
        self.setSizePolicy ( QSizePolicy.Fixed, QSizePolicy.Expanding )
        self.setLayout(controls_bar_l)

class QCanvasOperator(QWidget):
    def __init__(self, cbar):
        super().__init__()
        self.cbar = cbar

        self.set_cbar_disabled()

        self.cbar.btn_poly_color_red_act.triggered.connect ( lambda : self.set_color_scheme_id(0) )
        self.cbar.btn_poly_color_green_act.triggered.connect ( lambda : self.set_color_scheme_id(1) )
        self.cbar.btn_poly_color_blue_act.triggered.connect ( lambda : self.set_color_scheme_id(2) )
        self.cbar.btn_view_baked_mask_act.triggered.connect ( lambda : self.set_op_mode(OpMode.VIEW_BAKED) )
        self.cbar.btn_view_xseg_mask_act.triggered.connect ( lambda : self.set_op_mode(OpMode.VIEW_XSEG_MASK) )

        self.cbar.btn_view_xseg_overlay_mask_act.toggled.connect ( lambda is_checked: self.update() )

        self.cbar.btn_poly_type_include_act.triggered.connect ( lambda : self.set_poly_include_type(SegIEPolyType.INCLUDE) )
        self.cbar.btn_poly_type_exclude_act.triggered.connect ( lambda : self.set_poly_include_type(SegIEPolyType.EXCLUDE) )

        self.cbar.btn_undo_pt_act.triggered.connect ( lambda : self.action_undo_pt() )
        self.cbar.btn_redo_pt_act.triggered.connect ( lambda : self.action_redo_pt() )

        self.cbar.btn_delete_poly_act.triggered.connect ( lambda : self.action_delete_poly() )

        self.cbar.btn_pt_edit_mode_act.toggled.connect ( lambda is_checked: self.set_pt_edit_mode( PTEditMode.ADD_DEL if is_checked else PTEditMode.MOVE ) )
        self.cbar.btn_view_lock_center_act.toggled.connect ( lambda is_checked: self.set_view_lock( ViewLock.CENTER if is_checked else ViewLock.NONE ) )

        self.mouse_in_widget = False

        QXMainWindow.inst.add_keyPressEvent_listener ( self.on_keyPressEvent )
        QXMainWindow.inst.add_keyReleaseEvent_listener ( self.on_keyReleaseEvent )

        self.qp = QPainter()
        self.initialized = False
        self.last_state = None

    def initialize(self, img, img_look_pt=None, view_scale=None, ie_polys=None, xseg_mask=None, canvas_config=None ):
        q_img = self.q_img = QImage_from_np(img)
        self.img_pixmap = QPixmap.fromImage(q_img)

        self.xseg_mask_pixmap = None
        self.xseg_overlay_mask_pixmap = None
        if xseg_mask is not None:
            h,w,c = img.shape
            xseg_mask = cv2.resize(xseg_mask, (w,h), interpolation=cv2.INTER_CUBIC)
            xseg_mask = imagelib.normalize_channels(xseg_mask, 1)
            xseg_img = img.astype(np.float32)/255.0
            xseg_overlay_mask = xseg_img*(1-xseg_mask)*0.5 + xseg_img*xseg_mask
            xseg_overlay_mask = np.clip(xseg_overlay_mask*255, 0, 255).astype(np.uint8)
            xseg_mask = np.clip(xseg_mask*255, 0, 255).astype(np.uint8)
            self.xseg_mask_pixmap = QPixmap.fromImage(QImage_from_np(xseg_mask))
            self.xseg_overlay_mask_pixmap = QPixmap.fromImage(QImage_from_np(xseg_overlay_mask))

        self.img_size = QSize_to_np (self.img_pixmap.size())

        self.img_look_pt = img_look_pt
        self.view_scale = view_scale

        if ie_polys is None:
            ie_polys = SegIEPolys()
        self.ie_polys = ie_polys

        if canvas_config is None:
            canvas_config = CanvasConfig()
        self.canvas_config = canvas_config

        # UI init
        self.set_cbar_disabled()
        self.cbar.btn_poly_color_act_grp.setDisabled(False)
        self.cbar.btn_view_xseg_overlay_mask_act.setDisabled(False)
        self.cbar.btn_poly_type_act_grp.setDisabled(False)

        # Initial vars
        self.current_cursor = None
        self.mouse_hull_poly = None
        self.mouse_wire_poly = None
        self.drag_type = DragType.NONE
        self.mouse_cli_pt = np.zeros((2,), np.float32 )

        # Initial state
        self.set_op_mode(OpMode.NONE)
        self.set_color_scheme_id(1)
        self.set_poly_include_type(SegIEPolyType.INCLUDE)
        self.set_pt_edit_mode(PTEditMode.MOVE)
        self.set_view_lock(ViewLock.NONE)

        # Apply last state
        if self.last_state is not None:
            self.set_color_scheme_id(self.last_state.color_scheme_id)
            if self.last_state.op_mode is not None:
                self.set_op_mode(self.last_state.op_mode)

        self.initialized = True

        self.setMouseTracking(True)
        self.update_cursor()
        self.update()


    def finalize(self):
        if self.initialized:
            if self.op_mode == OpMode.DRAW_PTS:
                self.set_op_mode(OpMode.EDIT_PTS)

            self.last_state = sn(op_mode = self.op_mode if self.op_mode in [OpMode.VIEW_BAKED, OpMode.VIEW_XSEG_MASK] else None,
                                 color_scheme_id = self.color_scheme_id)

            self.img_pixmap = None
            self.update_cursor(is_finalize=True)
            self.setMouseTracking(False)
            self.setFocusPolicy(Qt.NoFocus)
            self.set_cbar_disabled()
            self.initialized = False
            self.update()

    # ====================================================================================
    # ====================================================================================
    # ====================================== GETTERS =====================================
    # ====================================================================================
    # ====================================================================================
    def is_initialized(self):
        return self.initialized

    def get_ie_polys(self):
        return self.ie_polys

    def get_cli_center_pt(self):
        return np.round(QSize_to_np(self.size())/2.0)

    def get_img_look_pt(self):
        img_look_pt = self.img_look_pt
        if img_look_pt is None:
            img_look_pt = self.img_size / 2
        return img_look_pt

    def get_view_scale(self):
        view_scale = self.view_scale
        if view_scale is None:
            # Calc as scale to fit
            min_cli_size = np.min(QSize_to_np(self.size()))
            max_img_size = np.max(self.img_size)
            view_scale =  min_cli_size / max_img_size

        return view_scale

    def get_current_color_scheme(self):
        return self.canvas_config.color_schemes[self.color_scheme_id]

    def get_poly_pt_id_under_pt(self, poly, cli_pt):
        w = np.argwhere ( npla.norm ( cli_pt - self.img_to_cli_pt( poly.get_pts() ), axis=1 )  <= self.canvas_config.pt_select_radius )
        return None if len(w) == 0 else w[-1][0]

    def get_poly_edge_id_pt_under_pt(self, poly, cli_pt):
        cli_pts = self.img_to_cli_pt(poly.get_pts())
        if len(cli_pts) >= 3:
            edge_dists, projs = sd.dist_to_edges(cli_pts, cli_pt, is_closed=True)
            edge_id = np.argmin(edge_dists)
            dist = edge_dists[edge_id]
            pt = projs[edge_id]
            if dist <= self.canvas_config.pt_select_radius:
                return edge_id, pt
        return None, None

    def get_poly_by_pt_near_wire(self, cli_pt):
        pt_select_radius = self.canvas_config.pt_select_radius

        for poly in reversed(self.ie_polys.get_polys()):
            pts = poly.get_pts()
            if len(pts) >= 3:
                cli_pts = self.img_to_cli_pt(pts)

                edge_dists, _ = sd.dist_to_edges(cli_pts, cli_pt, is_closed=True)

                if np.min(edge_dists) <= pt_select_radius or \
                   any( npla.norm ( cli_pt - cli_pts, axis=1 ) <= pt_select_radius ):
                    return poly
        return None

    def get_poly_by_pt_in_hull(self, cli_pos):
        img_pos = self.cli_to_img_pt(cli_pos)

        for poly in reversed(self.ie_polys.get_polys()):
            pts = poly.get_pts()
            if len(pts) >= 3:
                if cv2.pointPolygonTest( pts, tuple(img_pos), False) >= 0:
                    return poly

        return None

    def img_to_cli_pt(self, p):
        return (p - self.get_img_look_pt()) * self.get_view_scale() + self.get_cli_center_pt()# QSize_to_np(self.size())/2.0

    def cli_to_img_pt(self, p):
        return (p - self.get_cli_center_pt() ) / self.get_view_scale() + self.get_img_look_pt()

    def img_to_cli_rect(self, rect):
        tl = QPoint_to_np(rect.topLeft())
        xy = self.img_to_cli_pt(tl)
        xy2 = self.img_to_cli_pt(tl + QSize_to_np(rect.size()) ) - xy
        return QRect ( *xy.astype(np.int), *xy2.astype(np.int) )

    # ====================================================================================
    # ====================================================================================
    # ====================================== SETTERS =====================================
    # ====================================================================================
    # ====================================================================================
    def set_op_mode(self, op_mode, op_poly=None):
        if not hasattr(self,'op_mode'):
            self.op_mode = None
            self.op_poly = None

        if self.op_mode != op_mode:
            # Finalize prev mode
            if self.op_mode == OpMode.NONE:
                self.cbar.btn_poly_type_act_grp.setDisabled(True)
            elif self.op_mode == OpMode.DRAW_PTS:
                self.cbar.btn_undo_pt_act.setDisabled(True)
                self.cbar.btn_redo_pt_act.setDisabled(True)
                self.cbar.btn_view_lock_center_act.setDisabled(True)
                # Reset view_lock when exit from DRAW_PTS
                self.set_view_lock(ViewLock.NONE)
                # Remove unfinished poly
                if self.op_poly.get_pts_count() < 3:
                    self.ie_polys.remove_poly(self.op_poly)

            elif self.op_mode == OpMode.EDIT_PTS:
                self.cbar.btn_pt_edit_mode_act.setDisabled(True)
                self.cbar.btn_delete_poly_act.setDisabled(True)
                # Reset pt_edit_move when exit from EDIT_PTS
                self.set_pt_edit_mode(PTEditMode.MOVE)
            elif self.op_mode == OpMode.VIEW_BAKED:
                self.cbar.btn_view_baked_mask_act.setChecked(False)
            elif self.op_mode == OpMode.VIEW_XSEG_MASK:
                self.cbar.btn_view_xseg_mask_act.setChecked(False)

            self.op_mode = op_mode

            # Initialize new mode
            if op_mode == OpMode.NONE:
                self.cbar.btn_poly_type_act_grp.setDisabled(False)
            elif op_mode == OpMode.DRAW_PTS:
                self.cbar.btn_undo_pt_act.setDisabled(False)
                self.cbar.btn_redo_pt_act.setDisabled(False)
                self.cbar.btn_view_lock_center_act.setDisabled(False)
            elif op_mode == OpMode.EDIT_PTS:
                self.cbar.btn_pt_edit_mode_act.setDisabled(False)
                self.cbar.btn_delete_poly_act.setDisabled(False)
            elif op_mode == OpMode.VIEW_BAKED:
                self.cbar.btn_view_baked_mask_act.setChecked(True )
                n = QImage_to_np ( self.q_img ).astype(np.float32) / 255.0
                h,w,c = n.shape
                mask = np.zeros( (h,w,1), dtype=np.float32 )
                self.ie_polys.overlay_mask(mask)
                n = (mask*255).astype(np.uint8)
                self.img_baked_pixmap = QPixmap.fromImage(QImage_from_np(n))
            elif op_mode == OpMode.VIEW_XSEG_MASK:
                self.cbar.btn_view_xseg_mask_act.setChecked(True)

            if op_mode in [OpMode.DRAW_PTS, OpMode.EDIT_PTS]:
                self.mouse_op_poly_pt_id = None
                self.mouse_op_poly_edge_id = None
                self.mouse_op_poly_edge_id_pt = None

            self.op_poly = op_poly
            if op_poly is not None:
                self.update_mouse_info()

            self.update_cursor()
            self.update()

    def set_pt_edit_mode(self, pt_edit_mode):
        if not hasattr(self, 'pt_edit_mode') or self.pt_edit_mode != pt_edit_mode:
            self.pt_edit_mode = pt_edit_mode
            self.update_cursor()
            self.update()
        self.cbar.btn_pt_edit_mode_act.setChecked( self.pt_edit_mode == PTEditMode.ADD_DEL )

    def set_view_lock(self, view_lock):
        if not hasattr(self, 'view_lock') or self.view_lock != view_lock:
            if hasattr(self, 'view_lock') and self.view_lock != view_lock:
                if view_lock == ViewLock.CENTER:
                    self.img_look_pt = self.mouse_img_pt
                    QCursor.setPos ( self.mapToGlobal( QPoint_from_np(self.img_to_cli_pt(self.img_look_pt)) ))

            self.view_lock = view_lock
            self.update()
        self.cbar.btn_view_lock_center_act.setChecked( self.view_lock == ViewLock.CENTER )

    def set_cbar_disabled(self):
        self.cbar.btn_delete_poly_act.setDisabled(True)
        self.cbar.btn_undo_pt_act.setDisabled(True)
        self.cbar.btn_redo_pt_act.setDisabled(True)
        self.cbar.btn_pt_edit_mode_act.setDisabled(True)
        self.cbar.btn_view_lock_center_act.setDisabled(True)
        self.cbar.btn_poly_color_act_grp.setDisabled(True)
        self.cbar.btn_view_xseg_overlay_mask_act.setDisabled(True)
        self.cbar.btn_poly_type_act_grp.setDisabled(True)


    def set_color_scheme_id(self, id):
        if self.op_mode == OpMode.VIEW_BAKED or self.op_mode == OpMode.VIEW_XSEG_MASK:
            self.set_op_mode(OpMode.NONE)

        if not hasattr(self, 'color_scheme_id') or self.color_scheme_id != id:
            self.color_scheme_id = id
            self.update_cursor()
            self.update()

        if self.color_scheme_id == 0:
            self.cbar.btn_poly_color_red_act.setChecked( True )
        elif self.color_scheme_id == 1:
            self.cbar.btn_poly_color_green_act.setChecked( True )
        elif self.color_scheme_id == 2:
            self.cbar.btn_poly_color_blue_act.setChecked( True )

    def set_poly_include_type(self, poly_include_type):
        if not hasattr(self, 'poly_include_type' ) or \
           ( self.poly_include_type != poly_include_type and \
             self.op_mode in [OpMode.NONE, OpMode.EDIT_PTS] ):
            self.poly_include_type = poly_include_type
            self.update()
        self.cbar.btn_poly_type_include_act.setChecked(self.poly_include_type == SegIEPolyType.INCLUDE)
        self.cbar.btn_poly_type_exclude_act.setChecked(self.poly_include_type == SegIEPolyType.EXCLUDE)

    # ====================================================================================
    # ====================================================================================
    # ====================================== METHODS =====================================
    # ====================================================================================
    # ====================================================================================

    def update_cursor(self, is_finalize=False):
        if not self.initialized:
            return

        if not self.mouse_in_widget or is_finalize:
            if self.current_cursor is not None:
                QApplication.restoreOverrideCursor()
                self.current_cursor = None
        else:
            color_cc = self.get_current_color_scheme().cross_cursor
            nc = Qt.ArrowCursor

            if self.drag_type == DragType.IMAGE_LOOK:
                nc = Qt.ClosedHandCursor
            else:

                if self.op_mode == OpMode.NONE:
                    nc = color_cc
                    if self.mouse_wire_poly is not None:
                        nc = Qt.PointingHandCursor

                elif self.op_mode == OpMode.DRAW_PTS:
                    nc = color_cc
                elif self.op_mode == OpMode.EDIT_PTS:
                    nc = Qt.ArrowCursor

                    if self.mouse_op_poly_pt_id is not None:
                        nc = Qt.PointingHandCursor

                    if self.pt_edit_mode == PTEditMode.ADD_DEL:

                        if self.mouse_op_poly_edge_id is not None and \
                        self.mouse_op_poly_pt_id is None:
                            nc = color_cc
            if self.current_cursor != nc:
                if self.current_cursor is None:
                    QApplication.setOverrideCursor(nc)
                else:
                    QApplication.changeOverrideCursor(nc)
                self.current_cursor = nc

    def update_mouse_info(self, mouse_cli_pt=None):
        """
        Update selected polys/edges/points by given mouse position
        """
        if mouse_cli_pt is not None:
            self.mouse_cli_pt = mouse_cli_pt.astype(np.float32)

        self.mouse_img_pt = self.cli_to_img_pt(self.mouse_cli_pt)

        new_mouse_hull_poly = self.get_poly_by_pt_in_hull(self.mouse_cli_pt)

        if self.mouse_hull_poly != new_mouse_hull_poly:
            self.mouse_hull_poly = new_mouse_hull_poly
            self.update_cursor()
            self.update()

        new_mouse_wire_poly = self.get_poly_by_pt_near_wire(self.mouse_cli_pt)

        if self.mouse_wire_poly != new_mouse_wire_poly:
            self.mouse_wire_poly = new_mouse_wire_poly
            self.update_cursor()
            self.update()

        if self.op_mode in [OpMode.DRAW_PTS, OpMode.EDIT_PTS]:
            new_mouse_op_poly_pt_id = self.get_poly_pt_id_under_pt (self.op_poly, self.mouse_cli_pt)
            if self.mouse_op_poly_pt_id != new_mouse_op_poly_pt_id:
                self.mouse_op_poly_pt_id = new_mouse_op_poly_pt_id
                self.update_cursor()
                self.update()

            new_mouse_op_poly_edge_id,\
            new_mouse_op_poly_edge_id_pt = self.get_poly_edge_id_pt_under_pt (self.op_poly, self.mouse_cli_pt)
            if self.mouse_op_poly_edge_id != new_mouse_op_poly_edge_id:
                self.mouse_op_poly_edge_id = new_mouse_op_poly_edge_id
                self.update_cursor()
                self.update()

            if (self.mouse_op_poly_edge_id_pt.__class__ != new_mouse_op_poly_edge_id_pt.__class__) or \
               (isinstance(self.mouse_op_poly_edge_id_pt, np.ndarray) and \
                all(self.mouse_op_poly_edge_id_pt != new_mouse_op_poly_edge_id_pt)):

                self.mouse_op_poly_edge_id_pt = new_mouse_op_poly_edge_id_pt
                self.update_cursor()
                self.update()


    def action_undo_pt(self):
        if self.drag_type == DragType.NONE:
            if self.op_mode == OpMode.DRAW_PTS:
                if self.op_poly.undo() == 0:
                    self.ie_polys.remove_poly (self.op_poly)
                    self.set_op_mode(OpMode.NONE)
                self.update()

    def action_redo_pt(self):
        if self.drag_type == DragType.NONE:
            if self.op_mode == OpMode.DRAW_PTS:
                self.op_poly.redo()
                self.update()

    def action_delete_poly(self):
        if self.op_mode == OpMode.EDIT_PTS and \
            self.drag_type == DragType.NONE and \
            self.pt_edit_mode == PTEditMode.MOVE:
            # Delete current poly
            self.ie_polys.remove_poly (self.op_poly)
            self.set_op_mode(OpMode.NONE)

    # ====================================================================================
    # ====================================================================================
    # ================================== OVERRIDE QT METHODS =============================
    # ====================================================================================
    # ====================================================================================
    def on_keyPressEvent(self, ev):
        if not self.initialized:
            return
        key = ev.key()
        key_mods = int(ev.modifiers())
        if self.op_mode == OpMode.DRAW_PTS:
            self.set_view_lock(ViewLock.CENTER if key_mods == Qt.ShiftModifier else ViewLock.NONE )
        elif self.op_mode == OpMode.EDIT_PTS:
            self.set_pt_edit_mode(PTEditMode.ADD_DEL if key_mods == Qt.ControlModifier else PTEditMode.MOVE )

    def on_keyReleaseEvent(self, ev):
        if not self.initialized:
            return
        key = ev.key()
        key_mods = int(ev.modifiers())
        if self.op_mode == OpMode.DRAW_PTS:
            self.set_view_lock(ViewLock.CENTER if key_mods == Qt.ShiftModifier else ViewLock.NONE )
        elif self.op_mode == OpMode.EDIT_PTS:
            self.set_pt_edit_mode(PTEditMode.ADD_DEL if key_mods == Qt.ControlModifier else PTEditMode.MOVE )

    def enterEvent(self, ev):
        super().enterEvent(ev)
        self.mouse_in_widget = True
        self.update_cursor()

    def leaveEvent(self, ev):
        super().leaveEvent(ev)
        self.mouse_in_widget = False
        self.update_cursor()

    def mousePressEvent(self, ev):
        super().mousePressEvent(ev)
        if not self.initialized:
            return

        self.update_mouse_info(QPoint_to_np(ev.pos()))

        btn = ev.button()

        if btn == Qt.LeftButton:
            if self.op_mode == OpMode.NONE:
                # Clicking in NO OPERATION mode
                if self.mouse_wire_poly is not None:
                    # Click on wire on any poly -> switch to EDIT_MODE
                    self.set_op_mode(OpMode.EDIT_PTS, op_poly=self.mouse_wire_poly)
                else:
                    # Click on empty space -> create new poly with one point
                    new_poly = self.ie_polys.add_poly(self.poly_include_type)
                    self.ie_polys.sort()
                    new_poly.add_pt(*self.mouse_img_pt)
                    self.set_op_mode(OpMode.DRAW_PTS, op_poly=new_poly )

            elif self.op_mode == OpMode.DRAW_PTS:
                # Clicking in DRAW_PTS mode
                if len(self.op_poly.get_pts()) >= 3 and self.mouse_op_poly_pt_id == 0:
                    # Click on first point -> close poly and switch to edit mode
                    self.set_op_mode(OpMode.EDIT_PTS, op_poly=self.op_poly)
                else:
                    # Click on empty space -> add point to current poly
                    self.op_poly.add_pt(*self.mouse_img_pt)
                    self.update()

            elif self.op_mode == OpMode.EDIT_PTS:
                # Clicking in EDIT_PTS mode

                if self.mouse_op_poly_pt_id is not None:
                    # Click on point of op_poly
                    if self.pt_edit_mode == PTEditMode.ADD_DEL:
                        # in mode 'delete point'
                        self.op_poly.remove_pt(self.mouse_op_poly_pt_id)
                        if self.op_poly.get_pts_count() < 3:
                            # not enough points after delete -> remove poly
                            self.ie_polys.remove_poly (self.op_poly)
                            self.set_op_mode(OpMode.NONE)
                        self.update()

                    elif self.drag_type == DragType.NONE:
                        # otherwise -> start drag
                        self.drag_type = DragType.POLY_PT
                        self.drag_cli_pt     = self.mouse_cli_pt
                        self.drag_poly_pt_id = self.mouse_op_poly_pt_id
                        self.drag_poly_pt    = self.op_poly.get_pts()[ self.drag_poly_pt_id ]
                elif self.mouse_op_poly_edge_id is not None:
                    # Click on edge of op_poly
                    if self.pt_edit_mode == PTEditMode.ADD_DEL:
                        # in mode 'insert new point'
                        edge_img_pt = self.cli_to_img_pt(self.mouse_op_poly_edge_id_pt)
                        self.op_poly.insert_pt (self.mouse_op_poly_edge_id+1, edge_img_pt)
                        self.update()
                    else:
                        # Otherwise do nothing
                        pass
                else:
                    # other cases -> unselect poly
                    self.set_op_mode(OpMode.NONE)

        elif btn == Qt.MiddleButton:
            if self.drag_type == DragType.NONE:
                # Start image drag
                self.drag_type = DragType.IMAGE_LOOK
                self.drag_cli_pt = self.mouse_cli_pt
                self.drag_img_look_pt = self.get_img_look_pt()
                self.update_cursor()


    def mouseReleaseEvent(self, ev):
        super().mouseReleaseEvent(ev)
        if not self.initialized:
            return

        self.update_mouse_info(QPoint_to_np(ev.pos()))

        btn = ev.button()

        if btn == Qt.LeftButton:
            if self.op_mode == OpMode.EDIT_PTS:
                if self.drag_type == DragType.POLY_PT:
                    self.drag_type = DragType.NONE
                    self.update()

        elif btn == Qt.MiddleButton:
            if self.drag_type == DragType.IMAGE_LOOK:
                self.drag_type = DragType.NONE
                self.update_cursor()
                self.update()

    def mouseMoveEvent(self, ev):
        super().mouseMoveEvent(ev)
        if not self.initialized:
            return

        prev_mouse_cli_pt = self.mouse_cli_pt
        self.update_mouse_info(QPoint_to_np(ev.pos()))

        if self.view_lock == ViewLock.CENTER:
            if npla.norm(self.mouse_cli_pt - prev_mouse_cli_pt) >= 1:
                self.img_look_pt = self.mouse_img_pt
                QCursor.setPos ( self.mapToGlobal( QPoint_from_np(self.img_to_cli_pt(self.img_look_pt)) ))
            self.update()

        if self.drag_type == DragType.IMAGE_LOOK:
            delta_pt = self.cli_to_img_pt(self.mouse_cli_pt) - self.cli_to_img_pt(self.drag_cli_pt)
            self.img_look_pt = self.drag_img_look_pt - delta_pt
            self.update()

        if self.op_mode == OpMode.DRAW_PTS:
            self.update()
        elif self.op_mode == OpMode.EDIT_PTS:
            if self.drag_type == DragType.POLY_PT:
                delta_pt = self.cli_to_img_pt(self.mouse_cli_pt) - self.cli_to_img_pt(self.drag_cli_pt)
                self.op_poly.set_point(self.drag_poly_pt_id, self.drag_poly_pt + delta_pt)
                self.update()

    def wheelEvent(self, ev):
        super().wheelEvent(ev)

        if not self.initialized:
            return

        mods = int(ev.modifiers())
        delta = ev.angleDelta()

        cli_pt = QPoint_to_np(ev.pos())

        if self.drag_type == DragType.NONE:
            sign = np.sign( delta.y() )
            prev_img_pos = self.cli_to_img_pt (cli_pt)
            delta_scale = sign*0.2 + sign * self.get_view_scale() / 10.0
            self.view_scale = np.clip(self.get_view_scale() + delta_scale, 1.0, 20.0)
            new_img_pos = self.cli_to_img_pt (cli_pt)
            if sign > 0:
                self.img_look_pt = self.get_img_look_pt() + (prev_img_pos-new_img_pos)#*1.5
            else:
                QCursor.setPos ( self.mapToGlobal(QPoint_from_np(self.img_to_cli_pt(prev_img_pos))) )
            self.update()

    def paintEvent(self, event):
        super().paintEvent(event)
        if not self.initialized:
            return

        qp = self.qp
        qp.begin(self)
        qp.setRenderHint(QPainter.Antialiasing)
        qp.setRenderHint(QPainter.HighQualityAntialiasing)
        qp.setRenderHint(QPainter.SmoothPixmapTransform)

        src_rect = QRect(0, 0, *self.img_size)
        dst_rect = self.img_to_cli_rect( src_rect )

        if self.op_mode == OpMode.VIEW_BAKED:
            qp.drawPixmap(dst_rect, self.img_baked_pixmap, src_rect)
        elif self.op_mode == OpMode.VIEW_XSEG_MASK:
            if self.xseg_mask_pixmap is not None:
                qp.drawPixmap(dst_rect, self.xseg_mask_pixmap, src_rect)
        else:
            if self.cbar.btn_view_xseg_overlay_mask_act.isChecked() and \
               self.xseg_overlay_mask_pixmap is not None:
                qp.drawPixmap(dst_rect, self.xseg_overlay_mask_pixmap, src_rect)
            elif self.img_pixmap is not None:
                qp.drawPixmap(dst_rect, self.img_pixmap, src_rect)

            polys = self.ie_polys.get_polys()
            polys_len = len(polys)

            color_scheme = self.get_current_color_scheme()

            pt_rad = self.canvas_config.pt_radius
            pt_rad_x2 = pt_rad*2

            pt_select_radius = self.canvas_config.pt_select_radius

            op_mode = self.op_mode
            op_poly = self.op_poly

            for i,poly in enumerate(polys):

                selected_pt_path = QPainterPath()
                poly_line_path = QPainterPath()
                pts_line_path = QPainterPath()

                pt_remove_cli_pt = None
                poly_pts = poly.get_pts()
                for pt_id, img_pt in enumerate(poly_pts):
                    cli_pt = self.img_to_cli_pt(img_pt)
                    q_cli_pt = QPoint_from_np(cli_pt)

                    if pt_id == 0:
                        poly_line_path.moveTo(q_cli_pt)
                    else:
                        poly_line_path.lineTo(q_cli_pt)


                    if poly == op_poly:
                        if self.op_mode == OpMode.DRAW_PTS or \
                            (self.op_mode == OpMode.EDIT_PTS and \
                            (self.pt_edit_mode == PTEditMode.MOVE) or \
                            (self.pt_edit_mode == PTEditMode.ADD_DEL and self.mouse_op_poly_pt_id == pt_id) \
                            ):
                            pts_line_path.moveTo( QPoint_from_np(cli_pt + np.float32([0,-pt_rad])) )
                            pts_line_path.lineTo( QPoint_from_np(cli_pt + np.float32([0,pt_rad])) )
                            pts_line_path.moveTo( QPoint_from_np(cli_pt + np.float32([-pt_rad,0])) )
                            pts_line_path.lineTo( QPoint_from_np(cli_pt + np.float32([pt_rad,0])) )

                        if (self.op_mode == OpMode.EDIT_PTS and \
                            self.pt_edit_mode == PTEditMode.ADD_DEL and \
                            self.mouse_op_poly_pt_id == pt_id):
                            pt_remove_cli_pt = cli_pt

                        if self.op_mode == OpMode.DRAW_PTS and \
                            len(op_poly.get_pts()) >= 3 and pt_id == 0 and self.mouse_op_poly_pt_id == pt_id:
                            # Circle around poly point
                            selected_pt_path.addEllipse(q_cli_pt, pt_rad_x2, pt_rad_x2)


                if poly == op_poly:
                    if op_mode == OpMode.DRAW_PTS:
                        # Line from last point to mouse
                        poly_line_path.lineTo( QPoint_from_np(self.mouse_cli_pt) )

                    if self.mouse_op_poly_pt_id is not None:
                        pass

                    if self.mouse_op_poly_edge_id_pt is not None:
                        if self.pt_edit_mode == PTEditMode.ADD_DEL and self.mouse_op_poly_pt_id is None:
                            # Ready to insert point on edge
                            m_cli_pt = self.mouse_op_poly_edge_id_pt
                            pts_line_path.moveTo( QPoint_from_np(m_cli_pt + np.float32([0,-pt_rad])) )
                            pts_line_path.lineTo( QPoint_from_np(m_cli_pt + np.float32([0,pt_rad])) )
                            pts_line_path.moveTo( QPoint_from_np(m_cli_pt + np.float32([-pt_rad,0])) )
                            pts_line_path.lineTo( QPoint_from_np(m_cli_pt + np.float32([pt_rad,0])) )

                if len(poly_pts) >= 2:
                    # Closing poly line
                    poly_line_path.lineTo( QPoint_from_np(self.img_to_cli_pt(poly_pts[0])) )

                # Draw calls
                qp.setPen(color_scheme.pt_outline_pen)
                qp.setBrush(QBrush())
                qp.drawPath(selected_pt_path)

                qp.setPen(color_scheme.poly_outline_solid_pen)
                qp.setBrush(QBrush())
                qp.drawPath(pts_line_path)

                if poly.get_type() == SegIEPolyType.INCLUDE:
                    qp.setPen(color_scheme.poly_outline_solid_pen)
                else:
                    qp.setPen(color_scheme.poly_outline_dot_pen)

                qp.setBrush(color_scheme.poly_unselected_brush)
                if op_mode == OpMode.NONE:
                    if poly == self.mouse_wire_poly:
                        qp.setBrush(color_scheme.poly_selected_brush)
                #else:
                #    if poly == op_poly:
                #        qp.setBrush(color_scheme.poly_selected_brush)

                qp.drawPath(poly_line_path)

                if pt_remove_cli_pt is not None:
                    qp.setPen(color_scheme.poly_outline_solid_pen)
                    qp.setBrush(QBrush())

                    qp.drawLine( *(pt_remove_cli_pt + np.float32([-pt_rad_x2,-pt_rad_x2])), *(pt_remove_cli_pt + np.float32([pt_rad_x2,pt_rad_x2])) )
                    qp.drawLine( *(pt_remove_cli_pt + np.float32([-pt_rad_x2,pt_rad_x2])), *(pt_remove_cli_pt + np.float32([pt_rad_x2,-pt_rad_x2])) )

        qp.end()

class QCanvas(QFrame):
    def __init__(self):
        super().__init__()

        self.canvas_control_left_bar = QCanvasControlsLeftBar()
        self.canvas_control_right_bar = QCanvasControlsRightBar()

        cbar = sn( btn_poly_color_red_act   = self.canvas_control_right_bar.btn_poly_color_red_act,
                   btn_poly_color_green_act = self.canvas_control_right_bar.btn_poly_color_green_act,
                   btn_poly_color_blue_act  = self.canvas_control_right_bar.btn_poly_color_blue_act,
                   btn_view_baked_mask_act  = self.canvas_control_right_bar.btn_view_baked_mask_act,
                   btn_view_xseg_mask_act  = self.canvas_control_right_bar.btn_view_xseg_mask_act,
                   btn_view_xseg_overlay_mask_act  = self.canvas_control_right_bar.btn_view_xseg_overlay_mask_act,
                   btn_poly_color_act_grp = self.canvas_control_right_bar.btn_poly_color_act_grp,
                   btn_view_lock_center_act = self.canvas_control_right_bar.btn_view_lock_center_act,

                   btn_poly_type_include_act = self.canvas_control_left_bar.btn_poly_type_include_act,
                   btn_poly_type_exclude_act = self.canvas_control_left_bar.btn_poly_type_exclude_act,
                   btn_poly_type_act_grp = self.canvas_control_left_bar.btn_poly_type_act_grp,
                   btn_undo_pt_act = self.canvas_control_left_bar.btn_undo_pt_act,
                   btn_redo_pt_act = self.canvas_control_left_bar.btn_redo_pt_act,
                   btn_delete_poly_act = self.canvas_control_left_bar.btn_delete_poly_act,
                   btn_pt_edit_mode_act = self.canvas_control_left_bar.btn_pt_edit_mode_act )

        self.op = QCanvasOperator(cbar)
        self.l = QHBoxLayout()
        self.l.setContentsMargins(0,0,0,0)
        self.l.addWidget(self.canvas_control_left_bar)
        self.l.addWidget(self.op)
        self.l.addWidget(self.canvas_control_right_bar)
        self.setLayout(self.l)

class LoaderQSubprocessor(QSubprocessor):
    def __init__(self, image_paths, q_label, q_progressbar, on_finish_func ):

        self.image_paths = image_paths
        self.image_paths_len = len(image_paths)
        self.idxs = [*range(self.image_paths_len)]

        self.filtered_image_paths = self.image_paths.copy()

        self.image_paths_has_ie_polys = { image_path : False for image_path in self.image_paths }

        self.q_label = q_label
        self.q_progressbar = q_progressbar
        self.q_progressbar.setRange(0, self.image_paths_len)
        self.q_progressbar.setValue(0)
        self.q_progressbar.update()
        self.on_finish_func = on_finish_func
        self.done_count = 0
        super().__init__('LoaderQSubprocessor', LoaderQSubprocessor.Cli, 60)

    def get_data(self, host_dict):
        if len (self.idxs) > 0:
            idx = self.idxs.pop(0)
            image_path = self.image_paths[idx]
            self.q_label.setText(f'{QStringDB.loading_tip}... {image_path.name}')

            return idx, image_path

        return None

    def on_clients_finalized(self):
        self.on_finish_func([x for x in self.filtered_image_paths if x is not None], self.image_paths_has_ie_polys)

    def on_data_return (self, host_dict, data):
        self.idxs.insert(0, data[0])

    def on_result (self, host_dict, data, result):
        idx, has_dflimg, has_ie_polys = result

        if not has_dflimg:
            self.filtered_image_paths[idx] = None
        self.image_paths_has_ie_polys[self.image_paths[idx]] = has_ie_polys

        self.done_count += 1
        if self.q_progressbar is not None:
            self.q_progressbar.setValue(self.done_count)

    class Cli(QSubprocessor.Cli):
        def process_data(self, data):
            idx, filename = data
            dflimg = DFLIMG.load(filename)
            if dflimg is not None and dflimg.has_data():
                ie_polys = dflimg.get_seg_ie_polys()

                return idx, True, ie_polys.has_polys()
            return idx, False, False

class MainWindow(QXMainWindow):

    def __init__(self, input_dirpath, cfg_root_path):
        self.loading_frame = None
        self.help_frame = None

        super().__init__()

        self.input_dirpath = input_dirpath
        self.trash_dirpath = input_dirpath.parent / (input_dirpath.name + '_trash')
        self.cfg_root_path = cfg_root_path

        self.cfg_path = cfg_root_path / 'MainWindow_cfg.dat'
        self.cfg_dict = pickle.loads(self.cfg_path.read_bytes()) if self.cfg_path.exists() else {}

        self.cached_images = {}
        self.cached_has_ie_polys = {}

        self.initialize_ui()

        # Loader
        self.loading_frame = QFrame(self.main_canvas_frame)
        self.loading_frame.setAutoFillBackground(True)
        self.loading_frame.setFrameShape(QFrame.StyledPanel)
        self.loader_label = QLabel()
        self.loader_progress_bar = QProgressBar()

        intro_image = QLabel()
        intro_image.setPixmap( QPixmap.fromImage(QImageDB.intro) )

        intro_image_frame_l = QVBoxLayout()
        intro_image_frame_l.addWidget(intro_image, alignment=Qt.AlignCenter)
        intro_image_frame = QFrame()
        intro_image_frame.setSizePolicy (QSizePolicy.Expanding, QSizePolicy.Expanding)
        intro_image_frame.setLayout(intro_image_frame_l)

        loading_frame_l = QVBoxLayout()
        loading_frame_l.addWidget (intro_image_frame)
        loading_frame_l.addWidget (self.loader_label)
        loading_frame_l.addWidget (self.loader_progress_bar)
        self.loading_frame.setLayout(loading_frame_l)

        self.loader_subprocessor = LoaderQSubprocessor( image_paths=pathex.get_image_paths(input_dirpath, return_Path_class=True),
                                                        q_label=self.loader_label,
                                                        q_progressbar=self.loader_progress_bar,
                                                        on_finish_func=self.on_loader_finish )


    def on_loader_finish(self, image_paths, image_paths_has_ie_polys):
        self.image_paths_done = []
        self.image_paths = image_paths
        self.image_paths_has_ie_polys = image_paths_has_ie_polys
        self.set_has_ie_polys_count ( len([ 1 for x in self.image_paths_has_ie_polys if self.image_paths_has_ie_polys[x] == True]) )
        self.loading_frame.hide()
        self.loading_frame = None

        self.process_next_image(first_initialization=True)

    def closeEvent(self, ev):
        self.cfg_dict['geometry'] = self.saveGeometry().data()
        self.cfg_path.write_bytes( pickle.dumps(self.cfg_dict) )


    def update_cached_images (self, count=5):
        d = self.cached_images

        for image_path in self.image_paths_done[:-count]+self.image_paths[count:]:
            if image_path in d:
                del d[image_path]

        for image_path in self.image_paths[:count]+self.image_paths_done[-count:]:
            if image_path not in d:
                img = cv2_imread(image_path)
                if img is not None:
                    d[image_path] = img

    def load_image(self, image_path):
        try:
            img = self.cached_images.get(image_path, None)
            if img is None:
                img = cv2_imread(image_path)
                self.cached_images[image_path] = img
            if img is None:
                io.log_err(f'Unable to load {image_path}')
        except:
            img = None

        return img

    def update_preview_bar(self):
        count = self.image_bar.get_preview_images_count()
        d = self.cached_images
        prev_imgs = [ d.get(image_path, None) for image_path in self.image_paths_done[-1:-count:-1] ]
        next_imgs = [ d.get(image_path, None) for image_path in self.image_paths[:count] ]
        self.image_bar.update_images(prev_imgs, next_imgs)


    def canvas_initialize(self, image_path, only_has_polys=False):
        if only_has_polys and not self.image_paths_has_ie_polys[image_path]:
            return False

        dflimg = DFLIMG.load(image_path)
        if not dflimg or not dflimg.has_data():
            return False

        ie_polys = dflimg.get_seg_ie_polys()
        xseg_mask = dflimg.get_xseg_mask()
        img = self.load_image(image_path)
        if img is None:
            return False

        self.canvas.op.initialize ( img,  ie_polys=ie_polys, xseg_mask=xseg_mask )

        self.filename_label.setText(f"{image_path.name}")

        return True

    def canvas_finalize(self, image_path):
        self.canvas.op.finalize()

        if image_path.exists():
            dflimg = DFLIMG.load(image_path)
            ie_polys = dflimg.get_seg_ie_polys()
            new_ie_polys = self.canvas.op.get_ie_polys()

            if not new_ie_polys.identical(ie_polys):
                prev_has_polys = self.image_paths_has_ie_polys[image_path]
                self.image_paths_has_ie_polys[image_path] = new_ie_polys.has_polys()
                new_has_polys = self.image_paths_has_ie_polys[image_path]

                if not prev_has_polys and new_has_polys:
                    self.set_has_ie_polys_count ( self.get_has_ie_polys_count() +1)
                elif prev_has_polys and not new_has_polys:
                    self.set_has_ie_polys_count ( self.get_has_ie_polys_count() -1)

                dflimg.set_seg_ie_polys( new_ie_polys )
                dflimg.save()

        self.filename_label.setText(f"")

    def process_prev_image(self):
        key_mods = QApplication.keyboardModifiers()
        step = 5 if key_mods == Qt.ShiftModifier else 1
        only_has_polys = key_mods == Qt.ControlModifier

        if self.canvas.op.is_initialized():
            self.canvas_finalize(self.image_paths[0])

        while True:
            for _ in range(step):
                if len(self.image_paths_done) != 0:
                    self.image_paths.insert (0, self.image_paths_done.pop(-1))
                else:
                    break
            if len(self.image_paths) == 0:
                break

            ret = self.canvas_initialize(self.image_paths[0], len(self.image_paths_done) != 0 and only_has_polys)

            if ret or len(self.image_paths_done) == 0:
                break

        self.update_cached_images()
        self.update_preview_bar()

    def process_next_image(self, first_initialization=False):
        key_mods = QApplication.keyboardModifiers()

        step = 0 if first_initialization else 5 if key_mods == Qt.ShiftModifier else 1
        only_has_polys = False if first_initialization else key_mods == Qt.ControlModifier

        if self.canvas.op.is_initialized():
            self.canvas_finalize(self.image_paths[0])

        while True:
            for _ in range(step):
                if len(self.image_paths) != 0:
                    self.image_paths_done.append(self.image_paths.pop(0))
                else:
                    break
            if len(self.image_paths) == 0:
                break
            if self.canvas_initialize(self.image_paths[0], only_has_polys):
                break

        self.update_cached_images()
        self.update_preview_bar()
        
    def trash_current_image(self):
        self.process_next_image()
        
        img_path = self.image_paths_done.pop(-1)
        img_path = Path(img_path)
        self.trash_dirpath.mkdir(parents=True, exist_ok=True)
        img_path.rename( self.trash_dirpath / img_path.name )
        
        self.update_cached_images()
        self.update_preview_bar()
        
    def initialize_ui(self):

        self.canvas = QCanvas()

        image_bar = self.image_bar = ImagePreviewSequenceBar(preview_images_count=9, icon_size=QUIConfig.preview_bar_icon_q_size.width())
        image_bar.setSizePolicy ( QSizePolicy.Fixed, QSizePolicy.Fixed )


        btn_prev_image = QXIconButton(QIconDB.left, QStringDB.btn_prev_image_tip, shortcut='A', click_func=self.process_prev_image)
        btn_prev_image.setIconSize(QUIConfig.preview_bar_icon_q_size)

        btn_next_image = QXIconButton(QIconDB.right, QStringDB.btn_next_image_tip, shortcut='D', click_func=self.process_next_image)
        btn_next_image.setIconSize(QUIConfig.preview_bar_icon_q_size)

        btn_delete_image = QXIconButton(QIconDB.trashcan, QStringDB.btn_delete_image_tip, shortcut='X', click_func=self.trash_current_image)
        btn_delete_image.setIconSize(QUIConfig.preview_bar_icon_q_size)
    
        pad_image = QWidget()
        pad_image.setFixedSize(QUIConfig.preview_bar_icon_q_size)
        
        preview_image_bar_frame_l = QHBoxLayout()
        preview_image_bar_frame_l.setContentsMargins(0,0,0,0)
        preview_image_bar_frame_l.addWidget ( pad_image, alignment=Qt.AlignCenter)
        preview_image_bar_frame_l.addWidget ( btn_prev_image, alignment=Qt.AlignCenter)
        preview_image_bar_frame_l.addWidget ( image_bar)
        preview_image_bar_frame_l.addWidget ( btn_next_image, alignment=Qt.AlignCenter)
        #preview_image_bar_frame_l.addWidget ( btn_delete_image, alignment=Qt.AlignCenter)

        preview_image_bar_frame = QFrame()
        preview_image_bar_frame.setSizePolicy ( QSizePolicy.Fixed, QSizePolicy.Fixed )
        preview_image_bar_frame.setLayout(preview_image_bar_frame_l)

        preview_image_bar_frame2_l = QHBoxLayout()
        preview_image_bar_frame2_l.setContentsMargins(0,0,0,0)
        preview_image_bar_frame2_l.addWidget ( btn_delete_image, alignment=Qt.AlignCenter)

        preview_image_bar_frame2 = QFrame()
        preview_image_bar_frame2.setSizePolicy ( QSizePolicy.Fixed, QSizePolicy.Fixed )
        preview_image_bar_frame2.setLayout(preview_image_bar_frame2_l)
        
        preview_image_bar_l = QHBoxLayout()
        preview_image_bar_l.addWidget (preview_image_bar_frame, alignment=Qt.AlignCenter)
        preview_image_bar_l.addWidget (preview_image_bar_frame2)
        
        preview_image_bar = QFrame()
        preview_image_bar.setFrameShape(QFrame.StyledPanel)
        preview_image_bar.setSizePolicy ( QSizePolicy.Expanding, QSizePolicy.Fixed )
        preview_image_bar.setLayout(preview_image_bar_l)

        label_font = QFont('Courier New')
        self.filename_label = QLabel()
        self.filename_label.setFont(label_font)

        self.has_ie_polys_count_label = QLabel()

        status_frame_l = QHBoxLayout()
        status_frame_l.setContentsMargins(0,0,0,0)
        status_frame_l.addWidget ( QLabel(), alignment=Qt.AlignCenter)
        status_frame_l.addWidget (self.filename_label, alignment=Qt.AlignCenter)
        status_frame_l.addWidget (self.has_ie_polys_count_label, alignment=Qt.AlignCenter)
        status_frame = QFrame()
        status_frame.setLayout(status_frame_l)

        main_canvas_l = QVBoxLayout()
        main_canvas_l.setContentsMargins(0,0,0,0)
        main_canvas_l.addWidget (self.canvas)
        main_canvas_l.addWidget (status_frame)
        main_canvas_l.addWidget (preview_image_bar)

        self.main_canvas_frame = QFrame()
        self.main_canvas_frame.setLayout(main_canvas_l)

        self.main_l = QHBoxLayout()
        self.main_l.setContentsMargins(0,0,0,0)
        self.main_l.addWidget (self.main_canvas_frame)

        self.setLayout(self.main_l)

        geometry = self.cfg_dict.get('geometry', None)
        if geometry is not None:
            self.restoreGeometry(geometry)
        else:
            self.move( QPoint(0,0))

    def get_has_ie_polys_count(self):
        return self.has_ie_polys_count

    def set_has_ie_polys_count(self, c):
        self.has_ie_polys_count = c
        self.has_ie_polys_count_label.setText(f"{c} {QStringDB.labeled_tip}")

    def resizeEvent(self, ev):
        if self.loading_frame is not None:
            self.loading_frame.resize( ev.size() )
        if self.help_frame is not None:
            self.help_frame.resize( ev.size() )

def start(input_dirpath):
    """
    returns exit_code
    """
    io.log_info("Running XSeg editor.")

    if PackedFaceset.path_contains(input_dirpath):
        io.log_info (f'\n{input_dirpath} contains packed faceset! Unpack it first.\n')
        return 1

    root_path = Path(__file__).parent
    cfg_root_path = Path(tempfile.gettempdir())

    QApplication.setAttribute(Qt.AA_EnableHighDpiScaling, True)
    QApplication.setAttribute(Qt.AA_UseHighDpiPixmaps, True)

    app = QApplication([])
    app.setApplicationName("XSegEditor")
    app.setStyle('Fusion')

    QFontDatabase.addApplicationFont( str(root_path / 'gfx' / 'fonts' / 'NotoSans-Medium.ttf') )

    app.setFont( QFont('NotoSans'))

    QUIConfig.initialize()
    QStringDB.initialize()

    QIconDB.initialize( root_path / 'gfx' / 'icons' )
    QCursorDB.initialize( root_path / 'gfx' / 'cursors' )
    QImageDB.initialize( root_path / 'gfx' / 'images' )

    app.setWindowIcon(QIconDB.app_icon)
    app.setPalette( QDarkPalette() )

    win = MainWindow( input_dirpath=input_dirpath, cfg_root_path=cfg_root_path)

    win.show()
    win.raise_()

    app.exec_()
    return 0

XSegEditor/XSegEditor.py
oo92/TorchFaceLab
import cv2
import numpy as np
from pathlib import Path
from core.interact import interact as io
from core import imagelib 
import traceback

def cv2_imread(filename, flags=cv2.IMREAD_UNCHANGED, loader_func=None, verbose=True):
    """
    allows to open non-english characters path
    """
    try:
        if loader_func is not None:
            bytes = bytearray(loader_func(filename))
        else:
            with open(filename, "rb") as stream:
                bytes = bytearray(stream.read())
        numpyarray = np.asarray(bytes, dtype=np.uint8)
        return cv2.imdecode(numpyarray, flags)
    except:
        if verbose:
            io.log_err(f"Exception occured in cv2_imread : {traceback.format_exc()}")
        return None

def cv2_imwrite(filename, img, *args):
    ret, buf = cv2.imencode( Path(filename).suffix, img, *args)
    if ret == True:
        try:
            with open(filename, "wb") as stream:
                stream.write( buf )
        except:
            pass

def cv2_resize(x, *args, **kwargs):
    h,w,c = x.shape
    x = cv2.resize(x, *args, **kwargs)
    
    x = imagelib.normalize_channels(x, c)
    return x
    
core/cv2ex.py
oo92/TorchFaceLab
import numpy as np
import cv2
from enum import IntEnum


class SegIEPolyType(IntEnum):
    EXCLUDE = 0
    INCLUDE = 1
    


class SegIEPoly():
    def __init__(self, type=None, pts=None, **kwargs):
        self.type = type
        
        if pts is None:
            pts = np.empty( (0,2), dtype=np.float32 )
        else:
            pts = np.float32(pts)
        self.pts = pts
        self.n_max = self.n = len(pts)

    def dump(self):
        return {'type': int(self.type),
                'pts' : self.get_pts(),
               }
    
    def identical(self, b):
        if self.n != b.n:
            return False            
        return (self.pts[0:self.n] == b.pts[0:b.n]).all()        
        
    def get_type(self):
        return self.type

    def add_pt(self, x, y):
        self.pts = np.append(self.pts[0:self.n], [ ( float(x), float(y) ) ], axis=0).astype(np.float32)
        self.n_max = self.n = self.n + 1

    def undo(self):
        self.n = max(0, self.n-1)
        return self.n

    def redo(self):
        self.n = min(len(self.pts), self.n+1)
        return self.n

    def redo_clip(self):
        self.pts = self.pts[0:self.n]
        self.n_max = self.n

    def insert_pt(self, n, pt):
        if n < 0 or n > self.n:
            raise ValueError("insert_pt out of range")
        self.pts = np.concatenate( (self.pts[0:n], pt[None,...].astype(np.float32), self.pts[n:]), axis=0)
        self.n_max = self.n = self.n+1
        
    def remove_pt(self, n):
        if n < 0 or n >= self.n:
            raise ValueError("remove_pt out of range")
        self.pts = np.concatenate( (self.pts[0:n], self.pts[n+1:]), axis=0)
        self.n_max = self.n = self.n-1
        
    def get_last_point(self):
        return self.pts[self.n-1].copy()

    def get_pts(self):
        return self.pts[0:self.n].copy()
        
    def get_pts_count(self):
        return self.n

    def set_point(self, id, pt):
        self.pts[id] = pt
        
    def set_points(self, pts):
        self.pts = np.array(pts)
        self.n_max = self.n = len(pts)
        
    def mult_points(self, val):
        self.pts *= val       
    
        

class SegIEPolys():
    def __init__(self):
        self.polys = []

    def identical(self, b):
        polys_len = len(self.polys)
        o_polys_len = len(b.polys)
        if polys_len != o_polys_len:
            return False
        
        return all ([ a_poly.identical(b_poly) for a_poly, b_poly in zip(self.polys, b.polys) ])
        
    def add_poly(self, ie_poly_type):       
        poly = SegIEPoly(ie_poly_type)
        self.polys.append (poly)
        return poly

    def remove_poly(self, poly):
        if poly in self.polys:
            self.polys.remove(poly)

    def has_polys(self):
        return len(self.polys) != 0
        
    def get_poly(self, id):
        return self.polys[id]

    def get_polys(self):
        return self.polys
        
    def get_pts_count(self):
        return sum([poly.get_pts_count() for poly in self.polys])
        
    def sort(self):
        poly_by_type = { SegIEPolyType.EXCLUDE : [], SegIEPolyType.INCLUDE : [] }

        for poly in self.polys:
            poly_by_type[poly.type].append(poly)
            
        self.polys = poly_by_type[SegIEPolyType.INCLUDE] + poly_by_type[SegIEPolyType.EXCLUDE]

    def __iter__(self):
        for poly in self.polys:
            yield poly

    def overlay_mask(self, mask):
        h,w,c = mask.shape
        white = (1,)*c
        black = (0,)*c
        for poly in self.polys:
            pts = poly.get_pts().astype(np.int32)
            if len(pts) != 0:
                cv2.fillPoly(mask, [pts], white if poly.type == SegIEPolyType.INCLUDE else black )

    def dump(self):
        return {'polys' : [ poly.dump() for poly in self.polys ] }
    
    def mult_points(self, val):
        for poly in self.polys:
            poly.mult_points(val)
        
    @staticmethod
    def load(data=None):
        ie_polys = SegIEPolys()
        if data is not None:
            if isinstance(data, list):
                # Backward comp
                ie_polys.polys = [ SegIEPoly(type=type, pts=pts) for (type, pts) in data ]
            elif isinstance(data, dict):                
                ie_polys.polys = [ SegIEPoly(**poly_cfg) for poly_cfg in data['polys'] ]
                
        ie_polys.sort()   
                
        return ie_polys
core/imagelib/SegIEPolys.py
oo92/TorchFaceLab
from .estimate_sharpness import estimate_sharpness

from .equalize_and_stack_square import equalize_and_stack_square

from .text import get_text_image, get_draw_text_lines

from .draw import draw_polygon, draw_rect

from .morph import morph_by_points

from .warp import gen_warp_params, warp_by_params

from .reduce_colors import reduce_colors

from .color_transfer import color_transfer, color_transfer_mix, color_transfer_sot, color_transfer_mkl, color_transfer_idt, color_hist_match, reinhard_color_transfer, linear_color_transfer

from .common import random_crop, normalize_channels, cut_odd_image, overlay_alpha_image

from .SegIEPolys import *

from .blursharpen import LinearMotionBlur, blursharpen

from .filters import apply_random_rgb_levels, \
                     apply_random_overlay_triangle, \
                     apply_random_hsv_shift, \
                     apply_random_sharpen, \
                     apply_random_motion_blur, \
                     apply_random_gaussian_blur, \
                     apply_random_nearest_resize, \
                     apply_random_bilinear_resize, \
                     apply_random_jpeg_compress, \
                     apply_random_relight

core/imagelib/__init__.py
oo92/TorchFaceLab
import cv2
import numpy as np

def LinearMotionBlur(image, size, angle):
    k = np.zeros((size, size), dtype=np.float32)
    k[ (size-1)// 2 , :] = np.ones(size, dtype=np.float32)
    k = cv2.warpAffine(k, cv2.getRotationMatrix2D( (size / 2 -0.5 , size / 2 -0.5 ) , angle, 1.0), (size, size) )
    k = k * ( 1.0 / np.sum(k) )
    return cv2.filter2D(image, -1, k)
    
def blursharpen (img, sharpen_mode=0, kernel_size=3, amount=100):
    if kernel_size % 2 == 0:
        kernel_size += 1
    if amount > 0:
        if sharpen_mode == 1: #box
            kernel = np.zeros( (kernel_size, kernel_size), dtype=np.float32)
            kernel[ kernel_size//2, kernel_size//2] = 1.0
            box_filter = np.ones( (kernel_size, kernel_size), dtype=np.float32) / (kernel_size**2)
            kernel = kernel + (kernel - box_filter) * amount
            return cv2.filter2D(img, -1, kernel)
        elif sharpen_mode == 2: #gaussian
            blur = cv2.GaussianBlur(img, (kernel_size, kernel_size) , 0)
            img = cv2.addWeighted(img, 1.0 + (0.5 * amount), blur, -(0.5 * amount), 0)
            return img
    elif amount < 0:
        n = -amount
        while n > 0:

            img_blur = cv2.medianBlur(img, 5)
            if int(n / 10) != 0:
                img = img_blur
            else:
                pass_power = (n % 10) / 10.0
                img = img*(1.0-pass_power)+img_blur*pass_power
            n = max(n-10,0)

        return img
    return img
core/imagelib/blursharpen.py
oo92/TorchFaceLab
import cv2
import numexpr as ne
import numpy as np
import scipy as sp
from numpy import linalg as npla


def color_transfer_sot(src,trg, steps=10, batch_size=5, reg_sigmaXY=16.0, reg_sigmaV=5.0):
    """
    Color Transform via Sliced Optimal Transfer
    ported by @iperov from https://github.com/dcoeurjo/OTColorTransfer

    src         - any float range any channel image
    dst         - any float range any channel image, same shape as src
    steps       - number of solver steps
    batch_size  - solver batch size
    reg_sigmaXY - apply regularization and sigmaXY of filter, otherwise set to 0.0
    reg_sigmaV  - sigmaV of filter

    return value - clip it manually
    """
    if not np.issubdtype(src.dtype, np.floating):
        raise ValueError("src value must be float")
    if not np.issubdtype(trg.dtype, np.floating):
        raise ValueError("trg value must be float")

    if len(src.shape) != 3:
        raise ValueError("src shape must have rank 3 (h,w,c)")

    if src.shape != trg.shape:
        raise ValueError("src and trg shapes must be equal")

    src_dtype = src.dtype
    h,w,c = src.shape
    new_src = src.copy()

    advect = np.empty ( (h*w,c), dtype=src_dtype )
    for step in range (steps):
        advect.fill(0)
        for batch in range (batch_size):
            dir = np.random.normal(size=c).astype(src_dtype)
            dir /= npla.norm(dir)

            projsource = np.sum( new_src*dir, axis=-1).reshape ((h*w))
            projtarget = np.sum( trg*dir, axis=-1).reshape ((h*w))

            idSource = np.argsort (projsource)
            idTarget = np.argsort (projtarget)

            a = projtarget[idTarget]-projsource[idSource]
            for i_c in range(c):
                advect[idSource,i_c] += a * dir[i_c]
        new_src += advect.reshape( (h,w,c) ) / batch_size

    if reg_sigmaXY != 0.0:
        src_diff = new_src-src
        src_diff_filt = cv2.bilateralFilter (src_diff, 0, reg_sigmaV, reg_sigmaXY )
        if len(src_diff_filt.shape) == 2:
            src_diff_filt = src_diff_filt[...,None]
        new_src = src + src_diff_filt
    return new_src

def color_transfer_mkl(x0, x1):
    eps = np.finfo(float).eps

    h,w,c = x0.shape
    h1,w1,c1 = x1.shape

    x0 = x0.reshape ( (h*w,c) )
    x1 = x1.reshape ( (h1*w1,c1) )

    a = np.cov(x0.T)
    b = np.cov(x1.T)

    Da2, Ua = np.linalg.eig(a)
    Da = np.diag(np.sqrt(Da2.clip(eps, None)))

    C = np.dot(np.dot(np.dot(np.dot(Da, Ua.T), b), Ua), Da)

    Dc2, Uc = np.linalg.eig(C)
    Dc = np.diag(np.sqrt(Dc2.clip(eps, None)))

    Da_inv = np.diag(1./(np.diag(Da)))

    t = np.dot(np.dot(np.dot(np.dot(np.dot(np.dot(Ua, Da_inv), Uc), Dc), Uc.T), Da_inv), Ua.T)

    mx0 = np.mean(x0, axis=0)
    mx1 = np.mean(x1, axis=0)

    result = np.dot(x0-mx0, t) + mx1
    return np.clip ( result.reshape ( (h,w,c) ).astype(x0.dtype), 0, 1)

def color_transfer_idt(i0, i1, bins=256, n_rot=20):
    import scipy.stats
    
    relaxation = 1 / n_rot
    h,w,c = i0.shape
    h1,w1,c1 = i1.shape

    i0 = i0.reshape ( (h*w,c) )
    i1 = i1.reshape ( (h1*w1,c1) )

    n_dims = c

    d0 = i0.T
    d1 = i1.T

    for i in range(n_rot):

        r = sp.stats.special_ortho_group.rvs(n_dims).astype(np.float32)

        d0r = np.dot(r, d0)
        d1r = np.dot(r, d1)
        d_r = np.empty_like(d0)

        for j in range(n_dims):

            lo = min(d0r[j].min(), d1r[j].min())
            hi = max(d0r[j].max(), d1r[j].max())

            p0r, edges = np.histogram(d0r[j], bins=bins, range=[lo, hi])
            p1r, _     = np.histogram(d1r[j], bins=bins, range=[lo, hi])

            cp0r = p0r.cumsum().astype(np.float32)
            cp0r /= cp0r[-1]

            cp1r = p1r.cumsum().astype(np.float32)
            cp1r /= cp1r[-1]

            f = np.interp(cp0r, cp1r, edges[1:])

            d_r[j] = np.interp(d0r[j], edges[1:], f, left=0, right=bins)

        d0 = relaxation * np.linalg.solve(r, (d_r - d0r)) + d0

    return np.clip ( d0.T.reshape ( (h,w,c) ).astype(i0.dtype) , 0, 1)

def reinhard_color_transfer(target : np.ndarray, source : np.ndarray, target_mask : np.ndarray = None, source_mask : np.ndarray = None, mask_cutoff=0.5) -> np.ndarray:
    """
    Transfer color using rct method.

        target      np.ndarray H W 3C   (BGR)   np.float32
        source      np.ndarray H W 3C   (BGR)   np.float32

        target_mask(None)   np.ndarray H W 1C  np.float32
        source_mask(None)   np.ndarray H W 1C  np.float32
        
        mask_cutoff(0.5)    float

    masks are used to limit the space where color statistics will be computed to adjust the target

    reference: Color Transfer between Images https://www.cs.tau.ac.il/~turkel/imagepapers/ColorTransfer.pdf
    """
    source = cv2.cvtColor(source, cv2.COLOR_BGR2LAB)
    target = cv2.cvtColor(target, cv2.COLOR_BGR2LAB)

    source_input = source
    if source_mask is not None:
        source_input = source_input.copy()
        source_input[source_mask[...,0] < mask_cutoff] = [0,0,0]
    
    target_input = target
    if target_mask is not None:
        target_input = target_input.copy()
        target_input[target_mask[...,0] < mask_cutoff] = [0,0,0]

    target_l_mean, target_l_std, target_a_mean, target_a_std, target_b_mean, target_b_std, \
        = target_input[...,0].mean(), target_input[...,0].std(), target_input[...,1].mean(), target_input[...,1].std(), target_input[...,2].mean(), target_input[...,2].std()
    
    source_l_mean, source_l_std, source_a_mean, source_a_std, source_b_mean, source_b_std, \
        = source_input[...,0].mean(), source_input[...,0].std(), source_input[...,1].mean(), source_input[...,1].std(), source_input[...,2].mean(), source_input[...,2].std()
    
    # not as in the paper: scale by the standard deviations using reciprocal of paper proposed factor
    target_l = target[...,0]
    target_l = ne.evaluate('(target_l - target_l_mean) * source_l_std / target_l_std + source_l_mean')

    target_a = target[...,1]
    target_a = ne.evaluate('(target_a - target_a_mean) * source_a_std / target_a_std + source_a_mean')
    
    target_b = target[...,2]
    target_b = ne.evaluate('(target_b - target_b_mean) * source_b_std / target_b_std + source_b_mean')

    np.clip(target_l,    0, 100, out=target_l)
    np.clip(target_a, -127, 127, out=target_a)
    np.clip(target_b, -127, 127, out=target_b)

    return cv2.cvtColor(np.stack([target_l,target_a,target_b], -1), cv2.COLOR_LAB2BGR)


def linear_color_transfer(target_img, source_img, mode='pca', eps=1e-5):
    '''
    Matches the colour distribution of the target image to that of the source image
    using a linear transform.
    Images are expected to be of form (w,h,c) and float in [0,1].
    Modes are chol, pca or sym for different choices of basis.
    '''
    mu_t = target_img.mean(0).mean(0)
    t = target_img - mu_t
    t = t.transpose(2,0,1).reshape( t.shape[-1],-1)
    Ct = t.dot(t.T) / t.shape[1] + eps * np.eye(t.shape[0])
    mu_s = source_img.mean(0).mean(0)
    s = source_img - mu_s
    s = s.transpose(2,0,1).reshape( s.shape[-1],-1)
    Cs = s.dot(s.T) / s.shape[1] + eps * np.eye(s.shape[0])
    if mode == 'chol':
        chol_t = np.linalg.cholesky(Ct)
        chol_s = np.linalg.cholesky(Cs)
        ts = chol_s.dot(np.linalg.inv(chol_t)).dot(t)
    if mode == 'pca':
        eva_t, eve_t = np.linalg.eigh(Ct)
        Qt = eve_t.dot(np.sqrt(np.diag(eva_t))).dot(eve_t.T)
        eva_s, eve_s = np.linalg.eigh(Cs)
        Qs = eve_s.dot(np.sqrt(np.diag(eva_s))).dot(eve_s.T)
        ts = Qs.dot(np.linalg.inv(Qt)).dot(t)
    if mode == 'sym':
        eva_t, eve_t = np.linalg.eigh(Ct)
        Qt = eve_t.dot(np.sqrt(np.diag(eva_t))).dot(eve_t.T)
        Qt_Cs_Qt = Qt.dot(Cs).dot(Qt)
        eva_QtCsQt, eve_QtCsQt = np.linalg.eigh(Qt_Cs_Qt)
        QtCsQt = eve_QtCsQt.dot(np.sqrt(np.diag(eva_QtCsQt))).dot(eve_QtCsQt.T)
        ts = np.linalg.inv(Qt).dot(QtCsQt).dot(np.linalg.inv(Qt)).dot(t)
    matched_img = ts.reshape(*target_img.transpose(2,0,1).shape).transpose(1,2,0)
    matched_img += mu_s
    matched_img[matched_img>1] = 1
    matched_img[matched_img<0] = 0
    return np.clip(matched_img.astype(source_img.dtype), 0, 1)

def lab_image_stats(image):
    # compute the mean and standard deviation of each channel
    (l, a, b) = cv2.split(image)
    (lMean, lStd) = (l.mean(), l.std())
    (aMean, aStd) = (a.mean(), a.std())
    (bMean, bStd) = (b.mean(), b.std())

    # return the color statistics
    return (lMean, lStd, aMean, aStd, bMean, bStd)

def _scale_array(arr, clip=True):
    if clip:
        return np.clip(arr, 0, 255)

    mn = arr.min()
    mx = arr.max()
    scale_range = (max([mn, 0]), min([mx, 255]))

    if mn < scale_range[0] or mx > scale_range[1]:
        return (scale_range[1] - scale_range[0]) * (arr - mn) / (mx - mn) + scale_range[0]

    return arr

def channel_hist_match(source, template, hist_match_threshold=255, mask=None):
    # Code borrowed from:
    # https://stackoverflow.com/questions/32655686/histogram-matching-of-two-images-in-python-2-x
    masked_source = source
    masked_template = template

    if mask is not None:
        masked_source = source * mask
        masked_template = template * mask

    oldshape = source.shape
    source = source.ravel()
    template = template.ravel()
    masked_source = masked_source.ravel()
    masked_template = masked_template.ravel()
    s_values, bin_idx, s_counts = np.unique(source, return_inverse=True,
                                            return_counts=True)
    t_values, t_counts = np.unique(template, return_counts=True)

    s_quantiles = np.cumsum(s_counts).astype(np.float64)
    s_quantiles = hist_match_threshold * s_quantiles / s_quantiles[-1]
    t_quantiles = np.cumsum(t_counts).astype(np.float64)
    t_quantiles = 255 * t_quantiles / t_quantiles[-1]
    interp_t_values = np.interp(s_quantiles, t_quantiles, t_values)

    return interp_t_values[bin_idx].reshape(oldshape)

def color_hist_match(src_im, tar_im, hist_match_threshold=255):
    h,w,c = src_im.shape
    matched_R = channel_hist_match(src_im[:,:,0], tar_im[:,:,0], hist_match_threshold, None)
    matched_G = channel_hist_match(src_im[:,:,1], tar_im[:,:,1], hist_match_threshold, None)
    matched_B = channel_hist_match(src_im[:,:,2], tar_im[:,:,2], hist_match_threshold, None)

    to_stack = (matched_R, matched_G, matched_B)
    for i in range(3, c):
        to_stack += ( src_im[:,:,i],)


    matched = np.stack(to_stack, axis=-1).astype(src_im.dtype)
    return matched

def color_transfer_mix(img_src,img_trg):
    img_src = np.clip(img_src*255.0, 0, 255).astype(np.uint8)
    img_trg = np.clip(img_trg*255.0, 0, 255).astype(np.uint8)

    img_src_lab = cv2.cvtColor(img_src, cv2.COLOR_BGR2LAB)
    img_trg_lab = cv2.cvtColor(img_trg, cv2.COLOR_BGR2LAB)

    rct_light = np.clip ( linear_color_transfer(img_src_lab[...,0:1].astype(np.float32)/255.0,
                                                img_trg_lab[...,0:1].astype(np.float32)/255.0 )[...,0]*255.0,
                          0, 255).astype(np.uint8)

    img_src_lab[...,0] = (np.ones_like (rct_light)*100).astype(np.uint8)
    img_src_lab = cv2.cvtColor(img_src_lab, cv2.COLOR_LAB2BGR)

    img_trg_lab[...,0] = (np.ones_like (rct_light)*100).astype(np.uint8)
    img_trg_lab = cv2.cvtColor(img_trg_lab, cv2.COLOR_LAB2BGR)

    img_rct = color_transfer_sot( img_src_lab.astype(np.float32), img_trg_lab.astype(np.float32) )
    img_rct = np.clip(img_rct, 0, 255).astype(np.uint8)

    img_rct = cv2.cvtColor(img_rct, cv2.COLOR_BGR2LAB)
    img_rct[...,0] = rct_light
    img_rct = cv2.cvtColor(img_rct, cv2.COLOR_LAB2BGR)


    return (img_rct / 255.0).astype(np.float32)

def color_transfer(ct_mode, img_src, img_trg):
    """
    color transfer for [0,1] float32 inputs
    """
    if ct_mode == 'lct':
        out = linear_color_transfer (img_src, img_trg)
    elif ct_mode == 'rct':
        out = reinhard_color_transfer(img_src, img_trg)
    elif ct_mode == 'mkl':
        out = color_transfer_mkl (img_src, img_trg)
    elif ct_mode == 'idt':
        out = color_transfer_idt (img_src, img_trg)
    elif ct_mode == 'sot':
        out = color_transfer_sot (img_src, img_trg)
        out = np.clip( out, 0.0, 1.0)
    else:
        raise ValueError(f"unknown ct_mode {ct_mode}")
    return out

core/imagelib/color_transfer.py
oo92/TorchFaceLab
import numpy as np

def random_crop(img, w, h):
    height, width = img.shape[:2]
    
    h_rnd = height - h
    w_rnd = width - w
    
    y = np.random.randint(0, h_rnd) if h_rnd > 0 else 0
    x = np.random.randint(0, w_rnd) if w_rnd > 0 else 0
    
    return img[y:y+height, x:x+width]
                        
def normalize_channels(img, target_channels):
    img_shape_len = len(img.shape)
    if img_shape_len == 2:
        h, w = img.shape
        c = 0
    elif img_shape_len == 3:
        h, w, c = img.shape
    else:
        raise ValueError("normalize: incorrect image dimensions.")

    if c == 0 and target_channels > 0:
        img = img[...,np.newaxis]
        c = 1

    if c == 1 and target_channels > 1:
        img = np.repeat (img, target_channels, -1)
        c = target_channels

    if c > target_channels:
        img = img[...,0:target_channels]
        c = target_channels

    return img

def cut_odd_image(img):
    h, w, c = img.shape
    wm, hm = w % 2, h % 2
    if wm + hm != 0:
        img = img[0:h-hm,0:w-wm,:]
    return img

def overlay_alpha_image(img_target, img_source, xy_offset=(0,0) ):
    (h,w,c) = img_source.shape
    if c != 4:
        raise ValueError("overlay_alpha_image, img_source must have 4 channels")

    x1, x2 = xy_offset[0], xy_offset[0] + w
    y1, y2 = xy_offset[1], xy_offset[1] + h

    alpha_s = img_source[:, :, 3] / 255.0
    alpha_l = 1.0 - alpha_s

    for c in range(0, 3):
        img_target[y1:y2, x1:x2, c] = (alpha_s * img_source[:, :, c] +
                                        alpha_l * img_target[y1:y2, x1:x2, c])
core/imagelib/common.py
oo92/TorchFaceLab
import numpy as np
import cv2

def draw_polygon (image, points, color, thickness = 1):
    points_len = len(points)
    for i in range (0, points_len):
        p0 = tuple( points[i] )
        p1 = tuple( points[ (i+1) % points_len] )
        cv2.line (image, p0, p1, color, thickness=thickness)

def draw_rect(image, rect, color, thickness=1):
    l,t,r,b = rect
    draw_polygon (image, [ (l,t), (r,t), (r,b), (l,b ) ], color, thickness)

core/imagelib/draw.py
oo92/TorchFaceLab
import numpy as np
import cv2

def equalize_and_stack_square (images, axis=1):
    max_c = max ([ 1 if len(image.shape) == 2 else image.shape[2]  for image in images ] )

    target_wh = 99999
    for i,image in enumerate(images):
        if len(image.shape) == 2:
            h,w = image.shape
            c = 1
        else:
            h,w,c = image.shape

        if h < target_wh:
            target_wh = h

        if w < target_wh:
            target_wh = w

    for i,image in enumerate(images):
        if len(image.shape) == 2:
            h,w = image.shape
            c = 1
        else:
            h,w,c = image.shape

        if c < max_c:
            if c == 1:
                if len(image.shape) == 2:
                    image = np.expand_dims ( image, -1 )
                image = np.concatenate ( (image,)*max_c, -1 )
            elif c == 2: #GA
                image = np.expand_dims ( image[...,0], -1 )
                image = np.concatenate ( (image,)*max_c, -1 )
            else:
                image = np.concatenate ( (image, np.ones((h,w,max_c - c))), -1 )

        if h != target_wh or w != target_wh:
            image = cv2.resize ( image, (target_wh, target_wh) )
            h,w,c = image.shape

        images[i] = image

    return np.concatenate ( images, axis = 1 )
core/imagelib/equalize_and_stack_square.py
oo92/TorchFaceLab
"""
Copyright (c) 2009-2010 Arizona Board of Regents.  All Rights Reserved.
 Contact: Lina Karam (karam@asu.edu) and Niranjan Narvekar (nnarveka@asu.edu)
 Image, Video, and Usabilty (IVU) Lab, http://ivulab.asu.edu , Arizona State University
 This copyright statement may not be removed from any file containing it or from modifications to these files.
 This copyright notice must also be included in any file or product that is derived from the source files.

 Redistribution and use of this code in source and binary forms,  with or without modification, are permitted provided that the
 following conditions are met:
 - Redistribution's of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
 - Redistribution's in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer
in the documentation and/or other materials provided with the distribution.
 - The Image, Video, and Usability Laboratory (IVU Lab, http://ivulab.asu.edu) is acknowledged in any publication that
 reports research results using this code, copies of this code, or modifications of this code.
 The code and our papers are to be cited in the bibliography as:

N. D. Narvekar and L. J. Karam, "CPBD Sharpness Metric Software", http://ivulab.asu.edu/Quality/CPBD

N. D. Narvekar and L. J. Karam, "A No-Reference Image Blur Metric Based on the Cumulative
Probability of Blur Detection (CPBD)," accepted and to appear in the IEEE Transactions on Image Processing,  2011.

N. D. Narvekar and L. J. Karam, "An Improved No-Reference Sharpness Metric Based on the Probability of Blur Detection," International Workshop on Video Processing and Quality Metrics for Consumer Electronics (VPQM), January 2010, http://www.vpqm.org (pdf)

N. D. Narvekar and L. J. Karam, "A No Reference Perceptual Quality Metric based on Cumulative Probability of Blur Detection," First International Workshop on the Quality of Multimedia Experience (QoMEX), pp. 87-91, July 2009.

 DISCLAIMER:
 This software is provided by the copyright holders and contributors "as is" and any express or implied warranties, including, but not limited to, the implied warranties of merchantability and fitness for a particular purpose are disclaimed. In no event shall the Arizona Board of Regents, Arizona State University, IVU Lab members, authors or contributors be liable for any direct, indirect, incidental, special, exemplary, or consequential damages (including, but not limited to, procurement of substitute
goods or services; loss of use, data, or profits; or business interruption) however caused and on any theory of liability, whether in contract, strict liability, or tort (including negligence or otherwise) arising in any way out of the use of this software, even if advised of the possibility of such damage.
"""

import numpy as np
import cv2
from math import atan2, pi


def sobel(image):
    # type: (numpy.ndarray) -> numpy.ndarray
    """
    Find edges using the Sobel approximation to the derivatives.

    Inspired by the [Octave implementation](https://sourceforge.net/p/octave/image/ci/default/tree/inst/edge.m#l196).
    """
    from skimage.filters.edges import HSOBEL_WEIGHTS
    h1 = np.array(HSOBEL_WEIGHTS)
    h1 /= np.sum(abs(h1))  # normalize h1
    
    from scipy.ndimage import convolve
    strength2 = np.square(convolve(image, h1.T))

    # Note: https://sourceforge.net/p/octave/image/ci/default/tree/inst/edge.m#l59
    thresh2 = 2 * np.sqrt(np.mean(strength2))

    strength2[strength2 <= thresh2] = 0
    return _simple_thinning(strength2)


def _simple_thinning(strength):
    # type: (numpy.ndarray) -> numpy.ndarray
    """
    Perform a very simple thinning.

    Inspired by the [Octave implementation](https://sourceforge.net/p/octave/image/ci/default/tree/inst/edge.m#l512).
    """
    num_rows, num_cols = strength.shape

    zero_column = np.zeros((num_rows, 1))
    zero_row = np.zeros((1, num_cols))

    x = (
        (strength > np.c_[zero_column, strength[:, :-1]]) &
        (strength > np.c_[strength[:, 1:], zero_column])
    )

    y = (
        (strength > np.r_[zero_row, strength[:-1, :]]) &
        (strength > np.r_[strength[1:, :], zero_row])
    )

    return x | y





# threshold to characterize blocks as edge/non-edge blocks
THRESHOLD = 0.002
# fitting parameter
BETA = 3.6
# block size
BLOCK_HEIGHT, BLOCK_WIDTH = (64, 64)
# just noticeable widths based on the perceptual experiments
WIDTH_JNB = np.concatenate([5*np.ones(51), 3*np.ones(205)])


def compute(image):
    # type: (numpy.ndarray) -> float
    """Compute the sharpness metric for the given data."""

    # convert the image to double for further processing
    image = image.astype(np.float64)

    # edge detection using canny and sobel canny edge detection is done to
    # classify the blocks as edge or non-edge blocks and sobel edge
    # detection is done for the purpose of edge width measurement.
    from skimage.feature import canny
    canny_edges = canny(image)
    sobel_edges = sobel(image)

    # edge width calculation
    marziliano_widths = marziliano_method(sobel_edges, image)

    # sharpness metric calculation
    return _calculate_sharpness_metric(image, canny_edges, marziliano_widths)


def marziliano_method(edges, image):
    # type: (numpy.ndarray, numpy.ndarray) -> numpy.ndarray
    """
    Calculate the widths of the given edges.

    :return: A matrix with the same dimensions as the given image with 0's at
        non-edge locations and edge-widths at the edge locations.
    """

    # `edge_widths` consists of zero and non-zero values. A zero value
    # indicates that there is no edge at that position and a non-zero value
    # indicates that there is an edge at that position and the value itself
    # gives the edge width.
    edge_widths = np.zeros(image.shape)

    # find the gradient for the image
    gradient_y, gradient_x = np.gradient(image)

    # dimensions of the image
    img_height, img_width = image.shape

    # holds the angle information of the edges
    edge_angles = np.zeros(image.shape)

    # calculate the angle of the edges
    for row in range(img_height):
        for col in range(img_width):
            if gradient_x[row, col] != 0:
                edge_angles[row, col] = atan2(gradient_y[row, col], gradient_x[row, col]) * (180 / pi)
            elif gradient_x[row, col] == 0 and gradient_y[row, col] == 0:
                edge_angles[row,col] = 0
            elif gradient_x[row, col] == 0 and gradient_y[row, col] == pi/2:
                edge_angles[row, col] = 90


    if np.any(edge_angles):

        # quantize the angle
        quantized_angles = 45 * np.round(edge_angles / 45)

        for row in range(1, img_height - 1):
            for col in range(1, img_width - 1):
                if edges[row, col] == 1:

                    # gradient angle = 180 or -180
                    if quantized_angles[row, col] == 180 or quantized_angles[row, col] == -180:
                        for margin in range(100 + 1):
                            inner_border = (col - 1) - margin
                            outer_border = (col - 2) - margin

                            # outside image or intensity increasing from left to right
                            if outer_border < 0 or (image[row, outer_border] - image[row, inner_border]) <= 0:
                                break

                        width_left = margin + 1

                        for margin in range(100 + 1):
                            inner_border = (col + 1) + margin
                            outer_border = (col + 2) + margin

                            # outside image or intensity increasing from left to right
                            if outer_border >= img_width or (image[row, outer_border] - image[row, inner_border]) >= 0:
                                break

                        width_right = margin + 1

                        edge_widths[row, col] = width_left + width_right


                    # gradient angle = 0
                    if quantized_angles[row, col] == 0:
                        for margin in range(100 + 1):
                            inner_border = (col - 1) - margin
                            outer_border = (col - 2) - margin

                            # outside image or intensity decreasing from left to right
                            if outer_border < 0 or (image[row, outer_border] - image[row, inner_border]) >= 0:
                                break

                        width_left = margin + 1

                        for margin in range(100 + 1):
                            inner_border = (col + 1) + margin
                            outer_border = (col + 2) + margin

                            # outside image or intensity decreasing from left to right
                            if outer_border >= img_width or (image[row, outer_border] - image[row, inner_border]) <= 0:
                                break

                        width_right = margin + 1

                        edge_widths[row, col] = width_right + width_left

    return edge_widths


def _calculate_sharpness_metric(image, edges, edge_widths):
    # type: (numpy.array, numpy.array, numpy.array) -> numpy.float64

    # get the size of image
    img_height, img_width = image.shape

    total_num_edges = 0
    hist_pblur = np.zeros(101)

    # maximum block indices
    num_blocks_vertically = int(img_height / BLOCK_HEIGHT)
    num_blocks_horizontally = int(img_width / BLOCK_WIDTH)

    #  loop over the blocks
    for i in range(num_blocks_vertically):
        for j in range(num_blocks_horizontally):

            # get the row and col indices for the block pixel positions
            rows = slice(BLOCK_HEIGHT * i, BLOCK_HEIGHT * (i + 1))
            cols = slice(BLOCK_WIDTH * j, BLOCK_WIDTH * (j + 1))

            if is_edge_block(edges[rows, cols], THRESHOLD):
                block_widths = edge_widths[rows, cols]
                # rotate block to simulate column-major boolean indexing
                block_widths = np.rot90(np.flipud(block_widths), 3)
                block_widths = block_widths[block_widths != 0]

                block_contrast = get_block_contrast(image[rows, cols])
                block_jnb = WIDTH_JNB[block_contrast]

                # calculate the probability of blur detection at the edges
                # detected in the block
                prob_blur_detection = 1 - np.exp(-abs(block_widths/block_jnb) ** BETA)

                # update the statistics using the block information
                for probability in prob_blur_detection:
                    bucket = int(round(probability * 100))
                    hist_pblur[bucket] += 1
                    total_num_edges += 1

    # normalize the pdf
    if total_num_edges > 0:
        hist_pblur = hist_pblur / total_num_edges

    # calculate the sharpness metric
    return np.sum(hist_pblur[:64])


def is_edge_block(block, threshold):
    # type: (numpy.ndarray, float) -> bool
    """Decide whether the given block is an edge block."""
    return np.count_nonzero(block) > (block.size * threshold)


def get_block_contrast(block):
    # type: (numpy.ndarray) -> int
    return int(np.max(block) - np.min(block))


def estimate_sharpness(image):
    if image.ndim == 3:
        if image.shape[2] > 1:
            image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        else:
            image = image[...,0]
        
    return compute(image)

core/imagelib/estimate_sharpness.py
oo92/TorchFaceLab
import numpy as np
from .blursharpen import LinearMotionBlur, blursharpen
import cv2

def apply_random_rgb_levels(img, mask=None, rnd_state=None):
    if rnd_state is None:
        rnd_state = np.random
    np_rnd = rnd_state.rand

    inBlack  = np.array([np_rnd()*0.25    , np_rnd()*0.25    , np_rnd()*0.25], dtype=np.float32)
    inWhite  = np.array([1.0-np_rnd()*0.25, 1.0-np_rnd()*0.25, 1.0-np_rnd()*0.25], dtype=np.float32)
    inGamma  = np.array([0.5+np_rnd(), 0.5+np_rnd(), 0.5+np_rnd()], dtype=np.float32)

    outBlack  = np.array([np_rnd()*0.25    , np_rnd()*0.25    , np_rnd()*0.25], dtype=np.float32)
    outWhite  = np.array([1.0-np_rnd()*0.25, 1.0-np_rnd()*0.25, 1.0-np_rnd()*0.25], dtype=np.float32)

    result = np.clip( (img - inBlack) / (inWhite - inBlack), 0, 1 )
    result = ( result ** (1/inGamma) ) *  (outWhite - outBlack) + outBlack
    result = np.clip(result, 0, 1)

    if mask is not None:
        result = img*(1-mask) + result*mask

    return result

def apply_random_hsv_shift(img, mask=None, rnd_state=None):
    if rnd_state is None:
        rnd_state = np.random

    h, s, v = cv2.split(cv2.cvtColor(img, cv2.COLOR_BGR2HSV))
    h = ( h + rnd_state.randint(360) ) % 360
    s = np.clip ( s + rnd_state.random()-0.5, 0, 1 )
    v = np.clip ( v + rnd_state.random()-0.5, 0, 1 )

    result = np.clip( cv2.cvtColor(cv2.merge([h, s, v]), cv2.COLOR_HSV2BGR) , 0, 1 )
    if mask is not None:
        result = img*(1-mask) + result*mask

    return result

def apply_random_sharpen( img, chance, kernel_max_size, mask=None, rnd_state=None ):
    if rnd_state is None:
        rnd_state = np.random

    sharp_rnd_kernel = rnd_state.randint(kernel_max_size)+1

    result = img
    if rnd_state.randint(100) < np.clip(chance, 0, 100):
        if rnd_state.randint(2) == 0:
            result = blursharpen(result, 1, sharp_rnd_kernel, rnd_state.randint(10) )
        else:
            result = blursharpen(result, 2, sharp_rnd_kernel, rnd_state.randint(50) )
                
        if mask is not None:
            result = img*(1-mask) + result*mask

    return result

def apply_random_motion_blur( img, chance, mb_max_size, mask=None, rnd_state=None ):
    if rnd_state is None:
        rnd_state = np.random

    mblur_rnd_kernel = rnd_state.randint(mb_max_size)+1
    mblur_rnd_deg    = rnd_state.randint(360)

    result = img
    if rnd_state.randint(100) < np.clip(chance, 0, 100):
        result = LinearMotionBlur (result, mblur_rnd_kernel, mblur_rnd_deg )
        if mask is not None:
            result = img*(1-mask) + result*mask

    return result

def apply_random_gaussian_blur( img, chance, kernel_max_size, mask=None, rnd_state=None ):
    if rnd_state is None:
        rnd_state = np.random

    result = img
    if rnd_state.randint(100) < np.clip(chance, 0, 100):
        gblur_rnd_kernel = rnd_state.randint(kernel_max_size)*2+1
        result = cv2.GaussianBlur(result, (gblur_rnd_kernel,)*2 , 0)
        if mask is not None:
            result = img*(1-mask) + result*mask

    return result

def apply_random_resize( img, chance, max_size_per, interpolation=cv2.INTER_LINEAR, mask=None, rnd_state=None ):
    if rnd_state is None:
        rnd_state = np.random

    result = img
    if rnd_state.randint(100) < np.clip(chance, 0, 100):
        h,w,c = result.shape

        trg = rnd_state.rand()
        rw = w - int( trg * int(w*(max_size_per/100.0)) )
        rh = h - int( trg * int(h*(max_size_per/100.0)) )

        result = cv2.resize (result, (rw,rh), interpolation=interpolation )
        result = cv2.resize (result, (w,h), interpolation=interpolation )
        if mask is not None:
            result = img*(1-mask) + result*mask

    return result

def apply_random_nearest_resize( img, chance, max_size_per, mask=None, rnd_state=None ):
    return apply_random_resize( img, chance, max_size_per, interpolation=cv2.INTER_NEAREST, mask=mask, rnd_state=rnd_state )

def apply_random_bilinear_resize( img, chance, max_size_per, mask=None, rnd_state=None ):
    return apply_random_resize( img, chance, max_size_per, interpolation=cv2.INTER_LINEAR, mask=mask, rnd_state=rnd_state )

def apply_random_jpeg_compress( img, chance, mask=None, rnd_state=None ):
    if rnd_state is None:
        rnd_state = np.random

    result = img
    if rnd_state.randint(100) < np.clip(chance, 0, 100):
        h,w,c = result.shape

        quality = rnd_state.randint(10,101)

        ret, result = cv2.imencode('.jpg', np.clip(img*255, 0,255).astype(np.uint8), [int(cv2.IMWRITE_JPEG_QUALITY), quality] )
        if ret == True:
            result = cv2.imdecode(result, flags=cv2.IMREAD_UNCHANGED)
            result = result.astype(np.float32) / 255.0
            if mask is not None:
                result = img*(1-mask) + result*mask

    return result
    
def apply_random_overlay_triangle( img, max_alpha, mask=None, rnd_state=None ):
    if rnd_state is None:
        rnd_state = np.random

    h,w,c = img.shape
    pt1 = [rnd_state.randint(w), rnd_state.randint(h) ]
    pt2 = [rnd_state.randint(w), rnd_state.randint(h) ]
    pt3 = [rnd_state.randint(w), rnd_state.randint(h) ]
    
    alpha = rnd_state.uniform()*max_alpha
    
    tri_mask = cv2.fillPoly( np.zeros_like(img), [ np.array([pt1,pt2,pt3], np.int32) ], (alpha,)*c )
    
    if rnd_state.randint(2) == 0:
        result = np.clip(img+tri_mask, 0, 1)
    else:
        result = np.clip(img-tri_mask, 0, 1)
    
    if mask is not None:
        result = img*(1-mask) + result*mask

    return result
    
def _min_resize(x, m):
    if x.shape[0] < x.shape[1]:
        s0 = m
        s1 = int(float(m) / float(x.shape[0]) * float(x.shape[1]))
    else:
        s0 = int(float(m) / float(x.shape[1]) * float(x.shape[0]))
        s1 = m
    new_max = min(s1, s0)
    raw_max = min(x.shape[0], x.shape[1])
    return cv2.resize(x, (s1, s0), interpolation=cv2.INTER_LANCZOS4)
    
def _d_resize(x, d, fac=1.0):
    new_min = min(int(d[1] * fac), int(d[0] * fac))
    raw_min = min(x.shape[0], x.shape[1])
    if new_min < raw_min:
        interpolation = cv2.INTER_AREA
    else:
        interpolation = cv2.INTER_LANCZOS4
    y = cv2.resize(x, (int(d[1] * fac), int(d[0] * fac)), interpolation=interpolation)
    return y
    
def _get_image_gradient(dist):
    cols = cv2.filter2D(dist, cv2.CV_32F, np.array([[-1, 0, +1], [-2, 0, +2], [-1, 0, +1]]))
    rows = cv2.filter2D(dist, cv2.CV_32F, np.array([[-1, -2, -1], [0, 0, 0], [+1, +2, +1]]))
    return cols, rows

def _generate_lighting_effects(content):
    h512 = content
    h256 = cv2.pyrDown(h512)
    h128 = cv2.pyrDown(h256)
    h64 = cv2.pyrDown(h128)
    h32 = cv2.pyrDown(h64)
    h16 = cv2.pyrDown(h32)
    c512, r512 = _get_image_gradient(h512)
    c256, r256 = _get_image_gradient(h256)
    c128, r128 = _get_image_gradient(h128)
    c64, r64 = _get_image_gradient(h64)
    c32, r32 = _get_image_gradient(h32)
    c16, r16 = _get_image_gradient(h16)
    c = c16
    c = _d_resize(cv2.pyrUp(c), c32.shape) * 4.0 + c32
    c = _d_resize(cv2.pyrUp(c), c64.shape) * 4.0 + c64
    c = _d_resize(cv2.pyrUp(c), c128.shape) * 4.0 + c128
    c = _d_resize(cv2.pyrUp(c), c256.shape) * 4.0 + c256
    c = _d_resize(cv2.pyrUp(c), c512.shape) * 4.0 + c512
    r = r16
    r = _d_resize(cv2.pyrUp(r), r32.shape) * 4.0 + r32
    r = _d_resize(cv2.pyrUp(r), r64.shape) * 4.0 + r64
    r = _d_resize(cv2.pyrUp(r), r128.shape) * 4.0 + r128
    r = _d_resize(cv2.pyrUp(r), r256.shape) * 4.0 + r256
    r = _d_resize(cv2.pyrUp(r), r512.shape) * 4.0 + r512
    coarse_effect_cols = c
    coarse_effect_rows = r
    EPS = 1e-10

    max_effect = np.max((coarse_effect_cols**2 + coarse_effect_rows**2)**0.5, axis=0, keepdims=True, ).max(1, keepdims=True)
    coarse_effect_cols = (coarse_effect_cols + EPS) / (max_effect + EPS)
    coarse_effect_rows = (coarse_effect_rows + EPS) / (max_effect + EPS)

    return np.stack([ np.zeros_like(coarse_effect_rows), coarse_effect_rows, coarse_effect_cols], axis=-1)
    
def apply_random_relight(img, mask=None, rnd_state=None):
    if rnd_state is None:
        rnd_state = np.random
        
    def_img = img
        
    if rnd_state.randint(2) == 0:
        light_pos_y = 1.0 if rnd_state.randint(2) == 0 else -1.0
        light_pos_x = rnd_state.uniform()*2-1.0
    else:
        light_pos_y = rnd_state.uniform()*2-1.0
        light_pos_x = 1.0 if rnd_state.randint(2) == 0 else -1.0
                    
    light_source_height = 0.3*rnd_state.uniform()*0.7
    light_intensity = 1.0+rnd_state.uniform()
    ambient_intensity = 0.5
    
    light_source_location = np.array([[[light_source_height, light_pos_y, light_pos_x ]]], dtype=np.float32)
    light_source_direction = light_source_location / np.sqrt(np.sum(np.square(light_source_location)))

    lighting_effect = _generate_lighting_effects(img)
    lighting_effect = np.sum(lighting_effect * light_source_direction, axis=-1).clip(0, 1)
    lighting_effect = np.mean(lighting_effect, axis=-1, keepdims=True)

    result = def_img * (ambient_intensity + lighting_effect * light_intensity) #light_source_color
    result = np.clip(result, 0, 1)
    
    if mask is not None:
        result = def_img*(1-mask) + result*mask
    
    return result
core/imagelib/filters.py
oo92/TorchFaceLab
import numpy as np
import cv2
from scipy.spatial import Delaunay


def applyAffineTransform(src, srcTri, dstTri, size) :
    warpMat = cv2.getAffineTransform( np.float32(srcTri), np.float32(dstTri) )
    return cv2.warpAffine( src, warpMat, (size[0], size[1]), None, flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT_101 )

def morphTriangle(dst_img, src_img, st, dt) :
    (h,w,c) = dst_img.shape
    sr = np.array( cv2.boundingRect(np.float32(st)) )
    dr = np.array( cv2.boundingRect(np.float32(dt)) )
    sRect = st - sr[0:2]
    dRect = dt - dr[0:2]
    d_mask = np.zeros((dr[3], dr[2], c), dtype = np.float32)
    cv2.fillConvexPoly(d_mask, np.int32(dRect), (1.0,)*c, 8, 0);
    imgRect = src_img[sr[1]:sr[1] + sr[3], sr[0]:sr[0] + sr[2]]
    size = (dr[2], dr[3])
    warpImage1 = applyAffineTransform(imgRect, sRect, dRect, size)

    if c == 1:
        warpImage1 = np.expand_dims( warpImage1, -1 )

    dst_img[dr[1]:dr[1]+dr[3], dr[0]:dr[0]+dr[2]] = dst_img[dr[1]:dr[1]+dr[3], dr[0]:dr[0]+dr[2]]*(1-d_mask) + warpImage1 * d_mask

def morph_by_points (image, sp, dp):
    if sp.shape != dp.shape:
        raise ValueError ('morph_by_points() sp.shape != dp.shape')
    (h,w,c) = image.shape

    result_image = np.zeros(image.shape, dtype = image.dtype)

    for tri in Delaunay(dp).simplices:
        morphTriangle(result_image, image, sp[tri], dp[tri])

    return result_image
core/imagelib/morph.py
oo92/TorchFaceLab
import numpy as np
import cv2
from PIL import Image

#n_colors = [0..256]
def reduce_colors (img_bgr, n_colors):
    img_rgb = (img_bgr[...,::-1] * 255.0).astype(np.uint8)
    img_rgb_pil = Image.fromarray(img_rgb)
    img_rgb_pil_p = img_rgb_pil.convert('P', palette=Image.ADAPTIVE, colors=n_colors)

    img_rgb_p = img_rgb_pil_p.convert('RGB')
    img_bgr = cv2.cvtColor( np.array(img_rgb_p, dtype=np.float32) / 255.0, cv2.COLOR_RGB2BGR )

    return img_bgr

core/imagelib/reduce_colors.py
oo92/TorchFaceLab
from .draw import circle_faded, random_circle_faded, bezier, random_bezier_split_faded, random_faded
from .calc import *
core/imagelib/sd/__init__.py
oo92/TorchFaceLab
import numpy as np
import numpy.linalg as npla

def dist_to_edges(pts, pt, is_closed=False):
    """
    returns array of dist from pt to edge and projection pt to edges
    """
    if is_closed:
        a = pts
        b = np.concatenate( (pts[1:,:], pts[0:1,:]), axis=0 )
    else:
        a = pts[:-1,:]
        b = pts[1:,:]

    pa = pt-a
    ba = b-a
    
    div = np.einsum('ij,ij->i', ba, ba)
    div[div==0]=1
    h = np.clip( np.einsum('ij,ij->i', pa, ba) / div, 0, 1 )
    
    x = npla.norm ( pa - ba*h[...,None], axis=1 )
    
    return x, a+ba*h[...,None]
    

core/imagelib/sd/calc.py
oo92/TorchFaceLab
"""
Signed distance drawing functions using numpy.
"""
import math

import numpy as np
from numpy import linalg as npla


def vector2_dot(a,b):
    return a[...,0]*b[...,0]+a[...,1]*b[...,1]
    
def vector2_dot2(a):
    return a[...,0]*a[...,0]+a[...,1]*a[...,1]
    
def vector2_cross(a,b):
    return a[...,0]*b[...,1]-a[...,1]*b[...,0]


def circle_faded( wh, center, fade_dists ):
    """
    returns drawn circle in [h,w,1] output range [0..1.0] float32
    
    wh         = [w,h]                      resolution
    center     = [x,y]                      center of circle
    fade_dists = [fade_start, fade_end]     fade values
    """
    w,h = wh      
    
    pts = np.empty( (h,w,2), dtype=np.float32 )
    pts[...,0] = np.arange(w)[:,None]
    pts[...,1] = np.arange(h)[None,:]
    
    pts = pts.reshape ( (h*w, -1) )

    pts_dists = np.abs ( npla.norm(pts-center, axis=-1) )
    
    if fade_dists[1] == 0:
        fade_dists[1] = 1
        
    pts_dists = ( pts_dists - fade_dists[0] ) / fade_dists[1]
        
    pts_dists = np.clip( 1-pts_dists, 0, 1)
    
    return pts_dists.reshape ( (h,w,1) ).astype(np.float32)


def bezier( wh, A, B, C ):
    """
    returns drawn bezier in [h,w,1] output range float32, 
    every pixel contains signed distance to bezier line

        wh      [w,h]       resolution
        A,B,C   points [x,y]
    """
    
    width,height = wh
    
    A = np.float32(A)
    B = np.float32(B)
    C = np.float32(C)
    

    pos = np.empty( (height,width,2), dtype=np.float32 )
    pos[...,0] = np.arange(width)[:,None]
    pos[...,1] = np.arange(height)[None,:]
    

    a = B-A
    b = A - 2.0*B + C
    c = a * 2.0
    d = A - pos

    b_dot = vector2_dot(b,b)
    if b_dot == 0.0:
        return np.zeros( (height,width), dtype=np.float32 )
    
    kk = 1.0 / b_dot

    kx = kk * vector2_dot(a,b)
    ky = kk * (2.0*vector2_dot(a,a)+vector2_dot(d,b))/3.0;
    kz = kk * vector2_dot(d,a);

    res = 0.0;
    sgn = 0.0;

    p = ky - kx*kx;
    
    p3 = p*p*p;
    q = kx*(2.0*kx*kx - 3.0*ky) + kz;
    h = q*q + 4.0*p3;
    
    hp_sel = h >= 0.0
    
    hp_p = h[hp_sel]
    hp_p = np.sqrt(hp_p)
    
    hp_x = ( np.stack( (hp_p,-hp_p), -1) -q[hp_sel,None] ) / 2.0
    hp_uv = np.sign(hp_x) * np.power( np.abs(hp_x), [1.0/3.0, 1.0/3.0] )
    hp_t = np.clip( hp_uv[...,0] + hp_uv[...,1] - kx, 0.0, 1.0 )
    
    hp_t = hp_t[...,None]
    hp_q = d[hp_sel]+(c+b*hp_t)*hp_t
    hp_res = vector2_dot2(hp_q)
    hp_sgn = vector2_cross(c+2.0*b*hp_t,hp_q)
    
    hl_sel = h < 0.0
    
    hl_q = q[hl_sel]
    hl_p = p[hl_sel]
    hl_z = np.sqrt(-hl_p)
    hl_v = np.arccos( hl_q / (hl_p*hl_z*2.0)) / 3.0
    
    hl_m = np.cos(hl_v)
    hl_n = np.sin(hl_v)*1.732050808;
    
    hl_t = np.clip( np.stack( (hl_m+hl_m,-hl_n-hl_m,hl_n-hl_m), -1)*hl_z[...,None]-kx, 0.0, 1.0 );
    
    hl_d = d[hl_sel]
    
    hl_qx = hl_d+(c+b*hl_t[...,0:1])*hl_t[...,0:1]
    
    hl_dx = vector2_dot2(hl_qx)
    hl_sx = vector2_cross(c+2.0*b*hl_t[...,0:1], hl_qx)
    
    hl_qy = hl_d+(c+b*hl_t[...,1:2])*hl_t[...,1:2]
    hl_dy = vector2_dot2(hl_qy)
    hl_sy = vector2_cross(c+2.0*b*hl_t[...,1:2],hl_qy);
    
    hl_dx_l_dy = hl_dx<hl_dy
    hl_dx_ge_dy = hl_dx>=hl_dy
    
    hl_res = np.empty_like(hl_dx)
    hl_res[hl_dx_l_dy] = hl_dx[hl_dx_l_dy]
    hl_res[hl_dx_ge_dy] = hl_dy[hl_dx_ge_dy]
    
    hl_sgn = np.empty_like(hl_sx)
    hl_sgn[hl_dx_l_dy] = hl_sx[hl_dx_l_dy]
    hl_sgn[hl_dx_ge_dy] = hl_sy[hl_dx_ge_dy]
    
    res = np.empty( (height, width), np.float32 )
    res[hp_sel] = hp_res
    res[hl_sel] = hl_res
    
    sgn = np.empty( (height, width), np.float32 )
    sgn[hp_sel] = hp_sgn
    sgn[hl_sel] = hl_sgn
    
    sgn = np.sign(sgn)
    res = np.sqrt(res)*sgn
    
    return res[...,None]
        
def random_faded(wh):
    """
    apply one of them:
     random_circle_faded
     random_bezier_split_faded
    """
    rnd = np.random.randint(2)
    if rnd == 0:
        return random_circle_faded(wh)
    elif rnd == 1:
        return random_bezier_split_faded(wh)
    
def random_circle_faded ( wh, rnd_state=None ):
    if rnd_state is None:
        rnd_state = np.random
        
    w,h = wh
    wh_max = max(w,h)
    fade_start = rnd_state.randint(wh_max)
    fade_end = fade_start + rnd_state.randint(wh_max- fade_start)
    
    return circle_faded (wh, [ rnd_state.randint(h), rnd_state.randint(w) ], 
                             [fade_start, fade_end] ) 
                             
def random_bezier_split_faded( wh ):
    width, height = wh
    
    degA = np.random.randint(360)
    degB = np.random.randint(360)
    degC = np.random.randint(360)
    
    deg_2_rad = math.pi / 180.0
    
    center = np.float32([width / 2.0, height / 2.0])
    
    radius = max(width, height)
    
    A = center + radius*np.float32([ math.sin( degA * deg_2_rad), math.cos( degA * deg_2_rad) ] ) 
    B = center + np.random.randint(radius)*np.float32([ math.sin( degB * deg_2_rad), math.cos( degB * deg_2_rad) ] ) 
    C = center + radius*np.float32([ math.sin( degC * deg_2_rad), math.cos( degC * deg_2_rad) ] ) 
    
    x = bezier( (width,height), A, B, C )
    
    x = x / (1+np.random.randint(radius)) + 0.5
    
    x = np.clip(x, 0, 1)
    return x

core/imagelib/sd/draw.py
oo92/TorchFaceLab
import localization
import numpy as np
from PIL import Image, ImageDraw, ImageFont

pil_fonts = {}
def _get_pil_font (font, size):
    global pil_fonts
    try:
        font_str_id = '%s_%d' % (font, size)
        if font_str_id not in pil_fonts.keys():
            pil_fonts[font_str_id] = ImageFont.truetype(font + ".ttf", size=size, encoding="unic")
        pil_font = pil_fonts[font_str_id]
        return pil_font
    except:
        return ImageFont.load_default()

def get_text_image( shape, text, color=(1,1,1), border=0.2, font=None):
    h,w,c = shape
    try:
        pil_font = _get_pil_font( localization.get_default_ttf_font_name() , h-2)

        canvas = Image.new('RGB', (w,h) , (0,0,0) )
        draw = ImageDraw.Draw(canvas)
        offset = ( 0, 0)
        draw.text(offset, text, font=pil_font, fill=tuple((np.array(color)*255).astype(np.int)) )

        result = np.asarray(canvas) / 255

        if c > 3:
            result = np.concatenate ( (result, np.ones ((h,w,c-3)) ), axis=-1 )
        elif c < 3:
            result = result[...,0:c]
        return result
    except:
        return np.zeros ( (h,w,c) )

def draw_text( image, rect, text, color=(1,1,1), border=0.2, font=None):
    h,w,c = image.shape

    l,t,r,b = rect
    l = np.clip (l, 0, w-1)
    r = np.clip (r, 0, w-1)
    t = np.clip (t, 0, h-1)
    b = np.clip (b, 0, h-1)

    image[t:b, l:r] += get_text_image (  (b-t,r-l,c) , text, color, border, font )


def draw_text_lines (image, rect, text_lines, color=(1,1,1), border=0.2, font=None):
    text_lines_len = len(text_lines)
    if text_lines_len == 0:
        return

    l,t,r,b = rect
    h = b-t
    h_per_line = h // text_lines_len

    for i in range(0, text_lines_len):
        draw_text (image, (l, i*h_per_line, r, (i+1)*h_per_line), text_lines[i], color, border, font)

def get_draw_text_lines ( image, rect, text_lines, color=(1,1,1), border=0.2, font=None):
    image = np.zeros ( image.shape, dtype=np.float )
    draw_text_lines ( image, rect, text_lines, color, border, font)
    return image

core/imagelib/text.py
oo92/TorchFaceLab
from .interact import interact

core/interact/__init__.py
oo92/TorchFaceLab
import multiprocessing
import os
import sys
import threading
import time
import types

import colorama
import cv2
import numpy as np
from tqdm import tqdm

from core import stdex

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

yn_str = {True:'y',False:'n'}

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
        self.focus_wnd_name = None
        self.error_log_line_prefix = '/!\\ '

        self.process_messages_callbacks = {}

    def is_support_windows(self):
        return False

    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_destroy_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'):
        if self.pg_bar is not None:
            print ("\n")
        print (msg, end=end)

    def log_err(self, msg, end='\n'):
        if self.pg_bar is not None:
            print ("\n")
        print (f'{self.error_log_line_prefix}{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
            self.focus_wnd_name = wnd_name
        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 = {}
            self.focus_wnd_name = None

    def destroy_window(self, wnd_name):
        if wnd_name in self.named_windows:
            self.on_destroy_window(wnd_name)
            self.named_windows.pop(wnd_name)

            if wnd_name == self.focus_wnd_name:
                self.focus_wnd_name = list(self.named_windows.keys())[-1] if len( self.named_windows ) != 0 else None

            if wnd_name in self.capture_mouse_windows:
                self.capture_mouse_windows.pop(wnd_name)

            if wnd_name in self.capture_keys_windows:
                self.capture_keys_windows.pop(wnd_name)

            if wnd_name in self.mouse_events:
                self.mouse_events.pop(wnd_name)

            if wnd_name in self.key_events:
                self.key_events.pop(wnd_name)

    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, initial=0):
        if self.pg_bar is None:
            self.pg_bar = tqdm( total=total, desc=desc, leave=leave, ascii=True, initial=initial )
        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=None, leave=True, initial=0):
        self.pg_bar = tqdm( data, desc=desc, leave=leave, ascii=True, initial=initial )
        for x in self.pg_bar:
            yield x
        self.pg_bar.close()
        self.pg_bar = None

    def add_process_messages_callback(self, func ):
        tid = threading.get_ident()
        callbacks = self.process_messages_callbacks.get(tid, None)
        if callbacks is None:
            callbacks = []
            self.process_messages_callbacks[tid] = callbacks

        callbacks.append ( func )

    def process_messages(self, sleep_time=0):
        callbacks = self.process_messages_callbacks.get(threading.get_ident(), None)
        if callbacks is not None:
            for func in callbacks:
                func()

        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, ord_key, ctrl_pressed, alt_pressed, shift_pressed):
        if wnd_name not in self.key_events:
            self.key_events[wnd_name] = []
        self.key_events[wnd_name] += [ (ord_key, chr(ord_key) if ord_key <= 255 else chr(0), ctrl_pressed, alt_pressed, shift_pressed) ]

    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(self, s):
        return input(s)

    def input_number(self, s, default_value, valid_list=None, show_default_value=True, add_info=None, help_message=None):
        if show_default_value and default_value is not None:
            s = f"[{default_value}] {s}"

        if add_info is not None or \
           help_message is not None:
            s += " ("

        if add_info is not None:
            s += f" {add_info}"
        if help_message is not None:
            s += " ?:help"

        if add_info is not None or \
           help_message is not None:
            s += " )"

        s += " : "

        while True:
            try:
                inp = input(s)
                if len(inp) == 0:
                    result = default_value
                    break

                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):
                    result = default_value
                    break
                result = i
                break
            except:
                result = default_value
                break

        print(result)
        return result

    def input_int(self, s, default_value, valid_range=None, valid_list=None, add_info=None, show_default_value=True, help_message=None):
        if show_default_value:
            if len(s) != 0:
                s = f"[{default_value}] {s}"
            else:
                s = f"[{default_value}]"

        if add_info is not None or \
           valid_range is not None or \
           help_message is not None:
            s += " ("

        if valid_range is not None:
            s += f" {valid_range[0]}-{valid_range[1]}"

        if add_info is not None:
            s += f" {add_info}"

        if help_message is not None:
            s += " ?:help"

        if add_info is not None or \
           valid_range is not None or \
           help_message is not None:
            s += " )"

        s += " : "

        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_range is not None:
                    i = int(np.clip(i, valid_range[0], valid_range[1]))

                if (valid_list is not None) and (i not in valid_list):
                    i = default_value

                result = i
                break
            except:
                result = default_value
                break
        print (result)
        return result

    def input_bool(self, s, default_value, help_message=None):
        s = f"[{yn_str[default_value]}] {s} ( y/n"

        if help_message is not None:
            s += " ?:help"
        s += " ) : "

        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}.get(inp.lower(), default_value) )
            except:
                print ( "y" if default_value else "n" )
                return default_value

    def input_str(self, s, default_value=None, valid_list=None, show_default_value=True, help_message=None):
        if show_default_value and default_value is not None:
            s = f"[{default_value}] {s}"

        if valid_list is not None or \
           help_message is not None:
            s += " ("

        if valid_list is not None:
            s += " " + "/".join(valid_list)

        if help_message is not None:
            s += " ?:help"

        if valid_list is not None or \
           help_message is not None:
            s += " )"

        s += " : "


        while True:
            try:
                inp = input(s)

                if len(inp) == 0:
                    if default_value is None:
                        print("")
                        return None
                    result = default_value
                    break

                if help_message is not None and inp == '?':
                    print(help_message)
                    continue

                if valid_list is not None:
                    if inp.lower() in valid_list:
                        result = inp.lower()
                        break
                    if inp in valid_list:
                        result = inp
                        break
                    continue

                result = inp
                break
            except:
                result = default_value
                break

        print(result)
        return result

    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.daemon = True
        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()
        p.join()

        old_stdin = sys.stdin
        sys.stdin = os.fdopen( os.dup(sys.stdin.fileno()) )
        old_stdin.close()
        return inp

    def input_process_skip_pending(self, stdin_fd):
        sys.stdin = os.fdopen(stdin_fd)
        while True:
            try:
                if sys.stdin.isatty():
                    sys.stdin.read()
            except:
                pass

    def input_skip_pending(self):
        if is_colab:
            # currently it does not work on Colab
            return
        """
        skips unnecessary inputs between the dialogs
        """
        p = multiprocessing.Process(target=self.input_process_skip_pending, args=( sys.stdin.fileno(), ))
        p.daemon = True
        p.start()
        time.sleep(0.5)
        p.terminate()
        p.join()
        sys.stdin = os.fdopen( sys.stdin.fileno() )


class InteractDesktop(InteractBase):
    def __init__(self):
        colorama.init()
        super().__init__()

    def color_red(self):
        pass


    def is_support_windows(self):
        return True

    def on_destroy_all_windows(self):
        #cv2.destroyAllWindows()
        pass

    def on_create_window (self, wnd_name):
        #cv2.namedWindow(wnd_name)
        pass

    def on_destroy_window (self, wnd_name):
        #cv2.destroyWindow(wnd_name)
        pass

    def on_show_image (self, wnd_name, img):
        #cv2.imshow (wnd_name, img)
        pass

    def on_capture_mouse (self, wnd_name):
        self.last_xy = (0,0)

        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
                x,y = self.last_xy #fix opencv bug when window size more than screen size
            else: ev = 0

            self.last_xy = (x,y)
            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) )
            ord_key = cv2.waitKeyEx(wait_key_time)
            
            shift_pressed = False
            if ord_key != -1:
                chr_key = chr(ord_key) if ord_key <= 255 else chr(0)

                if chr_key >= 'A' and chr_key <= 'Z':
                    shift_pressed = True
                    ord_key += 32
                elif chr_key == '?':
                    shift_pressed = True
                    ord_key = ord('/')
                elif chr_key == '<':
                    shift_pressed = True
                    ord_key = ord(',')
                elif chr_key == '>':
                    shift_pressed = True
                    ord_key = ord('.')
        else:
            if sleep_time != 0:
                time.sleep(sleep_time)

        if has_capture_keys and ord_key != -1:
            self.add_key_event ( self.focus_wnd_name, ord_key, False, False, shift_pressed)

    def on_wait_any_key(self):
        cv2.waitKey(0)

class InteractColab(InteractBase):

    def is_support_windows(self):
        return False

    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_destroy_window (self, wnd_name):
        pass

    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()

core/interact/interact.py
oo92/TorchFaceLab
import multiprocessing
from core.interact import interact as io

class MPClassFuncOnDemand():
    def __init__(self, class_handle, class_func_name, **class_kwargs):
        self.class_handle = class_handle
        self.class_func_name = class_func_name
        self.class_kwargs = class_kwargs
   
        self.class_func = None
        
        self.s2c = multiprocessing.Queue()
        self.c2s = multiprocessing.Queue()
        self.lock = multiprocessing.Lock()
        
        io.add_process_messages_callback(self.io_callback)

    def io_callback(self):        
        while not self.c2s.empty():
            func_args, func_kwargs = self.c2s.get()
            if self.class_func is None:
                self.class_func = getattr( self.class_handle(**self.class_kwargs), self.class_func_name)
            self.s2c.put ( self.class_func (*func_args, **func_kwargs) )

    def __call__(self, *args, **kwargs):
        with self.lock:
            self.c2s.put ( (args, kwargs) )
            return self.s2c.get()

    def __getstate__(self):
        return {'s2c':self.s2c, 'c2s':self.c2s, 'lock':self.lock}


core/joblib/MPClassFuncOnDemand.py
oo92/TorchFaceLab
import multiprocessing
from core.interact import interact as io

class MPFunc():
    def __init__(self, func):
        self.func = func
        
        self.s2c = multiprocessing.Queue()
        self.c2s = multiprocessing.Queue()
        self.lock = multiprocessing.Lock()
        
        io.add_process_messages_callback(self.io_callback)

    def io_callback(self):        
        while not self.c2s.empty():
            func_args, func_kwargs = self.c2s.get()
            self.s2c.put ( self.func (*func_args, **func_kwargs) )

    def __call__(self, *args, **kwargs):
        with self.lock:
            self.c2s.put ( (args, kwargs) )
            return self.s2c.get()

    def __getstate__(self):
        return {'s2c':self.s2c, 'c2s':self.c2s, 'lock':self.lock}

core/joblib/MPFunc.py
oo92/TorchFaceLab
import multiprocessing
import queue as Queue
import threading
import time


class SubprocessGenerator(object):
    
    @staticmethod
    def launch_thread(generator): 
        generator._start()
        
    @staticmethod
    def start_in_parallel( generator_list ):
        """
        Start list of generators in parallel
        """
        for generator in generator_list:
            thread = threading.Thread(target=SubprocessGenerator.launch_thread, args=(generator,) )
            thread.daemon = True
            thread.start()

        while not all ([generator._is_started() for generator in generator_list]):
            time.sleep(0.005)
    
    def __init__(self, generator_func, user_param=None, prefetch=2, start_now=True):
        super().__init__()
        self.prefetch = prefetch
        self.generator_func = generator_func
        self.user_param = user_param
        self.sc_queue = multiprocessing.Queue()
        self.cs_queue = multiprocessing.Queue()
        self.p = None
        if start_now:
            self._start()

    def _start(self):
        if self.p == None:
            user_param = self.user_param
            self.user_param = None
            p = multiprocessing.Process(target=self.process_func, args=(user_param,) )
            p.daemon = True
            p.start()
            self.p = p
            
    def _is_started(self):
        return self.p is not None
        
    def process_func(self, user_param):
        self.generator_func = self.generator_func(user_param)
        while True:
            while self.prefetch > -1:
                try:
                    gen_data = next (self.generator_func)
                except StopIteration:
                    self.cs_queue.put (None)
                    return
                self.cs_queue.put (gen_data)
                self.prefetch -= 1
            self.sc_queue.get()
            self.prefetch += 1

    def __iter__(self):
        return self

    def __getstate__(self):
        self_dict = self.__dict__.copy()
        del self_dict['p']
        return self_dict

    def __next__(self):
        self._start()
        gen_data = self.cs_queue.get()
        if gen_data is None:
            self.p.terminate()
            self.p.join()
            raise StopIteration()
        self.sc_queue.put (1)
        return gen_data

core/joblib/SubprocessGenerator.py
oo92/TorchFaceLab
import traceback
import multiprocessing
import time
import sys
from core.interact import interact as io


class Subprocessor(object):

    class SilenceException(Exception):
        pass

    class Cli(object):
        def __init__ ( self, client_dict ):
            s2c = multiprocessing.Queue()
            c2s = multiprocessing.Queue()
            self.p = multiprocessing.Process(target=self._subprocess_run, args=(client_dict,s2c,c2s) )
            self.s2c = s2c
            self.c2s = c2s
            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, s2c, c2s):
            self.c2s = c2s
            data = None
            is_error = False
            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 Subprocessor.SilenceException as e:
                c2s.put ( {'op': 'error', 'data' : data} )
            except Exception as e:
                err_msg = traceback.format_exc()
                c2s.put ( {'op': 'error', 'data' : data, 'err_msg' : err_msg} )

            c2s.close()
            s2c.close()
            self.c2s = None

        # disable pickling
        def __getstate__(self):
            return dict()
        def __setstate__(self, d):
            self.__dict__.update(d)

    #overridable
    def __init__(self, name, SubprocessorCli_class, no_response_time_sec = 0, io_loop_sleep_time=0.005, initialize_subprocesses_in_serial=False):
        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
        self.io_loop_sleep_time = io_loop_sleep_time
        self.initialize_subprocesses_in_serial = initialize_subprocesses_in_serial

    #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()
        return None

    #overridable
    def on_tick(self):
        #tick in main loop
        #return True if system can be finalized when no data in get_data, orelse False
        return True

    #overridable
    def on_check_run(self):
        return True

    def run(self):
        if not self.on_check_run():
            return self.get_result()

        self.clis = []

        def cli_init_dispatcher(cli):
            while not cli.c2s.empty():
                obj = cli.c2s.get()
                op = obj.get('op','')
                if op == 'init_ok':
                    cli.state = 0
                elif op == 'log_info':
                    io.log_info(obj['msg'])
                elif op == 'log_err':
                    io.log_err(obj['msg'])
                elif op == 'error':
                    err_msg = obj.get('err_msg', None)
                    if err_msg is not None:
                        io.log_info(f'Error while subprocess initialization: {err_msg}')
                    cli.kill()
                    self.clis.remove(cli)
                    break

        #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 = 0
                cli.sent_data = None
                cli.name = name
                cli.host_dict = host_dict

                self.clis.append (cli)

                if self.initialize_subprocesses_in_serial:
                    while True:
                        cli_init_dispatcher(cli)
                        if cli.state == 0:
                            break
                        io.process_messages(0.005)
            except:
                raise Exception (f"Unable to start subprocess {name}. Error: {traceback.format_exc()}")

        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[:]:
                cli_init_dispatcher(cli)
            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
                        err_msg = obj.get('err_msg', None)
                        if err_msg is not None:
                            io.log_info(f'Error while processing data: {err_msg}')
                            
                        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 == 1:
                    if cli.sent_time != 0 and 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)

            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

            if self.io_loop_sleep_time != 0:
                io.process_messages(self.io_loop_sleep_time)

            if self.on_tick() and all ([cli.state == 0 for cli in self.clis]):
                #all subprocesses free and no more data available to process, ending loop
                break



        #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 (time.time() - cli.sent_time) > 30:
                    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()

core/joblib/SubprocessorBase.py
oo92/TorchFaceLab
class ThisThreadGenerator(object):
    def __init__(self, generator_func, user_param=None):
        super().__init__()
        self.generator_func = generator_func
        self.user_param = user_param
        self.initialized = False

    def __iter__(self):
        return self

    def __next__(self):
        if not self.initialized:
            self.initialized = True
            self.generator_func = self.generator_func(self.user_param)

        return next(self.generator_func)
core/joblib/ThisThreadGenerator.py
oo92/TorchFaceLab
from .SubprocessorBase import Subprocessor
from .ThisThreadGenerator import ThisThreadGenerator
from .SubprocessGenerator import SubprocessGenerator
from .MPFunc import MPFunc
from .MPClassFuncOnDemand import MPClassFuncOnDemand
core/joblib/__init__.py
oo92/TorchFaceLab
import torch.nn as nn

class ArchiBase(nn.Module):
    
    def __init__(self, *args, name=None, **kwargs):
        super(ArchiBase, self).__init__()
        self.name = name
        
    # Overridable 
    def flow(self, *args, **kwargs):
        raise Exception("This architecture does not support flow. Use model classes directly.")
    
    # Overridable
    def get_weights(self):
        pass

    
nn.ArchiBase = ArchiBase
core/leras/archis/ArchiBase.py
oo92/TorchFaceLab
from core.leras import nn
tf = nn.tf

class DeepFakeArchi(nn.ArchiBase):

    def __init__(self, resolution, use_fp16=False, mod=None, opts=None):
        super().__init__()

        if opts is None:
            opts = ''

        conv_dtype = tf.float16 if use_fp16 else tf.float32
        
        if 'c' in opts:
            def act(x, alpha=0.1):
                return x*tf.cos(x)
        else:
            def act(x, alpha=0.1):
                return tf.nn.leaky_relu(x, alpha)
                
        if mod is None:
            class Downscale(nn.ModelBase):
                def __init__(self, in_ch, out_ch, kernel_size=5, *kwargs ):
                    self.in_ch = in_ch
                    self.out_ch = out_ch
                    self.kernel_size = kernel_size
                    super().__init__(*kwargs)

                def on_build(self, *args, **kwargs ):
                    self.conv1 = nn.Conv2D( self.in_ch, self.out_ch, kernel_size=self.kernel_size, strides=2, padding='SAME', dtype=conv_dtype)

                def forward(self, x):
                    x = self.conv1(x)
                    x = act(x, 0.1)
                    return x

                def get_out_ch(self):
                    return self.out_ch

            class DownscaleBlock(nn.ModelBase):
                def on_build(self, in_ch, ch, n_downscales, kernel_size):
                    self.downs = []

                    last_ch = in_ch
                    for i in range(n_downscales):
                        cur_ch = ch*( min(2**i, 8)  )
                        self.downs.append ( Downscale(last_ch, cur_ch, kernel_size=kernel_size))
                        last_ch = self.downs[-1].get_out_ch()

                def forward(self, inp):
                    x = inp
                    for down in self.downs:
                        x = down(x)
                    return x

            class Upscale(nn.ModelBase):
                def on_build(self, in_ch, out_ch, kernel_size=3):
                    self.conv1 = nn.Conv2D( in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype)

                def forward(self, x):
                    x = self.conv1(x)
                    x = act(x, 0.1)
                    x = nn.depth_to_space(x, 2)
                    return x

            class ResidualBlock(nn.ModelBase):
                def on_build(self, ch, kernel_size=3):
                    self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype)
                    self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype)

                def forward(self, inp):
                    x = self.conv1(inp)
                    x = act(x, 0.2)
                    x = self.conv2(x)
                    x = act(inp + x, 0.2)
                    return x

            class Encoder(nn.ModelBase):
                def __init__(self, in_ch, e_ch, **kwargs ):
                    self.in_ch = in_ch
                    self.e_ch = e_ch
                    super().__init__(**kwargs)

                def on_build(self):
                    if 't' in opts:
                        self.down1 = Downscale(self.in_ch, self.e_ch, kernel_size=5)
                        self.res1 = ResidualBlock(self.e_ch)
                        self.down2 = Downscale(self.e_ch, self.e_ch*2, kernel_size=5)
                        self.down3 = Downscale(self.e_ch*2, self.e_ch*4, kernel_size=5)
                        self.down4 = Downscale(self.e_ch*4, self.e_ch*8, kernel_size=5)
                        self.down5 = Downscale(self.e_ch*8, self.e_ch*8, kernel_size=5)
                        self.res5 = ResidualBlock(self.e_ch*8)
                    else:
                        self.down1 = DownscaleBlock(self.in_ch, self.e_ch, n_downscales=4 if 't' not in opts else 5, kernel_size=5)

                def forward(self, x):
                    if use_fp16:
                        x = tf.cast(x, tf.float16)

                    if 't' in opts:
                        x = self.down1(x)
                        x = self.res1(x)
                        x = self.down2(x)
                        x = self.down3(x)
                        x = self.down4(x)
                        x = self.down5(x)
                        x = self.res5(x)
                    else:
                        x = self.down1(x)
                    x = nn.flatten(x)
                    if 'u' in opts:
                        x = nn.pixel_norm(x, axes=-1)

                    if use_fp16:
                        x = tf.cast(x, tf.float32)
                    return x

                def get_out_res(self, res):
                    return res // ( (2**4) if 't' not in opts else (2**5) )

                def get_out_ch(self):
                    return self.e_ch * 8

            lowest_dense_res = resolution // (32 if 'd' in opts else 16)

            class Inter(nn.ModelBase):
                def __init__(self, in_ch, ae_ch, ae_out_ch, **kwargs):
                    self.in_ch, self.ae_ch, self.ae_out_ch = in_ch, ae_ch, ae_out_ch
                    super().__init__(**kwargs)

                def on_build(self):
                    in_ch, ae_ch, ae_out_ch = self.in_ch, self.ae_ch, self.ae_out_ch

                    self.dense1 = nn.Dense( in_ch, ae_ch )
                    self.dense2 = nn.Dense( ae_ch, lowest_dense_res * lowest_dense_res * ae_out_ch )
                    if 't' not in opts:
                        self.upscale1 = Upscale(ae_out_ch, ae_out_ch)

                def forward(self, inp):
                    x = inp
                    x = self.dense1(x)
                    x = self.dense2(x)
                    x = nn.reshape_4D (x, lowest_dense_res, lowest_dense_res, self.ae_out_ch)

                    if use_fp16:
                        x = tf.cast(x, tf.float16)

                    if 't' not in opts:
                        x = self.upscale1(x)

                    return x

                def get_out_res(self):
                    return lowest_dense_res * 2 if 't' not in opts else lowest_dense_res

                def get_out_ch(self):
                    return self.ae_out_ch

            class Decoder(nn.ModelBase):
                def on_build(self, in_ch, d_ch, d_mask_ch):
                    if 't' not in opts:
                        self.upscale0 = Upscale(in_ch, d_ch*8, kernel_size=3)
                        self.upscale1 = Upscale(d_ch*8, d_ch*4, kernel_size=3)
                        self.upscale2 = Upscale(d_ch*4, d_ch*2, kernel_size=3)
                        self.res0 = ResidualBlock(d_ch*8, kernel_size=3)
                        self.res1 = ResidualBlock(d_ch*4, kernel_size=3)
                        self.res2 = ResidualBlock(d_ch*2, kernel_size=3)

                        self.upscalem0 = Upscale(in_ch, d_mask_ch*8, kernel_size=3)
                        self.upscalem1 = Upscale(d_mask_ch*8, d_mask_ch*4, kernel_size=3)
                        self.upscalem2 = Upscale(d_mask_ch*4, d_mask_ch*2, kernel_size=3)

                        self.out_conv  = nn.Conv2D( d_ch*2, 3, kernel_size=1, padding='SAME', dtype=conv_dtype)

                        if 'd' in opts:
                            self.out_conv1 = nn.Conv2D( d_ch*2, 3, kernel_size=3, padding='SAME', dtype=conv_dtype)
                            self.out_conv2 = nn.Conv2D( d_ch*2, 3, kernel_size=3, padding='SAME', dtype=conv_dtype)
                            self.out_conv3 = nn.Conv2D( d_ch*2, 3, kernel_size=3, padding='SAME', dtype=conv_dtype)
                            self.upscalem3 = Upscale(d_mask_ch*2, d_mask_ch*1, kernel_size=3)
                            self.out_convm = nn.Conv2D( d_mask_ch*1, 1, kernel_size=1, padding='SAME', dtype=conv_dtype)
                        else:
                            self.out_convm = nn.Conv2D( d_mask_ch*2, 1, kernel_size=1, padding='SAME', dtype=conv_dtype)
                    else:
                        self.upscale0 = Upscale(in_ch, d_ch*8, kernel_size=3)
                        self.upscale1 = Upscale(d_ch*8, d_ch*8, kernel_size=3)
                        self.upscale2 = Upscale(d_ch*8, d_ch*4, kernel_size=3)
                        self.upscale3 = Upscale(d_ch*4, d_ch*2, kernel_size=3)
                        self.res0 = ResidualBlock(d_ch*8, kernel_size=3)
                        self.res1 = ResidualBlock(d_ch*8, kernel_size=3)
                        self.res2 = ResidualBlock(d_ch*4, kernel_size=3)
                        self.res3 = ResidualBlock(d_ch*2, kernel_size=3)

                        self.upscalem0 = Upscale(in_ch, d_mask_ch*8, kernel_size=3)
                        self.upscalem1 = Upscale(d_mask_ch*8, d_mask_ch*8, kernel_size=3)
                        self.upscalem2 = Upscale(d_mask_ch*8, d_mask_ch*4, kernel_size=3)
                        self.upscalem3 = Upscale(d_mask_ch*4, d_mask_ch*2, kernel_size=3)
                        self.out_conv  = nn.Conv2D( d_ch*2, 3, kernel_size=1, padding='SAME', dtype=conv_dtype)

                        if 'd' in opts:
                            self.out_conv1 = nn.Conv2D( d_ch*2, 3, kernel_size=3, padding='SAME', dtype=conv_dtype)
                            self.out_conv2 = nn.Conv2D( d_ch*2, 3, kernel_size=3, padding='SAME', dtype=conv_dtype)
                            self.out_conv3 = nn.Conv2D( d_ch*2, 3, kernel_size=3, padding='SAME', dtype=conv_dtype)
                            self.upscalem4 = Upscale(d_mask_ch*2, d_mask_ch*1, kernel_size=3)
                            self.out_convm = nn.Conv2D( d_mask_ch*1, 1, kernel_size=1, padding='SAME', dtype=conv_dtype)
                        else:
                            self.out_convm = nn.Conv2D( d_mask_ch*2, 1, kernel_size=1, padding='SAME', dtype=conv_dtype)
                    
                def forward(self, z):
                    x = self.upscale0(z)
                    x = self.res0(x)
                    x = self.upscale1(x)
                    x = self.res1(x)
                    x = self.upscale2(x)
                    x = self.res2(x)

                    if 't' in opts:
                        x = self.upscale3(x)
                        x = self.res3(x)

                    if 'd' in opts:
                        x = tf.nn.sigmoid( nn.depth_to_space(tf.concat( (self.out_conv(x), self.out_conv1(x), self.out_conv2(x),self.out_conv3(x)), nn.conv2d_ch_axis), 2) )
                    else:
                        x = tf.nn.sigmoid(self.out_conv(x))

                    m = self.upscalem0(z)
                    m = self.upscalem1(m)
                    m = self.upscalem2(m)

                    if 't' in opts:
                        m = self.upscalem3(m)
                        if 'd' in opts:
                            m = self.upscalem4(m)
                    else:
                        if 'd' in opts:
                            m = self.upscalem3(m)

                    m = tf.nn.sigmoid(self.out_convm(m))

                    if use_fp16:
                        x = tf.cast(x, tf.float32)
                        m = tf.cast(m, tf.float32)

                    return x, m

        self.Encoder = Encoder
        self.Inter = Inter
        self.Decoder = Decoder

nn.DeepFakeArchi = DeepFakeArchi
core/leras/archis/DeepFakeArchi.py
oo92/TorchFaceLab
from .ArchiBase import *
from .DeepFakeArchi import *
core/leras/archis/__init__.py
oo92/TorchFaceLab
import sys, ctypes, os, multiprocessing, json, time, torch
from pathlib import Path
from core.interact import interact as io

class Device(object):
    def __init__(self, index, pytorch_dev_type, name, total_mem, free_mem):
        self.index = index
        self.pytorch_dev_type = pytorch_dev_type
        self.name = name
        
        self.total_mem = total_mem
        self.total_mem_gb = total_mem / 1024**3
        self.free_mem = free_mem
        self.free_mem_gb = free_mem / 1024**3

    def __str__(self):
        return f"[{self.index}]:[{self.name}][{self.free_mem_gb:.3}/{self.total_mem_gb :.3}]"

class Devices(object):
    all_devices = None

    def __init__(self, devices):
        self.devices = devices

    def __len__(self):
        return len(self.devices)

    def __getitem__(self, key):
        result = self.devices[key]
        if isinstance(key, slice):
            return Devices(result)
        return result

    def __iter__(self):
        for device in self.devices:
            yield device

    def get_best_device(self):
        result = None
        idx_mem = 0
        for device in self.devices:
            mem = device.total_mem
            if mem > idx_mem:
                result = device
                idx_mem = mem
        return result

    def get_worst_device(self):
        result = None
        idx_mem = sys.maxsize
        for device in self.devices:
            mem = device.total_mem
            if mem < idx_mem:
                result = device
                idx_mem = mem
        return result

    def get_device_by_index(self, idx):
        for device in self.devices:
            if device.index == idx:
                return device
        return None

    def get_devices_from_index_list(self, idx_list):
        result = []
        for device in self.devices:
            if device.index in idx_list:
                result += [device]
        return Devices(result)

    def get_equal_devices(self, device):
        device_name = device.name
        result = []
        for device in self.devices:
            if device.name == device_name:
                result.append(device)
        return Devices(result)

    def get_devices_at_least_mem(self, totalmemsize_gb):
        result = []
        for device in self.devices:
            if device.total_mem >= totalmemsize_gb * (1024**3):
                result.append(device)
        return Devices(result)

    @staticmethod
    def initialize_main_env():
        if int(os.environ.get("NN_DEVICES_INITIALIZED", 0)) != 0:
            return
            
        visible_devices = []
        
        device_count = torch.cuda.device_count()
        for i in range(device_count):
            device_name = torch.cuda.get_device_name(i)
            device_type = 'CUDA'  # PyTorch primarily supports CUDA for GPUs
            total_mem = torch.cuda.get_device_properties(i).total_memory
            free_mem = total_mem  # PyTorch does not provide a straightforward way to get free memory.
            visible_devices.append((device_type, device_name, total_mem))
        
        os.environ['NN_DEVICES_INITIALIZED'] = '1'
        os.environ['NN_DEVICES_COUNT'] = str(len(visible_devices))
        
        for i, (dev_type, name, total_mem) in enumerate(visible_devices):
            os.environ[f'NN_DEVICE_{i}_PYTORCH_DEV_TYPE'] = dev_type
            os.environ[f'NN_DEVICE_{i}_NAME'] = name
            os.environ[f'NN_DEVICE_{i}_TOTAL_MEM'] = str(total_mem)
            os.environ[f'NN_DEVICE_{i}_FREE_MEM'] = str(total_mem)

    @staticmethod
    def getDevices():
        if Devices.all_devices is None:
            if int(os.environ.get("NN_DEVICES_INITIALIZED", 0)) != 1:
                raise Exception("nn devices are not initialized. Run initialize_main_env() in main process.")
            devices = []
            for i in range(int(os.environ['NN_DEVICES_COUNT'])):
                devices.append(Device(index=i,
                                      pytorch_dev_type=os.environ[f'NN_DEVICE_{i}_PYTORCH_DEV_TYPE'],
                                      name=os.environ[f'NN_DEVICE_{i}_NAME'],
                                      total_mem=int(os.environ[f'NN_DEVICE_{i}_TOTAL_MEM']),
                                      free_mem=int(os.environ[f'NN_DEVICE_{i}_FREE_MEM'])))
            Devices.all_devices = Devices(devices)

        return Devices.all_devices

core/leras/device.py
oo92/TorchFaceLab
import multiprocessing
from core.joblib import Subprocessor
import numpy as np

class CAInitializerSubprocessor(Subprocessor):
    @staticmethod
    def generate(shape, dtype=np.float32, eps_std=0.05):
        """
        Super fast implementation of Convolution Aware Initialization for 4D shapes
        Convolution Aware Initialization https://arxiv.org/abs/1702.06295
        """
        if len(shape) != 4:
            raise ValueError("only shape with rank 4 supported.")

        row, column, stack_size, filters_size = shape

        fan_in = stack_size * (row * column)

        kernel_shape = (row, column)

        kernel_fft_shape = np.fft.rfft2(np.zeros(kernel_shape)).shape

        basis_size = np.prod(kernel_fft_shape)
        if basis_size == 1:
            x = np.random.normal( 0.0, eps_std, (filters_size, stack_size, basis_size) )
        else:
            nbb = stack_size // basis_size + 1
            x = np.random.normal(0.0, 1.0, (filters_size, nbb, basis_size, basis_size))
            x = x + np.transpose(x, (0,1,3,2) ) * (1-np.eye(basis_size))
            u, _, v = np.linalg.svd(x)
            x = np.transpose(u, (0,1,3,2) )
            x = np.reshape(x, (filters_size, -1, basis_size) )
            x = x[:,:stack_size,:]

        x = np.reshape(x, ( (filters_size,stack_size,) + kernel_fft_shape ) )

        x = np.fft.irfft2( x, kernel_shape ) \
            + np.random.normal(0, eps_std, (filters_size,stack_size,)+kernel_shape)

        x = x * np.sqrt( (2/fan_in) / np.var(x) )
        x = np.transpose( x, (2, 3, 1, 0) )
        return x.astype(dtype)

    class Cli(Subprocessor.Cli):
        #override
        def process_data(self, data):
            idx, shape, dtype = data
            weights = CAInitializerSubprocessor.generate (shape, dtype)
            return idx, weights

    #override
    def __init__(self, data_list):
        self.data_list = data_list
        self.data_list_idxs = [*range(len(data_list))]
        self.result = [None]*len(data_list)
        super().__init__('CAInitializerSubprocessor', CAInitializerSubprocessor.Cli)

    #override
    def process_info_generator(self):
        for i in range( min(multiprocessing.cpu_count(), len(self.data_list)) ):
            yield 'CPU%d' % (i), {}, {}

    #override
    def get_data(self, host_dict):
        if len (self.data_list_idxs) > 0:
            idx = self.data_list_idxs.pop(0)
            shape, dtype = self.data_list[idx]
            return idx, shape, dtype
        return None

    #override
    def on_data_return (self, host_dict, data):
        self.data_list_idxs.insert(0, data)

    #override
    def on_result (self, host_dict, data, result):
        idx, weights = result
        self.result[idx] = weights

    #override
    def get_result(self):
        return self.result

core/leras/initializers/CA.py
oo92/TorchFaceLab
import numpy as np
import torch
from torch import nn

from .CA import CAInitializerSubprocessor

class Initializers(nn.Module):
    def __init__(self):
        super(Initializers, self).__init__()

    def forward(self, shape, dtype=None):
        return torch.zeros(shape, dtype=dtype)

    @staticmethod
    def generate_batch(data_list, eps_std=0.05):
        # list of (shape, torch.dtype)
        return CAInitializerSubprocessor(data_list).run()

core/leras/initializers/__init__.py
oo92/TorchFaceLab
import torch
import torch.nn as nn

class AdaIN(nn.Module):
    def __init__(self, in_ch, mlp_ch, kernel_initializer=None):
        super(AdaIN, self).__init__()

        self.in_ch = in_ch
        self.mlp_ch = mlp_ch

        if kernel_initializer is None:
            kernel_initializer = nn.init.kaiming_normal_

        # Define weight and bias using nn.Parameter so they get registered as learnable parameters automatically
        self.weight1 = nn.Parameter(torch.empty(mlp_ch, in_ch))
        self.bias1 = nn.Parameter(torch.zeros(in_ch))

        self.weight2 = nn.Parameter(torch.empty(mlp_ch, in_ch))
        self.bias2 = nn.Parameter(torch.zeros(in_ch))

        # Initialize weights using provided initializer
        kernel_initializer(self.weight1)
        kernel_initializer(self.weight2)

    def forward(self, x, mlp):
        gamma = torch.matmul(mlp, self.weight1) + self.bias1
        beta = torch.matmul(mlp, self.weight2) + self.bias2

        # Reshape gamma and beta for broadcasting
        gamma = gamma.view(-1, self.in_ch, 1, 1)
        beta = beta.view(-1, self.in_ch, 1, 1)

        # Calculate mean and standard deviation
        x_mean = x.mean(dim=(2, 3), keepdim=True)
        x_std = x.std(dim=(2, 3), keepdim=True) + 1e-5

        # AdaIN transformation
        x = (x - x_mean) / x_std
        x *= gamma
        x += beta

        return x

nn.AdaIN = AdaIN
core/leras/layers/AdaIN.py
oo92/TorchFaceLab
import torch.nn as nn

class BatchNorm2D(nn.Module):
    """
    Custom BatchNorm2D layer using PyTorch.
    Currently not for training.
    """
    def __init__(self, dim, eps=1e-05, momentum=0.1, dtype=None):
        super(BatchNorm2D, self).__init__()
        self.dim = dim
        self.eps = eps
        self.momentum = momentum
        if dtype is None:
            dtype = torch.float32  # Use PyTorch's float32 as default
        self.dtype = dtype
        
        # Define weights
        self.weight = nn.Parameter(torch.ones(self.dim, dtype=self.dtype))
        self.bias = nn.Parameter(torch.zeros(self.dim, dtype=self.dtype))
        # These are non-trainable parameters
        self.register_buffer('running_mean', torch.zeros(self.dim, dtype=self.dtype))
        self.register_buffer('running_var', torch.ones(self.dim, dtype=self.dtype))

    def forward(self, x):
        if x.shape[1] == self.dim:  # PyTorch uses NCHW by default
            shape = (1, self.dim, 1, 1)
        else:
            shape = (1, 1, 1, self.dim)  # This might not be needed, but keeping for consistency

        weight = self.weight.view(shape)
        bias = self.bias.view(shape)
        running_mean = self.running_mean.view(shape)
        running_var = self.running_var.view(shape)

        x = (x - running_mean) / torch.sqrt(running_var + self.eps)
        x *= weight
        x += bias
        return x

nn.BatchNorm2D = BatchNorm2D
core/leras/layers/BatchNorm2D.py
oo92/TorchFaceLab
import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F

class BlurPool(nn.Module):
    def __init__(self, filt_size=3, stride=2):
        super(BlurPool, self).__init__()

        self.stride = stride
        pad = [int(1. * (filt_size - 1) / 2), int(np.ceil(1. * (filt_size - 1) / 2))]
        self.padding = pad + pad  # Double the pad list for 2D padding (height and width)

        if filt_size == 1:
            a = np.array([1.,])
        elif filt_size == 2:
            a = np.array([1., 1.])
        elif filt_size == 3:
            a = np.array([1., 2., 1.])
        elif filt_size == 4:
            a = np.array([1., 3., 3., 1.])
        elif filt_size == 5:
            a = np.array([1., 4., 6., 4., 1.])
        elif filt_size == 6:
            a = np.array([1., 5., 10., 10., 5., 1.])
        elif filt_size == 7:
            a = np.array([1., 6., 15., 20., 15., 6., 1.])

        a = a[:, None] * a[None, :]
        a = a / np.sum(a)
        a = a[:, :, None, None].astype(np.float32)  # Convert to float32 for torch tensor compatibility

        self.register_buffer('k', torch.tensor(a))

    def forward(self, x):
        k = self.k.expand(1, x.size(1), self.k.size(2), self.k.size(3))  # Copy the filter for each channel in the input
        x = F.pad(x, self.padding)
        x = F.conv2d(x, k, stride=self.stride, groups=x.size(1))  # Use depthwise convolution by setting groups to number of input channels

        return x

nn.BluePool = BlurPool
core/leras/layers/BlurPool.py
oo92/TorchFaceLab
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.init as init

class Conv2D(nn.Module):
    def __init__(self, in_ch, out_ch, kernel_size, strides=1, padding='SAME', dilations=1, use_bias=True, use_wscale=False, kernel_initializer=None, bias_initializer=None, trainable=True, dtype=None, **kwargs):
        super(Conv2D, self).__init__()
        
        # Determine padding type
        if isinstance(padding, str):
            if padding == "SAME":
                self.padding = ( (kernel_size - 1) * dilations + 1 ) // 2
            elif padding == "VALID":
                self.padding = 0
            else:
                raise ValueError ("Wrong padding type. Should be VALID, SAME or INT or 4x INTs")
        else:
            self.padding = int(padding)

        # Use weight scaling or not
        self.use_wscale = use_wscale
        if self.use_wscale:
            gain = 1.0 if kernel_size == 1 else np.sqrt(2)
            fan_in = kernel_size * kernel_size * in_ch
            he_std = gain / np.sqrt(fan_in)
            self.wscale = torch.tensor(he_std, dtype=dtype or torch.float32)

        # Create convolution layer
        self.conv = nn.Conv2d(in_ch, out_ch, kernel_size, stride=strides, padding=self.padding, dilation=dilations, bias=use_bias)

        # Weight initialization
        if kernel_initializer is None:
            if self.use_wscale:
                init.normal_(self.conv.weight, mean=0, std=1.0)
            else:
                init.xavier_uniform_(self.conv.weight)
        else:
            # Apply custom kernel initializer if provided
            kernel_initializer(self.conv.weight)

        # Bias initialization
        if use_bias:
            if bias_initializer is None:
                init.zeros_(self.conv.bias)
            else:
                # Apply custom bias initializer if provided
                bias_initializer(self.conv.bias)

    def forward(self, x):
        if self.use_wscale:
            return self.conv(x * self.wscale)
        else:
            return self.conv(x)

    def __str__(self):
        return f"{self.__class__.__name__} : in_ch:{self.conv.in_channels} out_ch:{self.conv.out_channels}"

nn.Conv2D = Conv2D
core/leras/layers/Conv2D.py
oo92/TorchFaceLab
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.init as init

class Conv2DTranspose(nn.Module):
    def __init__(self, in_ch, out_ch, kernel_size, strides=2, padding='SAME', use_bias=True, use_wscale=False, kernel_initializer=None, bias_initializer=None, trainable=True, dtype=None, **kwargs):
        super(Conv2DTranspose, self).__init__()
        
        if not isinstance(strides, int):
            raise ValueError("strides must be an int type")
        kernel_size = int(kernel_size)

        # Compute padding
        if padding == 'SAME':
            self.padding = kernel_size // 2
        elif padding == 'VALID':
            self.padding = 0
        else:
            raise ValueError("Unsupported padding type")

        self.use_wscale = use_wscale
        if self.use_wscale:
            gain = 1.0 if kernel_size == 1 else np.sqrt(2)
            fan_in = kernel_size * kernel_size * in_ch
            he_std = gain / np.sqrt(fan_in)
            self.wscale = torch.tensor(he_std, dtype=dtype or torch.float32)
        
        # Create transposed convolution layer
        self.conv_transpose = nn.ConvTranspose2d(in_ch, out_ch, kernel_size, stride=strides, padding=self.padding, bias=use_bias)

        # Weight initialization
        if kernel_initializer is None:
            if self.use_wscale:
                init.normal_(self.conv_transpose.weight, mean=0, std=1.0)
            else:
                init.xavier_uniform_(self.conv_transpose.weight)
        else:
            kernel_initializer(self.conv_transpose.weight)

        # Bias initialization
        if use_bias:
            if bias_initializer is None:
                init.zeros_(self.conv_transpose.bias)
            else:
                bias_initializer(self.conv_transpose.bias)
                
    def forward(self, x):
        if self.use_wscale:
            return self.conv_transpose(x * self.wscale)
        else:
            return self.conv_transpose(x)

    def __str__(self):
        return f"{self.__class__.__name__} : in_ch:{self.conv_transpose.in_channels} out_ch:{self.conv_transpose.out_channels}"

nn.Conv2DTranspose = Conv2DTranspose
core/leras/layers/Conv2DTranspose.py
oo92/TorchFaceLab
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.init as init

class Dense(nn.Module):
    def __init__(self, in_ch, out_ch, use_bias=True, use_wscale=False, maxout_ch=0, kernel_initializer=None, bias_initializer=None, trainable=True, dtype=None, **kwargs):
        super(Dense, self).__init__()
        
        self.in_ch = in_ch
        self.out_ch = out_ch
        self.use_bias = use_bias
        self.use_wscale = use_wscale
        self.maxout_ch = maxout_ch
        self.trainable = trainable

        # Compute the weight shape considering maxout
        if self.maxout_ch > 1:
            weight_shape = (in_ch, out_ch * maxout_ch)
        else:
            weight_shape = (in_ch, out_ch)

        # Initialize weights
        self.weight = nn.Parameter(torch.empty(*weight_shape, dtype=dtype or torch.float32), requires_grad=trainable)
        
        # Weight scaling
        if self.use_wscale:
            gain = 1.0
            fan_in = np.prod(weight_shape[:-1])
            he_std = gain / np.sqrt(fan_in)
            self.wscale = torch.tensor(he_std, dtype=dtype or torch.float32)
            if kernel_initializer is None:
                init.normal_(self.weight, mean=0, std=1.0)
        else:
            if kernel_initializer is None:
                init.xavier_uniform_(self.weight)
            else:
                kernel_initializer(self.weight)
        
        # Initialize bias
        if self.use_bias:
            self.bias = nn.Parameter(torch.empty(out_ch, dtype=dtype or torch.float32), requires_grad=trainable)
            if bias_initializer is None:
                init.zeros_(self.bias)
            else:
                bias_initializer(self.bias)
                
    def forward(self, x):
        if self.use_wscale:
            x = torch.matmul(x, self.weight * self.wscale)
        else:
            x = torch.matmul(x, self.weight)

        if self.maxout_ch > 1:
            x = x.view(-1, self.out_ch, self.maxout_ch)
            x, _ = torch.max(x, dim=-1)

        if self.use_bias:
            x += self.bias.unsqueeze(0)

        return x

nn.Dense = Dense

core/leras/layers/Dense.py
oo92/TorchFaceLab
import torch
import torch.nn as nn

class DenseNorm(nn.Module):
    def __init__(self, dense=False, eps=1e-06, dtype=None):
        super(DenseNorm, self).__init__()
        self.dense = dense
        
        if dtype is None:
            dtype = torch.float32
        self.eps = torch.tensor(eps, dtype=dtype)

    def forward(self, x):
        return x * (torch.mean(x.pow(2), dim=-1, keepdim=True) + self.eps).rsqrt()

nn.DenseNorm = DenseNorm
core/leras/layers/DenseNorm.py
oo92/TorchFaceLab
import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F

class DepthwiseConv2D(nn.Module):
    def __init__(self, in_ch, kernel_size, strides=1, padding='SAME', depth_multiplier=1, 
                 dilations=1, use_bias=True, use_wscale=False, kernel_initializer=None, 
                 bias_initializer=None, trainable=True, dtype=None):
        
        super(DepthwiseConv2D, self).__init__()
        
        if dtype is None:
            dtype = torch.float32
        
        if padding == "SAME":
            padding = ( (kernel_size - 1) * dilations + 1 ) // 2
        elif padding == "VALID":
            padding = 0
        else:
            raise ValueError("Wrong padding type. Should be VALID, SAME, or int.")

        self.use_bias = use_bias
        self.use_wscale = use_wscale
        self.in_ch = in_ch
        self.depth_multiplier = depth_multiplier
        self.kernel_size = kernel_size
        self.strides = strides
        self.padding = padding
        self.dilations = dilations
        self.kernel_initializer = kernel_initializer
        self.bias_initializer = bias_initializer

        if self.use_wscale:
            gain = 1.0 if self.kernel_size == 1 else np.sqrt(2)
            fan_in = self.kernel_size * self.kernel_size * self.in_ch
            he_std = gain / np.sqrt(fan_in)
            self.wscale = torch.tensor(he_std, dtype=dtype)

        self.weight = nn.Parameter(torch.empty(self.kernel_size, self.kernel_size, self.in_ch, self.depth_multiplier, dtype=dtype))
        if self.kernel_initializer:
            nn.init.xavier_uniform_(self.weight)
        
        if self.use_bias:
            self.bias = nn.Parameter(torch.zeros(self.in_ch * self.depth_multiplier, dtype=dtype))
            if self.bias_initializer:
                nn.init.zeros_(self.bias)
        else:
            self.register_parameter('bias', None)

    def forward(self, x):
        weight = self.weight
        if self.use_wscale:
            weight = weight * self.wscale
        
        x = F.conv2d(x, weight, self.bias, stride=self.strides, padding=self.padding, dilation=self.dilations, groups=self.in_ch)
        return x

    def __str__(self):
        return f"{self.__class__.__name__} : in_ch:{self.in_ch} depth_multiplier:{self.depth_multiplier}"

nn.DepthwiseConv2D = DepthwiseConv2D
core/leras/layers/DepthwiseConv2D.py
oo92/TorchFaceLab
import torch
import torch.nn as nn

class FRNorm2D(nn.Module):
    """
    PyTorch implementation of
    Filter Response Normalization Layer: Eliminating Batch Dependence in theTraining of Deep Neural Networks
    https://arxiv.org/pdf/1911.09737.pdf
    """
    def __init__(self, in_ch, dtype=None):
        super(FRNorm2D, self).__init__()
        self.in_ch = in_ch
        
        if dtype is None:
            dtype = torch.float32
        self.dtype = dtype

        # In PyTorch, the preferred way to declare weights and biases is using nn.Parameter.
        self.weight = nn.Parameter(torch.ones(self.in_ch, dtype=self.dtype))
        self.bias = nn.Parameter(torch.zeros(self.in_ch, dtype=self.dtype))
        self.eps = nn.Parameter(torch.tensor(1e-6, dtype=self.dtype))

    def forward(self, x):
        # Assuming that the input tensor x has the shape (N, C, H, W), which is PyTorch's default format.

        # Compute nu2 as the mean square of x along spatial dimensions
        nu2 = x.pow(2).mean(dim=[2, 3], keepdim=True)
        x_normalized = x * (1.0 / torch.sqrt(nu2 + self.eps.abs()))

        # Scale and shift the normalized tensor
        x_out = self.weight.view(1, self.in_ch, 1, 1) * x_normalized + self.bias.view(1, self.in_ch, 1, 1)

        return x_out

nn.FRNorm2D = FRNorm2D
core/leras/layers/FRNorm2D.py
oo92/TorchFaceLab
import torch
import torch.nn as nn

class InstanceNorm2D(nn.Module):
    def __init__(self, in_ch, dtype=None):
        super(InstanceNorm2D, self).__init__()
        self.in_ch = in_ch
        if dtype is None:
            dtype = torch.float32
        self.dtype = dtype
        
        # Instead of using tf.get_variable(), we can use nn.Parameter to declare weights
        # Note that in PyTorch, the preferred way to declare weights and biases is using nn.Parameter.
        self.weight = nn.Parameter(torch.randn(self.in_ch, dtype=self.dtype))
        self.bias = nn.Parameter(torch.zeros(self.in_ch, dtype=self.dtype))
        
        # Initialize the weights
        nn.init.xavier_uniform_(self.weight)

    def forward(self, x):
        # Assuming that data_format "NHWC" corresponds to shape (N, H, W, C) in TensorFlow,
        # this is equivalent to (N, C, H, W) in PyTorch, which is PyTorch's default format.
        
        # Computing mean and std dev along the spatial dimensions
        x_mean = x.mean(dim=[2,3], keepdim=True)
        x_std = x.std(dim=[2,3], keepdim=True) + 1e-5
        
        # Perform instance normalization
        x_normalized = (x - x_mean) / x_std
        x_out = self.weight.view(1, self.in_ch, 1, 1) * x_normalized + self.bias.view(1, self.in_ch, 1, 1)
        
        return x_out

nn.InstanceNorm2D = InstanceNorm2D
core/leras/layers/InstanceNorm2D.py
oo92/TorchFaceLab
import torch.nn as nn

class LayerBase():
    #override
    def build_weights(self):
        pass
    
    #override
    def forward(self, *args, **kwargs):
        pass
    
    def __call__(self, *args, **kwargs):
        return self.forward(*args, **kwargs)

nn.LayerBase = LayerBase
core/leras/layers/LayerBase.py
oo92/TorchFaceLab
import pickle, torch
import numpy as np
from pathlib import Path
from core import pathex
import torch.nn as nn

class Saveable(nn.Module):
    def __init__(self, name=None):
        super(Saveable, self).__init__()
        self.name = name

    def get_weights(self):
        # Return PyTorch tensors that should be initialized/loaded/saved
        return [param for param in self.parameters()]

    def get_weights_np(self):
        weights = self.get_weights()
        return [w.detach().numpy() for w in weights]

    def set_weights(self, new_weights):
        weights = self.get_weights()
        if len(weights) != len(new_weights):
            raise ValueError('len of lists mismatch')

        for w, new_w in zip(weights, new_weights):
            w_shape = list(w.shape)
            if w_shape != list(new_w.shape):
                new_w = new_w.reshape(w_shape)
            w.data.copy_(torch.from_numpy(new_w))

    def save_weights(self, filename, force_dtype=None):
        d = {}
        weights = self.get_weights()

        if self.name is None:
            raise Exception("name must be defined.")

        name = self.name

        for w in weights:
            w_val = w.detach().numpy().copy()
            w_name = w.name if hasattr(w, 'name') else str(id(w))
            w_name_split = w_name.split('/', 1)
            if name != w_name_split[0]:
                raise Exception("weight first name != Saveable.name")

            if force_dtype is not None:
                w_val = w_val.astype(force_dtype)

            d[w_name_split[1]] = w_val

        d_dumped = pickle.dumps(d, 4)
        pathex.write_bytes_safe(Path(filename), d_dumped)

    def load_weights(self, filename):
        filepath = Path(filename)
        if not filepath.exists():
            return False

        d_dumped = filepath.read_bytes()
        d = pickle.loads(d_dumped)

        weights = self.get_weights()

        if self.name is None:
            raise Exception("name must be defined.")

        try:
            for w in weights:
                w_name = w.name if hasattr(w, 'name') else str(id(w))
                w_name_split = w_name.split('/')
                if self.name != w_name_split[0]:
                    raise Exception("weight first name != Saveable.name")

                sub_w_name = "/".join(w_name_split[1:])
                w_val = d.get(sub_w_name, None)

                if w_val is None:
                    torch.nn.init.xavier_uniform_(w)
                else:
                    w_val = np.reshape(w_val, list(w.shape))
                    w.data.copy_(torch.from_numpy(w_val))
        except:
            return False

        return True

nn.Saveable = Saveable
core/leras/layers/Saveable.py
oo92/TorchFaceLab
import torch
import torch.nn as nn

class ScaleAdd(nn.Module):
    def __init__(self, ch, dtype=None):
        super(ScaleAdd, self).__init__()
        
        if dtype is None:
            dtype = torch.float32
        self.dtype = dtype
        self.ch = ch
        
        # Initialize weights
        self.weight = nn.Parameter(torch.zeros(self.ch, dtype=self.dtype))
        
    def forward(self, inputs):
        # Depending on the data format, choose the shape for reshaping the weights
        # Here, I'm assuming a default of "NCHW" for PyTorch. If you need "NHWC", you'll have to adjust.
        shape = (1, self.ch, 1, 1)
        
        # Reshape the weights
        weight = self.weight.view(shape)
        
        # Extract inputs
        x0, x1 = inputs
        
        # Compute the output
        x = x0 + x1 * weight
        
        return x
    
nn.ScaleAdd = ScaleAdd
core/leras/layers/ScaleAdd.py
oo92/TorchFaceLab
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.init as init

class TLU(nn.Module):
    """
    PyTorch implementation of
    Filter Response Normalization Layer: Eliminating Batch Dependence in the Training of Deep Neural Networks
    https://arxiv.org/pdf/1911.09737.pdf
    """
    def __init__(self, in_ch, dtype=None):
        super(TLU, self).__init__()
        
        self.in_ch = in_ch
        self.tau = nn.Parameter(torch.zeros(in_ch, dtype=dtype or torch.float32), requires_grad=True)
                
    def forward(self, x):
        tau = self.tau.view(1, self.in_ch, 1, 1)
        return torch.max(x, tau)

nn.TLU = TLU
core/leras/layers/TLU.py
oo92/TorchFaceLab
import torch
import torch.nn.functional as F
import numpy as np
import torch.nn as nn


class TanhPolar(nn.Module):
    def __init__(self, width, height, angular_offset_deg=270):
        super(TanhPolar, self).__init__()

        self.width = width
        self.height = height

        warp_gridx, warp_gridy = self._get_tanh_polar_warp_grids(width, height, angular_offset_deg=angular_offset_deg)
        restore_gridx, restore_gridy = self._get_tanh_polar_restore_grids(width, height, angular_offset_deg=angular_offset_deg)

        self.register_buffer('warp_gridx', torch.tensor(warp_gridx).unsqueeze(0))
        self.register_buffer('warp_gridy', torch.tensor(warp_gridy).unsqueeze(0))
        self.register_buffer('restore_gridx', torch.tensor(restore_gridx).unsqueeze(0))
        self.register_buffer('restore_gridy', torch.tensor(restore_gridy).unsqueeze(0))

    def warp(self, inp):
        batch_size = inp.shape[0]
        warp_gridx = self.warp_gridx.repeat(batch_size, 1, 1)
        warp_gridy = self.warp_gridy.repeat(batch_size, 1, 1)
        grid = torch.stack((warp_gridy, warp_gridx), dim=3)
        out = F.grid_sample(inp, grid, mode='bilinear', padding_mode='zeros', align_corners=True)
        return out

    def restore(self, inp):
        batch_size = inp.shape[0]
        restore_gridx = self.restore_gridx.repeat(batch_size, 1, 1)
        restore_gridy = self.restore_gridy.repeat(batch_size, 1, 1)
        inp = F.pad(inp, (0, 0, 1, 0, 1, 1), mode='reflect')
        grid = torch.stack((restore_gridy, restore_gridx), dim=3)
        out = F.grid_sample(inp, grid, mode='bilinear', padding_mode='zeros', align_corners=True)
        return out

    @staticmethod
    def _get_tanh_polar_warp_grids(W, H, angular_offset_deg):
        angular_offset_pi = angular_offset_deg * np.pi / 180.0
        roi_center = np.array([W // 2, H // 2], np.float32)
        roi_radii = np.array([W, H], np.float32) / np.pi ** 0.5
        cos_offset, sin_offset = np.cos(angular_offset_pi), np.sin(angular_offset_pi)
        normalised_dest_indices = np.stack(np.meshgrid(np.arange(0.0, 1.0, 1.0 / W), np.arange(0.0, 2.0 * np.pi, 2.0 * np.pi / H)), axis=-1)
        radii = normalised_dest_indices[..., 0]
        orientation_x = np.cos(normalised_dest_indices[..., 1])
        orientation_y = np.sin(normalised_dest_indices[..., 1])
        src_radii = np.arctanh(radii) * (roi_radii[0] * roi_radii[1] / np.sqrt(roi_radii[1] ** 2 * orientation_x ** 2 + roi_radii[0] ** 2 * orientation_y ** 2))
        src_x_indices = src_radii * orientation_x
        src_y_indices = src_radii * orientation_y
        src_x_indices, src_y_indices = (roi_center[0] + cos_offset * src_x_indices - sin_offset * src_y_indices,
                                        roi_center[1] + cos_offset * src_y_indices + sin_offset * src_x_indices)
        return src_x_indices.astype(np.float32), src_y_indices.astype(np.float32)

    @staticmethod
    def _get_tanh_polar_restore_grids(W, H, angular_offset_deg):
        angular_offset_pi = angular_offset_deg * np.pi / 180.0
        roi_center = np.array([W // 2, H // 2], np.float32)
        roi_radii = np.array([W, H], np.float32) / np.pi ** 0.5
        cos_offset, sin_offset = np.cos(angular_offset_pi), np.sin(angular_offset_pi)
        dest_indices = np.stack(np.meshgrid(np.arange(W), np.arange(H)), axis=-1).astype(float)
        normalised_dest_indices = np.matmul(dest_indices - roi_center, np.array([[cos_offset, -sin_offset],
                                                                                [sin_offset, cos_offset]]))
        radii = np.linalg.norm(normalised_dest_indices, axis=-1)
        normalised_dest_indices[..., 0] /= np.clip(radii, 1e-9, None)
        normalised_dest_indices[..., 1] /= np.clip(radii, 1e-9, None)
        radii *= np.sqrt(roi_radii[1] ** 2 * normalised_dest_indices[..., 0] ** 2 +
                         roi_radii[0] ** 2 * normalised_dest_indices[..., 1] ** 2) / roi_radii[0] / roi_radii[1]
        src_radii = np.tanh(radii)
        src_x_indices = src_radii * W + 1.0
        src_y_indices = np.mod((np.arctan2(normalised_dest_indices[..., 1], normalised_dest_indices[..., 0]) /
                                2.0 / np.pi) * H, H) + 1.0
        return src_x_indices.astype(np.float32), src_y_indices.astype(np.float32)

nn.TanhPolar = TanhPolar
core/leras/layers/TanhPolar.py
oo92/TorchFaceLab
from .Saveable import *
from .LayerBase import *

from .Conv2D import *
from .Conv2DTranspose import *
from .DepthwiseConv2D import *
from .Dense import *
from .BlurPool import *

from .BatchNorm2D import *
from .InstanceNorm2D import *
from .FRNorm2D import *

from .TLU import *
from .ScaleAdd import *
from .DenseNorm import *
from .AdaIN import *
from .TanhPolar import *
core/leras/layers/__init__.py
oo92/TorchFaceLab
from core.leras import nn
tf = nn.tf

class CodeDiscriminator(nn.ModelBase):
    def on_build(self, in_ch, code_res, ch=256, conv_kernel_initializer=None):            
        n_downscales = 1 + code_res // 8

        self.convs = []
        prev_ch = in_ch
        for i in range(n_downscales):
            cur_ch = ch * min( (2**i), 8 )
            self.convs.append ( nn.Conv2D( prev_ch, cur_ch, kernel_size=4 if i == 0 else 3, strides=2, padding='SAME', kernel_initializer=conv_kernel_initializer) )
            prev_ch = cur_ch

        self.out_conv =  nn.Conv2D( prev_ch, 1, kernel_size=1, padding='VALID', kernel_initializer=conv_kernel_initializer)

    def forward(self, x):
        for conv in self.convs:
            x = tf.nn.leaky_relu( conv(x), 0.1 )
        return self.out_conv(x)
        
nn.CodeDiscriminator = CodeDiscriminator
core/leras/models/CodeDiscriminator.py
oo92/TorchFaceLab
import torch
import numpy as np

from core.interact import interact as io
import torch.nn as nn
import core.leras.ops
import core.leras.layers.LayerBase as LayerBase
import core.leras.initializers
import core.leras.optimizers
import core.leras.models.ModelBase as ModelBase
import core.leras.archis

class ModelBase(nn.Module):
    def __init__(self, *args, name=None, **kwargs):
        super().__init__(name=name)
        self.layers = []
        self.layers_by_name = {}
        self.built = False
        self.args = args
        self.kwargs = kwargs
        self.run_placeholders = None

    def _build_sub(self, layer, name):
        # Simplified _build_sub for PyTorch since there's no direct equivalent of tf.variable_scope
        # Just handling naming and adding layers
        if isinstance(layer, (LayerBase, ModelBase)):
            layer.name = name
            self.layers.append(layer)
            self.layers_by_name[layer.name] = layer

    def xor_list(self, lst1, lst2):
        return [value for value in lst1 + lst2 if (value not in lst1) or (value not in lst2)]

    def build(self):
        generator = self.on_build(*self.args, **self.kwargs)
        if not isinstance(generator, type(iter([]))):
            generator = None

        if generator is not None:
            for _ in generator:
                pass

        current_vars = list(vars(self))
        for name in current_vars:
            self._build_sub(vars(self)[name], name)

        self.built = True

    def get_weights(self):
        if not self.built:
            self.build()
        return [layer.parameters() for layer in self.layers]

    def get_layer_by_name(self, name):
        return self.layers_by_name.get(name, None)

    def get_layers(self):
        if not self.built:
            self.build()
        return self.layers

    # Intended to be overridden by child classes
    def on_build(self, *args, **kwargs):
        pass

    # Intended to be overridden by child classes
    def forward(self, *args, **kwargs):
        pass

    def __call__(self, *args, **kwargs):
        if not self.built:
            self.build()
        return self.forward(*args, **kwargs)

    def build_for_run(self, shapes_list):
        if not isinstance(shapes_list, list):
            raise ValueError("shapes_list must be a list.")
        self.run_placeholders = [torch.zeros(*sh) for _, sh in shapes_list]
        self.run_output = self.__call__(self.run_placeholders)

    def run(self, inputs):
        if self.run_placeholders is None:
            raise Exception("Model didn't build for run.")

        if len(inputs) != len(self.run_placeholders):
            raise ValueError("len(inputs) != self.run_placeholders")

        # Direct forward pass in PyTorch
        with torch.no_grad():
            return self.forward(*inputs)

    # Similar summary function as in original code
    def summary(self):
        layers = self.get_layers()
        layers_names = []
        layers_params = []

        max_len_str = 0
        max_len_param_str = 0
        delim_str = "-"

        total_params = 0

        # Get layers names and str length for delim
        for layer in layers:
            name = str(layer)
            layers_names.append(name.capitalize())
            max_len_str = max(max_len_str, len(name))

        # Get params for each layer
        for layer in layers:
            params_count = sum(p.numel() for p in layer.parameters())
            layers_params.append(params_count)
            max_len_param_str = max(max_len_param_str, len(str(params_count)))

        total_params = sum(layers_params)

        # Set delim
        delim_str = delim_str * (max_len_str + max_len_param_str + 3)
        output = "\n" + delim_str + "\n"

        # Format model name str
        model_name_str = "| " + self.name.capitalize()
        while len(model_name_str) < len(delim_str) - 2:
            model_name_str += " "
        model_name_str += " |"

        output += model_name_str + "\n"
        output += delim_str + "\n"

        # Format layers table
        for layer_name, layer_param in zip(layers_names, layers_params):
            output += delim_str + "\n"
            formatted_name = layer_name.ljust(max_len_str)
            formatted_param = str(layer_param).rjust(max_len_param_str)
            output += f"| {formatted_name} | {formatted_param} |\n"

        output += delim_str + "\n"

        # Format sum of params
        total_params_str = "| Total params count: " + str(total_params)
        while len(total_params_str) < len(delim_str) - 2:
            total_params_str += " "
        total_params_str += " |"

        output += total_params_str + "\n"
        output += delim_str + "\n"

        io.log_info(output)


nn.ModelBase = ModelBase

core/leras/models/ModelBase.py
oo92/TorchFaceLab
import numpy as np
from core.leras import nn
tf = nn.tf

patch_discriminator_kernels = \
    { 1  : (512, [ [1,1] ]),
      2  : (512, [ [2,1] ]),
      3  : (512, [ [2,1], [2,1] ]),
      4  : (512, [ [2,2], [2,2] ]),
      5  : (512, [ [3,2], [2,2] ]),
      6  : (512, [ [4,2], [2,2] ]),
      7  : (512, [ [3,2], [3,2] ]),
      8  : (512, [ [4,2], [3,2] ]),
      9  : (512, [ [3,2], [4,2] ]),
      10 : (512, [ [4,2], [4,2] ]),
      11 : (512, [ [3,2], [3,2], [2,1] ]),
      12 : (512, [ [4,2], [3,2], [2,1] ]),
      13 : (512, [ [3,2], [4,2], [2,1] ]),
      14 : (512, [ [4,2], [4,2], [2,1] ]),
      15 : (512, [ [3,2], [3,2], [3,1] ]),
      16 : (512, [ [4,2], [3,2], [3,1] ]),
      17 : (512, [ [3,2], [4,2], [3,1] ]),
      18 : (512, [ [4,2], [4,2], [3,1] ]),
      19 : (512, [ [3,2], [3,2], [4,1] ]),
      20 : (512, [ [4,2], [3,2], [4,1] ]),
      21 : (512, [ [3,2], [4,2], [4,1] ]),
      22 : (512, [ [4,2], [4,2], [4,1] ]),
      23 : (256, [ [3,2], [3,2], [3,2], [2,1] ]),
      24 : (256, [ [4,2], [3,2], [3,2], [2,1] ]),
      25 : (256, [ [3,2], [4,2], [3,2], [2,1] ]),
      26 : (256, [ [4,2], [4,2], [3,2], [2,1] ]),
      27 : (256, [ [3,2], [4,2], [4,2], [2,1] ]),
      28 : (256, [ [4,2], [3,2], [4,2], [2,1] ]),
      29 : (256, [ [3,2], [4,2], [4,2], [2,1] ]),
      30 : (256, [ [4,2], [4,2], [4,2], [2,1] ]),
      31 : (256, [ [3,2], [3,2], [3,2], [3,1] ]),
      32 : (256, [ [4,2], [3,2], [3,2], [3,1] ]),
      33 : (256, [ [3,2], [4,2], [3,2], [3,1] ]),
      34 : (256, [ [4,2], [4,2], [3,2], [3,1] ]),
      35 : (256, [ [3,2], [4,2], [4,2], [3,1] ]),
      36 : (256, [ [4,2], [3,2], [4,2], [3,1] ]),
      37 : (256, [ [3,2], [4,2], [4,2], [3,1] ]),
      38 : (256, [ [4,2], [4,2], [4,2], [3,1] ]),
      39 : (256, [ [3,2], [3,2], [3,2], [4,1] ]),
      40 : (256, [ [4,2], [3,2], [3,2], [4,1] ]),
      41 : (256, [ [3,2], [4,2], [3,2], [4,1] ]),
      42 : (256, [ [4,2], [4,2], [3,2], [4,1] ]),
      43 : (256, [ [3,2], [4,2], [4,2], [4,1] ]),
      44 : (256, [ [4,2], [3,2], [4,2], [4,1] ]),
      45 : (256, [ [3,2], [4,2], [4,2], [4,1] ]),
      46 : (256, [ [4,2], [4,2], [4,2], [4,1] ]),
    }


class PatchDiscriminator(nn.ModelBase):
    def on_build(self, patch_size, in_ch, base_ch=None, conv_kernel_initializer=None):
        suggested_base_ch, kernels_strides = patch_discriminator_kernels[patch_size]

        if base_ch is None:
            base_ch = suggested_base_ch

        prev_ch = in_ch
        self.convs = []
        for i, (kernel_size, strides) in enumerate(kernels_strides):
            cur_ch = base_ch * min( (2**i), 8 )

            self.convs.append ( nn.Conv2D( prev_ch, cur_ch, kernel_size=kernel_size, strides=strides, padding='SAME', kernel_initializer=conv_kernel_initializer) )
            prev_ch = cur_ch

        self.out_conv =  nn.Conv2D( prev_ch, 1, kernel_size=1, padding='VALID', kernel_initializer=conv_kernel_initializer)

    def forward(self, x):
        for conv in self.convs:
            x = tf.nn.leaky_relu( conv(x), 0.1 )
        return self.out_conv(x)

nn.PatchDiscriminator = PatchDiscriminator

class UNetPatchDiscriminator(nn.ModelBase):
    """
    Inspired by https://arxiv.org/abs/2002.12655 "A U-Net Based Discriminator for Generative Adversarial Networks"
    """
    def calc_receptive_field_size(self, layers):
        """
        result the same as https://fomoro.com/research/article/receptive-field-calculatorindex.html
        """
        rf = 0
        ts = 1
        for i, (k, s) in enumerate(layers):
            if i == 0:
                rf = k
            else:
                rf += (k-1)*ts
            ts *= s
        return rf

    def find_archi(self, target_patch_size, max_layers=9):
        """
        Find the best configuration of layers using only 3x3 convs for target patch size
        """
        s = {}
        for layers_count in range(1,max_layers+1):
            val = 1 << (layers_count-1)
            while True:
                val -= 1

                layers = []
                sum_st = 0
                layers.append ( [3, 2])
                sum_st += 2
                for i in range(layers_count-1):
                    st = 1 + (1 if val & (1 << i) !=0 else 0 )
                    layers.append ( [3, st ])
                    sum_st += st                

                rf = self.calc_receptive_field_size(layers)

                s_rf = s.get(rf, None)
                if s_rf is None:
                    s[rf] = (layers_count, sum_st, layers)
                else:
                    if layers_count < s_rf[0] or \
                    ( layers_count == s_rf[0] and sum_st > s_rf[1] ):
                        s[rf] = (layers_count, sum_st, layers)

                if val == 0:
                    break

        x = sorted(list(s.keys()))
        q=x[np.abs(np.array(x)-target_patch_size).argmin()]
        return s[q][2]

    def on_build(self, patch_size, in_ch, base_ch = 16, use_fp16 = False):
        self.use_fp16 = use_fp16
        conv_dtype = tf.float16 if use_fp16 else tf.float32 
        
        class ResidualBlock(nn.ModelBase):
            def on_build(self, ch, kernel_size=3 ):
                self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype)
                self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype)

            def forward(self, inp):
                x = self.conv1(inp)
                x = tf.nn.leaky_relu(x, 0.2)
                x = self.conv2(x)
                x = tf.nn.leaky_relu(inp + x, 0.2)
                return x

        prev_ch = in_ch
        self.convs = []
        self.upconvs = []
        layers = self.find_archi(patch_size)
        
        level_chs = { i-1:v for i,v in enumerate([ min( base_ch * (2**i), 512 ) for i in range(len(layers)+1)]) }

        self.in_conv = nn.Conv2D( in_ch, level_chs[-1], kernel_size=1, padding='VALID', dtype=conv_dtype)

        for i, (kernel_size, strides) in enumerate(layers):
            self.convs.append ( nn.Conv2D( level_chs[i-1], level_chs[i], kernel_size=kernel_size, strides=strides, padding='SAME', dtype=conv_dtype) )

            self.upconvs.insert (0, nn.Conv2DTranspose( level_chs[i]*(2 if i != len(layers)-1 else 1), level_chs[i-1], kernel_size=kernel_size, strides=strides, padding='SAME', dtype=conv_dtype) )

        self.out_conv = nn.Conv2D( level_chs[-1]*2, 1, kernel_size=1, padding='VALID', dtype=conv_dtype)

        self.center_out  =  nn.Conv2D( level_chs[len(layers)-1], 1, kernel_size=1, padding='VALID', dtype=conv_dtype)
        self.center_conv =  nn.Conv2D( level_chs[len(layers)-1], level_chs[len(layers)-1], kernel_size=1, padding='VALID', dtype=conv_dtype)


    def forward(self, x):
        if self.use_fp16:
            x = tf.cast(x, tf.float16)
            
        x = tf.nn.leaky_relu( self.in_conv(x), 0.2 )

        encs = []
        for conv in self.convs:
            encs.insert(0, x)
            x = tf.nn.leaky_relu( conv(x), 0.2 )
            
        center_out, x = self.center_out(x), tf.nn.leaky_relu( self.center_conv(x), 0.2 )

        for i, (upconv, enc) in enumerate(zip(self.upconvs, encs)):
            x = tf.nn.leaky_relu( upconv(x), 0.2 )
            x = tf.concat( [enc, x], axis=nn.conv2d_ch_axis)

        x = self.out_conv(x)
        
        if self.use_fp16:
            center_out = tf.cast(center_out, tf.float32)
            x = tf.cast(x, tf.float32)

        return center_out, x

nn.UNetPatchDiscriminator = UNetPatchDiscriminator

core/leras/models/PatchDiscriminator.py
oo92/TorchFaceLab
from core.leras import nn
tf = nn.tf

class XSeg(nn.ModelBase):
    
    def on_build (self, in_ch, base_ch, out_ch):
        
        class ConvBlock(nn.ModelBase):
            def on_build(self, in_ch, out_ch):              
                self.conv = nn.Conv2D (in_ch, out_ch, kernel_size=3, padding='SAME')
                self.frn = nn.FRNorm2D(out_ch)
                self.tlu = nn.TLU(out_ch)

            def forward(self, x):                
                x = self.conv(x)
                x = self.frn(x)
                x = self.tlu(x)
                return x

        class UpConvBlock(nn.ModelBase):
            def on_build(self, in_ch, out_ch):
                self.conv = nn.Conv2DTranspose (in_ch, out_ch, kernel_size=3, padding='SAME')
                self.frn = nn.FRNorm2D(out_ch)
                self.tlu = nn.TLU(out_ch)

            def forward(self, x):
                x = self.conv(x)
                x = self.frn(x)
                x = self.tlu(x)
                return x
                
        self.base_ch = base_ch

        self.conv01 = ConvBlock(in_ch, base_ch)
        self.conv02 = ConvBlock(base_ch, base_ch)
        self.bp0 = nn.BlurPool (filt_size=4)

        self.conv11 = ConvBlock(base_ch, base_ch*2)
        self.conv12 = ConvBlock(base_ch*2, base_ch*2)
        self.bp1 = nn.BlurPool (filt_size=3)

        self.conv21 = ConvBlock(base_ch*2, base_ch*4)
        self.conv22 = ConvBlock(base_ch*4, base_ch*4)
        self.bp2 = nn.BlurPool (filt_size=2)

        self.conv31 = ConvBlock(base_ch*4, base_ch*8)
        self.conv32 = ConvBlock(base_ch*8, base_ch*8)
        self.conv33 = ConvBlock(base_ch*8, base_ch*8)
        self.bp3 = nn.BlurPool (filt_size=2)

        self.conv41 = ConvBlock(base_ch*8, base_ch*8)
        self.conv42 = ConvBlock(base_ch*8, base_ch*8)
        self.conv43 = ConvBlock(base_ch*8, base_ch*8)
        self.bp4 = nn.BlurPool (filt_size=2)
        
        self.conv51 = ConvBlock(base_ch*8, base_ch*8)
        self.conv52 = ConvBlock(base_ch*8, base_ch*8)
        self.conv53 = ConvBlock(base_ch*8, base_ch*8)
        self.bp5 = nn.BlurPool (filt_size=2)
        
        self.dense1 = nn.Dense ( 4*4* base_ch*8, 512)
        self.dense2 = nn.Dense ( 512, 4*4* base_ch*8)
                
        self.up5 = UpConvBlock (base_ch*8, base_ch*4)
        self.uconv53 = ConvBlock(base_ch*12, base_ch*8)
        self.uconv52 = ConvBlock(base_ch*8, base_ch*8)
        self.uconv51 = ConvBlock(base_ch*8, base_ch*8)
        
        self.up4 = UpConvBlock (base_ch*8, base_ch*4)
        self.uconv43 = ConvBlock(base_ch*12, base_ch*8)
        self.uconv42 = ConvBlock(base_ch*8, base_ch*8)
        self.uconv41 = ConvBlock(base_ch*8, base_ch*8)

        self.up3 = UpConvBlock (base_ch*8, base_ch*4)
        self.uconv33 = ConvBlock(base_ch*12, base_ch*8)
        self.uconv32 = ConvBlock(base_ch*8, base_ch*8)
        self.uconv31 = ConvBlock(base_ch*8, base_ch*8)

        self.up2 = UpConvBlock (base_ch*8, base_ch*4)
        self.uconv22 = ConvBlock(base_ch*8, base_ch*4)
        self.uconv21 = ConvBlock(base_ch*4, base_ch*4)

        self.up1 = UpConvBlock (base_ch*4, base_ch*2)
        self.uconv12 = ConvBlock(base_ch*4, base_ch*2)
        self.uconv11 = ConvBlock(base_ch*2, base_ch*2)

        self.up0 = UpConvBlock (base_ch*2, base_ch)
        self.uconv02 = ConvBlock(base_ch*2, base_ch)
        self.uconv01 = ConvBlock(base_ch, base_ch)
        self.out_conv = nn.Conv2D (base_ch, out_ch, kernel_size=3, padding='SAME')
    
        
    def forward(self, inp, pretrain=False):
        x = inp

        x = self.conv01(x)
        x = x0 = self.conv02(x)
        x = self.bp0(x)

        x = self.conv11(x)
        x = x1 = self.conv12(x)
        x = self.bp1(x)

        x = self.conv21(x)
        x = x2 = self.conv22(x)
        x = self.bp2(x)

        x = self.conv31(x)
        x = self.conv32(x)
        x = x3 = self.conv33(x)
        x = self.bp3(x)

        x = self.conv41(x)
        x = self.conv42(x)
        x = x4 = self.conv43(x)
        x = self.bp4(x)

        x = self.conv51(x)
        x = self.conv52(x)
        x = x5 = self.conv53(x)
        x = self.bp5(x)
        
        x = nn.flatten(x)
        x = self.dense1(x)
        x = self.dense2(x)
        x = nn.reshape_4D (x, 4, 4, self.base_ch*8 )
                          
        x = self.up5(x)
        if pretrain:
            x5 = tf.zeros_like(x5)
        x = self.uconv53(tf.concat([x,x5],axis=nn.conv2d_ch_axis))
        x = self.uconv52(x)
        x = self.uconv51(x)
        
        x = self.up4(x)
        if pretrain:
            x4 = tf.zeros_like(x4)
        x = self.uconv43(tf.concat([x,x4],axis=nn.conv2d_ch_axis))
        x = self.uconv42(x)
        x = self.uconv41(x)

        x = self.up3(x)
        if pretrain:
            x3 = tf.zeros_like(x3)
        x = self.uconv33(tf.concat([x,x3],axis=nn.conv2d_ch_axis))
        x = self.uconv32(x)
        x = self.uconv31(x)

        x = self.up2(x)
        if pretrain:
            x2 = tf.zeros_like(x2)
        x = self.uconv22(tf.concat([x,x2],axis=nn.conv2d_ch_axis))
        x = self.uconv21(x)

        x = self.up1(x)
        if pretrain:
            x1 = tf.zeros_like(x1)
        x = self.uconv12(tf.concat([x,x1],axis=nn.conv2d_ch_axis))
        x = self.uconv11(x)

        x = self.up0(x)
        if pretrain:
            x0 = tf.zeros_like(x0)
        x = self.uconv02(tf.concat([x,x0],axis=nn.conv2d_ch_axis))
        x = self.uconv01(x)

        logits = self.out_conv(x)
        return logits, tf.nn.sigmoid(logits)

nn.XSeg = XSeg
core/leras/models/XSeg.py
oo92/TorchFaceLab
from .ModelBase import *
from .PatchDiscriminator import *
from .CodeDiscriminator import *
from .XSeg import *
core/leras/models/__init__.py
oo92/TorchFaceLab
# from pathlib import Path
# import numpy as np
# import torch
# from core.interact import interact as io
# from .device import Devices
# import core.leras.ops
# import core.leras.layers
# import core.leras.initializers
# import core.leras.optimizers
# import core.leras.models
# import core.leras.archis

# class nn():
#     current_DeviceConfig = None
#     device = None
#     data_format = "NCHW"  # Default format for PyTorch
#     conv2d_ch_axis = 1
#     conv2d_spatial_axes = [2, 3]
#     floatx = None

#     @staticmethod
#     def initialize(device_config=None, floatx="float32"):

#         if nn.device is None:
#             if device_config is None:
#                 device_config = nn.getCurrentDeviceConfig()
#             nn.setCurrentDeviceConfig(device_config)

#             if torch.cuda.is_available():
#                 io.log_info(f"Using GPU: {torch.cuda.get_device_name(0)}")
#                 nn.device = torch.device('cuda')
#             else:
#                 io.log_info("Using CPU")
#                 nn.device = torch.device('cpu')

#         if floatx == "float32":
#             floatx = torch.float32
#         elif floatx == "float16":
#             floatx = torch.float16
#         else:
#             raise ValueError(f"unsupported floatx {floatx}")
#         nn.set_floatx(floatx)

#     @staticmethod
#     def initialize_main_env():
#         Devices.initialize_main_env()

#     @staticmethod
#     def set_floatx(torch_dtype):
#         """Set default float type for all layers when dtype is None for them"""
#         nn.floatx = torch_dtype

#     # Rest of the data format functions remain unchanged since PyTorch uses NCHW by default

#     @staticmethod
#     def get4Dshape(w, h, c):
#         """Returns 4D shape based on current data_format"""
#         if nn.data_format == "NHWC":
#             return (None, h, w, c)
#         else:
#             return (None, c, h, w)

#     @staticmethod
#     def to_data_format(x, to_data_format, from_data_format):
#         if to_data_format == from_data_format:
#             return x

#         if to_data_format == "NHWC":
#             return np.transpose(x, (0, 2, 3, 1))
#         elif to_data_format == "NCHW":
#             return np.transpose(x, (0, 3, 1, 2))
#         else:
#             raise ValueError(f"unsupported to_data_format {to_data_format}")

#     @staticmethod
#     def getCurrentDeviceConfig():
#         if nn.current_DeviceConfig is None:
#             nn.current_DeviceConfig = DeviceConfig.BestGPU()
#         return nn.current_DeviceConfig

#     @staticmethod
#     def setCurrentDeviceConfig(device_config):
#         nn.current_DeviceConfig = device_config

#     @staticmethod
#     def ask_choose_device_idxs(choose_only_one=False, allow_cpu=True, suggest_best_multi_gpu=False, suggest_all_gpu=False):
#         devices = Devices.getDevices()
#         if len(devices) == 0:
#             return []

#         all_devices_indexes = [device.index for device in devices]

#         if choose_only_one:
#             suggest_best_multi_gpu = False
#             suggest_all_gpu = False

#         if suggest_all_gpu:
#             best_device_indexes = all_devices_indexes
#         elif suggest_best_multi_gpu:
#             best_device_indexes = [device.index for device in devices.get_equal_devices(devices.get_best_device()) ]
#         else:
#             best_device_indexes = [ devices.get_best_device().index ]
#         best_device_indexes = ",".join([str(x) for x in best_device_indexes])

#         io.log_info ("")
#         if choose_only_one:
#             io.log_info ("Choose one GPU idx.")
#         else:
#             io.log_info ("Choose one or several GPU idxs (separated by comma).")
#         io.log_info ("")

#         if allow_cpu:
#             io.log_info ("[CPU] : CPU")
#         for device in devices:
#             io.log_info (f"  [{device.index}] : {device.name}")

#         io.log_info ("")

#         while True:
#             try:
#                 if choose_only_one:
#                     choosed_idxs = io.input_str("Which GPU index to choose?", best_device_indexes)
#                 else:
#                     choosed_idxs = io.input_str("Which GPU indexes to choose?", best_device_indexes)

#                 if allow_cpu and choosed_idxs.lower() == "cpu":
#                     choosed_idxs = []
#                     break

#                 choosed_idxs = [ int(x) for x in choosed_idxs.split(',') ]

#                 if choose_only_one:
#                     if len(choosed_idxs) == 1:
#                         break
#                 else:
#                     if all( [idx in all_devices_indexes for idx in choosed_idxs] ):
#                         break
#             except:
#                 pass
#         io.log_info ("")

#         return choosed_idxs

#     class DeviceConfig():
#         @staticmethod
#         def ask_choose_device(*args, **kwargs):
#             return nn.DeviceConfig.GPUIndexes( nn.ask_choose_device_idxs(*args,**kwargs) )
        
#         def __init__ (self, devices=None):
#             devices = devices or []

#             if not isinstance(devices, Devices):
#                 devices = Devices(devices)

#             self.devices = devices
#             self.cpu_only = len(devices) == 0

#         @staticmethod
#         def BestGPU():
#             devices = Devices.getDevices()
#             if len(devices) == 0:
#                 return nn.DeviceConfig.CPU()

#             return nn.DeviceConfig([devices.get_best_device()])

#         @staticmethod
#         def WorstGPU():
#             devices = Devices.getDevices()
#             if len(devices) == 0:
#                 return nn.DeviceConfig.CPU()

#             return nn.DeviceConfig([devices.get_worst_device()])

#         @staticmethod
#         def GPUIndexes(indexes):
#             if len(indexes) != 0:
#                 devices = Devices.getDevices().get_devices_from_index_list(indexes)
#             else:
#                 devices = []

#             return nn.DeviceConfig(devices)

#         @staticmethod
#         def CPU():
#             return nn.DeviceConfig([])

core/leras/nn.py
oo92/TorchFaceLab
import multiprocessing
from core.joblib import Subprocessor
import numpy as np

class CAInitializerSubprocessor(Subprocessor):
    @staticmethod
    def generate(shape, dtype=np.float32, eps_std=0.05):
        """
        Super fast implementation of Convolution Aware Initialization for 4D shapes
        Convolution Aware Initialization https://arxiv.org/abs/1702.06295
        """
        if len(shape) != 4:
            raise ValueError("only shape with rank 4 supported.")

        row, column, stack_size, filters_size = shape

        fan_in = stack_size * (row * column)

        kernel_shape = (row, column)

        kernel_fft_shape = np.fft.rfft2(np.zeros(kernel_shape)).shape

        basis_size = np.prod(kernel_fft_shape)
        if basis_size == 1:
            x = np.random.normal( 0.0, eps_std, (filters_size, stack_size, basis_size) )
        else:
            nbb = stack_size // basis_size + 1
            x = np.random.normal(0.0, 1.0, (filters_size, nbb, basis_size, basis_size))
            x = x + np.transpose(x, (0,1,3,2) ) * (1-np.eye(basis_size))
            u, _, v = np.linalg.svd(x)
            x = np.transpose(u, (0,1,3,2) )
            x = np.reshape(x, (filters_size, -1, basis_size) )
            x = x[:,:stack_size,:]

        x = np.reshape(x, ( (filters_size,stack_size,) + kernel_fft_shape ) )

        x = np.fft.irfft2( x, kernel_shape ) \
            + np.random.normal(0, eps_std, (filters_size,stack_size,)+kernel_shape)

        x = x * np.sqrt( (2/fan_in) / np.var(x) )
        x = np.transpose( x, (2, 3, 1, 0) )
        return x.astype(dtype)

    class Cli(Subprocessor.Cli):
        #override
        def process_data(self, data):
            idx, shape, dtype = data
            weights = CAInitializerSubprocessor.generate (shape, dtype)
            return idx, weights

    #override
    def __init__(self, data_list):
        self.data_list = data_list
        self.data_list_idxs = [*range(len(data_list))]
        self.result = [None]*len(data_list)
        super().__init__('CAInitializerSubprocessor', CAInitializerSubprocessor.Cli)

    #override
    def process_info_generator(self):
        for i in range( min(multiprocessing.cpu_count(), len(self.data_list)) ):
            yield 'CPU%d' % (i), {}, {}

    #override
    def get_data(self, host_dict):
        if len (self.data_list_idxs) > 0:
            idx = self.data_list_idxs.pop(0)
            shape, dtype = self.data_list[idx]
            return idx, shape, dtype
        return None

    #override
    def on_data_return (self, host_dict, data):
        self.data_list_idxs.insert(0, data)

    #override
    def on_result (self, host_dict, data, result):
        idx, weights = result
        self.result[idx] = weights

    #override
    def get_result(self):
        return self.result

core/leras/ops/CA.py
oo92/TorchFaceLab
import numpy as np
import torch

# from core.leras import nn

from .CA import CAInitializerSubprocessor

class initializers():
    class ca:
        def __call__(self, shape, dtype=None):
            if dtype:
                return torch.zeros(shape, dtype=dtype)
            else:
                return torch.zeros(shape)

        @staticmethod
        def generate_batch(data_list, eps_std=0.05):
            # list of (shape, np.dtype)
            return CAInitializerSubprocessor(data_list).run()

# nn.initializers = initializers

core/leras/ops/__init__.py
oo92/TorchFaceLab
import numpy as np
from core.leras import nn
from tensorflow.python.ops import control_flow_ops, state_ops

tf = nn.tf

class AdaBelief(nn.OptimizerBase):
    def __init__(self, lr=0.001, beta_1=0.9, beta_2=0.999, lr_dropout=1.0, lr_cos=0, clipnorm=0.0, name=None, **kwargs):
        super().__init__(name=name)

        if name is None:
            raise ValueError('name must be defined.')

        self.lr = lr
        self.beta_1 = beta_1
        self.beta_2 = beta_2
        self.lr_dropout = lr_dropout
        self.lr_cos = lr_cos
        self.clipnorm = clipnorm

        with tf.device('/CPU:0') :
            with tf.variable_scope(self.name):
                self.iterations = tf.Variable(0, dtype=tf.int64, name='iters')

        self.ms_dict = {}
        self.vs_dict = {}
        self.lr_rnds_dict = {}

    def get_weights(self):
        return [self.iterations] + list(self.ms_dict.values()) + list(self.vs_dict.values())

    def initialize_variables(self, trainable_weights, vars_on_cpu=True, lr_dropout_on_cpu=False):
        # Initialize here all trainable variables used in training
        e = tf.device('/CPU:0') if vars_on_cpu else None
        if e: e.__enter__()
        with tf.variable_scope(self.name):
            ms = { v.name : tf.get_variable ( f'ms_{v.name}'.replace(':','_'), v.shape, dtype=v.dtype, initializer=tf.initializers.constant(0.0), trainable=False) for v in trainable_weights }
            vs = { v.name : tf.get_variable ( f'vs_{v.name}'.replace(':','_'), v.shape, dtype=v.dtype, initializer=tf.initializers.constant(0.0), trainable=False) for v in trainable_weights }
            self.ms_dict.update (ms)
            self.vs_dict.update (vs)
            
            if self.lr_dropout != 1.0:
                e = tf.device('/CPU:0') if lr_dropout_on_cpu else None
                if e: e.__enter__()                    
                lr_rnds = [ nn.random_binomial( v.shape, p=self.lr_dropout, dtype=v.dtype) for v in trainable_weights ]
                if e: e.__exit__(None, None, None)                
                self.lr_rnds_dict.update ( { v.name : rnd for v,rnd in zip(trainable_weights,lr_rnds) } )
        if e: e.__exit__(None, None, None)

    def get_update_op(self, grads_vars):
        updates = []

        if self.clipnorm > 0.0:
            norm = tf.sqrt( sum([tf.reduce_sum(tf.square(tf.cast(g, tf.float32))) for g,v in grads_vars]))
        updates += [ state_ops.assign_add( self.iterations, 1) ]
        for i, (g,v) in enumerate(grads_vars):
            if self.clipnorm > 0.0:
                g = self.tf_clip_norm(g, self.clipnorm, tf.cast(norm, g.dtype) )

            ms = self.ms_dict[ v.name ]
            vs = self.vs_dict[ v.name ]
            
            m_t = self.beta_1*ms + (1.0-self.beta_1) * g
            v_t = self.beta_2*vs + (1.0-self.beta_2) * tf.square(g-m_t)

            lr = tf.constant(self.lr, g.dtype)
            if self.lr_cos != 0:
                lr *= (tf.cos(  tf.cast(self.iterations, g.dtype) * (2*3.1415926535/ float(self.lr_cos) )  ) + 1.0) / 2.0

            v_diff = - lr * m_t / (tf.sqrt(v_t) + np.finfo( g.dtype.as_numpy_dtype ).resolution )
            if self.lr_dropout != 1.0:
                lr_rnd = self.lr_rnds_dict[v.name]
                v_diff *= lr_rnd
            new_v = v + v_diff

            updates.append (state_ops.assign(ms, m_t))
            updates.append (state_ops.assign(vs, v_t))
            updates.append (state_ops.assign(v, new_v))

        return control_flow_ops.group ( *updates, name=self.name+'_updates')
nn.AdaBelief = AdaBelief

core/leras/optimizers/AdaBelief.py
oo92/TorchFaceLab
import copy
from core.leras import nn
tf = nn.tf

class OptimizerBase(nn.Saveable):
    def __init__(self, name=None):
        super().__init__(name=name)

    def tf_clip_norm(self, g, c, n):
        """Clip the gradient `g` if the L2 norm `n` exceeds `c`.
        # Arguments
            g: Tensor, the gradient tensor
            c: float >= 0. Gradients will be clipped
                when their L2 norm exceeds this value.
            n: Tensor, actual norm of `g`.
        # Returns
            Tensor, the gradient clipped if required.
        """
        if c <= 0:  # if clipnorm == 0 no need to add ops to the graph
            return g

        condition = n >= c
        then_expression = tf.scalar_mul(c / n, g)
        else_expression = g

        # saving the shape to avoid converting sparse tensor to dense
        if isinstance(then_expression, tf.Tensor):
            g_shape = copy.copy(then_expression.get_shape())
        elif isinstance(then_expression, tf.IndexedSlices):
            g_shape = copy.copy(then_expression.dense_shape)
        if condition.dtype != tf.bool:
            condition = tf.cast(condition, 'bool')
        g = tf.cond(condition,
                    lambda: then_expression,
                    lambda: else_expression)
        if isinstance(then_expression, tf.Tensor):
            g.set_shape(g_shape)
        elif isinstance(then_expression, tf.IndexedSlices):
            g._dense_shape = g_shape

        return g
nn.OptimizerBase = OptimizerBase

core/leras/optimizers/OptimizerBase.py
oo92/TorchFaceLab
import numpy as np
import torch
from torch import nn as t_nn

class RMSprop(t_nn.Module):
    def __init__(self, trainable_params, lr=0.001, rho=0.9, lr_dropout=1.0, lr_cos=0, clipnorm=0.0, name=None):
        super(RMSprop, self).__init__()

        if name is None:
            raise ValueError('name must be defined.')

        self.lr_dropout = lr_dropout
        self.lr_cos = lr_cos
        self.lr = lr
        self.rho = rho
        self.clipnorm = clipnorm

        self.iterations = torch.tensor(0, dtype=torch.int64)
        self.accumulators_dict = {}
        self.lr_rnds_dict = {}
        
        for name, param in trainable_params:
            self.accumulators_dict[name] = torch.zeros_like(param.data)
            if self.lr_dropout != 1.0:
                self.lr_rnds_dict[name] = torch.bernoulli(torch.ones_like(param.data) * self.lr_dropout)

    def step(self, grads_vars):
        updates = []

        if self.clipnorm > 0.0:
            norm = torch.sqrt(sum([(g ** 2).sum() for g, v in grads_vars]))
        
        self.iterations += 1
        
        for i, (name, g, v) in enumerate(grads_vars):
            if self.clipnorm > 0.0:
                g = torch.nn.utils.clip_by_norm(g, self.clipnorm, norm_type=2) * (norm if norm.is_nonzero() else 1.0)

            a = self.accumulators_dict[name]

            new_a = self.rho * a + (1. - self.rho) * g ** 2

            lr = torch.tensor(self.lr, dtype=g.dtype)
            if self.lr_cos != 0:
                lr *= (torch.cos(self.iterations.to(dtype=g.dtype) * (2*np.pi/ float(self.lr_cos))) + 1.0) / 2.0

            v_diff = - lr * g / (torch.sqrt(new_a) + np.finfo(g.dtype).eps)
            
            if self.lr_dropout != 1.0:
                lr_rnd = self.lr_rnds_dict[name]
                v_diff *= lr_rnd
            v.data += v_diff

            self.accumulators_dict[name] = new_a

core/leras/optimizers/RMSprop.py
oo92/TorchFaceLab
from .OptimizerBase import *
from .RMSprop import *
from .AdaBelief import *
core/leras/optimizers/__init__.py
oo92/TorchFaceLab
import math

import cv2
import numpy as np
import numpy.linalg as npla

from .umeyama import umeyama


def get_power_of_two(x):
    i = 0
    while (1 << i) < x:
        i += 1
    return i

def rotationMatrixToEulerAngles(R) :
    sy = math.sqrt(R[0,0] * R[0,0] +  R[1,0] * R[1,0])
    singular = sy < 1e-6
    if  not singular :
        x = math.atan2(R[2,1] , R[2,2])
        y = math.atan2(-R[2,0], sy)
        z = math.atan2(R[1,0], R[0,0])
    else :
        x = math.atan2(-R[1,2], R[1,1])
        y = math.atan2(-R[2,0], sy)
        z = 0
    return np.array([x, y, z])

def polygon_area(x,y):
    return 0.5*np.abs(np.dot(x,np.roll(y,1))-np.dot(y,np.roll(x,1)))

def rotate_point(origin, point, deg):
    """
    Rotate a point counterclockwise by a given angle around a given origin.

    The angle should be given in radians.
    """
    ox, oy = origin
    px, py = point

    rad = deg * math.pi / 180.0
    qx = ox + math.cos(rad) * (px - ox) - math.sin(rad) * (py - oy)
    qy = oy + math.sin(rad) * (px - ox) + math.cos(rad) * (py - oy)
    return np.float32([qx, qy])
    
def transform_points(points, mat, invert=False):
    if invert:
        mat = cv2.invertAffineTransform (mat)
    points = np.expand_dims(points, axis=1)
    points = cv2.transform(points, mat, points.shape)
    points = np.squeeze(points)
    return points

    
def transform_mat(mat, res, tx, ty, rotation, scale):
    """
    transform mat in local space of res
    scale -> translate -> rotate
    
        tx,ty       float
        rotation    int degrees
        scale       float
    """
    
    
    lt, rt, lb, ct = transform_points (  np.float32([(0,0),(res,0),(0,res),(res / 2, res/2) ]),mat, True)
    
    hor_v = (rt-lt).astype(np.float32)
    hor_size = npla.norm(hor_v)
    hor_v /= hor_size
    
    ver_v = (lb-lt).astype(np.float32)
    ver_size = npla.norm(ver_v)
    ver_v /= ver_size
    
    bt_diag_vec = (rt-ct).astype(np.float32)
    half_diag_len = npla.norm(bt_diag_vec)
    bt_diag_vec /= half_diag_len
    
    tb_diag_vec = np.float32( [ -bt_diag_vec[1], bt_diag_vec[0] ]  )

    rt = ct + bt_diag_vec*half_diag_len*scale 
    lb = ct - bt_diag_vec*half_diag_len*scale
    lt = ct - tb_diag_vec*half_diag_len*scale
    
    rt[0] += tx*hor_size
    lb[0] += tx*hor_size
    lt[0] += tx*hor_size
    rt[1] += ty*ver_size
    lb[1] += ty*ver_size
    lt[1] += ty*ver_size
    
    rt = rotate_point(ct, rt, rotation)
    lb = rotate_point(ct, lb, rotation)
    lt = rotate_point(ct, lt, rotation)
    
    return cv2.getAffineTransform( np.float32([lt, rt, lb]), np.float32([ [0,0], [res,0], [0,res] ]) )

core/mathlib/__init__.py
oo92/TorchFaceLab
import multiprocessing
import pickle
import struct
from core.joblib import Subprocessor

class MPSharedList():
    """
    Provides read-only pickled list of constant objects via shared memory aka 'multiprocessing.Array'
    Thus no 4GB limit for subprocesses.

    supports list concat via + or sum()
    """

    def __init__(self, obj_list):
        if obj_list is None:
            self.obj_counts    = None
            self.table_offsets = None
            self.data_offsets  = None
            self.sh_bs         = None
        else:
            obj_count, table_offset, data_offset, sh_b = MPSharedList.bake_data(obj_list)

            self.obj_counts    = [obj_count]
            self.table_offsets = [table_offset]
            self.data_offsets  = [data_offset]
            self.sh_bs         = [sh_b]

    def __add__(self, o):
        if isinstance(o, MPSharedList):
            m = MPSharedList(None)
            m.obj_counts    = self.obj_counts    + o.obj_counts
            m.table_offsets = self.table_offsets + o.table_offsets
            m.data_offsets  = self.data_offsets  + o.data_offsets
            m.sh_bs         = self.sh_bs         + o.sh_bs
            return m
        elif isinstance(o, int):
            return self
        else:
            raise ValueError(f"MPSharedList object of class {o.__class__} is not supported for __add__ operator.")

    def __radd__(self, o):
        return self+o

    def __len__(self):
        return sum(self.obj_counts)

    def __getitem__(self, key):
        obj_count = sum(self.obj_counts)
        if key < 0:
            key = obj_count+key
        if key < 0 or key >= obj_count:
            raise ValueError("out of range")

        for i in range(len(self.obj_counts)):

            if key < self.obj_counts[i]:
                table_offset = self.table_offsets[i]
                data_offset = self.data_offsets[i]
                sh_b = self.sh_bs[i]
                break
            key -= self.obj_counts[i]

        sh_b = memoryview(sh_b).cast('B')

        offset_start, offset_end = struct.unpack('<QQ', sh_b[ table_offset + key*8     : table_offset + (key+2)*8].tobytes() )

        return pickle.loads( sh_b[ data_offset + offset_start : data_offset + offset_end ].tobytes() )

    def __iter__(self):
        for i in range(self.__len__()):
            yield self.__getitem__(i)

    @staticmethod
    def bake_data(obj_list):
        if not isinstance(obj_list, list):
            raise ValueError("MPSharedList: obj_list should be list type.")

        obj_count = len(obj_list)

        if obj_count != 0:
            obj_pickled_ar = [pickle.dumps(o, 4) for o in obj_list]

            table_offset = 0
            table_size   = (obj_count+1)*8
            data_offset  = table_offset + table_size
            data_size    = sum([len(x) for x in obj_pickled_ar])

            sh_b = multiprocessing.RawArray('B', table_size + data_size)
            #sh_b[0:8] = struct.pack('<Q', obj_count)
            sh_b_view = memoryview(sh_b).cast('B')

            offset = 0

            sh_b_table = bytes()
            offsets = []

            offset = 0
            for i in range(obj_count):
                offsets.append(offset)
                offset += len(obj_pickled_ar[i])
            offsets.append(offset)

            sh_b_view[table_offset:table_offset+table_size] = struct.pack( '<'+'Q'*len(offsets), *offsets )

            for i, obj_pickled in enumerate(obj_pickled_ar):
                offset = data_offset+offsets[i]
                sh_b_view[offset:offset+len(obj_pickled)] = obj_pickled_ar[i]

            return obj_count, table_offset, data_offset, sh_b
        return 0, 0, 0, None


core/mplib/MPSharedList.py
oo92/TorchFaceLab
from .MPSharedList import MPSharedList
import multiprocessing
import threading
import time

import numpy as np


class IndexHost():
    """
    Provides random shuffled indexes for multiprocesses
    """
    def __init__(self, indexes_count, rnd_seed=None):
        self.sq = multiprocessing.Queue()
        self.cqs = []
        self.clis = []
        self.thread = threading.Thread(target=self.host_thread, args=(indexes_count,rnd_seed) )
        self.thread.daemon = True
        self.thread.start()

    def host_thread(self, indexes_count, rnd_seed):
        rnd_state = np.random.RandomState(rnd_seed) if rnd_seed is not None else np.random

        idxs = [*range(indexes_count)]
        shuffle_idxs = []
        sq = self.sq

        while True:
            while not sq.empty():
                obj = sq.get()
                cq_id, count = obj[0], obj[1]

                result = []
                for i in range(count):
                    if len(shuffle_idxs) == 0:
                        shuffle_idxs = idxs.copy()
                        rnd_state.shuffle(shuffle_idxs)
                    result.append(shuffle_idxs.pop())
                self.cqs[cq_id].put (result)

            time.sleep(0.001)

    def create_cli(self):
        cq = multiprocessing.Queue()
        self.cqs.append ( cq )
        cq_id = len(self.cqs)-1
        return IndexHost.Cli(self.sq, cq, cq_id)

    # disable pickling
    def __getstate__(self):
        return dict()
    def __setstate__(self, d):
        self.__dict__.update(d)

    class Cli():
        def __init__(self, sq, cq, cq_id):
            self.sq = sq
            self.cq = cq
            self.cq_id = cq_id

        def multi_get(self, count):
            self.sq.put ( (self.cq_id,count) )

            while True:
                if not self.cq.empty():
                    return self.cq.get()
                time.sleep(0.001)

class Index2DHost():
    """
    Provides random shuffled indexes for multiprocesses
    """
    def __init__(self, indexes2D):
        self.sq = multiprocessing.Queue()
        self.cqs = []
        self.clis = []
        self.thread = threading.Thread(target=self.host_thread, args=(indexes2D,) )
        self.thread.daemon = True
        self.thread.start()

    def host_thread(self, indexes2D):
        indexes2D_len = len(indexes2D)

        idxs = [*range(indexes2D_len)]
        idxs_2D = [None]*indexes2D_len
        shuffle_idxs = []
        shuffle_idxs_2D = [None]*indexes2D_len
        for i in range(indexes2D_len):
            idxs_2D[i] = [*range(len(indexes2D[i]))]
            shuffle_idxs_2D[i] = []

        #print(idxs)
        #print(idxs_2D)
        sq = self.sq

        while True:
            while not sq.empty():
                obj = sq.get()
                cq_id, count = obj[0], obj[1]

                result = []
                for i in range(count):
                    if len(shuffle_idxs) == 0:
                        shuffle_idxs = idxs.copy()
                        np.random.shuffle(shuffle_idxs)

                    idx_1D = shuffle_idxs.pop()
                    
                    #print(f'idx_1D = {idx_1D}, len(shuffle_idxs_2D[idx_1D])= {len(shuffle_idxs_2D[idx_1D])}')
                    
                    if len(shuffle_idxs_2D[idx_1D]) == 0:
                        shuffle_idxs_2D[idx_1D] = idxs_2D[idx_1D].copy()
                        #print(f'new shuffle_idxs_2d for {idx_1D} = { shuffle_idxs_2D[idx_1D] }')
                        
                        #print(f'len(shuffle_idxs_2D[idx_1D])= {len(shuffle_idxs_2D[idx_1D])}')
                    
                        np.random.shuffle( shuffle_idxs_2D[idx_1D] )

                    idx_2D = shuffle_idxs_2D[idx_1D].pop()
                    
                    #print(f'len(shuffle_idxs_2D[idx_1D])= {len(shuffle_idxs_2D[idx_1D])}')
                    
                    #print(f'idx_2D = {idx_2D}')
                    

                    result.append( indexes2D[idx_1D][idx_2D])

                self.cqs[cq_id].put (result)

            time.sleep(0.001)

    def create_cli(self):
        cq = multiprocessing.Queue()
        self.cqs.append ( cq )
        cq_id = len(self.cqs)-1
        return Index2DHost.Cli(self.sq, cq, cq_id)

    # disable pickling
    def __getstate__(self):
        return dict()
    def __setstate__(self, d):
        self.__dict__.update(d)

    class Cli():
        def __init__(self, sq, cq, cq_id):
            self.sq = sq
            self.cq = cq
            self.cq_id = cq_id

        def multi_get(self, count):
            self.sq.put ( (self.cq_id,count) )

            while True:
                if not self.cq.empty():
                    return self.cq.get()
                time.sleep(0.001)

class ListHost():
    def __init__(self, list_):
        self.sq = multiprocessing.Queue()
        self.cqs = []
        self.clis = []
        self.m_list = list_
        self.thread = threading.Thread(target=self.host_thread)
        self.thread.daemon = True
        self.thread.start()

    def host_thread(self):
        sq = self.sq
        while True:
            while not sq.empty():
                obj = sq.get()
                cq_id, cmd = obj[0], obj[1]

                if cmd == 0:
                    self.cqs[cq_id].put ( len(self.m_list) )
                elif cmd == 1:
                    idx = obj[2]
                    item = self.m_list[idx ]
                    self.cqs[cq_id].put ( item )
                elif cmd == 2:
                    result = []
                    for item in obj[2]:
                        result.append ( self.m_list[item] )
                    self.cqs[cq_id].put ( result )
                elif cmd == 3:
                    self.m_list.insert(obj[2], obj[3])
                elif cmd == 4:
                    self.m_list.append(obj[2])
                elif cmd == 5:
                    self.m_list.extend(obj[2])

            time.sleep(0.005)

    def create_cli(self):
        cq = multiprocessing.Queue()
        self.cqs.append ( cq )
        cq_id = len(self.cqs)-1
        return ListHost.Cli(self.sq, cq, cq_id)

    def get_list(self):
        return self.list_

    # disable pickling
    def __getstate__(self):
        return dict()
    def __setstate__(self, d):
        self.__dict__.update(d)

    class Cli():
        def __init__(self, sq, cq, cq_id):
            self.sq = sq
            self.cq = cq
            self.cq_id = cq_id

        def __len__(self):
            self.sq.put ( (self.cq_id,0) )

            while True:
                if not self.cq.empty():
                    return self.cq.get()
                time.sleep(0.001)

        def __getitem__(self, key):
            self.sq.put ( (self.cq_id,1,key) )

            while True:
                if not self.cq.empty():
                    return self.cq.get()
                time.sleep(0.001)

        def multi_get(self, keys):
            self.sq.put ( (self.cq_id,2,keys) )

            while True:
                if not self.cq.empty():
                    return self.cq.get()
                time.sleep(0.001)

        def insert(self, index, item):
            self.sq.put ( (self.cq_id,3,index,item) )

        def append(self, item):
            self.sq.put ( (self.cq_id,4,item) )

        def extend(self, items):
            self.sq.put ( (self.cq_id,5,items) )



class DictHost():
    def __init__(self, d, num_users):
        self.sqs = [ multiprocessing.Queue() for _ in range(num_users) ]
        self.cqs = [ multiprocessing.Queue() for _ in range(num_users) ]

        self.thread = threading.Thread(target=self.host_thread, args=(d,) )
        self.thread.daemon = True
        self.thread.start()

        self.clis = [ DictHostCli(sq,cq) for sq, cq in zip(self.sqs, self.cqs) ]

    def host_thread(self, d):
        while True:
            for sq, cq in zip(self.sqs, self.cqs):
                if not sq.empty():
                    obj = sq.get()
                    cmd = obj[0]
                    if cmd == 0:
                        cq.put (d[ obj[1] ])
                    elif cmd == 1:
                        cq.put ( list(d.keys()) )

            time.sleep(0.005)


    def get_cli(self, n_user):
        return self.clis[n_user]

    # disable pickling
    def __getstate__(self):
        return dict()
    def __setstate__(self, d):
        self.__dict__.update(d)

class DictHostCli():
    def __init__(self, sq, cq):
        self.sq = sq
        self.cq = cq

    def __getitem__(self, key):
        self.sq.put ( (0,key) )

        while True:
            if not self.cq.empty():
                return self.cq.get()
            time.sleep(0.001)

    def keys(self):
        self.sq.put ( (1,) )
        while True:
            if not self.cq.empty():
                return self.cq.get()
            time.sleep(0.001)

core/mplib/__init__.py
oo92/TorchFaceLab
import os
import sys

if sys.platform[0:3] == 'win':
    from ctypes import windll
    from ctypes import wintypes

def set_process_lowest_prio():
    try:
        if sys.platform[0:3] == 'win':
            GetCurrentProcess = windll.kernel32.GetCurrentProcess
            GetCurrentProcess.restype = wintypes.HANDLE
            SetPriorityClass = windll.kernel32.SetPriorityClass
            SetPriorityClass.argtypes = (wintypes.HANDLE, wintypes.DWORD)
            SetPriorityClass ( GetCurrentProcess(), 0x00000040 )
        elif 'darwin' in sys.platform:
            os.nice(10)
        elif 'linux' in sys.platform:
            os.nice(20)
    except:
        print("Unable to set lowest process priority")

def set_process_dpi_aware():
    if sys.platform[0:3] == 'win':
        windll.user32.SetProcessDPIAware(True)

def get_screen_size():
    if sys.platform[0:3] == 'win':
        user32 = windll.user32
        return user32.GetSystemMetrics(0), user32.GetSystemMetrics(1)
    elif 'darwin' in sys.platform:
        pass
    elif 'linux' in sys.platform:
        pass
        
    return (1366, 768)
        
core/osex.py
oo92/TorchFaceLab
from pathlib import Path
from os import scandir

image_extensions = [".jpg", ".jpeg", ".png", ".tif", ".tiff"]

def write_bytes_safe(p, bytes_data):
    """
    writes to .tmp first and then rename to target filename
    """
    p_tmp = p.parent / (p.name + '.tmp')
    p_tmp.write_bytes(bytes_data)
    if p.exists():
        p.unlink()
    p_tmp.rename (p)

def scantree(path):
    """Recursively yield DirEntry objects for given directory."""
    for entry in scandir(path):
        if entry.is_dir(follow_symlinks=False):
            yield from scantree(entry.path)  # see below for Python 2.x
        else:
            yield entry

def get_image_paths(dir_path, image_extensions=image_extensions, subdirs=False, return_Path_class=False):
    dir_path = Path (dir_path)

    result = []
    if dir_path.exists():

        if subdirs:
            gen = scantree(str(dir_path))
        else:
            gen = scandir(str(dir_path))

        for x in list(gen):
            if any([x.name.lower().endswith(ext) for ext in image_extensions]):
                result.append( x.path if not return_Path_class else Path(x.path) )
    return sorted(result)

def get_image_unique_filestem_paths(dir_path, verbose_print_func=None):
    result = get_image_paths(dir_path)
    result_dup = set()

    for f in result[:]:
        f_stem = Path(f).stem
        if f_stem in result_dup:
            result.remove(f)
            if verbose_print_func is not None:
                verbose_print_func ("Duplicate filenames are not allowed, skipping: %s" % Path(f).name )
            continue
        result_dup.add(f_stem)

    return sorted(result)

def get_paths(dir_path):
    dir_path = Path (dir_path)

    if dir_path.exists():
        return [ Path(x) for x in sorted([ x.path for x in list(scandir(str(dir_path))) ]) ]
    else:
        return []
        
def get_file_paths(dir_path):
    dir_path = Path (dir_path)

    if dir_path.exists():
        return [ Path(x) for x in sorted([ x.path for x in list(scandir(str(dir_path))) if x.is_file() ]) ]
    else:
        return []

def get_all_dir_names (dir_path):
    dir_path = Path (dir_path)

    if dir_path.exists():
        return sorted([ x.name for x in list(scandir(str(dir_path))) if x.is_dir() ])
    else:
        return []

def get_all_dir_names_startswith (dir_path, startswith):
    dir_path = Path (dir_path)
    startswith = startswith.lower()

    result = []
    if dir_path.exists():
        for x in list(scandir(str(dir_path))):
            if x.name.lower().startswith(startswith):
                result.append ( x.name[len(startswith):] )
    return sorted(result)

def get_first_file_by_stem (dir_path, stem, exts=None):
    dir_path = Path (dir_path)
    stem = stem.lower()

    if dir_path.exists():
        for x in sorted(list(scandir(str(dir_path))), key=lambda x: x.name):
            if not x.is_file():
                continue
            xp = Path(x.path)
            if xp.stem.lower() == stem and (exts is None or xp.suffix.lower() in exts):
                return xp

    return None

def move_all_files (src_dir_path, dst_dir_path):
    paths = get_file_paths(src_dir_path)
    for p in paths:
        p = Path(p)
        p.rename ( Path(dst_dir_path) / p.name )

def delete_all_files (dir_path):
    paths = get_file_paths(dir_path)
    for p in paths:
        p = Path(p)
        p.unlink()

core/pathex.py
oo92/TorchFaceLab
import multiprocessing
import sys
import time
import traceback

from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *

from core.interact import interact as io

from .qtex import *

class QSubprocessor(object):
    """
    
    """

    class Cli(object):
        def __init__ ( self, client_dict ):
            s2c = multiprocessing.Queue()
            c2s = multiprocessing.Queue()
            self.p = multiprocessing.Process(target=self._subprocess_run, args=(client_dict,s2c,c2s) )
            self.s2c = s2c
            self.c2s = c2s
            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, s2c, c2s):
            self.c2s = c2s
            data = None
            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:
                c2s.put ( {'op': 'error', 'data' : data} )
                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.close()
            s2c.close()
            self.c2s = None

        # disable pickling
        def __getstate__(self):
            return dict()
        def __setstate__(self, d):
            self.__dict__.update(d)

    #overridable
    def __init__(self, name, SubprocessorCli_class, no_response_time_sec = 0, io_loop_sleep_time=0.005):
        if not issubclass(SubprocessorCli_class, QSubprocessor.Cli):
            raise ValueError("SubprocessorCli_class must be subclass of QSubprocessor.Cli")

        self.name = name
        self.SubprocessorCli_class = SubprocessorCli_class
        self.no_response_time_sec = no_response_time_sec
        self.io_loop_sleep_time = io_loop_sleep_time

        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 = 0
                cli.sent_data = None
                cli.name = name
                cli.host_dict = host_dict

                self.clis.append (cli)
            except:
                raise Exception (f"Unable to start subprocess {name}. Error: {traceback.format_exc()}")

        if len(self.clis) == 0:
            raise Exception ("Unable to start QSubprocessor '%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 == 'log_info':
                        io.log_info(obj['msg'])
                    elif op == 'log_err':
                        io.log_err(obj['msg'])
                    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()
                
        self.q_timer = QTimer()
        self.q_timer.timeout.connect(self.tick)
        self.q_timer.start(5)
        
    #overridable
    def process_info_generator(self):
        #yield per process (name, host_dict, client_dict)
        for i in range(min(multiprocessing.cpu_count(), 8) ):
            yield 'CPU%d' % (i), {}, {}

    #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

    def tick(self):
        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 == 1:
                if cli.sent_time != 0 and self.no_response_time_sec != 0 and (time.time() - cli.sent_time) > self.no_response_time_sec:
                    #subprocess busy too long
                    io.log_info ( '%s doesnt response, terminating it.' % (cli.name) )
                    self.on_data_return (cli.host_dict, cli.sent_data )
                    cli.kill()
                    self.clis.remove(cli)

        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

        if all ([cli.state == 0 for cli in self.clis]):
            #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 (time.time() - cli.sent_time) > 30:
                        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
            self.q_timer.stop()
            self.q_timer = None
            self.on_clients_finalized()
            

core/qtex/QSubprocessor.py
oo92/TorchFaceLab
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *

from localization import StringsDB
from .QXMainWindow import *

class QXIconButton(QPushButton):
    """
    Custom Icon button that works through keyEvent system, without shortcut of QAction
    works only with QXMainWindow as global window class
    currently works only with one-key shortcut
    """

    def __init__(self, icon, 
                    tooltip=None, 
                    shortcut=None,                    
                    click_func=None,                  
                    first_repeat_delay=300,
                    repeat_delay=20,
                    ):

        super().__init__(icon, "")

        self.setIcon(icon)
        
        if shortcut is not None:
            tooltip = f"{tooltip} ( {StringsDB['S_HOT_KEY'] }: {shortcut} )"
        
        self.setToolTip(tooltip)
            
        
        self.seq = QKeySequence(shortcut) if shortcut is not None else None
        
        QXMainWindow.inst.add_keyPressEvent_listener ( self.on_keyPressEvent )
        QXMainWindow.inst.add_keyReleaseEvent_listener ( self.on_keyReleaseEvent )
        
        self.click_func = click_func
        self.first_repeat_delay = first_repeat_delay
        self.repeat_delay = repeat_delay
        self.repeat_timer = None
        
        self.op_device = None
        
        self.pressed.connect( lambda : self.action(is_pressed=True)  )
        self.released.connect( lambda : self.action(is_pressed=False)  )
        
    def action(self, is_pressed=None, op_device=None):
        if self.click_func is None:
            return

        if is_pressed is not None:
            if is_pressed:
                if self.repeat_timer is None:
                    self.click_func()
                    self.repeat_timer = QTimer()
                    self.repeat_timer.timeout.connect(self.action)
                    self.repeat_timer.start(self.first_repeat_delay)
            else:
                if self.repeat_timer is not None:
                    self.repeat_timer.stop()
                    self.repeat_timer = None
        else:
            self.click_func()
            if self.repeat_timer is not None:
                self.repeat_timer.setInterval(self.repeat_delay)
        
    def on_keyPressEvent(self, ev):              
        key = ev.nativeVirtualKey()
        if ev.isAutoRepeat():
            return
            
        if self.seq is not None:
            if key == self.seq[0]:
                self.action(is_pressed=True)

    def on_keyReleaseEvent(self, ev):
        key = ev.nativeVirtualKey()
        if ev.isAutoRepeat():
            return
        if self.seq is not None:
            if key == self.seq[0]:
                self.action(is_pressed=False)

core/qtex/QXIconButton.py
oo92/TorchFaceLab
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *

class QXMainWindow(QWidget):
    """
    Custom mainwindow class that provides global single instance and event listeners
    """
    inst = None
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)        
        if QXMainWindow.inst is not None:
            raise Exception("QXMainWindow can only be one.")        
        QXMainWindow.inst = self
        
        self.keyPressEvent_listeners = []
        self.keyReleaseEvent_listeners = []
        self.setFocusPolicy(Qt.WheelFocus)
        
    def add_keyPressEvent_listener(self, func):
        self.keyPressEvent_listeners.append (func)
        
    def add_keyReleaseEvent_listener(self, func):
        self.keyReleaseEvent_listeners.append (func)
        
    def keyPressEvent(self, ev):
        super().keyPressEvent(ev)        
        for func in self.keyPressEvent_listeners:
            func(ev)
            
    def keyReleaseEvent(self, ev):
        super().keyReleaseEvent(ev)        
        for func in self.keyReleaseEvent_listeners:
            func(ev)
core/qtex/QXMainWindow.py
oo92/TorchFaceLab
from .qtex import *
from .QSubprocessor import *
from .QXIconButton import *
core/qtex/__init__.py
oo92/TorchFaceLab
import numpy as np
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *
from localization import StringsDB

from .QXMainWindow import *
            
                
class QActionEx(QAction):
    def __init__(self, icon, text, shortcut=None, trigger_func=None, shortcut_in_tooltip=False, is_checkable=False, is_auto_repeat=False ):
        super().__init__(icon, text)
        if shortcut is not None:
            self.setShortcut(shortcut)
            if shortcut_in_tooltip:
                
                self.setToolTip( f"{text} ( {StringsDB['S_HOT_KEY'] }: {shortcut} )")
                
        if trigger_func is not None:
            self.triggered.connect(trigger_func)
        if is_checkable:            
            self.setCheckable(True)
        self.setAutoRepeat(is_auto_repeat)
            
def QImage_from_np(img):
    if img.dtype != np.uint8:
        raise ValueError("img should be in np.uint8 format")
        
    h,w,c = img.shape
    if c == 1:
        fmt = QImage.Format_Grayscale8
    elif c == 3:
        fmt = QImage.Format_BGR888
    elif c == 4:
        fmt = QImage.Format_ARGB32
    else:
      raise ValueError("unsupported channel count")  
    
    return QImage(img.data, w, h, c*w, fmt )
        
def QImage_to_np(q_img, fmt=QImage.Format_BGR888):
    q_img = q_img.convertToFormat(fmt)

    width = q_img.width()
    height = q_img.height()
    
    b = q_img.constBits()
    b.setsize(height * width * 3)
    arr = np.frombuffer(b, np.uint8).reshape((height, width, 3))
    return arr#[::-1]
        
def QPixmap_from_np(img):    
    return QPixmap.fromImage(QImage_from_np(img))
    
def QPoint_from_np(n):
    return QPoint(*n.astype(np.int))
    
def QPoint_to_np(q):
    return np.int32( [q.x(), q.y()] )
    
def QSize_to_np(q):
    return np.int32( [q.width(), q.height()] )
    
class QDarkPalette(QPalette):
    def __init__(self):
        super().__init__()
        text_color = QColor(200,200,200)
        self.setColor(QPalette.Window, QColor(53, 53, 53))
        self.setColor(QPalette.WindowText, text_color )
        self.setColor(QPalette.Base, QColor(25, 25, 25))
        self.setColor(QPalette.AlternateBase, QColor(53, 53, 53))
        self.setColor(QPalette.ToolTipBase, text_color )
        self.setColor(QPalette.ToolTipText, text_color )
        self.setColor(QPalette.Text, text_color ) 
        self.setColor(QPalette.Button, QColor(53, 53, 53))
        self.setColor(QPalette.ButtonText, Qt.white)
        self.setColor(QPalette.BrightText, Qt.red)
        self.setColor(QPalette.Link, QColor(42, 130, 218))
        self.setColor(QPalette.Highlight, QColor(42, 130, 218))
        self.setColor(QPalette.HighlightedText, Qt.black)
core/qtex/qtex.py
oo92/TorchFaceLab
import numpy as np

def random_normal( size=(1,), trunc_val = 2.5, rnd_state=None ):
    if rnd_state is None:
        rnd_state = np.random
    len = np.array(size).prod()
    result = np.empty ( (len,) , dtype=np.float32)

    for i in range (len):
        while True:
            x = rnd_state.normal()
            if x >= -trunc_val and x <= trunc_val:
                break
        result[i] = (x / trunc_val)

    return result.reshape ( size )
core/randomex.py
oo92/TorchFaceLab
import os
import sys

class suppress_stdout_stderr(object):
    def __enter__(self):
        self.outnull_file = open(os.devnull, 'w')
        self.errnull_file = open(os.devnull, 'w')

        self.old_stdout_fileno_undup    = sys.stdout.fileno()
        self.old_stderr_fileno_undup    = sys.stderr.fileno()

        self.old_stdout_fileno = os.dup ( sys.stdout.fileno() )
        self.old_stderr_fileno = os.dup ( sys.stderr.fileno() )

        self.old_stdout = sys.stdout
        self.old_stderr = sys.stderr

        os.dup2 ( self.outnull_file.fileno(), self.old_stdout_fileno_undup )
        os.dup2 ( self.errnull_file.fileno(), self.old_stderr_fileno_undup )

        sys.stdout = self.outnull_file
        sys.stderr = self.errnull_file
        return self

    def __exit__(self, *_):
        sys.stdout = self.old_stdout
        sys.stderr = self.old_stderr

        os.dup2 ( self.old_stdout_fileno, self.old_stdout_fileno_undup )
        os.dup2 ( self.old_stderr_fileno, self.old_stderr_fileno_undup )

        os.close ( self.old_stdout_fileno )
        os.close ( self.old_stderr_fileno )

        self.outnull_file.close()
        self.errnull_file.close()

core/stdex.py
oo92/TorchFaceLab
import struct

def struct_unpack(data, counter, fmt):
    fmt_size = struct.calcsize(fmt)
    return (counter+fmt_size,) + struct.unpack (fmt, data[counter:counter+fmt_size])

core/structex.py
oo92/TorchFaceLab
import torch, cv2
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
from pathlib import Path
from facelib import FaceType, LandmarksProcessor


class FANExtractor(object):
    def __init__(self, landmarks_3D=False, place_model_on_cpu=False):

        def can_squeeze(t):
            shape_set = set(t.shape)
            if len(shape_set) == 2 and 1 in shape_set:
                return True
            return False

        def reshape_tensor(k, t):
            if k == 'fc7.weight':
                return t.reshape([1024, 1024, 1, 1])
            if k in ['conv4_3_norm.weight', 'conv5_3_norm.weight']:
                return t.reshape([1, 512, 1, 1])
            if k == 'conv3_3_norm.weight':
                return t.reshape([1, 256, 1, 1])
            if can_squeeze(t):
                return t.squeeze()
            else:
                return torch.permute(t, [3, 2, 1, 0])

        # def rename_key(key):
        #     new_key = key.split(":")[0] # discard :0 and similar
        #     key_elements = new_key.split("/")
        #     new_key = ".".join(key_elements) # replace every / with .
        #     return new_key
        
        model_path = Path(__file__).parent / ("2DFAN.npy" if not landmarks_3D else "3DFAN.npy")
        if not model_path.exists():
            raise Exception("Unable to load FANExtractor model")

        self.device = torch.device('cpu') if place_model_on_cpu else torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        class ConvBlock(nn.Module):
            def __init__(self, in_planes, out_planes):
                super(ConvBlock, self).__init__()

                self.bn1 = nn.BatchNorm2d(in_planes)
                self.conv1 = nn.Conv2d(in_planes, out_planes // 2, kernel_size=3, stride=1, padding=1, bias=False)

                self.bn2 = nn.BatchNorm2d(out_planes // 2)
                self.conv2 = nn.Conv2d(out_planes // 2, out_planes // 4, kernel_size=3, stride=1, padding=1, bias=False)

                self.bn3 = nn.BatchNorm2d(out_planes // 4)
                self.conv3 = nn.Conv2d(out_planes // 4, out_planes // 4, kernel_size=3, stride=1, padding=1, bias=False)

                self.downsample = None
                if in_planes != out_planes:
                    self.downsample = nn.Sequential(
                        nn.BatchNorm2d(in_planes),
                        nn.ReLU(),
                        nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=1, bias=False)
                    )

            def forward(self, x):
                residual = x

                out1 = F.relu(self.bn1(x))
                out1 = self.conv1(out1)

                out2 = F.relu(self.bn2(out1))
                out2 = self.conv2(out2)

                out3 = F.relu(self.bn3(out2))
                out3 = self.conv3(out3)

                out = torch.cat([out1, out2, out3], dim=1)

                if self.downsample:
                    residual = self.downsample(x)

                out += residual
                return out


        class HourGlass(nn.Module):
            def __init__(self, in_planes, depth):
                super(HourGlass, self).__init__()
                self.b1 = ConvBlock(in_planes, 256)
                self.b2 = ConvBlock(in_planes, 256)

                self.b2_plus = ConvBlock(256, 256) if depth == 1 else HourGlass(256, depth - 1)

                self.b3 = ConvBlock(256, 256)

            def forward(self, x):
                up1 = self.b1(x)

                low = F.avg_pool2d(x, kernel_size=2, stride=2)
                low1 = self.b2(low)

                low2 = self.b2_plus(low1)
                low3 = self.b3(low2)

                up2 = F.interpolate(low3, scale_factor=2, mode='bilinear', align_corners=True)

                return up1 + up2


        class FAN(nn.Module):
            def __init__(self):
                super(FAN, self).__init__()
                self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3)
                self.bn1 = nn.BatchNorm2d(64)

                self.conv2 = ConvBlock(64, 128)
                self.conv3 = ConvBlock(128, 128)
                self.conv4 = ConvBlock(128, 256)

                self.hgs = nn.ModuleList([HourGlass(256, 4) for _ in range(4)])
                self.residuals = nn.ModuleList([ConvBlock(256, 256) for _ in range(4)])
                self.lstm1s = nn.ModuleList([nn.Conv2d(256, 256, kernel_size=1, stride=1) for _ in range(4)])
                self.bn_end = nn.ModuleList([nn.BatchNorm2d(256) for _ in range(4)])
                self.lstm2s = nn.ModuleList([nn.Conv2d(256, 68, kernel_size=1, stride=1) for _ in range(4)])
                # self.bl = nn.ModuleList([nn.Conv2d(256, 256, kernel_size=1, stride=1) for _ in range(3)])
                # self.al = nn.ModuleList([nn.Conv2d(68, 256, kernel_size=1, stride=1) for _ in range(3)])
                self.bl = nn.ModuleList([nn.Conv2d(68, 256, kernel_size=1, stride=1) for _ in range(3)])
                self.al = nn.ModuleList([nn.Conv2d(68, 256, kernel_size=1, stride=1) for _ in range(3)])


            def forward(self, inp):
                x, = inp

                x = F.relu(self.bn1(self.conv1(x)))
                x = self.conv2(x)
                x = F.avg_pool2d(x, kernel_size=2, stride=2)
                x = self.conv3(x)
                x = self.conv4(x)

                previous = x
                for i in range(4):
                    ll = self.hgs[i](previous)
                    ll = self.residuals[i](ll)
                    ll = F.relu(self.bn_end[i](self.lstm1s[i](ll)))
                    tmp_out = self.lstm2s[i](ll)
                    
                    if i < 3:
                        ll = ll + self.bl[i](tmp_out) + self.al[i](tmp_out)
                        previous = ll


                x = tmp_out
                x = x.permute(0, 3, 1, 2)
                return x

            def build_for_run(self, shapes_list):
                if not isinstance(shapes_list, list):
                    raise ValueError("shapes_list must be a list.")
                self.run_placeholders = [torch.zeros(*sh) for _, sh in shapes_list]
                self.run_output = self.__call__(self.run_placeholders)

        def rename_key(key):
            # Replace m_0 with hgs.0
            return key.replace("m_0", "hgs.0")

        e = None
        if place_model_on_cpu:
            device = torch.device("cpu")
            self.model.to(device)

        if e is not None: e.__enter__()
        self.model = FAN()
        # self.model.load_weights(str(model_path))
        if e is not None: e.__exit__(None,None,None)

        self.model.build_for_run ([(torch.float32, (1, 3, 256, 256))])
        
        self.model = FAN().to(self.device)

        # Load the state dictionary from the numpy file
        state_dict_np = np.load(model_path, allow_pickle=True)

        # Convert numpy arrays within the state dictionary to PyTorch tensors
        state_dict_torch = {rename_key(k): reshape_tensor(rename_key(k), torch.tensor(v, dtype=torch.float32)).cpu() for k, v in state_dict_np.items()}

        # Load the converted state dictionary into the model
        self.model.load_state_dict(state_dict_torch, strict=False)

        # missing_keys, unexpected_keys = self.model.load_state_dict(state_dict_torch, strict=False)
        # print("Missing Keys:", missing_keys)
        # print("Unexpected Keys:", unexpected_keys)

    
    def extract(self, input_image, rects, second_pass_extractor=None, is_bgr=True, multi_sample=False):
        if len(rects) == 0:
            return []

        if is_bgr:
            input_image = input_image[:, :, ::-1]  # Convert to RGB

        h, w, ch = input_image.shape
        landmarks = []

        for (left, top, right, bottom) in rects:
            scale = (right - left + bottom - top) / 195.0
            center = np.array([(left + right) / 2.0, (top + bottom) / 2.0])
            centers = [center]

            if multi_sample:
                centers += [center + [-1, -1], center + [1, -1], center + [1, 1], center + [-1, 1]]

            images = []
            ptss = []

            try:
                for c in centers:
                    images.append(self.crop(input_image, c, scale))

                images = torch.stack([torch.from_numpy(img) for img in images])
                images = images.float().div_(255.0).permute(0, 3, 1, 2).to(self.device)  # Convert images to [B, C, H, W]

                with torch.no_grad():
                    predicted = self.model(images)
                
                predicted = [p.cpu().numpy() for p in predicted]

                for i, pred in enumerate(predicted):
                    ptss.append(self.get_pts_from_predict(pred, centers[i], scale))
                pts_img = np.mean(np.array(ptss), 0)

                landmarks.append(pts_img)
            except:
                landmarks.append(None)

            if second_pass_extractor is not None:
                for i, lmrks in enumerate(landmarks):
                    try:
                        if lmrks is not None:
                            image_to_face_mat = LandmarksProcessor.get_transform_mat(lmrks, 256, FaceType.FULL)
                            face_image = cv2.warpAffine(input_image, image_to_face_mat, (256, 256), cv2.INTER_CUBIC)

                            # Assuming second_pass_extractor's extract method has a similar signature:
                            rects2 = second_pass_extractor.extract(face_image, is_bgr=False)
                            if len(rects2) == 1: 
                                lmrks2 = self.extract(face_image, [rects2[0]], is_bgr=False, multi_sample=True)[0]
                                landmarks[i] = LandmarksProcessor.transform_points(lmrks2, image_to_face_mat, True)
                    except:
                        pass

        return landmarks


    def transform(self, point, center, scale, resolution):
        pt = np.array ( [point[0], point[1], 1.0] )
        h = 200.0 * scale
        m = np.eye(3)
        m[0,0] = resolution / h
        m[1,1] = resolution / h
        m[0,2] = resolution * ( -center[0] / h + 0.5 )
        m[1,2] = resolution * ( -center[1] / h + 0.5 )
        m = np.linalg.inv(m)
        return np.matmul (m, pt)[0:2]

    def crop(self, image, center, scale, resolution=256.0):
        ul = self.transform([1, 1], center, scale, resolution).astype( np.int )
        br = self.transform([resolution, resolution], center, scale, resolution).astype( np.int )

        if image.ndim > 2:
            newDim = np.array([br[1] - ul[1], br[0] - ul[0], image.shape[2]], dtype=np.int32)
            newImg = np.zeros(newDim, dtype=np.uint8)
        else:
            newDim = np.array([br[1] - ul[1], br[0] - ul[0]], dtype=np.int)
            newImg = np.zeros(newDim, dtype=np.uint8)
        ht = image.shape[0]
        wd = image.shape[1]
        newX = np.array([max(1, -ul[0] + 1), min(br[0], wd) - ul[0]], dtype=np.int32)
        newY = np.array([max(1, -ul[1] + 1), min(br[1], ht) - ul[1]], dtype=np.int32)
        oldX = np.array([max(1, ul[0] + 1), min(br[0], wd)], dtype=np.int32)
        oldY = np.array([max(1, ul[1] + 1), min(br[1], ht)], dtype=np.int32)
        newImg[newY[0] - 1:newY[1], newX[0] - 1:newX[1] ] = image[oldY[0] - 1:oldY[1], oldX[0] - 1:oldX[1], :]

        newImg = cv2.resize(newImg, dsize=(int(resolution), int(resolution)), interpolation=cv2.INTER_LINEAR)
        return newImg

    def get_pts_from_predict(self, a, center, scale):
        a_ch, a_h, a_w = a.shape

        b = a.reshape(a_ch, a_h * a_w)
        c = b.argmax(axis=1).repeat(2).reshape(a_ch, 2).astype(np.float32)
        c[:, 0] %= a_w
        c[:, 1] = np.floor(c[:, 1] / a_w)

        for i in range(a_ch):
            pX, pY = int(c[i, 0]), int(c[i, 1])
            if 0 < pX < a_w - 1 and 0 < pY < a_h - 1:
                diff = np.array([a[i, pY, pX+1] - a[i, pY, pX-1], 
                                a[i, pY+1, pX] - a[i, pY-1, pX]])
                c[i] += np.sign(diff) * 0.25

        c += 0.5
        return np.array([self.transform(c[i], center, scale, a_w) for i in range(a_ch)])
facelib/FANExtractor.py
oo92/TorchFaceLab
import operator
from pathlib import Path

import cv2
import numpy as np

from core.leras import nn

class FaceEnhancer(object):
    """
    x4 face enhancer
    """
    def __init__(self, place_model_on_cpu=False, run_on_cpu=False):
        nn.initialize(data_format="NHWC")
        tf = nn.tf

        class FaceEnhancer (nn.ModelBase):
            def __init__(self, name='FaceEnhancer'):
                super().__init__(name=name)

            def on_build(self):
                self.conv1 = nn.Conv2D (3, 64, kernel_size=3, strides=1, padding='SAME')

                self.dense1 = nn.Dense (1, 64, use_bias=False)
                self.dense2 = nn.Dense (1, 64, use_bias=False)

                self.e0_conv0 = nn.Conv2D (64, 64, kernel_size=3, strides=1, padding='SAME')
                self.e0_conv1 = nn.Conv2D (64, 64, kernel_size=3, strides=1, padding='SAME')

                self.e1_conv0 = nn.Conv2D (64, 112, kernel_size=3, strides=1, padding='SAME')
                self.e1_conv1 = nn.Conv2D (112, 112, kernel_size=3, strides=1, padding='SAME')

                self.e2_conv0 = nn.Conv2D (112, 192, kernel_size=3, strides=1, padding='SAME')
                self.e2_conv1 = nn.Conv2D (192, 192, kernel_size=3, strides=1, padding='SAME')

                self.e3_conv0 = nn.Conv2D (192, 336, kernel_size=3, strides=1, padding='SAME')
                self.e3_conv1 = nn.Conv2D (336, 336, kernel_size=3, strides=1, padding='SAME')

                self.e4_conv0 = nn.Conv2D (336, 512, kernel_size=3, strides=1, padding='SAME')
                self.e4_conv1 = nn.Conv2D (512, 512, kernel_size=3, strides=1, padding='SAME')

                self.center_conv0 = nn.Conv2D (512, 512, kernel_size=3, strides=1, padding='SAME')
                self.center_conv1 = nn.Conv2D (512, 512, kernel_size=3, strides=1, padding='SAME')
                self.center_conv2 = nn.Conv2D (512, 512, kernel_size=3, strides=1, padding='SAME')
                self.center_conv3 = nn.Conv2D (512, 512, kernel_size=3, strides=1, padding='SAME')

                self.d4_conv0 = nn.Conv2D (1024, 512, kernel_size=3, strides=1, padding='SAME')
                self.d4_conv1 = nn.Conv2D (512, 512, kernel_size=3, strides=1, padding='SAME')

                self.d3_conv0 = nn.Conv2D (848, 512, kernel_size=3, strides=1, padding='SAME')
                self.d3_conv1 = nn.Conv2D (512, 512, kernel_size=3, strides=1, padding='SAME')

                self.d2_conv0 = nn.Conv2D (704, 288, kernel_size=3, strides=1, padding='SAME')
                self.d2_conv1 = nn.Conv2D (288, 288, kernel_size=3, strides=1, padding='SAME')

                self.d1_conv0 = nn.Conv2D (400, 160, kernel_size=3, strides=1, padding='SAME')
                self.d1_conv1 = nn.Conv2D (160, 160, kernel_size=3, strides=1, padding='SAME')

                self.d0_conv0 = nn.Conv2D (224, 96, kernel_size=3, strides=1, padding='SAME')
                self.d0_conv1 = nn.Conv2D (96, 96, kernel_size=3, strides=1, padding='SAME')

                self.out1x_conv0 = nn.Conv2D (96, 48, kernel_size=3, strides=1, padding='SAME')
                self.out1x_conv1 = nn.Conv2D (48, 3, kernel_size=3, strides=1, padding='SAME')

                self.dec2x_conv0 = nn.Conv2D (96, 96, kernel_size=3, strides=1, padding='SAME')
                self.dec2x_conv1 = nn.Conv2D (96, 96, kernel_size=3, strides=1, padding='SAME')

                self.out2x_conv0 = nn.Conv2D (96, 48, kernel_size=3, strides=1, padding='SAME')
                self.out2x_conv1 = nn.Conv2D (48, 3, kernel_size=3, strides=1, padding='SAME')

                self.dec4x_conv0 = nn.Conv2D (96, 72, kernel_size=3, strides=1, padding='SAME')
                self.dec4x_conv1 = nn.Conv2D (72, 72, kernel_size=3, strides=1, padding='SAME')

                self.out4x_conv0 = nn.Conv2D (72, 36, kernel_size=3, strides=1, padding='SAME')
                self.out4x_conv1 = nn.Conv2D (36, 3 , kernel_size=3, strides=1, padding='SAME')

            def forward(self, inp):
                bgr, param, param1 = inp

                x = self.conv1(bgr)
                a = self.dense1(param)
                a = tf.reshape(a, (-1,1,1,64) )

                b = self.dense2(param1)
                b = tf.reshape(b, (-1,1,1,64) )

                x = tf.nn.leaky_relu(x+a+b, 0.1)

                x = tf.nn.leaky_relu(self.e0_conv0(x), 0.1)
                x = e0 = tf.nn.leaky_relu(self.e0_conv1(x), 0.1)

                x = tf.nn.avg_pool(x, [1,2,2,1], [1,2,2,1], "VALID")
                x = tf.nn.leaky_relu(self.e1_conv0(x), 0.1)
                x = e1 = tf.nn.leaky_relu(self.e1_conv1(x), 0.1)

                x = tf.nn.avg_pool(x, [1,2,2,1], [1,2,2,1], "VALID")
                x = tf.nn.leaky_relu(self.e2_conv0(x), 0.1)
                x = e2 = tf.nn.leaky_relu(self.e2_conv1(x), 0.1)

                x = tf.nn.avg_pool(x, [1,2,2,1], [1,2,2,1], "VALID")
                x = tf.nn.leaky_relu(self.e3_conv0(x), 0.1)
                x = e3 = tf.nn.leaky_relu(self.e3_conv1(x), 0.1)

                x = tf.nn.avg_pool(x, [1,2,2,1], [1,2,2,1], "VALID")
                x = tf.nn.leaky_relu(self.e4_conv0(x), 0.1)
                x = e4 = tf.nn.leaky_relu(self.e4_conv1(x), 0.1)

                x = tf.nn.avg_pool(x, [1,2,2,1], [1,2,2,1], "VALID")
                x = tf.nn.leaky_relu(self.center_conv0(x), 0.1)
                x = tf.nn.leaky_relu(self.center_conv1(x), 0.1)
                x = tf.nn.leaky_relu(self.center_conv2(x), 0.1)
                x = tf.nn.leaky_relu(self.center_conv3(x), 0.1)

                x = tf.concat( [nn.resize2d_bilinear(x), e4], -1 )
                x = tf.nn.leaky_relu(self.d4_conv0(x), 0.1)
                x = tf.nn.leaky_relu(self.d4_conv1(x), 0.1)

                x = tf.concat( [nn.resize2d_bilinear(x), e3], -1 )
                x = tf.nn.leaky_relu(self.d3_conv0(x), 0.1)
                x = tf.nn.leaky_relu(self.d3_conv1(x), 0.1)

                x = tf.concat( [nn.resize2d_bilinear(x), e2], -1 )
                x = tf.nn.leaky_relu(self.d2_conv0(x), 0.1)
                x = tf.nn.leaky_relu(self.d2_conv1(x), 0.1)

                x = tf.concat( [nn.resize2d_bilinear(x), e1], -1 )
                x = tf.nn.leaky_relu(self.d1_conv0(x), 0.1)
                x = tf.nn.leaky_relu(self.d1_conv1(x), 0.1)

                x = tf.concat( [nn.resize2d_bilinear(x), e0], -1 )
                x = tf.nn.leaky_relu(self.d0_conv0(x), 0.1)
                x = d0 = tf.nn.leaky_relu(self.d0_conv1(x), 0.1)

                x = tf.nn.leaky_relu(self.out1x_conv0(x), 0.1)
                x = self.out1x_conv1(x)
                out1x = bgr + tf.nn.tanh(x)

                x = d0
                x = tf.nn.leaky_relu(self.dec2x_conv0(x), 0.1)
                x = tf.nn.leaky_relu(self.dec2x_conv1(x), 0.1)
                x = d2x = nn.resize2d_bilinear(x)

                x = tf.nn.leaky_relu(self.out2x_conv0(x), 0.1)
                x = self.out2x_conv1(x)

                out2x = nn.resize2d_bilinear(out1x) + tf.nn.tanh(x)

                x = d2x
                x = tf.nn.leaky_relu(self.dec4x_conv0(x), 0.1)
                x = tf.nn.leaky_relu(self.dec4x_conv1(x), 0.1)
                x = d4x = nn.resize2d_bilinear(x)

                x = tf.nn.leaky_relu(self.out4x_conv0(x), 0.1)
                x = self.out4x_conv1(x)

                out4x = nn.resize2d_bilinear(out2x) + tf.nn.tanh(x)

                return out4x

        model_path = Path(__file__).parent / "FaceEnhancer.npy"
        if not model_path.exists():
            raise Exception("Unable to load FaceEnhancer.npy")

        with tf.device ('/CPU:0' if place_model_on_cpu else nn.tf_default_device_name):
            self.model = FaceEnhancer()
            self.model.load_weights (model_path)
        
        with tf.device ('/CPU:0' if run_on_cpu else nn.tf_default_device_name):
            self.model.build_for_run ([ (tf.float32, nn.get4Dshape (192,192,3) ),
                                        (tf.float32, (None,1,) ),
                                        (tf.float32, (None,1,) ),
                                    ])

    def enhance (self, inp_img, is_tanh=False, preserve_size=True):
        if not is_tanh:
            inp_img = np.clip( inp_img * 2 -1, -1, 1 )

        param = np.array([0.2])
        param1 = np.array([1.0])
        up_res = 4
        patch_size = 192
        patch_size_half = patch_size // 2

        ih,iw,ic = inp_img.shape
        h,w,c = ih,iw,ic

        th,tw = h*up_res, w*up_res

        t_padding = 0
        b_padding = 0
        l_padding = 0
        r_padding = 0

        if h < patch_size:
            t_padding = (patch_size-h)//2
            b_padding = (patch_size-h) - t_padding

        if w < patch_size:
            l_padding = (patch_size-w)//2
            r_padding = (patch_size-w) - l_padding

        if t_padding != 0:
            inp_img = np.concatenate ([ np.zeros ( (t_padding,w,c), dtype=np.float32 ), inp_img ], axis=0 )
            h,w,c = inp_img.shape

        if b_padding != 0:
            inp_img = np.concatenate ([ inp_img, np.zeros ( (b_padding,w,c), dtype=np.float32 ) ], axis=0 )
            h,w,c = inp_img.shape

        if l_padding != 0:
            inp_img = np.concatenate ([ np.zeros ( (h,l_padding,c), dtype=np.float32 ), inp_img ], axis=1 )
            h,w,c = inp_img.shape

        if r_padding != 0:
            inp_img = np.concatenate ([ inp_img, np.zeros ( (h,r_padding,c), dtype=np.float32 ) ], axis=1 )
            h,w,c = inp_img.shape


        i_max = w-patch_size+1
        j_max = h-patch_size+1

        final_img = np.zeros ( (h*up_res,w*up_res,c), dtype=np.float32 )
        final_img_div = np.zeros ( (h*up_res,w*up_res,1), dtype=np.float32 )

        x = np.concatenate ( [ np.linspace (0,1,patch_size_half*up_res), np.linspace (1,0,patch_size_half*up_res) ] )
        x,y = np.meshgrid(x,x)
        patch_mask = (x*y)[...,None]

        j=0
        while j < j_max:
            i = 0
            while i < i_max:
                patch_img = inp_img[j:j+patch_size, i:i+patch_size,:]
                x = self.model.run( [ patch_img[None,...], [param], [param1] ] )[0]
                final_img    [j*up_res:(j+patch_size)*up_res, i*up_res:(i+patch_size)*up_res,:] += x*patch_mask
                final_img_div[j*up_res:(j+patch_size)*up_res, i*up_res:(i+patch_size)*up_res,:] += patch_mask
                if i == i_max-1:
                    break
                i = min( i+patch_size_half, i_max-1)
            if j == j_max-1:
                break
            j = min( j+patch_size_half, j_max-1)

        final_img_div[final_img_div==0] = 1.0
        final_img /= final_img_div

        if t_padding+b_padding+l_padding+r_padding != 0:
            final_img = final_img [t_padding*up_res:(h-b_padding)*up_res, l_padding*up_res:(w-r_padding)*up_res,:]

        if preserve_size:
            final_img = cv2.resize (final_img, (iw,ih), interpolation=cv2.INTER_LANCZOS4)

        if not is_tanh:
            final_img = np.clip( final_img/2+0.5, 0, 1 )

        return final_img


"""

    def enhance (self, inp_img, is_tanh=False, preserve_size=True):
        if not is_tanh:
            inp_img = np.clip( inp_img * 2 -1, -1, 1 )

        param = np.array([0.2])
        param1 = np.array([1.0])
        up_res = 4
        patch_size = 192
        patch_size_half = patch_size // 2

        h,w,c = inp_img.shape

        th,tw = h*up_res, w*up_res

        preupscale_rate = 1.0

        if h < patch_size or w < patch_size:
            preupscale_rate = 1.0 / ( max(h,w) / patch_size )

        if preupscale_rate != 1.0:
            inp_img = cv2.resize (inp_img, ( int(w*preupscale_rate), int(h*preupscale_rate) ), interpolation=cv2.INTER_LANCZOS4)
            h,w,c = inp_img.shape

        i_max = w-patch_size+1
        j_max = h-patch_size+1

        final_img = np.zeros ( (h*up_res,w*up_res,c), dtype=np.float32 )
        final_img_div = np.zeros ( (h*up_res,w*up_res,1), dtype=np.float32 )

        x = np.concatenate ( [ np.linspace (0,1,patch_size_half*up_res), np.linspace (1,0,patch_size_half*up_res) ] )
        x,y = np.meshgrid(x,x)
        patch_mask = (x*y)[...,None]

        j=0
        while j < j_max:
            i = 0
            while i < i_max:
                patch_img = inp_img[j:j+patch_size, i:i+patch_size,:]
                x = self.model.run( [ patch_img[None,...], [param], [param1] ] )[0]
                final_img    [j*up_res:(j+patch_size)*up_res, i*up_res:(i+patch_size)*up_res,:] += x*patch_mask
                final_img_div[j*up_res:(j+patch_size)*up_res, i*up_res:(i+patch_size)*up_res,:] += patch_mask
                if i == i_max-1:
                    break
                i = min( i+patch_size_half, i_max-1)
            if j == j_max-1:
                break
            j = min( j+patch_size_half, j_max-1)

        final_img_div[final_img_div==0] = 1.0
        final_img /= final_img_div

        if preserve_size:
            final_img = cv2.resize (final_img, (w,h), interpolation=cv2.INTER_LANCZOS4)
        else:
            if preupscale_rate != 1.0:
                final_img = cv2.resize (final_img, (tw,th), interpolation=cv2.INTER_LANCZOS4)

        if not is_tanh:
            final_img = np.clip( final_img/2+0.5, 0, 1 )

        return final_img
"""
facelib/FaceEnhancer.py
oo92/TorchFaceLab
import colorsys
import math
from enum import IntEnum

import cv2
import numpy as np
import numpy.linalg as npla

from core import imagelib
from core import mathlib
from facelib import FaceType
from core.mathlib.umeyama import umeyama

landmarks_2D = np.array([
[ 0.000213256,  0.106454  ], #17
[ 0.0752622,    0.038915  ], #18
[ 0.18113,      0.0187482 ], #19
[ 0.29077,      0.0344891 ], #20
[ 0.393397,     0.0773906 ], #21
[ 0.586856,     0.0773906 ], #22
[ 0.689483,     0.0344891 ], #23
[ 0.799124,     0.0187482 ], #24
[ 0.904991,     0.038915  ], #25
[ 0.98004,      0.106454  ], #26
[ 0.490127,     0.203352  ], #27
[ 0.490127,     0.307009  ], #28
[ 0.490127,     0.409805  ], #29
[ 0.490127,     0.515625  ], #30
[ 0.36688,      0.587326  ], #31
[ 0.426036,     0.609345  ], #32
[ 0.490127,     0.628106  ], #33
[ 0.554217,     0.609345  ], #34
[ 0.613373,     0.587326  ], #35
[ 0.121737,     0.216423  ], #36
[ 0.187122,     0.178758  ], #37
[ 0.265825,     0.179852  ], #38
[ 0.334606,     0.231733  ], #39
[ 0.260918,     0.245099  ], #40
[ 0.182743,     0.244077  ], #41
[ 0.645647,     0.231733  ], #42
[ 0.714428,     0.179852  ], #43
[ 0.793132,     0.178758  ], #44
[ 0.858516,     0.216423  ], #45
[ 0.79751,      0.244077  ], #46
[ 0.719335,     0.245099  ], #47
[ 0.254149,     0.780233  ], #48
[ 0.340985,     0.745405  ], #49
[ 0.428858,     0.727388  ], #50
[ 0.490127,     0.742578  ], #51
[ 0.551395,     0.727388  ], #52
[ 0.639268,     0.745405  ], #53
[ 0.726104,     0.780233  ], #54
[ 0.642159,     0.864805  ], #55
[ 0.556721,     0.902192  ], #56
[ 0.490127,     0.909281  ], #57
[ 0.423532,     0.902192  ], #58
[ 0.338094,     0.864805  ], #59
[ 0.290379,     0.784792  ], #60
[ 0.428096,     0.778746  ], #61
[ 0.490127,     0.785343  ], #62
[ 0.552157,     0.778746  ], #63
[ 0.689874,     0.784792  ], #64
[ 0.553364,     0.824182  ], #65
[ 0.490127,     0.831803  ], #66
[ 0.42689 ,     0.824182  ]  #67
], dtype=np.float32)


landmarks_2D_new = np.array([
[ 0.000213256,  0.106454  ], #17
[ 0.0752622,    0.038915  ], #18
[ 0.18113,      0.0187482 ], #19
[ 0.29077,      0.0344891 ], #20
[ 0.393397,     0.0773906 ], #21
[ 0.586856,     0.0773906 ], #22
[ 0.689483,     0.0344891 ], #23
[ 0.799124,     0.0187482 ], #24
[ 0.904991,     0.038915  ], #25
[ 0.98004,      0.106454  ], #26
[ 0.490127,     0.203352  ], #27
[ 0.490127,     0.307009  ], #28
[ 0.490127,     0.409805  ], #29
[ 0.490127,     0.515625  ], #30
[ 0.36688,      0.587326  ], #31
[ 0.426036,     0.609345  ], #32
[ 0.490127,     0.628106  ], #33
[ 0.554217,     0.609345  ], #34
[ 0.613373,     0.587326  ], #35
[ 0.121737,     0.216423  ], #36
[ 0.187122,     0.178758  ], #37
[ 0.265825,     0.179852  ], #38
[ 0.334606,     0.231733  ], #39
[ 0.260918,     0.245099  ], #40
[ 0.182743,     0.244077  ], #41
[ 0.645647,     0.231733  ], #42
[ 0.714428,     0.179852  ], #43
[ 0.793132,     0.178758  ], #44
[ 0.858516,     0.216423  ], #45
[ 0.79751,      0.244077  ], #46
[ 0.719335,     0.245099  ], #47
[ 0.254149,     0.780233  ], #48
[ 0.726104,     0.780233  ], #54
], dtype=np.float32)

mouth_center_landmarks_2D = np.array([
    [-4.4202591e-07,  4.4916576e-01],  #48
    [ 1.8399176e-01,  3.7537053e-01],  #49
    [ 3.7018123e-01,  3.3719531e-01],  #50
    [ 5.0000089e-01,  3.6938059e-01],  #51
    [ 6.2981832e-01,  3.3719531e-01],  #52
    [ 8.1600773e-01,  3.7537053e-01],  #53
    [ 1.0000000e+00,  4.4916576e-01],  #54
    [ 8.2213330e-01,  6.2836081e-01],  #55
    [ 6.4110327e-01,  7.0757812e-01],  #56
    [ 5.0000089e-01,  7.2259867e-01],  #57
    [ 3.5889623e-01,  7.0757812e-01],  #58
    [ 1.7786618e-01,  6.2836081e-01],  #59
    [ 7.6765373e-02,  4.5882553e-01],  #60
    [ 3.6856663e-01,  4.4601500e-01],  #61
    [ 5.0000089e-01,  4.5999300e-01],  #62
    [ 6.3143289e-01,  4.4601500e-01],  #63
    [ 9.2323411e-01,  4.5882553e-01],  #64
    [ 6.3399029e-01,  5.4228687e-01],  #65
    [ 5.0000089e-01,  5.5843467e-01],  #66
    [ 3.6601129e-01,  5.4228687e-01]   #67
], dtype=np.float32)

# 68 point landmark definitions
landmarks_68_pt = { "mouth": (48,68),
                    "right_eyebrow": (17, 22),
                    "left_eyebrow": (22, 27),
                    "right_eye": (36, 42),
                    "left_eye": (42, 48),
                    "nose": (27, 36), # missed one point
                    "jaw": (0, 17) }

landmarks_68_3D = np.array( [
[-73.393523  , -29.801432   , 47.667532   ], #00
[-72.775014  , -10.949766   , 45.909403   ], #01
[-70.533638  , 7.929818     , 44.842580   ], #02
[-66.850058  , 26.074280    , 43.141114   ], #03
[-59.790187  , 42.564390    , 38.635298   ], #04
[-48.368973  , 56.481080    , 30.750622   ], #05
[-34.121101  , 67.246992    , 18.456453   ], #06
[-17.875411  , 75.056892    , 3.609035    ], #07
[0.098749    , 77.061286    , -0.881698   ], #08
[17.477031   , 74.758448    , 5.181201    ], #09
[32.648966   , 66.929021    , 19.176563   ], #10
[46.372358   , 56.311389    , 30.770570   ], #11
[57.343480   , 42.419126    , 37.628629   ], #12
[64.388482   , 25.455880    , 40.886309   ], #13
[68.212038   , 6.990805     , 42.281449   ], #14
[70.486405   , -11.666193   , 44.142567   ], #15
[71.375822   , -30.365191   , 47.140426   ], #16
[-61.119406  , -49.361602   , 14.254422   ], #17
[-51.287588  , -58.769795   , 7.268147    ], #18
[-37.804800  , -61.996155   , 0.442051    ], #19
[-24.022754  , -61.033399   , -6.606501   ], #20
[-11.635713  , -56.686759   , -11.967398  ], #21
[12.056636   , -57.391033   , -12.051204  ], #22
[25.106256   , -61.902186   , -7.315098   ], #23
[38.338588   , -62.777713   , -1.022953   ], #24
[51.191007   , -59.302347   , 5.349435    ], #25
[60.053851   , -50.190255   , 11.615746   ], #26
[0.653940    , -42.193790   , -13.380835  ], #27
[0.804809    , -30.993721   , -21.150853  ], #28
[0.992204    , -19.944596   , -29.284036  ], #29
[1.226783    , -8.414541    , -36.948060  ], #00
[-14.772472  , 2.598255     , -20.132003  ], #01
[-7.180239   , 4.751589     , -23.536684  ], #02
[0.555920    , 6.562900     , -25.944448  ], #03
[8.272499    , 4.661005     , -23.695741  ], #04
[15.214351   , 2.643046     , -20.858157  ], #05
[-46.047290  , -37.471411   , 7.037989    ], #06
[-37.674688  , -42.730510   , 3.021217    ], #07
[-27.883856  , -42.711517   , 1.353629    ], #08
[-19.648268  , -36.754742   , -0.111088   ], #09
[-28.272965  , -35.134493   , -0.147273   ], #10
[-38.082418  , -34.919043   , 1.476612    ], #11
[19.265868   , -37.032306   , -0.665746   ], #12
[27.894191   , -43.342445   , 0.247660    ], #13
[37.437529   , -43.110822   , 1.696435    ], #14
[45.170805   , -38.086515   , 4.894163    ], #15
[38.196454   , -35.532024   , 0.282961    ], #16
[28.764989   , -35.484289   , -1.172675   ], #17
[-28.916267  , 28.612716    , -2.240310   ], #18
[-17.533194  , 22.172187    , -15.934335  ], #19
[-6.684590   , 19.029051    , -22.611355  ], #20
[0.381001    , 20.721118    , -23.748437  ], #21
[8.375443    , 19.035460    , -22.721995  ], #22
[18.876618   , 22.394109    , -15.610679  ], #23
[28.794412   , 28.079924    , -3.217393   ], #24
[19.057574   , 36.298248    , -14.987997  ], #25
[8.956375    , 39.634575    , -22.554245  ], #26
[0.381549    , 40.395647    , -23.591626  ], #27
[-7.428895   , 39.836405    , -22.406106  ], #28
[-18.160634  , 36.677899    , -15.121907  ], #29
[-24.377490  , 28.677771    , -4.785684   ], #30
[-6.897633   , 25.475976    , -20.893742  ], #31
[0.340663    , 26.014269    , -22.220479  ], #32
[8.444722    , 25.326198    , -21.025520  ], #33
[24.474473   , 28.323008    , -5.712776   ], #34
[8.449166    , 30.596216    , -20.671489  ], #35
[0.205322    , 31.408738    , -21.903670  ], #36 
[-7.198266   , 30.844876    , -20.328022  ]  #37
], dtype=np.float32)

FaceType_to_padding_remove_align = {
    FaceType.HALF: (0.0, False),
    FaceType.MID_FULL: (0.0675, False),
    FaceType.FULL: (0.2109375, False),
    FaceType.FULL_NO_ALIGN: (0.2109375, True),
    FaceType.WHOLE_FACE: (0.40, False),
    FaceType.HEAD: (0.70, False),
    FaceType.HEAD_NO_ALIGN: (0.70, True),
}

def convert_98_to_68(lmrks):
    #jaw
    result = [ lmrks[0] ]
    for i in range(2,16,2):
        result += [ ( lmrks[i] + (lmrks[i-1]+lmrks[i+1])/2 ) / 2  ]
    result += [ lmrks[16] ]
    for i in range(18,32,2):
        result += [ ( lmrks[i] + (lmrks[i-1]+lmrks[i+1])/2 ) / 2  ]
    result += [ lmrks[32] ]

    #eyebrows averaging
    result += [ lmrks[33],
                (lmrks[34]+lmrks[41])/2,
                (lmrks[35]+lmrks[40])/2,
                (lmrks[36]+lmrks[39])/2,
                (lmrks[37]+lmrks[38])/2,
              ]

    result += [ (lmrks[42]+lmrks[50])/2,
                (lmrks[43]+lmrks[49])/2,
                (lmrks[44]+lmrks[48])/2,
                (lmrks[45]+lmrks[47])/2,
                lmrks[46]
              ]

    #nose
    result += list ( lmrks[51:60] )

    #left eye (from our view)
    result += [ lmrks[60],
                lmrks[61],
                lmrks[63],
                lmrks[64],
                lmrks[65],
                lmrks[67] ]

    #right eye
    result += [ lmrks[68],
                lmrks[69],
                lmrks[71],
                lmrks[72],
                lmrks[73],
                lmrks[75] ]

    #mouth
    result += list ( lmrks[76:96] )

    return np.concatenate (result).reshape ( (68,2) )

def transform_points(points, mat, invert=False):
    if invert:
        mat = cv2.invertAffineTransform (mat)
    points = np.expand_dims(points, axis=1)
    points = cv2.transform(points, mat, points.shape)
    points = np.squeeze(points)
    return points

def get_transform_mat (image_landmarks, output_size, face_type, scale=1.0):
    if not isinstance(image_landmarks, np.ndarray):
        image_landmarks = np.array (image_landmarks)


    # estimate landmarks transform from global space to local aligned space with bounds [0..1]
    mat = umeyama( np.concatenate ( [ image_landmarks[17:49] , image_landmarks[54:55] ] ) , landmarks_2D_new, True)[0:2]

    # get corner points in global space
    g_p = transform_points (  np.float32([(0,0),(1,0),(1,1),(0,1),(0.5,0.5) ]) , mat, True)
    g_c = g_p[4]

    # calc diagonal vectors between corners in global space
    tb_diag_vec = (g_p[2]-g_p[0]).astype(np.float32)
    tb_diag_vec /= npla.norm(tb_diag_vec)
    bt_diag_vec = (g_p[1]-g_p[3]).astype(np.float32)
    bt_diag_vec /= npla.norm(bt_diag_vec)

    # calc modifier of diagonal vectors for scale and padding value
    padding, remove_align = FaceType_to_padding_remove_align.get(face_type, 0.0)
    mod = (1.0 / scale)* ( npla.norm(g_p[0]-g_p[2])*(padding*np.sqrt(2.0) + 0.5) )

    if face_type == FaceType.WHOLE_FACE:
        # adjust vertical offset for WHOLE_FACE, 7% below in order to cover more forehead
        vec = (g_p[0]-g_p[3]).astype(np.float32)
        vec_len = npla.norm(vec)
        vec /= vec_len
        g_c += vec*vec_len*0.07

    elif face_type == FaceType.HEAD:
        # assuming image_landmarks are 3D_Landmarks extracted for HEAD,
        # adjust horizontal offset according to estimated yaw        
        yaw = estimate_averaged_yaw(transform_points (image_landmarks, mat, False))
        
        hvec = (g_p[0]-g_p[1]).astype(np.float32)
        hvec_len = npla.norm(hvec)
        hvec /= hvec_len        

        yaw *= np.abs(math.tanh(yaw*2)) # Damp near zero
        
        g_c -= hvec * (yaw * hvec_len / 2.0)                     

        # adjust vertical offset for HEAD, 50% below
        vvec = (g_p[0]-g_p[3]).astype(np.float32)
        vvec_len = npla.norm(vvec)
        vvec /= vvec_len
        g_c += vvec*vvec_len*0.50

    # calc 3 points in global space to estimate 2d affine transform
    if not remove_align:
        l_t = np.array( [ g_c - tb_diag_vec*mod,
                          g_c + bt_diag_vec*mod,
                          g_c + tb_diag_vec*mod ] )
    else:
        # remove_align - face will be centered in the frame but not aligned
        l_t = np.array( [ g_c - tb_diag_vec*mod,
                          g_c + bt_diag_vec*mod,
                          g_c + tb_diag_vec*mod,
                          g_c - bt_diag_vec*mod,
                         ] )

        # get area of face square in global space
        area = mathlib.polygon_area(l_t[:,0], l_t[:,1] )

        # calc side of square
        side = np.float32(math.sqrt(area) / 2)

        # calc 3 points with unrotated square
        l_t = np.array( [ g_c + [-side,-side],
                          g_c + [ side,-side],
                          g_c + [ side, side] ] )

    # calc affine transform from 3 global space points to 3 local space points size of 'output_size'
    pts2 = np.float32(( (0,0),(output_size,0),(output_size,output_size) ))
    mat = cv2.getAffineTransform(l_t,pts2)
    return mat

def get_rect_from_landmarks(image_landmarks):
    mat = get_transform_mat(image_landmarks, 256, FaceType.FULL_NO_ALIGN)

    g_p = transform_points (  np.float32([(0,0),(255,255) ]) , mat, True)

    (l,t,r,b) = g_p[0][0], g_p[0][1], g_p[1][0], g_p[1][1]

    return (l,t,r,b)

def expand_eyebrows(lmrks, eyebrows_expand_mod=1.0):
    if len(lmrks) != 68:
        raise Exception('works only with 68 landmarks')
    lmrks = np.array( lmrks.copy(), dtype=np.int )

    # #nose
    ml_pnt = (lmrks[36] + lmrks[0]) // 2
    mr_pnt = (lmrks[16] + lmrks[45]) // 2

    # mid points between the mid points and eye
    ql_pnt = (lmrks[36] + ml_pnt) // 2
    qr_pnt = (lmrks[45] + mr_pnt) // 2

    # Top of the eye arrays
    bot_l = np.array((ql_pnt, lmrks[36], lmrks[37], lmrks[38], lmrks[39]))
    bot_r = np.array((lmrks[42], lmrks[43], lmrks[44], lmrks[45], qr_pnt))

    # Eyebrow arrays
    top_l = lmrks[17:22]
    top_r = lmrks[22:27]

    # Adjust eyebrow arrays
    lmrks[17:22] = top_l + eyebrows_expand_mod * 0.5 * (top_l - bot_l)
    lmrks[22:27] = top_r + eyebrows_expand_mod * 0.5 * (top_r - bot_r)
    return lmrks




def get_image_hull_mask (image_shape, image_landmarks, eyebrows_expand_mod=1.0 ):
    hull_mask = np.zeros(image_shape[0:2]+(1,),dtype=np.float32)

    lmrks = expand_eyebrows(image_landmarks, eyebrows_expand_mod)

    r_jaw = (lmrks[0:9], lmrks[17:18])
    l_jaw = (lmrks[8:17], lmrks[26:27])
    r_cheek = (lmrks[17:20], lmrks[8:9])
    l_cheek = (lmrks[24:27], lmrks[8:9])
    nose_ridge = (lmrks[19:25], lmrks[8:9],)
    r_eye = (lmrks[17:22], lmrks[27:28], lmrks[31:36], lmrks[8:9])
    l_eye = (lmrks[22:27], lmrks[27:28], lmrks[31:36], lmrks[8:9])
    nose = (lmrks[27:31], lmrks[31:36])
    parts = [r_jaw, l_jaw, r_cheek, l_cheek, nose_ridge, r_eye, l_eye, nose]

    for item in parts:
        merged = np.concatenate(item)
        cv2.fillConvexPoly(hull_mask, cv2.convexHull(merged), (1,) )

    return hull_mask

def get_image_eye_mask (image_shape, image_landmarks):
    if len(image_landmarks) != 68:
        raise Exception('get_image_eye_mask works only with 68 landmarks')

    h,w,c = image_shape

    hull_mask = np.zeros( (h,w,1),dtype=np.float32)

    image_landmarks = image_landmarks.astype(np.int)

    cv2.fillConvexPoly( hull_mask, cv2.convexHull( image_landmarks[36:42]), (1,) )
    cv2.fillConvexPoly( hull_mask, cv2.convexHull( image_landmarks[42:48]), (1,) )

    dilate = h // 32
    hull_mask = cv2.dilate(hull_mask, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(dilate,dilate)), iterations = 1 )

    blur = h // 16
    blur = blur + (1-blur % 2)
    hull_mask = cv2.GaussianBlur(hull_mask, (blur, blur) , 0)
    hull_mask = hull_mask[...,None]

    return hull_mask

def get_image_mouth_mask (image_shape, image_landmarks):
    if len(image_landmarks) != 68:
        raise Exception('get_image_eye_mask works only with 68 landmarks')

    h,w,c = image_shape

    hull_mask = np.zeros( (h,w,1),dtype=np.float32)

    image_landmarks = image_landmarks.astype(np.int)

    cv2.fillConvexPoly( hull_mask, cv2.convexHull( image_landmarks[60:]), (1,) )

    dilate = h // 32
    hull_mask = cv2.dilate(hull_mask, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(dilate,dilate)), iterations = 1 )

    blur = h // 16
    blur = blur + (1-blur % 2)
    hull_mask = cv2.GaussianBlur(hull_mask, (blur, blur) , 0)
    hull_mask = hull_mask[...,None]

    return hull_mask
    
def alpha_to_color (img_alpha, color):
    if len(img_alpha.shape) == 2:
        img_alpha = img_alpha[...,None]
    h,w,c = img_alpha.shape
    result = np.zeros( (h,w, len(color) ), dtype=np.float32 )
    result[:,:] = color

    return result * img_alpha



def get_cmask (image_shape, lmrks, eyebrows_expand_mod=1.0):
    h,w,c = image_shape

    hull = get_image_hull_mask (image_shape, lmrks, eyebrows_expand_mod )

    result = np.zeros( (h,w,3), dtype=np.float32 )



    def process(w,h, data ):
        d = {}
        cur_lc = 0
        all_lines = []
        for s, pts_loop_ar in data:
            lines = []
            for pts, loop in pts_loop_ar:
                pts_len = len(pts)
                lines.append ( [ [ pts[i], pts[(i+1) % pts_len ] ]  for i in range(pts_len - (0 if loop else 1) ) ] )
            lines = np.concatenate (lines)

            lc = lines.shape[0]
            all_lines.append(lines)
            d[s] = cur_lc, cur_lc+lc
            cur_lc += lc
        all_lines = np.concatenate (all_lines, 0)

        #calculate signed distance for all points and lines
        line_count = all_lines.shape[0]
        pts_count = w*h

        all_lines = np.repeat ( all_lines[None,...], pts_count, axis=0 ).reshape ( (pts_count*line_count,2,2) )

        pts = np.empty( (h,w,line_count,2), dtype=np.float32 )
        pts[...,1] = np.arange(h)[:,None,None]
        pts[...,0] = np.arange(w)[:,None]
        pts = pts.reshape ( (h*w*line_count, -1) )

        a = all_lines[:,0,:]
        b = all_lines[:,1,:]
        pa = pts-a
        ba = b-a
        ph = np.clip ( np.einsum('ij,ij->i', pa, ba) / np.einsum('ij,ij->i', ba, ba), 0, 1 )
        dists = npla.norm ( pa - ba*ph[...,None], axis=1).reshape ( (h,w,line_count) )

        def get_dists(name, thickness=0):
            s,e = d[name]
            result = dists[...,s:e]
            if thickness != 0:
                result = np.abs(result)-thickness
            return np.min (result, axis=-1)

        return get_dists

    l_eye = lmrks[42:48]
    r_eye = lmrks[36:42]
    l_brow = lmrks[22:27]
    r_brow = lmrks[17:22]
    mouth = lmrks[48:60]

    up_nose = np.concatenate( (lmrks[27:31], lmrks[33:34]) )
    down_nose = lmrks[31:36]
    nose = np.concatenate ( (up_nose, down_nose) )

    gdf = process ( w,h,
                         (
                          ('eyes',  ((l_eye, True), (r_eye, True)) ),
                          ('brows', ((l_brow, False), (r_brow,False)) ),
                          ('up_nose', ((up_nose, False),) ),
                          ('down_nose', ((down_nose, False),) ),
                          ('mouth', ((mouth, True),) ),
                         )
                        )

    eyes_fall_dist = w // 32
    eyes_thickness = max( w // 64, 1 )

    brows_fall_dist = w // 32
    brows_thickness = max( w // 256, 1 )

    nose_fall_dist = w / 12
    nose_thickness = max( w // 96, 1 )

    mouth_fall_dist = w // 32
    mouth_thickness = max( w // 64, 1 )

    eyes_mask = gdf('eyes',eyes_thickness)
    eyes_mask = 1-np.clip( eyes_mask/ eyes_fall_dist, 0, 1)
    #eyes_mask = np.clip ( 1- ( np.sqrt( np.maximum(eyes_mask,0) ) / eyes_fall_dist ), 0, 1)
    #eyes_mask = np.clip ( 1- ( np.cbrt( np.maximum(eyes_mask,0) ) / eyes_fall_dist ), 0, 1)

    brows_mask = gdf('brows', brows_thickness)
    brows_mask = 1-np.clip( brows_mask / brows_fall_dist, 0, 1)
    #brows_mask = np.clip ( 1- ( np.sqrt( np.maximum(brows_mask,0) ) / brows_fall_dist ), 0, 1)

    mouth_mask = gdf('mouth', mouth_thickness)
    mouth_mask = 1-np.clip( mouth_mask / mouth_fall_dist, 0, 1)
    #mouth_mask = np.clip ( 1- ( np.sqrt( np.maximum(mouth_mask,0) ) / mouth_fall_dist ), 0, 1)

    def blend(a,b,k):
        x = np.clip ( 0.5+0.5*(b-a)/k, 0.0, 1.0 )
        return (a-b)*x+b - k*x*(1.0-x)


    #nose_mask = (a-b)*x+b - k*x*(1.0-x)

    #nose_mask = np.minimum (up_nose_mask , down_nose_mask )
    #nose_mask = 1-np.clip( nose_mask / nose_fall_dist, 0, 1)

    nose_mask = blend ( gdf('up_nose', nose_thickness), gdf('down_nose', nose_thickness), nose_thickness*3 )
    nose_mask = 1-np.clip( nose_mask / nose_fall_dist, 0, 1)

    up_nose_mask = gdf('up_nose', nose_thickness)
    up_nose_mask = 1-np.clip( up_nose_mask / nose_fall_dist, 0, 1)
    #up_nose_mask = np.clip ( 1- ( np.cbrt( np.maximum(up_nose_mask,0) ) / nose_fall_dist ), 0, 1)

    down_nose_mask = gdf('down_nose', nose_thickness)
    down_nose_mask = 1-np.clip( down_nose_mask / nose_fall_dist, 0, 1)
    #down_nose_mask = np.clip ( 1- ( np.cbrt( np.maximum(down_nose_mask,0) ) / nose_fall_dist ), 0, 1)

    #nose_mask = np.clip( up_nose_mask + down_nose_mask, 0, 1 )
    #nose_mask /= np.max(nose_mask)
    #nose_mask = np.maximum (up_nose_mask , down_nose_mask )
    #nose_mask = down_nose_mask

    #nose_mask = np.zeros_like(nose_mask)

    eyes_mask = eyes_mask * (1-mouth_mask)
    nose_mask = nose_mask * (1-eyes_mask)

    hull_mask = hull[...,0].copy()
    hull_mask = hull_mask * (1-eyes_mask) * (1-brows_mask) * (1-nose_mask) * (1-mouth_mask)

    #eyes_mask = eyes_mask * (1-nose_mask)

    mouth_mask= mouth_mask * (1-nose_mask)

    brows_mask = brows_mask * (1-nose_mask)* (1-eyes_mask )

    hull_mask = alpha_to_color(hull_mask, (0,1,0) )
    eyes_mask = alpha_to_color(eyes_mask, (1,0,0) )
    brows_mask = alpha_to_color(brows_mask, (0,0,1) )
    nose_mask = alpha_to_color(nose_mask, (0,1,1) )
    mouth_mask = alpha_to_color(mouth_mask, (0,0,1) )

    #nose_mask = np.maximum( up_nose_mask, down_nose_mask )

    result = hull_mask + mouth_mask+ nose_mask + brows_mask  + eyes_mask
    result *= hull
    #result = np.clip (result, 0, 1)
    return result

def blur_image_hull_mask (hull_mask):

    maxregion = np.argwhere(hull_mask==1.0)
    miny,minx = maxregion.min(axis=0)[:2]
    maxy,maxx = maxregion.max(axis=0)[:2]
    lenx = maxx - minx;
    leny = maxy - miny;
    masky = int(minx+(lenx//2))
    maskx = int(miny+(leny//2))
    lowest_len = min (lenx, leny)
    ero = int( lowest_len * 0.085 )
    blur = int( lowest_len * 0.10 )

    hull_mask = cv2.erode(hull_mask, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(ero,ero)), iterations = 1 )
    hull_mask = cv2.blur(hull_mask, (blur, blur) )
    hull_mask = np.expand_dims (hull_mask,-1)

    return hull_mask

mirror_idxs = [
    [0,16],
    [1,15],
    [2,14],
    [3,13],
    [4,12],
    [5,11],
    [6,10],
    [7,9],

    [17,26],
    [18,25],
    [19,24],
    [20,23],
    [21,22],

    [36,45],
    [37,44],
    [38,43],
    [39,42],
    [40,47],
    [41,46],

    [31,35],
    [32,34],

    [50,52],
    [49,53],
    [48,54],
    [59,55],
    [58,56],
    [67,65],
    [60,64],
    [61,63] ]

def mirror_landmarks (landmarks, val):
    result = landmarks.copy()

    for idx in mirror_idxs:
        result [ idx ] = result [ idx[::-1] ]

    result[:,0] = val - result[:,0] - 1
    return result

def get_face_struct_mask (image_shape, image_landmarks, eyebrows_expand_mod=1.0, color=(1,) ):
    mask = np.zeros(image_shape[0:2]+( len(color),),dtype=np.float32)
    lmrks = expand_eyebrows(image_landmarks, eyebrows_expand_mod)
    draw_landmarks (mask, image_landmarks, color=color, draw_circles=False, thickness=2)
    return mask

def draw_landmarks (image, image_landmarks, color=(0,255,0), draw_circles=True, thickness=1, transparent_mask=False):
    if len(image_landmarks) != 68:
        raise Exception('get_image_eye_mask works only with 68 landmarks')

    int_lmrks = np.array(image_landmarks, dtype=np.int)

    jaw = int_lmrks[slice(*landmarks_68_pt["jaw"])]
    right_eyebrow = int_lmrks[slice(*landmarks_68_pt["right_eyebrow"])]
    left_eyebrow = int_lmrks[slice(*landmarks_68_pt["left_eyebrow"])]
    mouth = int_lmrks[slice(*landmarks_68_pt["mouth"])]
    right_eye = int_lmrks[slice(*landmarks_68_pt["right_eye"])]
    left_eye = int_lmrks[slice(*landmarks_68_pt["left_eye"])]
    nose = int_lmrks[slice(*landmarks_68_pt["nose"])]

    # open shapes
    cv2.polylines(image, tuple(np.array([v]) for v in ( right_eyebrow, jaw, left_eyebrow, np.concatenate((nose, [nose[-6]])) )),
                  False, color, thickness=thickness, lineType=cv2.LINE_AA)
    # closed shapes
    cv2.polylines(image, tuple(np.array([v]) for v in (right_eye, left_eye, mouth)),
                  True, color, thickness=thickness, lineType=cv2.LINE_AA)

    if draw_circles:
        # the rest of the cicles
        for x, y in np.concatenate((right_eyebrow, left_eyebrow, mouth, right_eye, left_eye, nose), axis=0):
            cv2.circle(image, (x, y), 1, color, 1, lineType=cv2.LINE_AA)
        # jaw big circles
        for x, y in jaw:
            cv2.circle(image, (x, y), 2, color, lineType=cv2.LINE_AA)

    if transparent_mask:
        mask = get_image_hull_mask (image.shape, image_landmarks)
        image[...] = ( image * (1-mask) + image * mask / 2 )[...]

def draw_rect_landmarks (image, rect, image_landmarks, face_type, face_size=256, transparent_mask=False, landmarks_color=(0,255,0)):
    draw_landmarks(image, image_landmarks, color=landmarks_color, transparent_mask=transparent_mask)
    imagelib.draw_rect (image, rect, (255,0,0), 2 )

    image_to_face_mat = get_transform_mat (image_landmarks, face_size, face_type)
    points = transform_points ( [ (0,0), (0,face_size-1), (face_size-1, face_size-1), (face_size-1,0) ], image_to_face_mat, True)
    imagelib.draw_polygon (image, points, (0,0,255), 2)

    points = transform_points ( [ ( int(face_size*0.05), 0), ( int(face_size*0.1), int(face_size*0.1) ), ( 0, int(face_size*0.1) ) ], image_to_face_mat, True)
    imagelib.draw_polygon (image, points, (0,0,255), 2)

def calc_face_pitch(landmarks):
    if not isinstance(landmarks, np.ndarray):
        landmarks = np.array (landmarks)
    t = ( (landmarks[6][1]-landmarks[8][1]) + (landmarks[10][1]-landmarks[8][1]) ) / 2.0
    b = landmarks[8][1]
    return float(b-t)

def estimate_averaged_yaw(landmarks):
    # Works much better than solvePnP if landmarks from "3DFAN"
    if not isinstance(landmarks, np.ndarray):
        landmarks = np.array (landmarks)
    l = ( (landmarks[27][0]-landmarks[0][0]) + (landmarks[28][0]-landmarks[1][0]) + (landmarks[29][0]-landmarks[2][0]) ) / 3.0   
    r = ( (landmarks[16][0]-landmarks[27][0]) + (landmarks[15][0]-landmarks[28][0]) + (landmarks[14][0]-landmarks[29][0]) ) / 3.0
    return float(r-l)
    
def estimate_pitch_yaw_roll(aligned_landmarks, size=256):
    """
    returns pitch,yaw,roll [-pi/2...+pi/2]
    """
    shape = (size,size)
    focal_length = shape[1]
    camera_center = (shape[1] / 2, shape[0] / 2)
    camera_matrix = np.array(
        [[focal_length, 0, camera_center[0]],
         [0, focal_length, camera_center[1]],
         [0, 0, 1]], dtype=np.float32)

    (_, rotation_vector, _) = cv2.solvePnP(
        np.concatenate( (landmarks_68_3D[:27],   landmarks_68_3D[30:36]) , axis=0) ,
        np.concatenate( (aligned_landmarks[:27], aligned_landmarks[30:36]) , axis=0).astype(np.float32),
        camera_matrix,
        np.zeros((4, 1)) )

    pitch, yaw, roll = mathlib.rotationMatrixToEulerAngles( cv2.Rodrigues(rotation_vector)[0] )
   
    half_pi = math.pi / 2.0
    pitch = np.clip ( pitch, -half_pi, half_pi )
    yaw   = np.clip ( yaw ,  -half_pi, half_pi )
    roll  = np.clip ( roll,  -half_pi, half_pi )

    return -pitch, yaw, roll

#if remove_align:
#    bbox = transform_points ( [ (0,0), (0,output_size), (output_size, output_size), (output_size,0) ], mat, True)
#    #import code
#    #code.interact(local=dict(globals(), **locals()))
#    area = mathlib.polygon_area(bbox[:,0], bbox[:,1] )
#    side = math.sqrt(area) / 2
#    center = transform_points ( [(output_size/2,output_size/2)], mat, True)
#    pts1 = np.float32(( center+[-side,-side], center+[side,-side], center+[side,-side] ))
#    pts2 = np.float32([[0,0],[output_size,0],[0,output_size]])
#    mat = cv2.getAffineTransform(pts1,pts2)
#if full_face_align_top and (face_type == FaceType.FULL or face_type == FaceType.FULL_NO_ALIGN):
#    #lmrks2 = expand_eyebrows(image_landmarks)
#    #lmrks2_ = transform_points( [ lmrks2[19], lmrks2[24] ], mat, False )
#    #y_diff = np.float32( (0,np.min(lmrks2_[:,1])) )
#    #y_diff = transform_points( [ np.float32( (0,0) ), y_diff], mat, True)
#    #y_diff = y_diff[1]-y_diff[0]
#
#    x_diff = np.float32((0,0))
#
#    lmrks2_ = transform_points( [ image_landmarks[0], image_landmarks[16] ], mat, False )
#    if lmrks2_[0,0] < 0:
#        x_diff = lmrks2_[0,0]
#        x_diff = transform_points( [ np.float32( (0,0) ), np.float32((x_diff,0)) ], mat, True)
#        x_diff = x_diff[1]-x_diff[0]
#    elif lmrks2_[1,0] >= output_size:
#        x_diff = lmrks2_[1,0]-(output_size-1)
#        x_diff = transform_points( [ np.float32( (0,0) ), np.float32((x_diff,0)) ], mat, True)
#        x_diff = x_diff[1]-x_diff[0]
#
#    mat = cv2.getAffineTransform( l_t+y_diff+x_diff ,pts2)


"""
def get_averaged_transform_mat (img_landmarks,
                                img_landmarks_prev,
                                img_landmarks_next,
                                average_frame_count,
                                average_center_frame_count,
                                output_size, face_type, scale=1.0):

    l_c_list = []
    tb_diag_vec_list = []
    bt_diag_vec_list = []
    mod_list = []

    count = max(average_frame_count,average_center_frame_count)
    for i in range ( -count, count+1, 1 ):
        if i < 0:
            lmrks = img_landmarks_prev[i] if -i < len(img_landmarks_prev) else None
        elif i > 0:
            lmrks = img_landmarks_next[i] if i < len(img_landmarks_next) else None
        else:
            lmrks = img_landmarks

        if lmrks is None:
            continue

        l_c, tb_diag_vec, bt_diag_vec, mod = get_transform_mat_data (lmrks, face_type, scale=scale)

        if i >= -average_frame_count and i <= average_frame_count:
            tb_diag_vec_list.append(tb_diag_vec)
            bt_diag_vec_list.append(bt_diag_vec)
            mod_list.append(mod)

        if i >= -average_center_frame_count and i <= average_center_frame_count:
            l_c_list.append(l_c)

    tb_diag_vec = np.mean( np.array(tb_diag_vec_list), axis=0 )
    bt_diag_vec = np.mean( np.array(bt_diag_vec_list), axis=0 )
    mod         = np.mean( np.array(mod_list), axis=0 )
    l_c         = np.mean( np.array(l_c_list), axis=0 )

    return get_transform_mat_by_data (l_c, tb_diag_vec, bt_diag_vec, mod, output_size, face_type)


def get_transform_mat (image_landmarks, output_size, face_type, scale=1.0):
    if not isinstance(image_landmarks, np.ndarray):
        image_landmarks = np.array (image_landmarks)

    # get face padding value for FaceType
    padding, remove_align = FaceType_to_padding_remove_align.get(face_type, 0.0)

    # estimate landmarks transform from global space to local aligned space with bounds [0..1]
    mat = umeyama( np.concatenate ( [ image_landmarks[17:49] , image_landmarks[54:55] ] ) , landmarks_2D_new, True)[0:2]

    # get corner points in global space
    l_p = transform_points (  np.float32([(0,0),(1,0),(1,1),(0,1),(0.5,0.5)]) , mat, True)
    l_c = l_p[4]

    # calc diagonal vectors between corners in global space
    tb_diag_vec = (l_p[2]-l_p[0]).astype(np.float32)
    tb_diag_vec /= npla.norm(tb_diag_vec)
    bt_diag_vec = (l_p[1]-l_p[3]).astype(np.float32)
    bt_diag_vec /= npla.norm(bt_diag_vec)

    # calc modifier of diagonal vectors for scale and padding value
    mod = (1.0 / scale)* ( npla.norm(l_p[0]-l_p[2])*(padding*np.sqrt(2.0) + 0.5) )

    # calc 3 points in global space to estimate 2d affine transform
    if not remove_align:
        l_t = np.array( [ np.round( l_c - tb_diag_vec*mod ),
                          np.round( l_c + bt_diag_vec*mod ),
                          np.round( l_c + tb_diag_vec*mod ) ] )
    else:
        # remove_align - face will be centered in the frame but not aligned
        l_t = np.array( [ np.round( l_c - tb_diag_vec*mod ),
                          np.round( l_c + bt_diag_vec*mod ),
                          np.round( l_c + tb_diag_vec*mod ),
                          np.round( l_c - bt_diag_vec*mod ),
                         ] )

        # get area of face square in global space
        area = mathlib.polygon_area(l_t[:,0], l_t[:,1] )

        # calc side of square
        side = np.float32(math.sqrt(area) / 2)

        # calc 3 points with unrotated square
        l_t = np.array( [ np.round( l_c + [-side,-side] ),
                          np.round( l_c + [ side,-side] ),
                          np.round( l_c + [ side, side] ) ] )

    # calc affine transform from 3 global space points to 3 local space points size of 'output_size'
    pts2 = np.float32(( (0,0),(output_size,0),(output_size,output_size) ))
    mat = cv2.getAffineTransform(l_t,pts2)

    return mat
"""
facelib/LandmarksProcessor.py
oo92/TorchFaceLab
from pathlib import Path

import cv2, torch, operator
import numpy as np
import torch.nn as nn
import torch.nn.functional as F

class L2Norm(nn.Module):
    def __init__(self, n_channels):
        super(L2Norm, self).__init__()
        self.n_channels = n_channels
        self.weight = nn.Parameter(torch.ones(1, self.n_channels, 1, 1))

    def forward(self, x):
        norm = torch.sqrt(torch.sum(x**2, dim=1, keepdim=True))
        return self.weight * x / (norm + 1e-10)

class S3FD(nn.Module):
    def __init__(self):
        super(S3FD, self).__init__()
        
        self.minus = torch.tensor([104,117,123], dtype=torch.float32 ).view(1,3,1,1)

        self.conv1_1 = nn.Conv2d(3, 64, 3, 1, 1)
        self.conv1_2 = nn.Conv2d(64, 64, 3, 1, 1)
        self.conv2_1 = nn.Conv2d(64, 128, 3, 1, 1)
        self.conv2_2 = nn.Conv2d(128, 128, 3, 1, 1)
        self.conv3_1 = nn.Conv2d(128, 256, 3, 1, 1)
        self.conv3_2 = nn.Conv2d(256, 256, 3, 1, 1)
        self.conv3_3 = nn.Conv2d(256, 256, 3, 1, 1)
        self.conv4_1 = nn.Conv2d(256, 512, 3, 1, 1)
        self.conv4_2 = nn.Conv2d(512, 512, 3, 1, 1)
        self.conv4_3 = nn.Conv2d(512, 512, 3, 1, 1)
        self.conv5_1 = nn.Conv2d(512, 512, 3, 1, 1)
        self.conv5_2 = nn.Conv2d(512, 512, 3, 1, 1)
        self.conv5_3 = nn.Conv2d(512, 512, 3, 1, 1)
        self.fc6 = nn.Conv2d(512, 1024, 3, 1, 3)
        self.fc7 = nn.Conv2d(1024, 1024, 1, 1, 0)
        self.conv6_1 = nn.Conv2d(1024, 256, 1, 1, 0)
        self.conv6_2 = nn.Conv2d(256, 512, 3, 2, 1)
        self.conv7_1 = nn.Conv2d(512, 128, 1, 1, 0)
        self.conv7_2 = nn.Conv2d(128, 256, 3, 2, 1)

        self.conv3_3_norm = L2Norm(256)
        self.conv4_3_norm = L2Norm(512)
        self.conv5_3_norm = L2Norm(512)

        self.conv3_3_norm_mbox_conf = nn.Conv2d(256, 4, 3, 1, 1)
        self.conv3_3_norm_mbox_loc = nn.Conv2d(256, 4, 3, 1, 1)

        self.conv4_3_norm_mbox_conf = nn.Conv2d(512, 2, 3, 1, 1)
        self.conv4_3_norm_mbox_loc = nn.Conv2d(512, 4, 3, 1, 1)

        self.conv5_3_norm_mbox_conf = nn.Conv2d(512, 2, 3, 1, 1)
        self.conv5_3_norm_mbox_loc = nn.Conv2d(512, 4, 3, 1, 1)

        self.fc7_mbox_conf = nn.Conv2d(1024, 2, 3, 1, 1)
        self.fc7_mbox_loc = nn.Conv2d(1024, 4, 3, 1, 1)

        self.conv6_2_mbox_conf = nn.Conv2d(512, 2, 3, 1, 1)
        self.conv6_2_mbox_loc = nn.Conv2d(512, 4, 3, 1, 1)

        self.conv7_2_mbox_conf = nn.Conv2d(256, 2, 3, 1, 1)
        self.conv7_2_mbox_loc = nn.Conv2d(256, 4, 3, 1, 1)

    def forward(self, x):
        x = x.cuda()
        self.minus = self.minus.cuda()
        x = x - self.minus

        # Initial layers
        h = F.relu(self.conv1_1(x))
        h = F.relu(self.conv1_2(h))
        h = F.max_pool2d(h, 2, 2)

        h = F.relu(self.conv2_1(h))
        h = F.relu(self.conv2_2(h))
        h = F.max_pool2d(h, 2, 2)

        h = F.relu(self.conv3_1(h))
        h = F.relu(self.conv3_2(h))
        h = F.relu(self.conv3_3(h))
        f3_3 = self.conv3_3_norm(h)
        h = F.max_pool2d(h, 2, 2)

        h = F.relu(self.conv4_1(h))
        h = F.relu(self.conv4_2(h))
        h = F.relu(self.conv4_3(h))
        f4_3 = self.conv4_3_norm(h)
        h = F.max_pool2d(h, 2, 2)

        h = F.relu(self.conv5_1(h))
        h = F.relu(self.conv5_2(h))
        h = F.relu(self.conv5_3(h))
        f5_3 = self.conv5_3_norm(h)
        h = F.max_pool2d(h, 2, 2)

        h = F.relu(self.fc6(h))
        h = F.relu(self.fc7(h))
        ffc7 = h 

        h = F.relu(self.conv6_1(h))
        h = F.relu(self.conv6_2(h))
        f6_2 = h

        h = F.relu(self.conv7_1(h))
        h = F.relu(self.conv7_2(h))
        f7_2 = h

        cls1 = F.softmax(self.conv3_3_norm_mbox_conf(f3_3).permute(0, 2, 3, 1).contiguous().view(x.size(0), -1, 4), dim=-1)
        reg1 = self.conv3_3_norm_mbox_loc(f3_3).permute(0, 2, 3, 1).contiguous().view(x.size(0), -1, 4)

        cls2 = F.softmax(self.conv4_3_norm_mbox_conf(f4_3).permute(0, 2, 3, 1).contiguous().view(x.size(0), -1, 2), dim=-1)
        reg2 = self.conv4_3_norm_mbox_loc(f4_3).permute(0, 2, 3, 1).contiguous().view(x.size(0), -1, 4)

        cls3 = F.softmax(self.conv5_3_norm_mbox_conf(f5_3).permute(0, 2, 3, 1).contiguous().view(x.size(0), -1, 2), dim=-1)
        reg3 = self.conv5_3_norm_mbox_loc(f5_3).permute(0, 2, 3, 1).contiguous().view(x.size(0), -1, 4)

        cls4 = F.softmax(self.fc7_mbox_conf(ffc7).permute(0, 2, 3, 1).contiguous().view(x.size(0), -1, 2), dim=-1)
        reg4 = self.fc7_mbox_loc(ffc7).permute(0, 2, 3, 1).contiguous().view(x.size(0), -1, 4)

        cls5 = F.softmax(self.conv6_2_mbox_conf(f6_2).permute(0, 2, 3, 1).contiguous().view(x.size(0), -1, 2), dim=-1)
        reg5 = self.conv6_2_mbox_loc(f6_2).permute(0, 2, 3, 1).contiguous().view(x.size(0), -1, 4)

        cls6 = F.softmax(self.conv7_2_mbox_conf(f7_2).permute(0, 2, 3, 1).contiguous().view(x.size(0), -1, 2), dim=-1)
        reg6 = self.conv7_2_mbox_loc(f7_2).permute(0, 2, 3, 1).contiguous().view(x.size(0), -1, 4)

        return [cls1, reg1, cls2, reg2, cls3, reg3, cls4, reg4, cls5, reg5, cls6, reg6]



class S3FDExtractor(object):
    def __init__(self, place_model_on_cpu=False):
        model_path = Path(__file__).parent / "S3FD.npy"
        if not model_path.exists():
            raise Exception("Unable to load S3FD.npy")

        self.model = S3FD()
        self.device = torch.device('cuda' if torch.cuda.is_available() and not place_model_on_cpu else 'cpu')
        self.model.to(self.device)
        
        def can_squeeze(t):
            shape_set = set(t.shape)
            if len(shape_set) == 2 and 1 in shape_set:
                return True
            return False

        def reshape_tensor(k, t):
            if k == 'fc7.weight':
                return t.reshape([1024, 1024, 1, 1])
            if k in ['conv4_3_norm.weight', 'conv5_3_norm.weight']:
                return t.reshape([1, 512, 1, 1])
            if k == 'conv3_3_norm.weight':
                return t.reshape([1, 256, 1, 1])
            if can_squeeze(t):
                return t.squeeze()
            else:
                return torch.permute(t, [3, 2, 1, 0])

        def rename_key(key):
            new_key = key.split(":")[0] # discard :0 and similar
            key_elements = new_key.split("/")
            new_key = ".".join(key_elements) # replace every / with .
            return new_key

        # Load the state dictionary from the numpy file
        state_dict_np = np.load(model_path, allow_pickle=True)

        # Convert numpy arrays within the state dictionary to PyTorch tensors
        state_dict_torch = {rename_key(k): reshape_tensor(rename_key(k), torch.tensor(v, dtype=torch.float32)).cpu() for k, v in state_dict_np.items()}

        # Load the converted state dictionary into the model
        self.model.load_state_dict(state_dict_torch)


    def __enter__(self):
        return self

    def __exit__(self, exc_type=None, exc_value=None, traceback=None):
        return False  # pass exception between __enter__ and __exit__ to outer level

    def extract(self, input_image, is_bgr=True, is_remove_intersects=False):
        # Convert BGR to RGB
        if is_bgr:
            input_image = input_image[:, :, ::-1]

        (h, w, ch) = input_image.shape

        # Scale logic
        d = max(w, h)
        scale_to = 640 if d >= 1280 else d / 2
        scale_to = max(64, scale_to)

        input_scale = d / scale_to
        input_image = cv2.resize(input_image, (int(w/input_scale), int(h/input_scale)), interpolation=cv2.INTER_LINEAR)

        # Convert to PyTorch tensor and normalize
        input_tensor = torch.from_numpy(input_image.transpose(2, 0, 1)).float().unsqueeze(0) / 255.0
        input_tensor = input_tensor.to(self.device)
        
        # Forward pass the tensor through the model
        outputs = self.model(input_tensor)

        detected_faces = []
        for output in outputs:
            for ltrb in output:
                print('ltrb: ', ltrb)
                l, t, r, b = [x*input_scale for x in ltrb] # <- Line causing the current issue
                if min(r-l, b-t) < 40:
                    continue
                detected_faces.append([int(x) for x in (l, t, r, b)])

        # Sort by largest area first
        detected_faces = [[(l, t, r, b), (r-l)*(b-t)] for (l, t, r, b) in detected_faces]
        detected_faces = sorted(detected_faces, key=operator.itemgetter(1), reverse=True)
        detected_faces = [x[0] for x in detected_faces]

        # Handling intersections
        if is_remove_intersects:
            for i in range(len(detected_faces)-1, 0, -1):
                l1, t1, r1, b1 = detected_faces[i]
                l0, t0, r0, b0 = detected_faces[i-1]
                dx = min(r0, r1) - max(l0, l1)
                dy = min(b0, b1) - max(t0, t1)
                if (dx >= 0) and (dy >= 0):
                    detected_faces.pop(i)

        return detected_faces


    def refine(self, olist):
        """
        Refine the output list to remove duplicates, apply confidence thresholding,
        and convert box coordinates if necessary.
        
        :param olist: A list containing raw model outputs.
        :return: A list of refined bounding boxes.
        """
        
        # Assuming each output in olist is a tuple of (confidence, bbox)
        confidences, boxes = zip(*olist)
        
        # Convert confidences and boxes to tensors for processing
        confidences = torch.tensor(confidences)
        boxes = torch.tensor(boxes)
        
        # Apply confidence thresholding (e.g., 0.5)
        mask = confidences > 0.5
        confidences = confidences[mask]
        boxes = boxes[mask]
        
        # Perform Non-Maximum Suppression
        # Here we're assuming you have an NMS function available.
        # PyTorch's torchvision library provides one such function.
        keep = nms(boxes, confidences, iou_threshold=0.5)
        
        # Convert boxes from relative to absolute coordinates if necessary
        # This step depends on how your model outputs bounding box coordinates.
        # If it's already in the desired format, you can skip this step.
        
        # width and height are assumed to be the dimensions of the input image.
        # If they are different for each image, you will need to pass them to this function.
        width, height = 300, 300  # Placeholder values
        boxes[keep, 0] *= width
        boxes[keep, 1] *= height
        boxes[keep, 2] *= width
        boxes[keep, 3] *= height
        
        return boxes[keep].tolist()


facelib/S3FDExtractor.py
oo92/TorchFaceLab
import os
import pickle
from functools import partial
from pathlib import Path

import cv2
import numpy as np

from core.interact import interact as io
from core.leras import nn


class XSegNet(object):
    VERSION = 1

    def __init__ (self, name, 
                        resolution=256, 
                        load_weights=True, 
                        weights_file_root=None, 
                        training=False, 
                        place_model_on_cpu=False, 
                        run_on_cpu=False, 
                        optimizer=None, 
                        data_format="NHWC",
                        raise_on_no_model_files=False):
                
        self.resolution = resolution
        self.weights_file_root = Path(weights_file_root) if weights_file_root is not None else Path(__file__).parent
        
        nn.initialize(data_format=data_format)
        tf = nn.tf
        
        model_name = f'{name}_{resolution}'
        self.model_filename_list = []
        
        with tf.device ('/CPU:0'):
            #Place holders on CPU
            self.input_t  = tf.placeholder (nn.floatx, nn.get4Dshape(resolution,resolution,3) )
            self.target_t = tf.placeholder (nn.floatx, nn.get4Dshape(resolution,resolution,1) )

        # Initializing model classes
        with tf.device ('/CPU:0' if place_model_on_cpu else nn.tf_default_device_name):
            self.model = nn.XSeg(3, 32, 1, name=name)
            self.model_weights = self.model.get_weights()
            if training:
                if optimizer is None:
                    raise ValueError("Optimizer should be provided for training mode.")                
                self.opt = optimizer              
                self.opt.initialize_variables (self.model_weights, vars_on_cpu=place_model_on_cpu)                    
                self.model_filename_list += [ [self.opt, f'{model_name}_opt.npy' ] ]
                
        
        self.model_filename_list += [ [self.model, f'{model_name}.npy'] ]

        if not training:
            with tf.device ('/CPU:0' if run_on_cpu else nn.tf_default_device_name):
                _, pred = self.model(self.input_t)

            def net_run(input_np):
                return nn.tf_sess.run ( [pred], feed_dict={self.input_t :input_np})[0]
            self.net_run = net_run

        self.initialized = True
        # Loading/initializing all models/optimizers weights
        for model, filename in self.model_filename_list:
            do_init = not load_weights

            if not do_init:
                model_file_path = self.weights_file_root / filename
                do_init = not model.load_weights( model_file_path )
                if do_init:
                    if raise_on_no_model_files:
                        raise Exception(f'{model_file_path} does not exists.')
                    if not training:
                        self.initialized = False
                        break

            if do_init:
                model.init_weights()
        
    def get_resolution(self):
        return self.resolution
        
    def flow(self, x, pretrain=False):
        return self.model(x, pretrain=pretrain)

    def get_weights(self):
        return self.model_weights

    def save_weights(self):
        for model, filename in io.progress_bar_generator(self.model_filename_list, "Saving", leave=False):
            model.save_weights( self.weights_file_root / filename )

    def extract (self, input_image):
        if not self.initialized:
            return 0.5*np.ones ( (self.resolution, self.resolution, 1), nn.floatx.as_numpy_dtype )
            
        input_shape_len = len(input_image.shape)            
        if input_shape_len == 3:
            input_image = input_image[None,...]

        result = np.clip ( self.net_run(input_image), 0, 1.0 )
        result[result < 0.1] = 0 #get rid of noise

        if input_shape_len == 3:
            result = result[0]

        return result
facelib/XSegNet.py
oo92/TorchFaceLab
from .FaceType import FaceType
from .S3FDExtractor import S3FDExtractor
from .FANExtractor import FANExtractor
from .FaceEnhancer import FaceEnhancer
from .XSegNet import XSegNet
facelib/__init__.py
oo92/TorchFaceLab
from .localization import StringsDB, system_language, get_default_ttf_font_name


localization/__init__.py
oo92/TorchFaceLab
﻿import sys
import locale

system_locale = locale.getdefaultlocale()[0]
# system_locale may be nil
system_language = system_locale[0:2] if system_locale is not None else "en"
if system_language not in ['en','ru','zh']:
    system_language = 'en'

windows_font_name_map = {
    'en' : 'cour',
    'ru' : 'cour',
    'zh' : 'simsun_01'
}

darwin_font_name_map = {
    'en' : 'cour',
    'ru' : 'cour',
    'zh' : 'Apple LiSung Light'
}

linux_font_name_map = {
    'en' : 'cour',
    'ru' : 'cour',
    'zh' : 'cour'
}

def get_default_ttf_font_name():
    platform = sys.platform
    if platform[0:3] == 'win': return windows_font_name_map.get(system_language, 'cour')
    elif platform == 'darwin': return darwin_font_name_map.get(system_language, 'cour')
    else: return linux_font_name_map.get(system_language, 'cour')

SID_HOT_KEY = 1

if system_language == 'en':
    StringsDB = {'S_HOT_KEY' : 'hot key'}
elif system_language == 'ru':
    StringsDB = {'S_HOT_KEY' : 'горячая клавиша'}    
elif system_language == 'zh':
    StringsDB = {'S_HOT_KEY' : '热键'}   
    
localization/localization.py
oo92/TorchFaceLab
if __name__ == "__main__":
    # Fix for linux
    import multiprocessing
    multiprocessing.set_start_method("spawn")

    from core.leras import nn
    # nn.initialize_main_env()
    import os
    import sys
    import time
    import argparse

    from core import pathex
    from core import osex
    from pathlib import Path
    from core.interact import interact as io
    from core.leras.device import Devices
    Devices.initialize_main_env()
    
    if sys.version_info[0] < 3 or (sys.version_info[0] == 3 and sys.version_info[1] < 6):
        raise Exception("This program requires at least Python 3.6")

    class fixPathAction(argparse.Action):
        def __call__(self, parser, namespace, values, option_string=None):
            setattr(namespace, self.dest, os.path.abspath(os.path.expanduser(values)))

    exit_code = 0
    
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers()

    def process_extract(arguments):
        osex.set_process_lowest_prio()
        from mainscripts import Extractor
        Extractor.main( detector                = arguments.detector,
                        input_path              = Path(arguments.input_dir),
                        output_path             = Path(arguments.output_dir),
                        output_debug            = arguments.output_debug,
                        manual_fix              = arguments.manual_fix,
                        manual_output_debug_fix = arguments.manual_output_debug_fix,
                        manual_window_size      = arguments.manual_window_size,
                        face_type               = arguments.face_type,
                        max_faces_from_image    = arguments.max_faces_from_image,
                        image_size              = arguments.image_size,
                        jpeg_quality            = arguments.jpeg_quality,
                        cpu_only                = arguments.cpu_only,
                        force_gpu_idxs          = [ int(x) for x in arguments.force_gpu_idxs.split(',') ] if arguments.force_gpu_idxs is not None else None,
                      )

    p = subparsers.add_parser( "extract", help="Extract the faces from a pictures.")
    p.add_argument('--detector', dest="detector", choices=['s3fd','manual'], default=None, help="Type of detector.")
    p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
    p.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the extracted files will be stored.")
    p.add_argument('--output-debug', action="store_true", dest="output_debug", default=None, help="Writes debug images to <output-dir>_debug\ directory.")
    p.add_argument('--no-output-debug', action="store_false", dest="output_debug", default=None, help="Don't writes debug images to <output-dir>_debug\ directory.")
    p.add_argument('--face-type', dest="face_type", choices=['half_face', 'full_face', 'whole_face', 'head', 'mark_only'], default=None)
    p.add_argument('--max-faces-from-image', type=int, dest="max_faces_from_image", default=None, help="Max faces from image.")    
    p.add_argument('--image-size', type=int, dest="image_size", default=None, help="Output image size.")
    p.add_argument('--jpeg-quality', type=int, dest="jpeg_quality", default=None, help="Jpeg quality.")    
    p.add_argument('--manual-fix', action="store_true", dest="manual_fix", default=False, help="Enables manual extract only frames where faces were not recognized.")
    p.add_argument('--manual-output-debug-fix', action="store_true", dest="manual_output_debug_fix", default=False, help="Performs manual reextract input-dir frames which were deleted from [output_dir]_debug\ dir.")
    p.add_argument('--manual-window-size', type=int, dest="manual_window_size", default=1368, help="Manual fix window size. Default: 1368.")
    p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Extract on CPU..")
    p.add_argument('--force-gpu-idxs', dest="force_gpu_idxs", default=None, help="Force to choose GPU indexes separated by comma.")

    p.set_defaults (func=process_extract)

    def process_sort(arguments):
        osex.set_process_lowest_prio()
        from mainscripts import Sorter
        Sorter.main (input_path=Path(arguments.input_dir), sort_by_method=arguments.sort_by_method)

    p = subparsers.add_parser( "sort", help="Sort faces in a directory.")
    p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
    p.add_argument('--by', dest="sort_by_method", default=None, choices=("blur", "motion-blur", "face-yaw", "face-pitch", "face-source-rect-size", "hist", "hist-dissim", "brightness", "hue", "black", "origname", "oneface", "final-by-blur", "final-by-size", "absdiff"), help="Method of sorting. 'origname' sort by original filename to recover original sequence." )
    p.set_defaults (func=process_sort)

    def process_util(arguments):
        osex.set_process_lowest_prio()
        from mainscripts import Util

        if arguments.add_landmarks_debug_images:
            Util.add_landmarks_debug_images (input_path=arguments.input_dir)

        if arguments.recover_original_aligned_filename:
            Util.recover_original_aligned_filename (input_path=arguments.input_dir)

        if arguments.save_faceset_metadata:
            Util.save_faceset_metadata_folder (input_path=arguments.input_dir)

        if arguments.restore_faceset_metadata:
            Util.restore_faceset_metadata_folder (input_path=arguments.input_dir)

        if arguments.pack_faceset:
            io.log_info ("Performing faceset packing...\r\n")
            from samplelib import PackedFaceset
            PackedFaceset.pack( Path(arguments.input_dir) )

        if arguments.unpack_faceset:
            io.log_info ("Performing faceset unpacking...\r\n")
            from samplelib import PackedFaceset
            PackedFaceset.unpack( Path(arguments.input_dir) )
            
        if arguments.export_faceset_mask:
            io.log_info ("Exporting faceset mask..\r\n")
            Util.export_faceset_mask( Path(arguments.input_dir) )

    p = subparsers.add_parser( "util", help="Utilities.")
    p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
    p.add_argument('--add-landmarks-debug-images', action="store_true", dest="add_landmarks_debug_images", default=False, help="Add landmarks debug image for aligned faces.")
    p.add_argument('--recover-original-aligned-filename', action="store_true", dest="recover_original_aligned_filename", default=False, help="Recover original aligned filename.")
    p.add_argument('--save-faceset-metadata', action="store_true", dest="save_faceset_metadata", default=False, help="Save faceset metadata to file.")
    p.add_argument('--restore-faceset-metadata', action="store_true", dest="restore_faceset_metadata", default=False, help="Restore faceset metadata to file. Image filenames must be the same as used with save.")
    p.add_argument('--pack-faceset', action="store_true", dest="pack_faceset", default=False, help="")
    p.add_argument('--unpack-faceset', action="store_true", dest="unpack_faceset", default=False, help="")
    p.add_argument('--export-faceset-mask', action="store_true", dest="export_faceset_mask", default=False, help="")

    p.set_defaults (func=process_util)

    def process_train(arguments):
        osex.set_process_lowest_prio()


        kwargs = {'model_class_name'         : arguments.model_name,
                  'saved_models_path'        : Path(arguments.model_dir),
                  'training_data_src_path'   : Path(arguments.training_data_src_dir),
                  'training_data_dst_path'   : Path(arguments.training_data_dst_dir),
                  'pretraining_data_path'    : Path(arguments.pretraining_data_dir) if arguments.pretraining_data_dir is not None else None,
                  'pretrained_model_path'    : Path(arguments.pretrained_model_dir) if arguments.pretrained_model_dir is not None else None,
                  'no_preview'               : arguments.no_preview,
                  'force_model_name'         : arguments.force_model_name,
                  'force_gpu_idxs'           : [ int(x) for x in arguments.force_gpu_idxs.split(',') ] if arguments.force_gpu_idxs is not None else None,
                  'cpu_only'                 : arguments.cpu_only,
                  'silent_start'             : arguments.silent_start,
                  'execute_programs'         : [ [int(x[0]), x[1] ] for x in arguments.execute_program ],
                  'debug'                    : arguments.debug,
                  }
        from mainscripts import Trainer
        Trainer.main(**kwargs)

    p = subparsers.add_parser( "train", help="Trainer")
    p.add_argument('--training-data-src-dir', required=True, action=fixPathAction, dest="training_data_src_dir", help="Dir of extracted SRC faceset.")
    p.add_argument('--training-data-dst-dir', required=True, action=fixPathAction, dest="training_data_dst_dir", help="Dir of extracted DST faceset.")
    p.add_argument('--pretraining-data-dir', action=fixPathAction, dest="pretraining_data_dir", default=None, help="Optional dir of extracted faceset that will be used in pretraining mode.")
    p.add_argument('--pretrained-model-dir', action=fixPathAction, dest="pretrained_model_dir", default=None, help="Optional dir of pretrain model files. (Currently only for Quick96).")
    p.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Saved models dir.")
    p.add_argument('--model', required=True, dest="model_name", choices=pathex.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Model class name.")
    p.add_argument('--debug', action="store_true", dest="debug", default=False, help="Debug samples.")
    p.add_argument('--no-preview', action="store_true", dest="no_preview", default=False, help="Disable preview window.")
    p.add_argument('--force-model-name', dest="force_model_name", default=None, help="Forcing to choose model name from model/ folder.")
    p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Train on CPU.")
    p.add_argument('--force-gpu-idxs', dest="force_gpu_idxs", default=None, help="Force to choose GPU indexes separated by comma.")
    p.add_argument('--silent-start', action="store_true", dest="silent_start", default=False, help="Silent start. Automatically chooses Best GPU and last used model.")
    
    p.add_argument('--execute-program', dest="execute_program", default=[], action='append', nargs='+')
    p.set_defaults (func=process_train)
    
    def process_exportdfm(arguments):
        osex.set_process_lowest_prio()
        from mainscripts import ExportDFM
        ExportDFM.main(model_class_name = arguments.model_name, saved_models_path = Path(arguments.model_dir))

    p = subparsers.add_parser( "exportdfm", help="Export model to use in DeepFaceLive.")
    p.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Saved models dir.")
    p.add_argument('--model', required=True, dest="model_name", choices=pathex.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Model class name.")
    p.set_defaults (func=process_exportdfm)

    def process_merge(arguments):
        osex.set_process_lowest_prio()
        from mainscripts import Merger
        Merger.main ( model_class_name       = arguments.model_name,
                      saved_models_path      = Path(arguments.model_dir),
                      force_model_name       = arguments.force_model_name,
                      input_path             = Path(arguments.input_dir),
                      output_path            = Path(arguments.output_dir),
                      output_mask_path       = Path(arguments.output_mask_dir),
                      aligned_path           = Path(arguments.aligned_dir) if arguments.aligned_dir is not None else None,
                      force_gpu_idxs         = arguments.force_gpu_idxs,
                      cpu_only               = arguments.cpu_only)

    p = subparsers.add_parser( "merge", help="Merger")
    p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
    p.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the merged files will be stored.")
    p.add_argument('--output-mask-dir', required=True, action=fixPathAction, dest="output_mask_dir", help="Output mask directory. This is where the mask files will be stored.")
    p.add_argument('--aligned-dir', action=fixPathAction, dest="aligned_dir", default=None, help="Aligned directory. This is where the extracted of dst faces stored.")
    p.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Model dir.")
    p.add_argument('--model', required=True, dest="model_name", choices=pathex.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Model class name.")
    p.add_argument('--force-model-name', dest="force_model_name", default=None, help="Forcing to choose model name from model/ folder.")
    p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Merge on CPU.")
    p.add_argument('--force-gpu-idxs', dest="force_gpu_idxs", default=None, help="Force to choose GPU indexes separated by comma.")
    p.set_defaults(func=process_merge)

    videoed_parser = subparsers.add_parser( "videoed", help="Video processing.").add_subparsers()

    def process_videoed_extract_video(arguments):
        osex.set_process_lowest_prio()
        from mainscripts import VideoEd
        VideoEd.extract_video (arguments.input_file, arguments.output_dir, arguments.output_ext, arguments.fps)
    p = videoed_parser.add_parser( "extract-video", help="Extract images from video file.")
    p.add_argument('--input-file', required=True, action=fixPathAction, dest="input_file", help="Input file to be processed. Specify .*-extension to find first file.")
    p.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the extracted images will be stored.")
    p.add_argument('--output-ext', dest="output_ext", default=None, help="Image format (extension) of output files.")
    p.add_argument('--fps', type=int, dest="fps", default=None, help="How many frames of every second of the video will be extracted. 0 - full fps.")
    p.set_defaults(func=process_videoed_extract_video)

    def process_videoed_cut_video(arguments):
        osex.set_process_lowest_prio()
        from mainscripts import VideoEd
        VideoEd.cut_video (arguments.input_file,
                           arguments.from_time,
                           arguments.to_time,
                           arguments.audio_track_id,
                           arguments.bitrate)
    p = videoed_parser.add_parser( "cut-video", help="Cut video file.")
    p.add_argument('--input-file', required=True, action=fixPathAction, dest="input_file", help="Input file to be processed. Specify .*-extension to find first file.")
    p.add_argument('--from-time', dest="from_time", default=None, help="From time, for example 00:00:00.000")
    p.add_argument('--to-time', dest="to_time", default=None, help="To time, for example 00:00:00.000")
    p.add_argument('--audio-track-id', type=int, dest="audio_track_id", default=None, help="Specify audio track id.")
    p.add_argument('--bitrate', type=int, dest="bitrate", default=None, help="Bitrate of output file in Megabits.")
    p.set_defaults(func=process_videoed_cut_video)

    def process_videoed_denoise_image_sequence(arguments):
        osex.set_process_lowest_prio()
        from mainscripts import VideoEd
        VideoEd.denoise_image_sequence (arguments.input_dir, arguments.factor)
    p = videoed_parser.add_parser( "denoise-image-sequence", help="Denoise sequence of images, keeping sharp edges. Helps to remove pixel shake from the predicted face.")
    p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory to be processed.")
    p.add_argument('--factor', type=int, dest="factor", default=None, help="Denoise factor (1-20).")
    p.set_defaults(func=process_videoed_denoise_image_sequence)

    def process_videoed_video_from_sequence(arguments):
        osex.set_process_lowest_prio()
        from mainscripts import VideoEd
        VideoEd.video_from_sequence (input_dir      = arguments.input_dir,
                                     output_file    = arguments.output_file,
                                     reference_file = arguments.reference_file,
                                     ext      = arguments.ext,
                                     fps      = arguments.fps,
                                     bitrate  = arguments.bitrate,
                                     include_audio = arguments.include_audio,
                                     lossless = arguments.lossless)

    p = videoed_parser.add_parser( "video-from-sequence", help="Make video from image sequence.")
    p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input file to be processed. Specify .*-extension to find first file.")
    p.add_argument('--output-file', required=True, action=fixPathAction, dest="output_file", help="Input file to be processed. Specify .*-extension to find first file.")
    p.add_argument('--reference-file', action=fixPathAction, dest="reference_file", help="Reference file used to determine proper FPS and transfer audio from it. Specify .*-extension to find first file.")
    p.add_argument('--ext', dest="ext", default='png', help="Image format (extension) of input files.")
    p.add_argument('--fps', type=int, dest="fps", default=None, help="FPS of output file. Overwritten by reference-file.")
    p.add_argument('--bitrate', type=int, dest="bitrate", default=None, help="Bitrate of output file in Megabits.")
    p.add_argument('--include-audio', action="store_true", dest="include_audio", default=False, help="Include audio from reference file.")
    p.add_argument('--lossless', action="store_true", dest="lossless", default=False, help="PNG codec.")

    p.set_defaults(func=process_videoed_video_from_sequence)

    facesettool_parser = subparsers.add_parser( "facesettool", help="Faceset tools.").add_subparsers()

    def process_faceset_enhancer(arguments):
        osex.set_process_lowest_prio()
        from mainscripts import FacesetEnhancer
        FacesetEnhancer.process_folder ( Path(arguments.input_dir),
                                         cpu_only=arguments.cpu_only,
                                         force_gpu_idxs=arguments.force_gpu_idxs
                                       )

    p = facesettool_parser.add_parser ("enhance", help="Enhance details in DFL faceset.")
    p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory of aligned faces.")
    p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Process on CPU.")
    p.add_argument('--force-gpu-idxs', dest="force_gpu_idxs", default=None, help="Force to choose GPU indexes separated by comma.")

    p.set_defaults(func=process_faceset_enhancer)
    
    
    p = facesettool_parser.add_parser ("resize", help="Resize DFL faceset.")
    p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory of aligned faces.")

    def process_faceset_resizer(arguments):
        osex.set_process_lowest_prio()
        from mainscripts import FacesetResizer
        FacesetResizer.process_folder ( Path(arguments.input_dir) )
    p.set_defaults(func=process_faceset_resizer)

    def process_dev_test(arguments):
        osex.set_process_lowest_prio()
        from mainscripts import dev_misc
        dev_misc.dev_gen_mask_files( arguments.input_dir )

    p = subparsers.add_parser( "dev_test", help="")
    p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
    p.set_defaults (func=process_dev_test)
    
    # ========== XSeg
    xseg_parser = subparsers.add_parser( "xseg", help="XSeg tools.").add_subparsers()
    
    p = xseg_parser.add_parser( "editor", help="XSeg editor.")

    def process_xsegeditor(arguments):
        osex.set_process_lowest_prio()
        from XSegEditor import XSegEditor
        global exit_code
        exit_code = XSegEditor.start (Path(arguments.input_dir))
        
    p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")

    p.set_defaults (func=process_xsegeditor)
  
    p = xseg_parser.add_parser( "apply", help="Apply trained XSeg model to the extracted faces.")

    def process_xsegapply(arguments):
        osex.set_process_lowest_prio()
        from mainscripts import XSegUtil
        XSegUtil.apply_xseg (Path(arguments.input_dir), Path(arguments.model_dir))
    p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
    p.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir")
    p.set_defaults (func=process_xsegapply)
    
    
    p = xseg_parser.add_parser( "remove", help="Remove applied XSeg masks from the extracted faces.")
    def process_xsegremove(arguments):
        osex.set_process_lowest_prio()
        from mainscripts import XSegUtil
        XSegUtil.remove_xseg (Path(arguments.input_dir) )
    p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
    p.set_defaults (func=process_xsegremove)
    
    
    p = xseg_parser.add_parser( "remove_labels", help="Remove XSeg labels from the extracted faces.")
    def process_xsegremovelabels(arguments):
        osex.set_process_lowest_prio()
        from mainscripts import XSegUtil
        XSegUtil.remove_xseg_labels (Path(arguments.input_dir) )
    p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
    p.set_defaults (func=process_xsegremovelabels)
    
    
    p = xseg_parser.add_parser( "fetch", help="Copies faces containing XSeg polygons in <input_dir>_xseg dir.")

    def process_xsegfetch(arguments):
        osex.set_process_lowest_prio()
        from mainscripts import XSegUtil
        XSegUtil.fetch_xseg (Path(arguments.input_dir) )
    p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
    p.set_defaults (func=process_xsegfetch)
    
    def bad_args(arguments):
        parser.print_help()
        exit(0)
    parser.set_defaults(func=bad_args)

    arguments = parser.parse_args()
    arguments.func(arguments)

    if exit_code == 0:
        print ("Done.")
        
    exit(exit_code)
    
'''
import code
code.interact(local=dict(globals(), **locals()))
'''

main.py
oo92/TorchFaceLab
import os
import sys
import traceback
import queue
import threading
import time
import numpy as np
import itertools
from pathlib import Path
from core import pathex
from core import imagelib
import cv2
import models
from core.interact import interact as io


def main(model_class_name, saved_models_path):
    model = models.import_model(model_class_name)(
                        is_exporting=True,
                        saved_models_path=saved_models_path,
                        cpu_only=True)
    model.export_dfm () 

mainscripts/ExportDFM.py
oo92/TorchFaceLab
﻿import traceback, math, multiprocessing, operator, os, shutil, sys, time, cv2, facelib
from pathlib import Path

import numpy as np

from numpy import linalg as npla
from core import imagelib
from core import mathlib
from facelib import FaceType, LandmarksProcessor
from core.interact import interact as io
from core.joblib import Subprocessor
from core.leras.device import Devices, Device
from core import pathex
from core.cv2ex import *
from DFLIMG import *
from core.leras.device import Devices

DEBUG = False

class DeviceConfig():
    current_DeviceConfig = None
    device = None

    def set_floatx(torch_dtype):
        """Set default float type for all layers when dtype is None for them"""
        floatx = torch_dtype

    @staticmethod
    def getCurrentDeviceConfig():
        if current_DeviceConfig is None:
            current_DeviceConfig = DeviceConfig.BestGPU()
        return current_DeviceConfig

    @staticmethod
    def setCurrentDeviceConfig(device_config):
        current_DeviceConfig = device_config

    @staticmethod
    def initialize(device_config=None, floatx="float32"):

        if device is None:
            if device_config is None:
                device_config = getCurrentDeviceConfig()
            setCurrentDeviceConfig(device_config)

            if torch.cuda.is_available():
                io.log_info(f"Using GPU: {torch.cuda.get_device_name(0)}")
                device = torch.device('cuda')
            else:
                io.log_info("Using CPU")
                device = torch.device('cpu')

        if floatx == "float32":
            floatx = torch.float32
        elif floatx == "float16":
            floatx = torch.float16
        else:
            raise ValueError(f"unsupported floatx {floatx}")
        set_floatx(floatx)

    @staticmethod
    def ask_choose_device_idxs(choose_only_one=False, allow_cpu=True, suggest_best_multi_gpu=False, suggest_all_gpu=False):
        devices = Devices.getDevices()
        if len(devices) == 0:
            return []

        all_devices_indexes = [device.index for device in devices]

        if choose_only_one:
            suggest_best_multi_gpu = False
            suggest_all_gpu = False

        if suggest_all_gpu:
            best_device_indexes = all_devices_indexes
        elif suggest_best_multi_gpu:
            best_device_indexes = [device.index for device in devices.get_equal_devices(devices.get_best_device()) ]
        else:
            best_device_indexes = [ devices.get_best_device().index ]
        best_device_indexes = ",".join([str(x) for x in best_device_indexes])

        io.log_info ("")
        if choose_only_one:
            io.log_info ("Choose one GPU idx.")
        else:
            io.log_info ("Choose one or several GPU idxs (separated by comma).")
        io.log_info ("")

        if allow_cpu:
            io.log_info ("[CPU] : CPU")
        for device in devices:
            io.log_info (f"  [{device.index}] : {device.name}")

        io.log_info ("")

        while True:
            try:
                if choose_only_one:
                    choosed_idxs = io.input_str("Which GPU index to choose?", best_device_indexes)
                else:
                    choosed_idxs = io.input_str("Which GPU indexes to choose?", best_device_indexes)

                if allow_cpu and choosed_idxs.lower() == "cpu":
                    choosed_idxs = []
                    break

                choosed_idxs = [ int(x) for x in choosed_idxs.split(',') ]

                if choose_only_one:
                    if len(choosed_idxs) == 1:
                        break
                else:
                    if all( [idx in all_devices_indexes for idx in choosed_idxs] ):
                        break
            except:
                pass
        io.log_info ("")

        return choosed_idxs

    @staticmethod
    def ask_choose_device(*args, **kwargs):
        return DeviceConfig.GPUIndexes(DeviceConfig.ask_choose_device_idxs(*args,**kwargs) )
    
    def __init__ (self, devices=None):
        devices = devices or []

        if not isinstance(devices, Devices):
            devices = Devices(devices)

        self.devices = devices
        self.cpu_only = len(devices) == 0

    @staticmethod
    def BestGPU():
        devices = Devices.getDevices()
        if len(devices) == 0:
            return DeviceConfig.CPU()

        return DeviceConfig([devices.get_best_device()])

    @staticmethod
    def WorstGPU():
        devices = Devices.getDevices()
        if len(devices) == 0:
            return DeviceConfig.CPU()

        return DeviceConfig([devices.get_worst_device()])

    @staticmethod
    def GPUIndexes(indexes):
        if len(indexes) != 0:
            devices = Devices.getDevices().get_devices_from_index_list(indexes)
        else:
            devices = []

        return DeviceConfig(devices)

    @staticmethod
    def CPU():
        return DeviceConfig([])

class ExtractSubprocessor(Subprocessor):
    class Data(object):
        def __init__(self, filepath=None, rects=None, landmarks = None, landmarks_accurate=True, manual=False, force_output_path=None, final_output_files = None):
            self.filepath = filepath
            self.rects = rects or []
            self.rects_rotation = 0
            self.landmarks_accurate = landmarks_accurate
            self.manual = manual
            self.landmarks = landmarks or []
            self.force_output_path = force_output_path
            self.final_output_files = final_output_files or []
            self.faces_detected = 0

    class Cli(Subprocessor.Cli):

        #override
        def on_initialize(self, client_dict):
            self.type                 = client_dict['type']
            self.image_size           = client_dict['image_size']
            self.jpeg_quality         = client_dict['jpeg_quality']
            self.face_type            = client_dict['face_type']
            self.max_faces_from_image = client_dict['max_faces_from_image']
            self.device_idx           = client_dict['device_idx']
            self.cpu_only             = client_dict['device_type'] == 'CPU'
            self.final_output_path    = client_dict['final_output_path']
            self.output_debug_path    = client_dict['output_debug_path']

            #transfer and set stdin in order to work code.interact in debug subprocess
            stdin_fd         = client_dict['stdin_fd']
            if stdin_fd is not None and DEBUG:
                sys.stdin = os.fdopen(stdin_fd)

            if self.cpu_only:
                device_config = DeviceConfig.CPU()
                place_model_on_cpu = True
            else:
                device_config = DeviceConfig.GPUIndexes ([self.device_idx])
                place_model_on_cpu = device_config.devices[0].total_mem_gb < 4

            if self.type == 'all' or 'rects' in self.type or 'landmarks' in self.type:
                Devices.initialize_main_env()

            self.log_info (f"Running on {client_dict['device_name'] }")

            if self.type == 'all' or self.type == 'rects-s3fd' or 'landmarks' in self.type:
                self.rects_extractor = facelib.S3FDExtractor(place_model_on_cpu=place_model_on_cpu)

            if self.type == 'all' or 'landmarks' in self.type:
                # for head type, extract "3D landmarks"
                self.landmarks_extractor = facelib.FANExtractor(landmarks_3D=self.face_type >= FaceType.HEAD,
                                                                place_model_on_cpu=place_model_on_cpu)

            self.cached_image = (None, None)

        #override
        def process_data(self, data):
            if 'landmarks' in self.type and len(data.rects) == 0:
                return data

            filepath = data.filepath
            cached_filepath, image = self.cached_image
            if cached_filepath != filepath:
                image = cv2_imread( filepath )
                if image is None:
                    self.log_err (f'Failed to open {filepath}, reason: cv2_imread() fail.')
                    return data
                image = imagelib.normalize_channels(image, 3)
                image = imagelib.cut_odd_image(image)
                self.cached_image = ( filepath, image )

            h, w, c = image.shape

            if 'rects' in self.type or self.type == 'all':
                data = ExtractSubprocessor.Cli.rects_stage (data=data,
                                                            image=image,
                                                            max_faces_from_image=self.max_faces_from_image,
                                                            rects_extractor=self.rects_extractor,
                                                            )

            if 'landmarks' in self.type or self.type == 'all':
                data = ExtractSubprocessor.Cli.landmarks_stage (data=data,
                                                                image=image,
                                                                landmarks_extractor=self.landmarks_extractor,
                                                                rects_extractor=self.rects_extractor,
                                                                )

            if self.type == 'final' or self.type == 'all':
                data = ExtractSubprocessor.Cli.final_stage(data=data,
                                                           image=image,
                                                           face_type=self.face_type,
                                                           image_size=self.image_size,
                                                           jpeg_quality=self.jpeg_quality,
                                                           output_debug_path=self.output_debug_path,
                                                           final_output_path=self.final_output_path,
                                                           )
            return data

        @staticmethod
        def rects_stage(data,
                        image,
                        max_faces_from_image,
                        rects_extractor,
                        ):
            h,w,c = image.shape
            if min(h,w) < 128:
                # Image is too small
                data.rects = []
            else:
                for rot in ([0, 90, 270, 180]):
                    if rot == 0:
                        rotated_image = image
                    elif rot == 90:
                        rotated_image = image.swapaxes( 0,1 )[:,::-1,:]
                    elif rot == 180:
                        rotated_image = image[::-1,::-1,:]
                    elif rot == 270:
                        rotated_image = image.swapaxes( 0,1 )[::-1,:,:]
                    rects = data.rects = rects_extractor.extract (rotated_image, is_bgr=True)
                    if len(rects) != 0:
                        data.rects_rotation = rot
                        break
                if max_faces_from_image is not None and \
                   max_faces_from_image > 0 and \
                   len(data.rects) > 0:
                    data.rects = data.rects[0:max_faces_from_image]
            return data


        @staticmethod
        def landmarks_stage(data,
                            image,
                            landmarks_extractor,
                            rects_extractor,
                            ):
            h, w, ch = image.shape

            if data.rects_rotation == 0:
                rotated_image = image
            elif data.rects_rotation == 90:
                rotated_image = image.swapaxes( 0,1 )[:,::-1,:]
            elif data.rects_rotation == 180:
                rotated_image = image[::-1,::-1,:]
            elif data.rects_rotation == 270:
                rotated_image = image.swapaxes( 0,1 )[::-1,:,:]

            data.landmarks = landmarks_extractor.extract (rotated_image, data.rects, rects_extractor if (data.landmarks_accurate) else None, is_bgr=True)
            if data.rects_rotation != 0:
                for i, (rect, lmrks) in enumerate(zip(data.rects, data.landmarks)):
                    new_rect, new_lmrks = rect, lmrks
                    (l,t,r,b) = rect
                    if data.rects_rotation == 90:
                        new_rect = ( t, h-l, b, h-r)
                        if lmrks is not None:
                            new_lmrks = lmrks[:,::-1].copy()
                            new_lmrks[:,1] = h - new_lmrks[:,1]
                    elif data.rects_rotation == 180:
                        if lmrks is not None:
                            new_rect = ( w-l, h-t, w-r, h-b)
                            new_lmrks = lmrks.copy()
                            new_lmrks[:,0] = w - new_lmrks[:,0]
                            new_lmrks[:,1] = h - new_lmrks[:,1]
                    elif data.rects_rotation == 270:
                        new_rect = ( w-b, l, w-t, r )
                        if lmrks is not None:
                            new_lmrks = lmrks[:,::-1].copy()
                            new_lmrks[:,0] = w - new_lmrks[:,0]
                    data.rects[i], data.landmarks[i] = new_rect, new_lmrks

            return data

        @staticmethod
        def final_stage(data,
                        image,
                        face_type,
                        image_size,
                        jpeg_quality,
                        output_debug_path=None,
                        final_output_path=None,
                        ):
            data.final_output_files = []
            filepath = data.filepath
            rects = data.rects
            landmarks = data.landmarks

            if output_debug_path is not None:
                debug_image = image.copy()

            face_idx = 0
            for rect, image_landmarks in zip( rects, landmarks ):
                if image_landmarks is None:
                    continue

                rect = np.array(rect)

                if face_type == FaceType.MARK_ONLY:
                    image_to_face_mat = None
                    face_image = image
                    face_image_landmarks = image_landmarks
                else:
                    image_to_face_mat = LandmarksProcessor.get_transform_mat (image_landmarks, image_size, face_type)

                    face_image = cv2.warpAffine(image, image_to_face_mat, (image_size, image_size), cv2.INTER_LANCZOS4)
                    face_image_landmarks = LandmarksProcessor.transform_points (image_landmarks, image_to_face_mat)

                    landmarks_bbox = LandmarksProcessor.transform_points ( [ (0,0), (0,image_size-1), (image_size-1, image_size-1), (image_size-1,0) ], image_to_face_mat, True)

                    rect_area      = mathlib.polygon_area(np.array(rect[[0,2,2,0]]).astype(np.float32), np.array(rect[[1,1,3,3]]).astype(np.float32))
                    landmarks_area = mathlib.polygon_area(landmarks_bbox[:,0].astype(np.float32), landmarks_bbox[:,1].astype(np.float32) )

                    if not data.manual and face_type <= FaceType.FULL_NO_ALIGN and landmarks_area > 4*rect_area: #get rid of faces which umeyama-landmark-area > 4*detector-rect-area
                        continue

                    if output_debug_path is not None:
                        LandmarksProcessor.draw_rect_landmarks (debug_image, rect, image_landmarks, face_type, image_size, transparent_mask=True)

                output_path = final_output_path
                if data.force_output_path is not None:
                    output_path = data.force_output_path

                output_filepath = output_path / f"{filepath.stem}_{face_idx}.jpg"
                cv2_imwrite(output_filepath, face_image, [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_quality ] )

                dflimg = DFLJPG.load(output_filepath)
                dflimg.set_face_type(FaceType.toString(face_type))
                dflimg.set_landmarks(face_image_landmarks.tolist())
                dflimg.set_source_filename(filepath.name)
                dflimg.set_source_rect(rect)
                dflimg.set_source_landmarks(image_landmarks.tolist())
                dflimg.set_image_to_face_mat(image_to_face_mat)
                dflimg.save()

                data.final_output_files.append (output_filepath)
                face_idx += 1
            data.faces_detected = face_idx

            if output_debug_path is not None:
                cv2_imwrite( output_debug_path / (filepath.stem+'.jpg'), debug_image, [int(cv2.IMWRITE_JPEG_QUALITY), 50] )

            return data

        #overridable
        def get_data_name (self, data):
            #return string identificator of your data
            return data.filepath

    @staticmethod
    def get_devices_for_config (type, device_config):
        devices = device_config.devices
        cpu_only = len(devices) == 0

        if 'rects'     in type or \
           'landmarks' in type or \
           'all'       in type:

            if not cpu_only:
                if type == 'landmarks-manual':
                    devices = [devices.get_best_device()]

                result = []

                for device in devices:
                    count = 1

                    if count == 1:
                        result += [ (device.index, 'GPU', device.name, device.total_mem_gb) ]
                    else:
                        for i in range(count):
                            result += [ (device.index, 'GPU', f"{device.name} #{i}", device.total_mem_gb) ]

                return result
            else:
                if type == 'landmarks-manual':
                    return [ (0, 'CPU', 'CPU', 0 ) ]
                else:
                    return [ (i, 'CPU', 'CPU%d' % (i), 0 ) for i in range( min(8, multiprocessing.cpu_count() // 2) ) ]

        elif type == 'final':
            return [ (i, 'CPU', 'CPU%d' % (i), 0 ) for i in (range(min(8, multiprocessing.cpu_count())) if not DEBUG else [0]) ]

    def __init__(self, input_data, type, image_size=None, jpeg_quality=None, face_type=None, output_debug_path=None, manual_window_size=0, max_faces_from_image=0, final_output_path=None, device_config=None):
        if type == 'landmarks-manual':
            for x in input_data:
                x.manual = True

        self.input_data = input_data

        self.type = type
        self.image_size = image_size
        self.jpeg_quality = jpeg_quality
        self.face_type = face_type
        self.output_debug_path = output_debug_path
        self.final_output_path = final_output_path
        self.manual_window_size = manual_window_size
        self.max_faces_from_image = max_faces_from_image
        self.result = []

        self.devices = ExtractSubprocessor.get_devices_for_config(self.type, device_config)

        super().__init__('Extractor', ExtractSubprocessor.Cli,
                             999999 if type == 'landmarks-manual' or DEBUG else 120)

    #override
    def on_clients_initialized(self):
        if self.type == 'landmarks-manual':
            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.hide_help = False
            self.landmarks_accurate = True
            self.force_landmarks = False

            self.landmarks = None
            self.x = 0
            self.y = 0
            self.rect_size = 100
            self.rect_locked = False
            self.extract_needed = True

            self.image = None
            self.image_filepath = None

        io.progress_bar (None, len (self.input_data))

    #override
    def on_clients_finalized(self):
        if self.type == 'landmarks-manual':
            io.destroy_all_windows()

        io.progress_bar_close()

    #override
    def process_info_generator(self):
        base_dict = {'type' : self.type,
                     'image_size': self.image_size,
                     'jpeg_quality' : self.jpeg_quality,
                     'face_type': self.face_type,
                     'max_faces_from_image':self.max_faces_from_image,
                     'output_debug_path': self.output_debug_path,
                     'final_output_path': self.final_output_path,
                     'stdin_fd': sys.stdin.fileno() }


        for (device_idx, device_type, device_name, device_total_vram_gb) in self.devices:
            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 self.type == 'landmarks-manual':
            need_remark_face = False
            while len (self.input_data) > 0:
                data = self.input_data[0]
                filepath, data_rects, data_landmarks = data.filepath, data.rects, data.landmarks
                is_frame_done = False

                if self.image_filepath != filepath:
                    self.image_filepath = filepath
                    if self.cache_original_image[0] == filepath:
                        self.original_image = self.cache_original_image[1]
                    else:
                        self.original_image = imagelib.normalize_channels( cv2_imread( filepath ), 3 )

                        self.cache_original_image = (filepath, 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,filepath):
                        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,filepath), 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 = (imagelib.get_draw_text_lines ( self.image, sh,
                                                        [   '[L Mouse click] - lock/unlock selection. [Mouse wheel] - change rect',
                                                            '[R Mouse Click] - manual face rectangle',
                                                            '[Enter] / [Space] - confirm / skip frame',
                                                            '[,] [.]- prev frame, next frame. [Q] - skip remaining frames',
                                                            '[a] - accuracy on/off (more fps)',
                                                            '[h] - hide this help'
                                                        ], (1, 1, 1) )*255).astype(np.uint8)

                        self.cache_text_lines_img = (sh, self.text_lines_img)

                if need_remark_face: # need remark image from input data that already has a marked face?
                    need_remark_face = False
                    if len(data_rects) != 0: # If there was already a face then lock the rectangle to it until the mouse is clicked
                        self.rect = data_rects.pop()
                        self.landmarks = data_landmarks.pop()
                        data_rects.clear()
                        data_landmarks.clear()

                        self.rect_locked = True
                        self.rect_size = ( self.rect[2] - self.rect[0] ) / 2
                        self.x = ( self.rect[0] + self.rect[2] ) / 2
                        self.y = ( self.rect[1] + self.rect[3] ) / 2
                        self.redraw()

                if len(data_rects) == 0:
                    (h,w,c) = self.image.shape
                    while True:
                        io.process_messages(0.0001)

                        if not self.force_landmarks:
                            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:
                                if self.force_landmarks:
                                    self.x = new_x
                                    self.y = new_y
                                    self.force_landmarks = False
                                    self.rect_locked = True
                                    self.redraw()
                                else:
                                    self.rect_locked = not self.rect_locked
                                    self.extract_needed = True
                            elif ev == io.EVENT_RBUTTONDOWN:
                                self.force_landmarks = not self.force_landmarks
                                if self.force_landmarks:
                                    self.rect_locked = False
                            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, chr_key, ctrl_pressed, alt_pressed, shift_pressed = key_events[-1] if len(key_events) > 0 else (0,0,False,False,False)

                        if key == ord('\r') or key == ord('\n'):
                            #confirm frame
                            is_frame_done = True
                            data_rects.append (self.rect)
                            data_landmarks.append (self.landmarks)
                            break
                        elif key == ord(' '):
                            #confirm skip frame
                            is_frame_done = True
                            break
                        elif key == ord(',')  and len(self.result) > 0:
                            #go prev frame

                            if self.rect_locked:
                                self.rect_locked = False
                                # Only save the face if the rect is still locked
                                data_rects.append (self.rect)
                                data_landmarks.append (self.landmarks)


                            self.input_data.insert(0, self.result.pop() )
                            io.progress_bar_inc(-1)
                            need_remark_face = True

                            break
                        elif key == ord('.'):
                            #go next frame

                            if self.rect_locked:
                                self.rect_locked = False
                                # Only save the face if the rect is still locked
                                data_rects.append (self.rect)
                                data_landmarks.append (self.landmarks)

                            need_remark_face = True
                            is_frame_done = True
                            break
                        elif key == ord('q'):
                            #skip remaining

                            if self.rect_locked:
                                self.rect_locked = False
                                data_rects.append (self.rect)
                                data_landmarks.append (self.landmarks)

                            while len(self.input_data) > 0:
                                self.result.append( self.input_data.pop(0) )
                                io.progress_bar_inc(1)

                            break

                        elif key == ord('h'):
                            self.hide_help = not self.hide_help
                            break
                        elif key == ord('a'):
                            self.landmarks_accurate = not self.landmarks_accurate
                            break

                        if self.force_landmarks:
                            pt2 = np.float32([new_x, new_y])
                            pt1 = np.float32([self.x, self.y])

                            pt_vec_len = npla.norm(pt2-pt1)
                            pt_vec = pt2-pt1
                            if pt_vec_len != 0:
                                pt_vec /= pt_vec_len

                            self.rect_size = pt_vec_len
                            self.rect = ( int(self.x-self.rect_size),
                                          int(self.y-self.rect_size),
                                          int(self.x+self.rect_size),
                                          int(self.y+self.rect_size) )

                            if pt_vec_len > 0:
                                lmrks = np.concatenate ( (np.zeros ((17,2), np.float32), LandmarksProcessor.landmarks_2D), axis=0 )
                                lmrks -= lmrks[30:31,:]
                                mat = cv2.getRotationMatrix2D( (0, 0), -np.arctan2( pt_vec[1], pt_vec[0] )*180/math.pi , pt_vec_len)
                                mat[:, 2] += (self.x, self.y)
                                self.landmarks = LandmarksProcessor.transform_points(lmrks, mat )


                            self.redraw()

                        elif self.x != new_x or \
                           self.y != new_y or \
                           self.rect_size != new_rect_size or \
                           self.extract_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 ExtractSubprocessor.Data (filepath, rects=[self.rect], landmarks_accurate=self.landmarks_accurate)

                else:
                    is_frame_done = True

                if is_frame_done:
                    self.result.append ( data )
                    self.input_data.pop(0)
                    io.progress_bar_inc(1)
                    self.extract_needed = True
                    self.rect_locked = False
        else:
            if len (self.input_data) > 0:
                return self.input_data.pop(0)

        return None

    #override
    def on_data_return (self, host_dict, data):
        if not self.type != 'landmarks-manual':
            self.input_data.insert(0, data)

    def redraw(self):
        (h,w,c) = self.image.shape

        if not self.hide_help:
            image = cv2.addWeighted (self.image,1.0,self.text_lines_img,1.0,0)
        else:
            image = self.image.copy()

        view_rect = (np.array(self.rect) * self.view_scale).astype(np.int).tolist()
        view_landmarks  = (np.array(self.landmarks) * self.view_scale).astype(np.int).tolist()

        if self.rect_size <= 40:
            scaled_rect_size = h // 3 if w > h else w // 3

            p1 = (self.x - self.rect_size, self.y - self.rect_size)
            p2 = (self.x + self.rect_size, self.y - self.rect_size)
            p3 = (self.x - self.rect_size, self.y + self.rect_size)

            wh = h if h < w else w
            np1 = (w / 2 - wh / 4, h / 2 - wh / 4)
            np2 = (w / 2 + wh / 4, h / 2 - wh / 4)
            np3 = (w / 2 - wh / 4, h / 2 + wh / 4)

            mat = cv2.getAffineTransform( np.float32([p1,p2,p3])*self.view_scale, np.float32([np1,np2,np3]) )
            image = cv2.warpAffine(image, mat,(w,h) )
            view_landmarks = LandmarksProcessor.transform_points (view_landmarks, mat)

        landmarks_color = (255,255,0) if self.rect_locked else (0,255,0)
        LandmarksProcessor.draw_rect_landmarks (image, view_rect, view_landmarks, self.face_type, self.image_size, landmarks_color=landmarks_color)
        self.extract_needed = False

        io.show_image (self.wnd_name, image)


    #override
    def on_result (self, host_dict, data, result):
        if self.type == 'landmarks-manual':
            filepath, landmarks = result.filepath, result.landmarks

            if len(landmarks) != 0 and landmarks[0] is not None:
                self.landmarks = landmarks[0]

            self.redraw()
        else:
            self.result.append ( result )
            io.progress_bar_inc(1)



    #override
    def get_result(self):
        return self.result


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
    def __init__(self, input_paths, debug_paths ):
        self.input_paths = input_paths
        self.debug_paths_stems = [ Path(d).stem for d in debug_paths]
        self.result = []
        super().__init__('DeletedFilesSearcherSubprocessor', DeletedFilesSearcherSubprocessor.Cli, 60)

    #override
    def process_info_generator(self):
        for i in range(min(multiprocessing.cpu_count(), 8)):
            yield 'CPU%d' % (i), {}, {'debug_paths_stems' : self.debug_paths_stems}

    #override
    def on_clients_initialized(self):
        io.progress_bar ("Searching deleted files", len (self.input_paths))

    #override
    def on_clients_finalized(self):
        io.progress_bar_close()

    #override
    def get_data(self, host_dict):
        if len (self.input_paths) > 0:
            return [self.input_paths.pop(0)]
        return None

    #override
    def on_data_return (self, host_dict, data):
        self.input_paths.insert(0, data[0])

    #override
    def on_result (self, host_dict, data, result):
        if result == False:
            self.result.append( data[0] )
        io.progress_bar_inc(1)

    #override
    def get_result(self):
        return self.result

def main(detector=None,
         input_path=None,
         output_path=None,
         output_debug=None,
         manual_fix=False,
         manual_output_debug_fix=False,
         manual_window_size=1368,
         face_type='full_face',
         max_faces_from_image=None,
         image_size=None,
         jpeg_quality=None,
         cpu_only = False,
         force_gpu_idxs = None,
         ):

    if not input_path.exists():
        io.log_err ('Input directory not found. Please ensure it exists.')
        return

    if not output_path.exists():
        output_path.mkdir(parents=True, exist_ok=True)

    if face_type is not None:
        face_type = FaceType.fromString(face_type)

    if face_type is None:
        if manual_output_debug_fix:
            files = pathex.get_image_paths(output_path)
            if len(files) != 0:
                dflimg = DFLIMG.load(Path(files[0]))
                if dflimg is not None and dflimg.has_data():
                     face_type = FaceType.fromString ( dflimg.get_face_type() )

    input_image_paths = pathex.get_image_unique_filestem_paths(input_path, verbose_print_func=io.log_info)
    output_images_paths = pathex.get_image_paths(output_path)
    output_debug_path = output_path.parent / (output_path.name + '_debug')

    continue_extraction = False
    if not manual_output_debug_fix and len(output_images_paths) > 0:
        if len(output_images_paths) > 128:
            continue_extraction = io.input_bool ("Continue extraction?", True, help_message="Extraction can be continued, but you must specify the same options again.")

        if len(output_images_paths) > 128 and continue_extraction:
            try:
                input_image_paths = input_image_paths[ [ Path(x).stem for x in input_image_paths ].index ( Path(output_images_paths[-128]).stem.split('_')[0] ) : ]
            except:
                io.log_err("Error in fetching the last index. Extraction cannot be continued.")
                return
        elif input_path != output_path:
                io.input(f"\n WARNING !!! \n {output_path} contains files! \n They will be deleted. \n Press enter to continue.\n")
                for filename in output_images_paths:
                    Path(filename).unlink()

    device_config = DeviceConfig.GPUIndexes(force_gpu_idxs or DeviceConfig.ask_choose_device_idxs(choose_only_one=detector=='manual', suggest_all_gpu=True) ) \
                    if not cpu_only else DeviceConfig.CPU()

    if face_type is None:
        face_type = io.input_str ("Face type", 'wf', ['f','wf','head'], help_message="Full face / whole face / head. 'Whole face' covers full area of face include forehead. 'head' covers full head, but requires XSeg for src and dst faceset.").lower()
        face_type = {'f'  : FaceType.FULL,
                     'wf' : FaceType.WHOLE_FACE,
                     'head' : FaceType.HEAD}[face_type]

    if max_faces_from_image is None:
        max_faces_from_image = io.input_int(f"Max number of faces from image", 0, help_message="If you extract a src faceset that has frames with a large number of faces, it is advisable to set max faces to 3 to speed up extraction. 0 - unlimited")

    if image_size is None:
        image_size = io.input_int(f"Image size", 512 if face_type < FaceType.HEAD else 768, valid_range=[256,2048], help_message="Output image size. The higher image size, the worse face-enhancer works. Use higher than 512 value only if the source image is sharp enough and the face does not need to be enhanced.")

    if jpeg_quality is None:
        jpeg_quality = io.input_int(f"Jpeg quality", 90, valid_range=[1,100], help_message="Jpeg quality. The higher jpeg quality the larger the output file size.")

    if detector is None:
        io.log_info ("Choose detector type.")
        io.log_info ("[0] S3FD")
        io.log_info ("[1] manual")
        detector = {0:'s3fd', 1:'manual'}[ io.input_int("", 0, [0,1]) ]


    if output_debug is None:
        output_debug = io.input_bool (f"Write debug images to {output_debug_path.name}?", False)

    if output_debug:
        output_debug_path.mkdir(parents=True, exist_ok=True)

    if manual_output_debug_fix:
        if not output_debug_path.exists():
            io.log_err(f'{output_debug_path} not found. Re-extract faces with "Write debug images" option.')
            return
        else:
            detector = 'manual'
            io.log_info('Performing re-extract frames which were deleted from _debug directory.')

            input_image_paths = DeletedFilesSearcherSubprocessor (input_image_paths, pathex.get_image_paths(output_debug_path) ).run()
            input_image_paths = sorted (input_image_paths)
            io.log_info('Found %d images.' % (len(input_image_paths)))
    else:
        if not continue_extraction and output_debug_path.exists():
            for filename in pathex.get_image_paths(output_debug_path):
                Path(filename).unlink()

    images_found = len(input_image_paths)
    faces_detected = 0
    if images_found != 0:
        if detector == 'manual':
            io.log_info ('Performing manual extract...')
            data = ExtractSubprocessor ([ ExtractSubprocessor.Data(Path(filename)) for filename in input_image_paths ], 'landmarks-manual', image_size, jpeg_quality, face_type, output_debug_path if output_debug else None, manual_window_size=manual_window_size, device_config=device_config).run()

            io.log_info ('Performing 3rd pass...')
            data = ExtractSubprocessor (data, 'final', image_size, jpeg_quality, face_type, output_debug_path if output_debug else None, final_output_path=output_path, device_config=device_config).run()

        else:
            io.log_info ('Extracting faces...')
            data = ExtractSubprocessor ([ ExtractSubprocessor.Data(Path(filename)) for filename in input_image_paths ],
                                         'all',
                                         image_size,
                                         jpeg_quality,
                                         face_type,
                                         output_debug_path if output_debug else None,
                                         max_faces_from_image=max_faces_from_image,
                                         final_output_path=output_path,
                                         device_config=device_config).run()

        faces_detected += sum([d.faces_detected for d in data])

        if manual_fix:
            if all ( np.array ( [ d.faces_detected > 0 for d in data] ) == True ):
                io.log_info ('All faces are detected, manual fix not needed.')
            else:
                fix_data = [ ExtractSubprocessor.Data(d.filepath) for d in data if d.faces_detected == 0 ]
                io.log_info ('Performing manual fix for %d images...' % (len(fix_data)) )
                fix_data = ExtractSubprocessor (fix_data, 'landmarks-manual', image_size, jpeg_quality, face_type, output_debug_path if output_debug else None, manual_window_size=manual_window_size, device_config=device_config).run()
                fix_data = ExtractSubprocessor (fix_data, 'final', image_size, jpeg_quality, face_type, output_debug_path if output_debug else None, final_output_path=output_path, device_config=device_config).run()
                faces_detected += sum([d.faces_detected for d in fix_data])


    io.log_info ('-------------------------')
    io.log_info ('Images found:        %d' % (images_found) )
    io.log_info ('Faces detected:      %d' % (faces_detected) )
    io.log_info ('-------------------------')

mainscripts/Extractor.py
oo92/TorchFaceLab
import multiprocessing
import shutil

from DFLIMG import *
from core.interact import interact as io
from core.joblib import Subprocessor
from core.leras import nn
from core import pathex
from core.cv2ex import *


class FacesetEnhancerSubprocessor(Subprocessor):

    #override
    def __init__(self, image_paths, output_dirpath, device_config):
        self.image_paths = image_paths
        self.output_dirpath = output_dirpath
        self.result = []
        self.nn_initialize_mp_lock = multiprocessing.Lock()
        self.devices = FacesetEnhancerSubprocessor.get_devices_for_config(device_config)

        super().__init__('FacesetEnhancer', FacesetEnhancerSubprocessor.Cli, 600)

    #override
    def on_clients_initialized(self):
        io.progress_bar (None, len (self.image_paths))

    #override
    def on_clients_finalized(self):
        io.progress_bar_close()

    #override
    def process_info_generator(self):
        base_dict = {'output_dirpath':self.output_dirpath,
                     'nn_initialize_mp_lock': self.nn_initialize_mp_lock,}

        for (device_idx, device_type, device_name, device_total_vram_gb) in self.devices:
            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 len (self.image_paths) > 0:
            return self.image_paths.pop(0)

    #override
    def on_data_return (self, host_dict, data):
        self.image_paths.insert(0, data)

    #override
    def on_result (self, host_dict, data, result):
        io.progress_bar_inc(1)
        if result[0] == 1:
            self.result +=[ (result[1], result[2]) ]

    #override
    def get_result(self):
        return self.result

    @staticmethod
    def get_devices_for_config (device_config):
        devices = device_config.devices
        cpu_only = len(devices) == 0

        if not cpu_only:
            return [ (device.index, 'GPU', device.name, device.total_mem_gb) for device in devices ]
        else:
            return [ (i, 'CPU', 'CPU%d' % (i), 0 ) for i in range( min(8, multiprocessing.cpu_count() // 2) ) ]

    class Cli(Subprocessor.Cli):

        #override
        def on_initialize(self, client_dict):
            device_idx   = client_dict['device_idx']
            cpu_only     = client_dict['device_type'] == 'CPU'
            self.output_dirpath = client_dict['output_dirpath']
            nn_initialize_mp_lock = client_dict['nn_initialize_mp_lock']

            if cpu_only:
                device_config = nn.DeviceConfig.CPU()
                device_vram = 99
            else:
                device_config = nn.DeviceConfig.GPUIndexes ([device_idx])
                device_vram = device_config.devices[0].total_mem_gb

            nn.initialize (device_config)

            intro_str = 'Running on %s.' % (client_dict['device_name'])

            self.log_info (intro_str)

            from facelib import FaceEnhancer
            self.fe = FaceEnhancer( place_model_on_cpu=(device_vram<=2 or cpu_only), run_on_cpu=cpu_only )

        #override
        def process_data(self, filepath):
            try:
                dflimg = DFLIMG.load (filepath)
                if dflimg is None or not dflimg.has_data():
                    self.log_err (f"{filepath.name} is not a dfl image file")
                else:
                    dfl_dict = dflimg.get_dict()

                    img = cv2_imread(filepath).astype(np.float32) / 255.0
                    img = self.fe.enhance(img)
                    img = np.clip (img*255, 0, 255).astype(np.uint8)

                    output_filepath = self.output_dirpath / filepath.name

                    cv2_imwrite ( str(output_filepath), img, [int(cv2.IMWRITE_JPEG_QUALITY), 100] )

                    dflimg = DFLIMG.load (output_filepath)
                    dflimg.set_dict(dfl_dict)
                    dflimg.save()

                    return (1, filepath, output_filepath)
            except:
                self.log_err (f"Exception occured while processing file {filepath}. Error: {traceback.format_exc()}")

            return (0, filepath, None)

def process_folder ( dirpath, cpu_only=False, force_gpu_idxs=None ):
    device_config = nn.DeviceConfig.GPUIndexes( force_gpu_idxs or nn.ask_choose_device_idxs(suggest_all_gpu=True) ) \
                    if not cpu_only else nn.DeviceConfig.CPU()

    output_dirpath = dirpath.parent / (dirpath.name + '_enhanced')
    output_dirpath.mkdir (exist_ok=True, parents=True)

    dirpath_parts = '/'.join( dirpath.parts[-2:])
    output_dirpath_parts = '/'.join( output_dirpath.parts[-2:] )
    io.log_info (f"Enhancing faceset in {dirpath_parts}")
    io.log_info ( f"Processing to {output_dirpath_parts}")

    output_images_paths = pathex.get_image_paths(output_dirpath)
    if len(output_images_paths) > 0:
        for filename in output_images_paths:
            Path(filename).unlink()

    image_paths = [Path(x) for x in pathex.get_image_paths( dirpath )]
    result = FacesetEnhancerSubprocessor ( image_paths, output_dirpath, device_config=device_config).run()

    is_merge = io.input_bool (f"\r\nMerge {output_dirpath_parts} to {dirpath_parts} ?", True)
    if is_merge:
        io.log_info (f"Copying processed files to {dirpath_parts}")

        for (filepath, output_filepath) in result:
            try:
                shutil.copy (output_filepath, filepath)
            except:
                pass

        io.log_info (f"Removing {output_dirpath_parts}")
        shutil.rmtree(output_dirpath)

mainscripts/FacesetEnhancer.py
oo92/TorchFaceLab
import multiprocessing
import shutil

import cv2
from core import pathex
from core.cv2ex import *
from core.interact import interact as io
from core.joblib import Subprocessor
from DFLIMG import *
from facelib import FaceType, LandmarksProcessor


class FacesetResizerSubprocessor(Subprocessor):

    #override
    def __init__(self, image_paths, output_dirpath, image_size, face_type=None):
        self.image_paths = image_paths
        self.output_dirpath = output_dirpath
        self.image_size = image_size
        self.face_type = face_type
        self.result = []

        super().__init__('FacesetResizer', FacesetResizerSubprocessor.Cli, 600)

    #override
    def on_clients_initialized(self):
        io.progress_bar (None, len (self.image_paths))

    #override
    def on_clients_finalized(self):
        io.progress_bar_close()

    #override
    def process_info_generator(self):
        base_dict = {'output_dirpath':self.output_dirpath, 'image_size':self.image_size, 'face_type':self.face_type}

        for device_idx in range( min(8, multiprocessing.cpu_count()) ):
            client_dict = base_dict.copy()
            device_name = f'CPU #{device_idx}'
            client_dict['device_name'] = device_name
            yield device_name, {}, client_dict

    #override
    def get_data(self, host_dict):
        if len (self.image_paths) > 0:
            return self.image_paths.pop(0)

    #override
    def on_data_return (self, host_dict, data):
        self.image_paths.insert(0, data)

    #override
    def on_result (self, host_dict, data, result):
        io.progress_bar_inc(1)
        if result[0] == 1:
            self.result +=[ (result[1], result[2]) ]

    #override
    def get_result(self):
        return self.result

    class Cli(Subprocessor.Cli):

        #override
        def on_initialize(self, client_dict):
            self.output_dirpath = client_dict['output_dirpath']
            self.image_size = client_dict['image_size']
            self.face_type = client_dict['face_type']
            self.log_info (f"Running on { client_dict['device_name'] }")

        #override
        def process_data(self, filepath):
            try:
                dflimg = DFLIMG.load (filepath)
                if dflimg is None or not dflimg.has_data():
                    self.log_err (f"{filepath.name} is not a dfl image file")
                else:
                    img = cv2_imread(filepath)
                    h,w = img.shape[:2]
                    if h != w:
                        raise Exception(f'w != h in {filepath}')
                    
                    image_size = self.image_size
                    face_type = self.face_type
                    output_filepath = self.output_dirpath / filepath.name
                    
                    if face_type is not None:
                        lmrks = dflimg.get_landmarks()
                        mat = LandmarksProcessor.get_transform_mat(lmrks, image_size, face_type)
                        
                        img = cv2.warpAffine(img, mat, (image_size, image_size), flags=cv2.INTER_LANCZOS4 )
                        img = np.clip(img, 0, 255).astype(np.uint8)
                        
                        cv2_imwrite ( str(output_filepath), img, [int(cv2.IMWRITE_JPEG_QUALITY), 100] )

                        dfl_dict = dflimg.get_dict()
                        dflimg = DFLIMG.load (output_filepath)
                        dflimg.set_dict(dfl_dict)
                        
                        xseg_mask = dflimg.get_xseg_mask()
                        if xseg_mask is not None:
                            xseg_res = 256
                            
                            xseg_lmrks = lmrks.copy()
                            xseg_lmrks *= (xseg_res / w)
                            xseg_mat = LandmarksProcessor.get_transform_mat(xseg_lmrks, xseg_res, face_type)
                            
                            xseg_mask = cv2.warpAffine(xseg_mask, xseg_mat, (xseg_res, xseg_res), flags=cv2.INTER_LANCZOS4 )
                            xseg_mask[xseg_mask < 0.5] = 0
                            xseg_mask[xseg_mask >= 0.5] = 1

                            dflimg.set_xseg_mask(xseg_mask)
                        
                        seg_ie_polys = dflimg.get_seg_ie_polys()
                        
                        for poly in seg_ie_polys.get_polys():
                            poly_pts = poly.get_pts()
                            poly_pts = LandmarksProcessor.transform_points(poly_pts, mat)
                            poly.set_points(poly_pts)
                            
                        dflimg.set_seg_ie_polys(seg_ie_polys)
                        
                        lmrks = LandmarksProcessor.transform_points(lmrks, mat)
                        dflimg.set_landmarks(lmrks)
    
                        image_to_face_mat = dflimg.get_image_to_face_mat()
                        if image_to_face_mat is not None:
                            image_to_face_mat = LandmarksProcessor.get_transform_mat ( dflimg.get_source_landmarks(), image_size, face_type )
                            dflimg.set_image_to_face_mat(image_to_face_mat)
                        dflimg.set_face_type( FaceType.toString(face_type) )
                        dflimg.save()
                        
                    else:
                        dfl_dict = dflimg.get_dict()
                         
                        scale = w / image_size
                        
                        img = cv2.resize(img, (image_size, image_size), interpolation=cv2.INTER_LANCZOS4)                    
                        
                        cv2_imwrite ( str(output_filepath), img, [int(cv2.IMWRITE_JPEG_QUALITY), 100] )

                        dflimg = DFLIMG.load (output_filepath)
                        dflimg.set_dict(dfl_dict)
                        
                        lmrks = dflimg.get_landmarks()                    
                        lmrks /= scale
                        dflimg.set_landmarks(lmrks)
                        
                        seg_ie_polys = dflimg.get_seg_ie_polys()
                        seg_ie_polys.mult_points( 1.0 / scale)
                        dflimg.set_seg_ie_polys(seg_ie_polys)
                        
                        image_to_face_mat = dflimg.get_image_to_face_mat()
    
                        if image_to_face_mat is not None:
                            face_type = FaceType.fromString ( dflimg.get_face_type() )
                            image_to_face_mat = LandmarksProcessor.get_transform_mat ( dflimg.get_source_landmarks(), image_size, face_type )
                            dflimg.set_image_to_face_mat(image_to_face_mat)
                        dflimg.save()

                    return (1, filepath, output_filepath)
            except:
                self.log_err (f"Exception occured while processing file {filepath}. Error: {traceback.format_exc()}")

            return (0, filepath, None)

def process_folder ( dirpath):
    
    image_size = io.input_int(f"New image size", 512, valid_range=[128,2048])
    
    face_type = io.input_str ("Change face type", 'same', ['h','mf','f','wf','head','same']).lower()
    if face_type == 'same':
        face_type = None
    else:
        face_type = {'h'  : FaceType.HALF,
                     'mf' : FaceType.MID_FULL,
                     'f'  : FaceType.FULL,
                     'wf' : FaceType.WHOLE_FACE,
                     'head' : FaceType.HEAD}[face_type]
                     

    output_dirpath = dirpath.parent / (dirpath.name + '_resized')
    output_dirpath.mkdir (exist_ok=True, parents=True)

    dirpath_parts = '/'.join( dirpath.parts[-2:])
    output_dirpath_parts = '/'.join( output_dirpath.parts[-2:] )
    io.log_info (f"Resizing faceset in {dirpath_parts}")
    io.log_info ( f"Processing to {output_dirpath_parts}")

    output_images_paths = pathex.get_image_paths(output_dirpath)
    if len(output_images_paths) > 0:
        for filename in output_images_paths:
            Path(filename).unlink()

    image_paths = [Path(x) for x in pathex.get_image_paths( dirpath )]
    result = FacesetResizerSubprocessor ( image_paths, output_dirpath, image_size, face_type).run()

    is_merge = io.input_bool (f"\r\nMerge {output_dirpath_parts} to {dirpath_parts} ?", True)
    if is_merge:
        io.log_info (f"Copying processed files to {dirpath_parts}")

        for (filepath, output_filepath) in result:
            try:
                shutil.copy (output_filepath, filepath)
            except:
                pass

        io.log_info (f"Removing {output_dirpath_parts}")
        shutil.rmtree(output_dirpath)

mainscripts/FacesetResizer.py
oo92/TorchFaceLab
﻿import math
import multiprocessing
import traceback
from pathlib import Path

import numpy as np
import numpy.linalg as npla

import samplelib
from core import pathex
from core.cv2ex import *
from core.interact import interact as io
from core.joblib import MPClassFuncOnDemand, MPFunc
from core.leras import nn
from DFLIMG import DFLIMG
from facelib import FaceEnhancer, FaceType, LandmarksProcessor, XSegNet
from merger import FrameInfo, InteractiveMergerSubprocessor, MergerConfig


def main (model_class_name=None,
          saved_models_path=None,
          training_data_src_path=None,
          force_model_name=None,
          input_path=None,
          output_path=None,
          output_mask_path=None,
          aligned_path=None,
          force_gpu_idxs=None,
          cpu_only=None):
    io.log_info ("Running merger.\r\n")

    try:
        if not input_path.exists():
            io.log_err('Input directory not found. Please ensure it exists.')
            return

        if not output_path.exists():
            output_path.mkdir(parents=True, exist_ok=True)

        if not output_mask_path.exists():
            output_mask_path.mkdir(parents=True, exist_ok=True)

        if not saved_models_path.exists():
            io.log_err('Model directory not found. Please ensure it exists.')
            return

        # Initialize model
        import models
        model = models.import_model(model_class_name)(is_training=False,
                                                      saved_models_path=saved_models_path,
                                                      force_gpu_idxs=force_gpu_idxs,
                                                      force_model_name=force_model_name,
                                                      cpu_only=cpu_only)

        predictor_func, predictor_input_shape, cfg = model.get_MergerConfig()

        # Preparing MP functions
        predictor_func = MPFunc(predictor_func)

        run_on_cpu = len(nn.getCurrentDeviceConfig().devices) == 0
        xseg_256_extract_func = MPClassFuncOnDemand(XSegNet, 'extract',
                                                    name='XSeg',
                                                    resolution=256,
                                                    weights_file_root=saved_models_path,
                                                    place_model_on_cpu=True,
                                                    run_on_cpu=run_on_cpu)

        face_enhancer_func = MPClassFuncOnDemand(FaceEnhancer, 'enhance',
                                                    place_model_on_cpu=True,
                                                    run_on_cpu=run_on_cpu)

        is_interactive = io.input_bool ("Use interactive merger?", True) if not io.is_colab() else False

        if not is_interactive:
            cfg.ask_settings()
            
        subprocess_count = io.input_int("Number of workers?", max(8, multiprocessing.cpu_count()), 
                                        valid_range=[1, multiprocessing.cpu_count()], help_message="Specify the number of threads to process. A low value may affect performance. A high value may result in memory error. The value may not be greater than CPU cores." )

        input_path_image_paths = pathex.get_image_paths(input_path)

        if cfg.type == MergerConfig.TYPE_MASKED:
            if not aligned_path.exists():
                io.log_err('Aligned directory not found. Please ensure it exists.')
                return

            packed_samples = None
            try:
                packed_samples = samplelib.PackedFaceset.load(aligned_path)
            except:
                io.log_err(f"Error occured while loading samplelib.PackedFaceset.load {str(aligned_path)}, {traceback.format_exc()}")


            if packed_samples is not None:
                io.log_info ("Using packed faceset.")
                def generator():
                    for sample in io.progress_bar_generator( packed_samples, "Collecting alignments"):
                        filepath = Path(sample.filename)
                        yield filepath, DFLIMG.load(filepath, loader_func=lambda x: sample.read_raw_file()  )
            else:
                def generator():
                    for filepath in io.progress_bar_generator( pathex.get_image_paths(aligned_path), "Collecting alignments"):
                        filepath = Path(filepath)
                        yield filepath, DFLIMG.load(filepath)

            alignments = {}
            multiple_faces_detected = False

            for filepath, dflimg in generator():
                if dflimg is None or not dflimg.has_data():
                    io.log_err (f"{filepath.name} is not a dfl image file")
                    continue

                source_filename = dflimg.get_source_filename()
                if source_filename is None:
                    continue

                source_filepath = Path(source_filename)
                source_filename_stem = source_filepath.stem

                if source_filename_stem not in alignments.keys():
                    alignments[ source_filename_stem ] = []

                alignments_ar = alignments[ source_filename_stem ]
                alignments_ar.append ( (dflimg.get_source_landmarks(), filepath, source_filepath ) )

                if len(alignments_ar) > 1:
                    multiple_faces_detected = True

            if multiple_faces_detected:
                io.log_info ("")
                io.log_info ("Warning: multiple faces detected. Only one alignment file should refer one source file.")
                io.log_info ("")

            for a_key in list(alignments.keys()):
                a_ar = alignments[a_key]
                if len(a_ar) > 1:
                    for _, filepath, source_filepath in a_ar:
                        io.log_info (f"alignment {filepath.name} refers to {source_filepath.name} ")
                    io.log_info ("")

                alignments[a_key] = [ a[0] for a in a_ar]

            if multiple_faces_detected:
                io.log_info ("It is strongly recommended to process the faces separatelly.")
                io.log_info ("Use 'recover original filename' to determine the exact duplicates.")
                io.log_info ("")

            frames = [ InteractiveMergerSubprocessor.Frame( frame_info=FrameInfo(filepath=Path(p),
                                                                     landmarks_list=alignments.get(Path(p).stem, None)
                                                                    )
                                              )
                       for p in input_path_image_paths ]

            if multiple_faces_detected:
                io.log_info ("Warning: multiple faces detected. Motion blur will not be used.")
                io.log_info ("")
            else:
                s = 256
                local_pts = [ (s//2-1, s//2-1), (s//2-1,0) ] #center+up
                frames_len = len(frames)
                for i in io.progress_bar_generator( range(len(frames)) , "Computing motion vectors"):
                    fi_prev = frames[max(0, i-1)].frame_info
                    fi      = frames[i].frame_info
                    fi_next = frames[min(i+1, frames_len-1)].frame_info
                    if len(fi_prev.landmarks_list) == 0 or \
                       len(fi.landmarks_list) == 0 or \
                       len(fi_next.landmarks_list) == 0:
                            continue

                    mat_prev = LandmarksProcessor.get_transform_mat ( fi_prev.landmarks_list[0], s, face_type=FaceType.FULL)
                    mat      = LandmarksProcessor.get_transform_mat ( fi.landmarks_list[0]     , s, face_type=FaceType.FULL)
                    mat_next = LandmarksProcessor.get_transform_mat ( fi_next.landmarks_list[0], s, face_type=FaceType.FULL)

                    pts_prev = LandmarksProcessor.transform_points (local_pts, mat_prev, True)
                    pts      = LandmarksProcessor.transform_points (local_pts, mat, True)
                    pts_next = LandmarksProcessor.transform_points (local_pts, mat_next, True)

                    prev_vector = pts[0]-pts_prev[0]
                    next_vector = pts_next[0]-pts[0]

                    motion_vector = pts_next[0] - pts_prev[0]
                    fi.motion_power = npla.norm(motion_vector)

                    motion_vector = motion_vector / fi.motion_power if fi.motion_power != 0 else np.array([0,0],dtype=np.float32)

                    fi.motion_deg = -math.atan2(motion_vector[1],motion_vector[0])*180 / math.pi


        if len(frames) == 0:
            io.log_info ("No frames to merge in input_dir.")
        else:
            if False:
                pass
            else:
                InteractiveMergerSubprocessor (
                            is_interactive         = is_interactive,
                            merger_session_filepath = model.get_strpath_storage_for_file('merger_session.dat'),
                            predictor_func         = predictor_func,
                            predictor_input_shape  = predictor_input_shape,
                            face_enhancer_func     = face_enhancer_func,
                            xseg_256_extract_func = xseg_256_extract_func,
                            merger_config          = cfg,
                            frames                 = frames,
                            frames_root_path       = input_path,
                            output_path            = output_path,
                            output_mask_path       = output_mask_path,
                            model_iter             = model.get_iter(),
                            subprocess_count       = subprocess_count,
                        ).run()

        model.finalize()

    except Exception as e:
        print ( traceback.format_exc() )


"""
elif cfg.type == MergerConfig.TYPE_FACE_AVATAR:
filesdata = []
for filepath in io.progress_bar_generator(input_path_image_paths, "Collecting info"):
    filepath = Path(filepath)

    dflimg = DFLIMG.x(filepath)
    if dflimg is None:
        io.log_err ("%s is not a dfl image file" % (filepath.name) )
        continue
    filesdata += [ ( FrameInfo(filepath=filepath, landmarks_list=[dflimg.get_landmarks()] ), dflimg.get_source_filename() ) ]

filesdata = sorted(filesdata, key=operator.itemgetter(1)) #sort by source_filename
frames = []
filesdata_len = len(filesdata)
for i in range(len(filesdata)):
    frame_info = filesdata[i][0]

    prev_temporal_frame_infos = []
    next_temporal_frame_infos = []

    for t in range (cfg.temporal_face_count):
        prev_frame_info = filesdata[ max(i -t, 0) ][0]
        next_frame_info = filesdata[ min(i +t, filesdata_len-1 )][0]

        prev_temporal_frame_infos.insert (0, prev_frame_info )
        next_temporal_frame_infos.append (   next_frame_info )

    frames.append ( InteractiveMergerSubprocessor.Frame(prev_temporal_frame_infos=prev_temporal_frame_infos,
                                                frame_info=frame_info,
                                                next_temporal_frame_infos=next_temporal_frame_infos) )
"""

#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)

mainscripts/Merger.py
oo92/TorchFaceLab
﻿import math
import multiprocessing
import operator
import os
import sys
import tempfile
from functools import cmp_to_key
from pathlib import Path

import cv2
import numpy as np
from numpy import linalg as npla

from core import imagelib, mathlib, pathex
from core.cv2ex import *
from core.imagelib import estimate_sharpness
from core.interact import interact as io
from core.joblib import Subprocessor
from core.leras import nn
from DFLIMG import *
from facelib import LandmarksProcessor


class BlurEstimatorSubprocessor(Subprocessor):
    class Cli(Subprocessor.Cli):
        def on_initialize(self, client_dict):
            self.estimate_motion_blur = client_dict['estimate_motion_blur']
        
        #override
        def process_data(self, data):
            filepath = Path( data[0] )
            dflimg = DFLIMG.load (filepath)

            if dflimg is None or not dflimg.has_data():
                self.log_err (f"{filepath.name} is not a dfl image file")
                return [ str(filepath), 0 ]
            else:
                image = cv2_imread( str(filepath) )
                
                face_mask = LandmarksProcessor.get_image_hull_mask (image.shape, dflimg.get_landmarks())
                image = (image*face_mask).astype(np.uint8)
                
                
                if self.estimate_motion_blur:
                    value = cv2.Laplacian(image, cv2.CV_64F, ksize=11).var()    
                else:
                    value = estimate_sharpness(image)
                
                return [ str(filepath), value ]


        #override
        def get_data_name (self, data):
            #return string identificator of your data
            return data[0]

    #override
    def __init__(self, input_data, estimate_motion_blur=False ):
        self.input_data = input_data
        self.estimate_motion_blur = estimate_motion_blur
        self.img_list = []
        self.trash_img_list = []
        super().__init__('BlurEstimator', BlurEstimatorSubprocessor.Cli, 60)

    #override
    def on_clients_initialized(self):
        io.progress_bar ("", len (self.input_data))

    #override
    def on_clients_finalized(self):
        io.progress_bar_close ()

    #override
    def process_info_generator(self):
        cpu_count = multiprocessing.cpu_count()
        io.log_info(f'Running on {cpu_count} CPUs')

        for i in range(cpu_count):
            yield 'CPU%d' % (i), {}, {'estimate_motion_blur':self.estimate_motion_blur}

    #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[1] == 0:
            self.trash_img_list.append ( result )
        else:
            self.img_list.append ( result )

        io.progress_bar_inc(1)

    #override
    def get_result(self):
        return self.img_list, self.trash_img_list


def sort_by_blur(input_path):
    io.log_info ("Sorting by blur...")

    img_list = [ (filename,[]) for filename in pathex.get_image_paths(input_path) ]
    img_list, trash_img_list = BlurEstimatorSubprocessor (img_list).run()

    io.log_info ("Sorting...")
    img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)

    return img_list, trash_img_list
    
def sort_by_motion_blur(input_path):
    io.log_info ("Sorting by motion blur...")

    img_list = [ (filename,[]) for filename in pathex.get_image_paths(input_path) ]
    img_list, trash_img_list = BlurEstimatorSubprocessor (img_list, estimate_motion_blur=True).run()

    io.log_info ("Sorting...")
    img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)

    return img_list, trash_img_list
    
def sort_by_face_yaw(input_path):
    io.log_info ("Sorting by face yaw...")
    img_list = []
    trash_img_list = []
    for filepath in io.progress_bar_generator( pathex.get_image_paths(input_path), "Loading"):
        filepath = Path(filepath)

        dflimg = DFLIMG.load (filepath)

        if dflimg is None or not dflimg.has_data():
            io.log_err (f"{filepath.name} is not a dfl image file")
            trash_img_list.append ( [str(filepath)] )
            continue

        pitch, yaw, roll = LandmarksProcessor.estimate_pitch_yaw_roll ( dflimg.get_landmarks(), size=dflimg.get_shape()[1] )

        img_list.append( [str(filepath), yaw ] )

    io.log_info ("Sorting...")
    img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)

    return img_list, trash_img_list

def sort_by_face_pitch(input_path):
    io.log_info ("Sorting by face pitch...")
    img_list = []
    trash_img_list = []
    for filepath in io.progress_bar_generator( pathex.get_image_paths(input_path), "Loading"):
        filepath = Path(filepath)

        dflimg = DFLIMG.load (filepath)

        if dflimg is None or not dflimg.has_data():
            io.log_err (f"{filepath.name} is not a dfl image file")
            trash_img_list.append ( [str(filepath)] )
            continue

        pitch, yaw, roll = LandmarksProcessor.estimate_pitch_yaw_roll ( dflimg.get_landmarks(), size=dflimg.get_shape()[1] )

        img_list.append( [str(filepath), pitch ] )

    io.log_info ("Sorting...")
    img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)

    return img_list, trash_img_list

def sort_by_face_source_rect_size(input_path):
    io.log_info ("Sorting by face rect size...")
    img_list = []
    trash_img_list = []
    for filepath in io.progress_bar_generator( pathex.get_image_paths(input_path), "Loading"):
        filepath = Path(filepath)

        dflimg = DFLIMG.load (filepath)

        if dflimg is None or not dflimg.has_data():
            io.log_err (f"{filepath.name} is not a dfl image file")
            trash_img_list.append ( [str(filepath)] )
            continue

        source_rect = dflimg.get_source_rect()
        rect_area = mathlib.polygon_area(np.array(source_rect[[0,2,2,0]]).astype(np.float32), np.array(source_rect[[1,1,3,3]]).astype(np.float32))

        img_list.append( [str(filepath), rect_area ] )

    io.log_info ("Sorting...")
    img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)

    return img_list, trash_img_list



class HistSsimSubprocessor(Subprocessor):
    class Cli(Subprocessor.Cli):
        #override
        def process_data(self, data):
            img_list = []
            for x in data:
                img = cv2_imread(x)
                img_list.append ([x, cv2.calcHist([img], [0], None, [256], [0, 256]),
                                     cv2.calcHist([img], [1], None, [256], [0, 256]),
                                     cv2.calcHist([img], [2], None, [256], [0, 256])
                                 ])

            img_list_len = len(img_list)
            for i in range(img_list_len-1):
                min_score = float("inf")
                j_min_score = i+1
                for j in range(i+1,len(img_list)):
                    score = cv2.compareHist(img_list[i][1], img_list[j][1], cv2.HISTCMP_BHATTACHARYYA) + \
                            cv2.compareHist(img_list[i][2], img_list[j][2], cv2.HISTCMP_BHATTACHARYYA) + \
                            cv2.compareHist(img_list[i][3], img_list[j][3], cv2.HISTCMP_BHATTACHARYYA)
                    if score < min_score:
                        min_score = score
                        j_min_score = j
                img_list[i+1], img_list[j_min_score] = img_list[j_min_score], img_list[i+1]

                self.progress_bar_inc(1)

            return img_list

        #override
        def get_data_name (self, data):
            return "Bunch of images"

    #override
    def __init__(self, img_list ):
        self.img_list = img_list
        self.img_list_len = len(img_list)

        slice_count = 20000
        sliced_count = self.img_list_len // slice_count

        if sliced_count > 12:
            sliced_count = 11.9
            slice_count = int(self.img_list_len / sliced_count)
            sliced_count = self.img_list_len // slice_count

        self.img_chunks_list = [ self.img_list[i*slice_count : (i+1)*slice_count] for i in range(sliced_count) ] + \
                               [ self.img_list[sliced_count*slice_count:] ]

        self.result = []
        super().__init__('HistSsim', HistSsimSubprocessor.Cli, 0)

    #override
    def process_info_generator(self):
        cpu_count = len(self.img_chunks_list)
        io.log_info(f'Running on {cpu_count} threads')
        for i in range(cpu_count):
            yield 'CPU%d' % (i), {'i':i}, {}

    #override
    def on_clients_initialized(self):
        io.progress_bar ("Sorting", len(self.img_list))
        io.progress_bar_inc(len(self.img_chunks_list))

    #override
    def on_clients_finalized(self):
        io.progress_bar_close()

    #override
    def get_data(self, host_dict):
        if len (self.img_chunks_list) > 0:
            return self.img_chunks_list.pop(0)
        return None

    #override
    def on_data_return (self, host_dict, data):
        raise Exception("Fail to process data. Decrease number of images and try again.")

    #override
    def on_result (self, host_dict, data, result):
        self.result += result
        return 0

    #override
    def get_result(self):
        return self.result

def sort_by_hist(input_path):
    io.log_info ("Sorting by histogram similarity...")
    img_list = HistSsimSubprocessor(pathex.get_image_paths(input_path)).run()
    return img_list, []

class HistDissimSubprocessor(Subprocessor):
    class Cli(Subprocessor.Cli):
        #override
        def on_initialize(self, client_dict):
            self.img_list = client_dict['img_list']
            self.img_list_len = len(self.img_list)

        #override
        def process_data(self, data):
            i = data[0]
            score_total = 0
            for j in range( 0, self.img_list_len):
                if i == j:
                    continue
                score_total += cv2.compareHist(self.img_list[i][1], self.img_list[j][1], cv2.HISTCMP_BHATTACHARYYA)

            return score_total

        #override
        def get_data_name (self, data):
            #return string identificator of your data
            return self.img_list[data[0]][0]

    #override
    def __init__(self, img_list ):
        self.img_list = img_list
        self.img_list_range = [i for i in range(0, len(img_list) )]
        self.result = []
        super().__init__('HistDissim', HistDissimSubprocessor.Cli, 60)

    #override
    def on_clients_initialized(self):
        io.progress_bar ("Sorting", len (self.img_list) )

    #override
    def on_clients_finalized(self):
        io.progress_bar_close()

    #override
    def process_info_generator(self):
        cpu_count = min(multiprocessing.cpu_count(), 8)
        io.log_info(f'Running on {cpu_count} CPUs')
        for i in range(cpu_count):
            yield 'CPU%d' % (i), {}, {'img_list' : self.img_list}

    #override
    def get_data(self, host_dict):
        if len (self.img_list_range) > 0:
            return [self.img_list_range.pop(0)]

        return None

    #override
    def on_data_return (self, host_dict, data):
        self.img_list_range.insert(0, data[0])

    #override
    def on_result (self, host_dict, data, result):
        self.img_list[data[0]][2] = result
        io.progress_bar_inc(1)

    #override
    def get_result(self):
        return self.img_list

def sort_by_hist_dissim(input_path):
    io.log_info ("Sorting by histogram dissimilarity...")

    img_list = []
    trash_img_list = []
    for filepath in io.progress_bar_generator( pathex.get_image_paths(input_path), "Loading"):
        filepath = Path(filepath)

        dflimg = DFLIMG.load (filepath)

        image = cv2_imread(str(filepath))

        if dflimg is not None and dflimg.has_data():
            face_mask = LandmarksProcessor.get_image_hull_mask (image.shape, dflimg.get_landmarks())
            image = (image*face_mask).astype(np.uint8)

        img_list.append ([str(filepath), cv2.calcHist([cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)], [0], None, [256], [0, 256]), 0 ])

    img_list = HistDissimSubprocessor(img_list).run()

    io.log_info ("Sorting...")
    img_list = sorted(img_list, key=operator.itemgetter(2), reverse=True)

    return img_list, trash_img_list

def sort_by_brightness(input_path):
    io.log_info ("Sorting by brightness...")
    img_list = [ [x, np.mean ( cv2.cvtColor(cv2_imread(x), cv2.COLOR_BGR2HSV)[...,2].flatten()  )] for x in io.progress_bar_generator( pathex.get_image_paths(input_path), "Loading") ]
    io.log_info ("Sorting...")
    img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)
    return img_list, []

def sort_by_hue(input_path):
    io.log_info ("Sorting by hue...")
    img_list = [ [x, np.mean ( cv2.cvtColor(cv2_imread(x), cv2.COLOR_BGR2HSV)[...,0].flatten()  )] for x in io.progress_bar_generator( pathex.get_image_paths(input_path), "Loading") ]
    io.log_info ("Sorting...")
    img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)
    return img_list, []

def sort_by_black(input_path):
    io.log_info ("Sorting by amount of black pixels...")

    img_list = []
    for x in io.progress_bar_generator( pathex.get_image_paths(input_path), "Loading"):
        img = cv2_imread(x)
        img_list.append ([x, img[(img == 0)].size ])

    io.log_info ("Sorting...")
    img_list = sorted(img_list, key=operator.itemgetter(1), reverse=False)

    return img_list, []

def sort_by_origname(input_path):
    io.log_info ("Sort by original filename...")

    img_list = []
    trash_img_list = []
    for filepath in io.progress_bar_generator( pathex.get_image_paths(input_path), "Loading"):
        filepath = Path(filepath)

        dflimg = DFLIMG.load (filepath)

        if dflimg is None or not dflimg.has_data():
            io.log_err (f"{filepath.name} is not a dfl image file")
            trash_img_list.append( [str(filepath)] )
            continue

        img_list.append( [str(filepath), dflimg.get_source_filename()] )

    io.log_info ("Sorting...")
    img_list = sorted(img_list, key=operator.itemgetter(1))
    return img_list, trash_img_list

def sort_by_oneface_in_image(input_path):
    io.log_info ("Sort by one face in images...")
    image_paths = pathex.get_image_paths(input_path)
    a = np.array ([ ( int(x[0]), int(x[1]) ) \
                      for x in [ Path(filepath).stem.split('_') for filepath in image_paths ] if len(x) == 2
                  ])
    if len(a) > 0:
        idxs = np.ndarray.flatten ( np.argwhere ( a[:,1] != 0 ) )
        idxs = np.unique ( a[idxs][:,0] )
        idxs = np.ndarray.flatten ( np.argwhere ( np.array([ x[0] in idxs for x in a ]) == True ) )
        if len(idxs) > 0:
            io.log_info ("Found %d images." % (len(idxs)) )
            img_list = [ (path,) for i,path in enumerate(image_paths) if i not in idxs ]
            trash_img_list = [ (image_paths[x],) for x in idxs ]
            return img_list, trash_img_list

    io.log_info ("Nothing found. Possible recover original filenames first.")
    return [], []

class FinalLoaderSubprocessor(Subprocessor):
    class Cli(Subprocessor.Cli):
        #override
        def on_initialize(self, client_dict):
            self.faster = client_dict['faster']

        #override
        def process_data(self, data):
            filepath = Path(data[0])

            try:
                dflimg = DFLIMG.load (filepath)

                if dflimg is None or not dflimg.has_data():
                    self.log_err (f"{filepath.name} is not a dfl image file")
                    return [ 1, [str(filepath)] ]

                bgr = cv2_imread(str(filepath))
                if bgr is None:
                    raise Exception ("Unable to load %s" % (filepath.name) )

                gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
                if self.faster:
                    source_rect = dflimg.get_source_rect()
                    sharpness = mathlib.polygon_area(np.array(source_rect[[0,2,2,0]]).astype(np.float32), np.array(source_rect[[1,1,3,3]]).astype(np.float32))
                else:
                    face_mask = LandmarksProcessor.get_image_hull_mask (gray.shape, dflimg.get_landmarks())     
                    sharpness = estimate_sharpness( (gray[...,None]*face_mask).astype(np.uint8) )

                pitch, yaw, roll = LandmarksProcessor.estimate_pitch_yaw_roll ( dflimg.get_landmarks(), size=dflimg.get_shape()[1] )

                hist = cv2.calcHist([gray], [0], None, [256], [0, 256])
            except Exception as e:
                self.log_err (e)
                return [ 1, [str(filepath)] ]

            return [ 0, [str(filepath), sharpness, hist, yaw, pitch ] ]

        #override
        def get_data_name (self, data):
            #return string identificator of your data
            return data[0]

    #override
    def __init__(self, img_list, faster ):
        self.img_list = img_list

        self.faster = faster
        self.result = []
        self.result_trash = []

        super().__init__('FinalLoader', FinalLoaderSubprocessor.Cli, 60)

    #override
    def on_clients_initialized(self):
        io.progress_bar ("Loading", len (self.img_list))

    #override
    def on_clients_finalized(self):
        io.progress_bar_close()

    #override
    def process_info_generator(self):
        cpu_count = min(multiprocessing.cpu_count(), 8)
        io.log_info(f'Running on {cpu_count} CPUs')

        for i in range(cpu_count):
            yield 'CPU%d' % (i), {}, {'faster': self.faster}

    #override
    def get_data(self, host_dict):
        if len (self.img_list) > 0:
            return [self.img_list.pop(0)]

        return None

    #override
    def on_data_return (self, host_dict, data):
        self.img_list.insert(0, data[0])

    #override
    def on_result (self, host_dict, data, result):
        if result[0] == 0:
            self.result.append (result[1])
        else:
            self.result_trash.append (result[1])
        io.progress_bar_inc(1)

    #override
    def get_result(self):
        return self.result, self.result_trash

class FinalHistDissimSubprocessor(Subprocessor):
    class Cli(Subprocessor.Cli):
        #override
        def process_data(self, data):
            idx, pitch_yaw_img_list = data

            for p in range ( len(pitch_yaw_img_list) ):

                img_list = pitch_yaw_img_list[p]
                if img_list is not None:
                    for i in range( len(img_list) ):
                        score_total = 0
                        for j in range( len(img_list) ):
                            if i == j:
                                continue
                            score_total += cv2.compareHist(img_list[i][2], img_list[j][2], cv2.HISTCMP_BHATTACHARYYA)
                        img_list[i][3] = score_total

                    pitch_yaw_img_list[p] = sorted(img_list, key=operator.itemgetter(3), reverse=True)

            return idx, pitch_yaw_img_list

        #override
        def get_data_name (self, data):
            return "Bunch of images"

    #override
    def __init__(self, pitch_yaw_sample_list ):
        self.pitch_yaw_sample_list = pitch_yaw_sample_list
        self.pitch_yaw_sample_list_len = len(pitch_yaw_sample_list)

        self.pitch_yaw_sample_list_idxs = [ i for i in range(self.pitch_yaw_sample_list_len) if self.pitch_yaw_sample_list[i] is not None ]
        self.result = [ None for _ in range(self.pitch_yaw_sample_list_len) ]
        super().__init__('FinalHistDissimSubprocessor', FinalHistDissimSubprocessor.Cli)

    #override
    def process_info_generator(self):
        cpu_count = min(multiprocessing.cpu_count(), 8)
        io.log_info(f'Running on {cpu_count} CPUs')
        for i in range(cpu_count):
            yield 'CPU%d' % (i), {}, {}

    #override
    def on_clients_initialized(self):
        io.progress_bar ("Sort by hist-dissim", len(self.pitch_yaw_sample_list_idxs) )

    #override
    def on_clients_finalized(self):
        io.progress_bar_close()

    #override
    def get_data(self, host_dict):
        if len (self.pitch_yaw_sample_list_idxs) > 0:
            idx = self.pitch_yaw_sample_list_idxs.pop(0)

            return idx, self.pitch_yaw_sample_list[idx]
        return None

    #override
    def on_data_return (self, host_dict, data):
        self.pitch_yaw_sample_list_idxs.insert(0, data[0])

    #override
    def on_result (self, host_dict, data, result):
        idx, yaws_sample_list = data
        self.result[idx] = yaws_sample_list
        io.progress_bar_inc(1)

    #override
    def get_result(self):
        return self.result

def sort_best_faster(input_path):
    return sort_best(input_path, faster=True)

def sort_best(input_path, faster=False):
    target_count = io.input_int ("Target number of faces?", 2000)

    io.log_info ("Performing sort by best faces.")
    if faster:
        io.log_info("Using faster algorithm. Faces will be sorted by source-rect-area instead of blur.")

    img_list, trash_img_list = FinalLoaderSubprocessor( pathex.get_image_paths(input_path), faster ).run()
    final_img_list = []

    grads = 128
    imgs_per_grad = round (target_count / grads)

    #instead of math.pi / 2, using -1.2,+1.2 because actually maximum yaw for 2DFAN landmarks are -1.2+1.2
    grads_space = np.linspace (-1.2, 1.2,grads)

    yaws_sample_list = [None]*grads
    for g in io.progress_bar_generator ( range(grads), "Sort by yaw"):
        yaw = grads_space[g]
        next_yaw = grads_space[g+1] if g < grads-1 else yaw

        yaw_samples = []
        for img in img_list:
            s_yaw = -img[3]
            if (g == 0          and s_yaw < next_yaw) or \
               (g < grads-1     and s_yaw >= yaw and s_yaw < next_yaw) or \
               (g == grads-1    and s_yaw >= yaw):
                yaw_samples += [ img ]
        if len(yaw_samples) > 0:
            yaws_sample_list[g] = yaw_samples

    total_lack = 0
    for g in io.progress_bar_generator ( range(grads), ""):
        img_list = yaws_sample_list[g]
        img_list_len = len(img_list) if img_list is not None else 0

        lack = imgs_per_grad - img_list_len
        total_lack += max(lack, 0)

    imgs_per_grad += total_lack // grads


    sharpned_imgs_per_grad = imgs_per_grad*10
    for g in io.progress_bar_generator ( range (grads), "Sort by blur"):
        img_list = yaws_sample_list[g]
        if img_list is None:
            continue

        img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)

        if len(img_list) > sharpned_imgs_per_grad:
            trash_img_list += img_list[sharpned_imgs_per_grad:]
            img_list = img_list[0:sharpned_imgs_per_grad]

        yaws_sample_list[g] = img_list


    yaw_pitch_sample_list = [None]*grads
    pitch_grads = imgs_per_grad

    for g in io.progress_bar_generator ( range (grads), "Sort by pitch"):
        img_list = yaws_sample_list[g]
        if img_list is None:
            continue

        pitch_sample_list = [None]*pitch_grads

        grads_space = np.linspace (-math.pi / 2,math.pi / 2, pitch_grads )

        for pg in range (pitch_grads):

            pitch = grads_space[pg]
            next_pitch = grads_space[pg+1] if pg < pitch_grads-1 else pitch

            pitch_samples = []
            for img in img_list:
                s_pitch = img[4]
                if (pg == 0                and s_pitch < next_pitch) or \
                   (pg < pitch_grads-1     and s_pitch >= pitch and s_pitch < next_pitch) or \
                   (pg == pitch_grads-1    and s_pitch >= pitch):
                    pitch_samples += [ img ]

            if len(pitch_samples) > 0:
                pitch_sample_list[pg] = pitch_samples
        yaw_pitch_sample_list[g] = pitch_sample_list

    yaw_pitch_sample_list = FinalHistDissimSubprocessor(yaw_pitch_sample_list).run()

    for g in io.progress_bar_generator (range (grads), "Fetching the best"):
        pitch_sample_list = yaw_pitch_sample_list[g]
        if pitch_sample_list is None:
            continue

        n = imgs_per_grad

        while n > 0:
            n_prev = n
            for pg in range(pitch_grads):
                img_list = pitch_sample_list[pg]
                if img_list is None:
                    continue
                final_img_list += [ img_list.pop(0) ]
                if len(img_list) == 0:
                    pitch_sample_list[pg] = None
                n -= 1
                if n == 0:
                    break
            if n_prev == n:
                break

        for pg in range(pitch_grads):
            img_list = pitch_sample_list[pg]
            if img_list is None:
                continue
            trash_img_list += img_list

    return final_img_list, trash_img_list

"""
def sort_by_vggface(input_path):
    io.log_info ("Sorting by face similarity using VGGFace model...")

    model = VGGFace()

    final_img_list = []
    trash_img_list = []

    image_paths = pathex.get_image_paths(input_path)
    img_list = [ (x,) for x in image_paths ]
    img_list_len = len(img_list)
    img_list_range = [*range(img_list_len)]

    feats = [None]*img_list_len
    for i in io.progress_bar_generator(img_list_range, "Loading"):
        img = cv2_imread( img_list[i][0] ).astype(np.float32)
        img = imagelib.normalize_channels (img, 3)
        img = cv2.resize (img, (224,224) )
        img = img[..., ::-1]
        img[..., 0] -= 93.5940
        img[..., 1] -= 104.7624
        img[..., 2] -= 129.1863
        feats[i] = model.predict( img[None,...] )[0]

    tmp = np.zeros( (img_list_len,) )
    float_inf = float("inf")
    for i in io.progress_bar_generator ( range(img_list_len-1), "Sorting" ):
        i_feat = feats[i]

        for j in img_list_range:
            tmp[j] = npla.norm(i_feat-feats[j]) if j >= i+1 else float_inf

        idx = np.argmin(tmp)

        img_list[i+1], img_list[idx] = img_list[idx], img_list[i+1]
        feats[i+1], feats[idx] = feats[idx], feats[i+1]

    return img_list, trash_img_list
"""

def sort_by_absdiff(input_path):
    io.log_info ("Sorting by absolute difference...")

    is_sim = io.input_bool ("Sort by similar?", True, help_message="Otherwise sort by dissimilar.")

    from core.leras import nn

    device_config = nn.DeviceConfig.ask_choose_device(choose_only_one=True)
    nn.initialize( device_config=device_config, data_format="NHWC" )
    tf = nn.tf

    image_paths = pathex.get_image_paths(input_path)
    image_paths_len = len(image_paths)

    batch_size = 512
    batch_size_remain = image_paths_len % batch_size

    i_t = tf.placeholder (tf.float32, (None,None,None,None) )
    j_t = tf.placeholder (tf.float32, (None,None,None,None) )

    outputs_full = []
    outputs_remain = []

    for i in range(batch_size):
        diff_t = tf.reduce_sum( tf.abs(i_t-j_t[i]), axis=[1,2,3] )
        outputs_full.append(diff_t)
        if i < batch_size_remain:
            outputs_remain.append(diff_t)

    def func_bs_full(i,j):
        return nn.tf_sess.run (outputs_full, feed_dict={i_t:i,j_t:j})

    def func_bs_remain(i,j):
        return nn.tf_sess.run (outputs_remain, feed_dict={i_t:i,j_t:j})

    import h5py
    db_file_path = Path(tempfile.gettempdir()) / 'sort_cache.hdf5'
    db_file = h5py.File( str(db_file_path), "w")
    db = db_file.create_dataset("results", (image_paths_len,image_paths_len), compression="gzip")

    pg_len = image_paths_len // batch_size
    if batch_size_remain != 0:
        pg_len += 1

    pg_len = int( (  pg_len*pg_len - pg_len ) / 2 + pg_len )

    io.progress_bar ("Computing", pg_len)
    j=0
    while j < image_paths_len:
        j_images = [ cv2_imread(x) for x in image_paths[j:j+batch_size] ]
        j_images_len = len(j_images)

        func = func_bs_remain if image_paths_len-j < batch_size else func_bs_full

        i=0
        while i < image_paths_len:
            if i >= j:
                i_images = [ cv2_imread(x) for x in image_paths[i:i+batch_size] ]
                i_images_len = len(i_images)
                result = func (i_images,j_images)
                db[j:j+j_images_len,i:i+i_images_len] = np.array(result)
                io.progress_bar_inc(1)

            i += batch_size
        db_file.flush()
        j += batch_size

    io.progress_bar_close()

    next_id = 0
    sorted = [next_id]
    for i in io.progress_bar_generator ( range(image_paths_len-1), "Sorting" ):
        id_ar = np.concatenate ( [ db[:next_id,next_id], db[next_id,next_id:] ] )
        id_ar = np.argsort(id_ar)


        next_id = np.setdiff1d(id_ar, sorted, True)[ 0 if is_sim else -1]
        sorted += [next_id]
    db_file.close()
    db_file_path.unlink()

    img_list = [ (image_paths[x],) for x in sorted]
    return img_list, []

def final_process(input_path, img_list, trash_img_list):
    if len(trash_img_list) != 0:
        parent_input_path = input_path.parent
        trash_path = parent_input_path / (input_path.stem + '_trash')
        trash_path.mkdir (exist_ok=True)

        io.log_info ("Trashing %d items to %s" % ( len(trash_img_list), str(trash_path) ) )

        for filename in pathex.get_image_paths(trash_path):
            Path(filename).unlink()

        for i in io.progress_bar_generator( range(len(trash_img_list)), "Moving trash", leave=False):
            src = Path (trash_img_list[i][0])
            dst = trash_path / src.name
            try:
                src.rename (dst)
            except:
                io.log_info ('fail to trashing %s' % (src.name) )

        io.log_info ("")

    if len(img_list) != 0:
        for i in io.progress_bar_generator( [*range(len(img_list))], "Renaming", leave=False):
            src = Path (img_list[i][0])
            dst = input_path / ('%.5d_%s' % (i, src.name ))
            try:
                src.rename (dst)
            except:
                io.log_info ('fail to rename %s' % (src.name) )

        for i in io.progress_bar_generator( [*range(len(img_list))], "Renaming"):
            src = Path (img_list[i][0])
            src = input_path / ('%.5d_%s' % (i, src.name))
            dst = input_path / ('%.5d%s' % (i, src.suffix))
            try:
                src.rename (dst)
            except:
                io.log_info ('fail to rename %s' % (src.name) )

sort_func_methods = {
    'blur':        ("blur", sort_by_blur),
    'motion-blur': ("motion_blur", sort_by_motion_blur),
    'face-yaw':    ("face yaw direction", sort_by_face_yaw),
    'face-pitch':  ("face pitch direction", sort_by_face_pitch),
    'face-source-rect-size' : ("face rect size in source image", sort_by_face_source_rect_size),
    'hist':        ("histogram similarity", sort_by_hist),
    'hist-dissim': ("histogram dissimilarity", sort_by_hist_dissim),
    'brightness':  ("brightness", sort_by_brightness),
    'hue':         ("hue", sort_by_hue),
    'black':       ("amount of black pixels", sort_by_black),
    'origname':    ("original filename", sort_by_origname),
    'oneface':     ("one face in image", sort_by_oneface_in_image),
    'absdiff':     ("absolute pixel difference", sort_by_absdiff),
    'final':       ("best faces", sort_best),
    'final-fast':  ("best faces faster", sort_best_faster),
}

def main (input_path, sort_by_method=None):
    io.log_info ("Running sort tool.\r\n")

    if sort_by_method is None:
        io.log_info(f"Choose sorting method:")

        key_list = list(sort_func_methods.keys())
        for i, key in enumerate(key_list):
            desc, func = sort_func_methods[key]
            io.log_info(f"[{i}] {desc}")

        io.log_info("")
        id = io.input_int("", 5, valid_list=[*range(len(key_list))] )

        sort_by_method = key_list[id]
    else:
        sort_by_method = sort_by_method.lower()

    desc, func = sort_func_methods[sort_by_method]
    img_list, trash_img_list = func(input_path)

    final_process (input_path, img_list, trash_img_list)

mainscripts/Sorter.py
oo92/TorchFaceLab
﻿import os
import sys
import traceback
import queue
import threading
import time
import numpy as np
import itertools
from pathlib import Path
from core import pathex
from core import imagelib
import cv2
import models
from core.interact import interact as io

def trainerThread (s2c, c2s, e,
                    model_class_name = None,
                    saved_models_path = None,
                    training_data_src_path = None,
                    training_data_dst_path = None,
                    pretraining_data_path = None,
                    pretrained_model_path = None,
                    no_preview=False,
                    force_model_name=None,
                    force_gpu_idxs=None,
                    cpu_only=None,
                    silent_start=False,
                    execute_programs = None,
                    debug=False,
                    **kwargs):
    while True:
        try:
            start_time = time.time()

            save_interval_min = 25

            if not training_data_src_path.exists():
                training_data_src_path.mkdir(exist_ok=True, parents=True)

            if not training_data_dst_path.exists():
                training_data_dst_path.mkdir(exist_ok=True, parents=True)

            if not saved_models_path.exists():
                saved_models_path.mkdir(exist_ok=True, parents=True)
                            
            model = models.import_model(model_class_name)(
                        is_training=True,
                        saved_models_path=saved_models_path,
                        training_data_src_path=training_data_src_path,
                        training_data_dst_path=training_data_dst_path,
                        pretraining_data_path=pretraining_data_path,
                        pretrained_model_path=pretrained_model_path,
                        no_preview=no_preview,
                        force_model_name=force_model_name,
                        force_gpu_idxs=force_gpu_idxs,
                        cpu_only=cpu_only,
                        silent_start=silent_start,
                        debug=debug)

            is_reached_goal = model.is_reached_iter_goal()

            shared_state = { 'after_save' : False }
            loss_string = ""
            save_iter =  model.get_iter()
            def model_save():
                if not debug and not is_reached_goal:
                    io.log_info ("Saving....", end='\r')
                    model.save()
                    shared_state['after_save'] = True
                    
            def model_backup():
                if not debug and not is_reached_goal:
                    model.create_backup()             

            def send_preview():
                if not debug:
                    previews = model.get_previews()
                    c2s.put ( {'op':'show', 'previews': previews, 'iter':model.get_iter(), 'loss_history': model.get_loss_history().copy() } )
                else:
                    previews = [( 'debug, press update for new', model.debug_one_iter())]
                    c2s.put ( {'op':'show', 'previews': previews} )
                e.set() #Set the GUI Thread as Ready

            if model.get_target_iter() != 0:
                if is_reached_goal:
                    io.log_info('Model already trained to target iteration. You can use preview.')
                else:
                    io.log_info('Starting. Target iteration: %d. Press "Enter" to stop training and save model.' % ( model.get_target_iter()  ) )
            else:
                io.log_info('Starting. Press "Enter" to stop training and save model.')

            last_save_time = time.time()

            execute_programs = [ [x[0], x[1], time.time() ] for x in execute_programs ]

            for i in itertools.count(0,1):
                if not debug:
                    cur_time = time.time()

                    for x in execute_programs:
                        prog_time, prog, last_time = x
                        exec_prog = False
                        if prog_time > 0 and (cur_time - start_time) >= prog_time:
                            x[0] = 0
                            exec_prog = True
                        elif prog_time < 0 and (cur_time - last_time)  >= -prog_time:
                            x[2] = cur_time
                            exec_prog = True

                        if exec_prog:
                            try:
                                exec(prog)
                            except Exception as e:
                                print("Unable to execute program: %s" % (prog) )

                    if not is_reached_goal:

                        if model.get_iter() == 0:
                            io.log_info("")
                            io.log_info("Trying to do the first iteration. If an error occurs, reduce the model parameters.")
                            io.log_info("")
                            
                            if sys.platform[0:3] == 'win':
                                io.log_info("!!!")
                                io.log_info("Windows 10 users IMPORTANT notice. You should set this setting in order to work correctly.")
                                io.log_info("https://i.imgur.com/B7cmDCB.jpg")
                                io.log_info("!!!")

                        iter, iter_time = model.train_one_iter()

                        loss_history = model.get_loss_history()
                        time_str = time.strftime("[%H:%M:%S]")
                        if iter_time >= 10:
                            loss_string = "{0}[#{1:06d}][{2:.5s}s]".format ( time_str, iter, '{:0.4f}'.format(iter_time) )
                        else:
                            loss_string = "{0}[#{1:06d}][{2:04d}ms]".format ( time_str, iter, int(iter_time*1000) )

                        if shared_state['after_save']:
                            shared_state['after_save'] = False
                            
                            mean_loss = np.mean ( loss_history[save_iter:iter], axis=0)

                            for loss_value in mean_loss:
                                loss_string += "[%.4f]" % (loss_value)

                            io.log_info (loss_string)

                            save_iter = iter
                        else:
                            for loss_value in loss_history[-1]:
                                loss_string += "[%.4f]" % (loss_value)

                            if io.is_colab():
                                io.log_info ('\r' + loss_string, end='')
                            else:
                                io.log_info (loss_string, end='\r')

                        if model.get_iter() == 1:
                            model_save()

                        if model.get_target_iter() != 0 and model.is_reached_iter_goal():
                            io.log_info ('Reached target iteration.')
                            model_save()
                            is_reached_goal = True
                            io.log_info ('You can use preview now.')
                
                need_save = False
                while time.time() - last_save_time >= save_interval_min*60:
                    last_save_time += save_interval_min*60
                    need_save = True
                
                if not is_reached_goal and need_save:
                    model_save()
                    send_preview()

                if i==0:
                    if is_reached_goal:
                        model.pass_one_iter()
                    send_preview()

                if debug:
                    time.sleep(0.005)

                while not s2c.empty():
                    input = s2c.get()
                    op = input['op']
                    if op == 'save':
                        model_save()
                    elif op == 'backup':
                        model_backup()
                    elif op == 'preview':
                        if is_reached_goal:
                            model.pass_one_iter()
                        send_preview()
                    elif op == 'close':
                        model_save()
                        i = -1
                        break

                if i == -1:
                    break



            model.finalize()

        except Exception as e:
            print ('Error: %s' % (str(e)))
            traceback.print_exc()
        break
    c2s.put ( {'op':'close'} )



def main(**kwargs):
    io.log_info ("Running trainer.\r\n")

    no_preview = kwargs.get('no_preview', False)

    s2c = queue.Queue()
    c2s = queue.Queue()

    e = threading.Event()
    thread = threading.Thread(target=trainerThread, args=(s2c, c2s, e), kwargs=kwargs )
    thread.start()

    e.wait() #Wait for inital load to occur.

    if no_preview:
        while True:
            if not c2s.empty():
                input = c2s.get()
                op = input.get('op','')
                if op == 'close':
                    break
            try:
                io.process_messages(0.1)
            except KeyboardInterrupt:
                s2c.put ( {'op': 'close'} )
    else:
        wnd_name = "Training preview"
        io.named_window(wnd_name)
        io.capture_keys(wnd_name)

        previews = None
        loss_history = None
        selected_preview = 0
        update_preview = False
        is_showing = False
        is_waiting_preview = False
        show_last_history_iters_count = 0
        iter = 0
        while True:
            if not c2s.empty():
                input = c2s.get()
                op = input['op']
                if op == 'show':
                    is_waiting_preview = False
                    loss_history = input['loss_history'] if 'loss_history' in input.keys() else None
                    previews = input['previews'] if 'previews' in input.keys() else None
                    iter = input['iter'] if 'iter' in input.keys() else 0
                    if previews is not None:
                        max_w = 0
                        max_h = 0
                        for (preview_name, preview_rgb) in previews:
                            (h, w, c) = preview_rgb.shape
                            max_h = max (max_h, h)
                            max_w = max (max_w, w)

                        max_size = 800
                        if max_h > max_size:
                            max_w = int( max_w / (max_h / max_size) )
                            max_h = max_size

                        #make all previews size equal
                        for preview in previews[:]:
                            (preview_name, preview_rgb) = preview
                            (h, w, c) = preview_rgb.shape
                            if h != max_h or w != max_w:
                                previews.remove(preview)
                                previews.append ( (preview_name, cv2.resize(preview_rgb, (max_w, max_h))) )
                        selected_preview = selected_preview % len(previews)
                        update_preview = True
                elif op == 'close':
                    break

            if update_preview:
                update_preview = False

                selected_preview_name = previews[selected_preview][0]
                selected_preview_rgb = previews[selected_preview][1]
                (h,w,c) = selected_preview_rgb.shape

                # HEAD
                head_lines = [
                    '[s]:save [b]:backup [enter]:exit',
                    '[p]:update [space]:next preview [l]:change history range',
                    'Preview: "%s" [%d/%d]' % (selected_preview_name,selected_preview+1, len(previews) )
                    ]
                head_line_height = 15
                head_height = len(head_lines) * head_line_height
                head = np.ones ( (head_height,w,c) ) * 0.1

                for i in range(0, len(head_lines)):
                    t = i*head_line_height
                    b = (i+1)*head_line_height
                    head[t:b, 0:w] += imagelib.get_text_image (  (head_line_height,w,c) , head_lines[i], color=[0.8]*c )

                final = head

                if loss_history is not None:
                    if show_last_history_iters_count == 0:
                        loss_history_to_show = loss_history
                    else:
                        loss_history_to_show = loss_history[-show_last_history_iters_count:]

                    lh_img = models.ModelBase.get_loss_history_preview(loss_history_to_show, iter, w, c)
                    final = np.concatenate ( [final, lh_img], axis=0 )

                final = np.concatenate ( [final, selected_preview_rgb], axis=0 )
                final = np.clip(final, 0, 1)

                io.show_image( wnd_name, (final*255).astype(np.uint8) )
                is_showing = True

            key_events = io.get_key_events(wnd_name)
            key, chr_key, ctrl_pressed, alt_pressed, shift_pressed = key_events[-1] if len(key_events) > 0 else (0,0,False,False,False)

            if key == ord('\n') or key == ord('\r'):
                s2c.put ( {'op': 'close'} )
            elif key == ord('s'):
                s2c.put ( {'op': 'save'} )
            elif key == ord('b'):
                s2c.put ( {'op': 'backup'} )
            elif key == ord('p'):
                if not is_waiting_preview:
                    is_waiting_preview = True
                    s2c.put ( {'op': 'preview'} )
            elif key == ord('l'):
                if show_last_history_iters_count == 0:
                    show_last_history_iters_count = 5000
                elif show_last_history_iters_count == 5000:
                    show_last_history_iters_count = 10000
                elif show_last_history_iters_count == 10000:
                    show_last_history_iters_count = 50000
                elif show_last_history_iters_count == 50000:
                    show_last_history_iters_count = 100000
                elif show_last_history_iters_count == 100000:
                    show_last_history_iters_count = 0
                update_preview = True
            elif key == ord(' '):
                selected_preview = (selected_preview + 1) % len(previews)
                update_preview = True

            try:
                io.process_messages(0.1)
            except KeyboardInterrupt:
                s2c.put ( {'op': 'close'} )

        io.destroy_all_windows()
mainscripts/Trainer.py
oo92/TorchFaceLab
﻿import pickle
from pathlib import Path

import cv2

from DFLIMG import *
from facelib import LandmarksProcessor, FaceType
from core.interact import interact as io
from core import pathex
from core.cv2ex import *


def save_faceset_metadata_folder(input_path):
    input_path = Path(input_path)

    metadata_filepath = input_path / 'meta.dat'

    io.log_info (f"Saving metadata to {str(metadata_filepath)}\r\n")

    d = {}
    for filepath in io.progress_bar_generator( pathex.get_image_paths(input_path), "Processing"):
        filepath = Path(filepath)
        dflimg = DFLIMG.load (filepath)
        if dflimg is None or not dflimg.has_data():
            io.log_info(f"{filepath} is not a dfl image file")
            continue
            
        dfl_dict = dflimg.get_dict()
        d[filepath.name] = ( dflimg.get_shape(), dfl_dict )

    try:
        with open(metadata_filepath, "wb") as f:
            f.write ( pickle.dumps(d) )
    except:
        raise Exception( 'cannot save %s' % (filename) )

    io.log_info("Now you can edit images.")
    io.log_info("!!! Keep same filenames in the folder.")
    io.log_info("You can change size of images, restoring process will downscale back to original size.")
    io.log_info("After that, use restore metadata.")

def restore_faceset_metadata_folder(input_path):
    input_path = Path(input_path)

    metadata_filepath = input_path / 'meta.dat'
    io.log_info (f"Restoring metadata from {str(metadata_filepath)}.\r\n")

    if not metadata_filepath.exists():
        io.log_err(f"Unable to find {str(metadata_filepath)}.")

    try:
        with open(metadata_filepath, "rb") as f:
            d = pickle.loads(f.read())
    except:
        raise FileNotFoundError(filename)

    for filepath in io.progress_bar_generator( pathex.get_image_paths(input_path, image_extensions=['.jpg'], return_Path_class=True), "Processing"):
        saved_data = d.get(filepath.name, None)
        if saved_data is None:
            io.log_info(f"No saved metadata for {filepath}")
            continue
        
        shape, dfl_dict = saved_data

        img = cv2_imread (filepath)
        if img.shape != shape:
            img = cv2.resize (img, (shape[1], shape[0]), interpolation=cv2.INTER_LANCZOS4 )

            cv2_imwrite (str(filepath), img, [int(cv2.IMWRITE_JPEG_QUALITY), 100] )

        if filepath.suffix == '.jpg':
            dflimg = DFLJPG.load(filepath)
            dflimg.set_dict(dfl_dict)
            dflimg.save()
        else:
            continue

    metadata_filepath.unlink()

def add_landmarks_debug_images(input_path):
    io.log_info ("Adding landmarks debug images...")

    for filepath in io.progress_bar_generator( pathex.get_image_paths(input_path), "Processing"):
        filepath = Path(filepath)

        img = cv2_imread(str(filepath))

        dflimg = DFLIMG.load (filepath)

        if dflimg is None or not dflimg.has_data():
            io.log_err (f"{filepath.name} is not a dfl image file")
            continue
        
        if img is not None:
            face_landmarks = dflimg.get_landmarks()
            face_type = FaceType.fromString ( dflimg.get_face_type() )
            
            if face_type == FaceType.MARK_ONLY:
                rect = dflimg.get_source_rect()
                LandmarksProcessor.draw_rect_landmarks(img, rect, face_landmarks, FaceType.FULL )
            else:
                LandmarksProcessor.draw_landmarks(img, face_landmarks, transparent_mask=True )
            
            
            
            output_file = '{}{}'.format( str(Path(str(input_path)) / filepath.stem),  '_debug.jpg')
            cv2_imwrite(output_file, img, [int(cv2.IMWRITE_JPEG_QUALITY), 50] )

def recover_original_aligned_filename(input_path):
    io.log_info ("Recovering original aligned filename...")

    files = []
    for filepath in io.progress_bar_generator( pathex.get_image_paths(input_path), "Processing"):
        filepath = Path(filepath)

        dflimg = DFLIMG.load (filepath)

        if dflimg is None or not dflimg.has_data():
            io.log_err (f"{filepath.name} is not a dfl image file")
            continue

        files += [ [filepath, None, dflimg.get_source_filename(), False] ]

    files_len = len(files)
    for i in io.progress_bar_generator( range(files_len), "Sorting" ):
        fp, _, sf, converted = files[i]

        if converted:
            continue

        sf_stem = Path(sf).stem

        files[i][1] = fp.parent / ( sf_stem + '_0' + fp.suffix )
        files[i][3] = True
        c = 1

        for j in range(i+1, files_len):
            fp_j, _, sf_j, converted_j = files[j]
            if converted_j:
                continue

            if sf_j == sf:
                files[j][1] = fp_j.parent / ( sf_stem + ('_%d' % (c)) + fp_j.suffix )
                files[j][3] = True
                c += 1

    for file in io.progress_bar_generator( files, "Renaming", leave=False ):
        fs, _, _, _ = file
        dst = fs.parent / ( fs.stem + '_tmp' + fs.suffix )
        try:
            fs.rename (dst)
        except:
            io.log_err ('fail to rename %s' % (fs.name) )

    for file in io.progress_bar_generator( files, "Renaming" ):
        fs, fd, _, _ = file
        fs = fs.parent / ( fs.stem + '_tmp' + fs.suffix )
        try:
            fs.rename (fd)
        except:
            io.log_err ('fail to rename %s' % (fs.name) )

def export_faceset_mask(input_dir):
    for filename in io.progress_bar_generator(pathex.get_image_paths (input_dir), "Processing"):
        filepath = Path(filename)
        
        if '_mask' in filepath.stem:
            continue
        
        mask_filepath = filepath.parent / (filepath.stem+'_mask'+filepath.suffix)
        
        dflimg = DFLJPG.load(filepath)
        
        H,W,C = dflimg.shape
        
        seg_ie_polys = dflimg.get_seg_ie_polys()
        
        if dflimg.has_xseg_mask():
            mask = dflimg.get_xseg_mask()
            mask[mask < 0.5] = 0.0
            mask[mask >= 0.5] = 1.0
        elif seg_ie_polys.has_polys():
            mask = np.zeros ((H,W,1), dtype=np.float32)
            seg_ie_polys.overlay_mask(mask)
        else:
            raise Exception(f'no mask in file {filepath}')
        
        
        cv2_imwrite(mask_filepath, (mask*255).astype(np.uint8), [int(cv2.IMWRITE_JPEG_QUALITY), 100] )
                
mainscripts/Util.py
oo92/TorchFaceLab
import subprocess
import numpy as np
import ffmpeg
from pathlib import Path
from core import pathex
from core.interact import interact as io

def extract_video(input_file, output_dir, output_ext=None, fps=None):
    input_file_path = Path(input_file)
    output_path = Path(output_dir)

    if not output_path.exists():
        output_path.mkdir(exist_ok=True)


    if input_file_path.suffix == '.*':
        input_file_path = pathex.get_first_file_by_stem (input_file_path.parent, input_file_path.stem)
    else:
        if not input_file_path.exists():
            input_file_path = None

    if input_file_path is None:
        io.log_err("input_file not found.")
        return

    if fps is None:
        fps = io.input_int ("Enter FPS", 0, help_message="How many frames of every second of the video will be extracted. 0 - full fps")

    if output_ext is None:
        output_ext = io.input_str ("Output image format", "png", ["png","jpg"], help_message="png is lossless, but extraction is x10 slower for HDD, requires x10 more disk space than jpg.")

    for filename in pathex.get_image_paths (output_path, ['.'+output_ext]):
        Path(filename).unlink()

    job = ffmpeg.input(str(input_file_path))

    kwargs = {'pix_fmt': 'rgb24'}
    if fps != 0:
        kwargs.update ({'r':str(fps)})

    if output_ext == 'jpg':
        kwargs.update ({'q:v':'2'}) #highest quality for jpg

    job = job.output( str (output_path / ('%5d.'+output_ext)), **kwargs )

    try:
        job = job.run()
    except:
        io.log_err ("ffmpeg fail, job commandline:" + str(job.compile()) )

def cut_video ( input_file, from_time=None, to_time=None, audio_track_id=None, bitrate=None):
    input_file_path = Path(input_file)
    if input_file_path is None:
        io.log_err("input_file not found.")
        return

    output_file_path = input_file_path.parent / (input_file_path.stem + "_cut" + input_file_path.suffix)

    if from_time is None:
        from_time = io.input_str ("From time", "00:00:00.000")

    if to_time is None:
        to_time = io.input_str ("To time", "00:00:00.000")

    if audio_track_id is None:
        audio_track_id = io.input_int ("Specify audio track id.", 0)

    if bitrate is None:
        bitrate = max (1, io.input_int ("Bitrate of output file in MB/s", 25) )

    kwargs = {"c:v": "libx264",
              "b:v": "%dM" %(bitrate),
              "pix_fmt": "yuv420p",
             }

    job = ffmpeg.input(str(input_file_path), ss=from_time, to=to_time)

    job_v = job['v:0']
    job_a = job['a:' + str(audio_track_id) + '?' ]

    job = ffmpeg.output(job_v, job_a, str(output_file_path), **kwargs).overwrite_output()

    try:
        job = job.run()
    except:
        io.log_err ("ffmpeg fail, job commandline:" + str(job.compile()) )

def denoise_image_sequence( input_dir, ext=None, factor=None ):
    input_path = Path(input_dir)

    if not input_path.exists():
        io.log_err("input_dir not found.")
        return

    image_paths = [ Path(filepath) for filepath in pathex.get_image_paths(input_path) ]

    # Check extension of all images
    image_paths_suffix = None
    for filepath in image_paths:
        if image_paths_suffix is None:
            image_paths_suffix = filepath.suffix
        else:
            if filepath.suffix != image_paths_suffix:
                io.log_err(f"All images in {input_path.name} should be with the same extension.")
                return

    if factor is None:
        factor = np.clip ( io.input_int ("Denoise factor?", 7, add_info="1-20"), 1, 20 )

    # Rename to temporary filenames
    for i,filepath in io.progress_bar_generator( enumerate(image_paths), "Renaming", leave=False):
        src = filepath
        dst = filepath.parent / ( f'{i+1:06}_{filepath.name}' )
        try:
            src.rename (dst)
        except:
            io.log_error ('fail to rename %s' % (src.name) )
            return

    # Rename to sequental filenames
    for i,filepath in io.progress_bar_generator( enumerate(image_paths), "Renaming", leave=False):

        src = filepath.parent / ( f'{i+1:06}_{filepath.name}' )
        dst = filepath.parent / ( f'{i+1:06}{filepath.suffix}' )
        try:
            src.rename (dst)
        except:
            io.log_error ('fail to rename %s' % (src.name) )
            return

    # Process image sequence in ffmpeg
    kwargs = {}
    if image_paths_suffix == '.jpg':
        kwargs.update ({'q:v':'2'})

    job = ( ffmpeg
            .input(str ( input_path / ('%6d'+image_paths_suffix) ) )
            .filter("hqdn3d", factor, factor, 5,5)
            .output(str ( input_path / ('%6d'+image_paths_suffix) ), **kwargs )
           )

    try:
        job = job.run()
    except:
        io.log_err ("ffmpeg fail, job commandline:" + str(job.compile()) )

    # Rename to temporary filenames
    for i,filepath in io.progress_bar_generator( enumerate(image_paths), "Renaming", leave=False):
        src = filepath.parent / ( f'{i+1:06}{filepath.suffix}' )
        dst = filepath.parent / ( f'{i+1:06}_{filepath.name}' )
        try:
            src.rename (dst)
        except:
            io.log_error ('fail to rename %s' % (src.name) )
            return

    # Rename to initial filenames
    for i,filepath in io.progress_bar_generator( enumerate(image_paths), "Renaming", leave=False):
        src = filepath.parent / ( f'{i+1:06}_{filepath.name}' )
        dst = filepath

        try:
            src.rename (dst)
        except:
            io.log_error ('fail to rename %s' % (src.name) )
            return

def video_from_sequence( input_dir, output_file, reference_file=None, ext=None, fps=None, bitrate=None, include_audio=False, lossless=None ):
    input_path = Path(input_dir)
    output_file_path = Path(output_file)
    reference_file_path = Path(reference_file) if reference_file is not None else None

    if not input_path.exists():
        io.log_err("input_dir not found.")
        return

    if not output_file_path.parent.exists():
        output_file_path.parent.mkdir(parents=True, exist_ok=True)
        return

    out_ext = output_file_path.suffix

    if ext is None:
        ext = io.input_str ("Input image format (extension)", "png")

    if lossless is None:
        lossless = io.input_bool ("Use lossless codec", False)

    video_id = None
    audio_id = None
    ref_in_a = None
    if reference_file_path is not None:
        if reference_file_path.suffix == '.*':
            reference_file_path = pathex.get_first_file_by_stem (reference_file_path.parent, reference_file_path.stem)
        else:
            if not reference_file_path.exists():
                reference_file_path = None

        if reference_file_path is None:
            io.log_err("reference_file not found.")
            return

        #probing reference file
        probe = ffmpeg.probe (str(reference_file_path))

        #getting first video and audio streams id with fps
        for stream in probe['streams']:
            if video_id is None and stream['codec_type'] == 'video':
                video_id = stream['index']
                fps = stream['r_frame_rate']

            if audio_id is None and stream['codec_type'] == 'audio':
                audio_id = stream['index']

        if audio_id is not None:
            #has audio track
            ref_in_a = ffmpeg.input (str(reference_file_path))[str(audio_id)]

    if fps is None:
        #if fps not specified and not overwritten by reference-file
        fps = max (1, io.input_int ("Enter FPS", 25) )

    if not lossless and bitrate is None:
        bitrate = max (1, io.input_int ("Bitrate of output file in MB/s", 16) )

    input_image_paths = pathex.get_image_paths(input_path)

    i_in = ffmpeg.input('pipe:', format='image2pipe', r=fps)

    output_args = [i_in]

    if include_audio and ref_in_a is not None:
        output_args += [ref_in_a]

    output_args += [str (output_file_path)]

    output_kwargs = {}

    if lossless:
        output_kwargs.update ({"c:v": "libx264",
                               "crf": "0",
                               "pix_fmt": "yuv420p",
                              })
    else:
        output_kwargs.update ({"c:v": "libx264",
                               "b:v": "%dM" %(bitrate),
                               "pix_fmt": "yuv420p",
                              })

    if include_audio and ref_in_a is not None:
        output_kwargs.update ({"c:a": "aac",
                               "b:a": "192k",
                               "ar" : "48000",
                               "strict": "experimental"
                               })

    job = ( ffmpeg.output(*output_args, **output_kwargs).overwrite_output() )

    try:
        job_run = job.run_async(pipe_stdin=True)

        for image_path in input_image_paths:
            with open (image_path, "rb") as f:
                image_bytes = f.read()
                job_run.stdin.write (image_bytes)

        job_run.stdin.close()
        job_run.wait()
    except:
        io.log_err ("ffmpeg fail, job commandline:" + str(job.compile()) )

mainscripts/VideoEd.py
oo92/TorchFaceLab
import json
import shutil
import traceback
from pathlib import Path

import numpy as np

from core import pathex
from core.cv2ex import *
from core.interact import interact as io
from core.leras import nn
from DFLIMG import *
from facelib import XSegNet, LandmarksProcessor, FaceType
import pickle

def apply_xseg(input_path, model_path):
    if not input_path.exists():
        raise ValueError(f'{input_path} not found. Please ensure it exists.')

    if not model_path.exists():
        raise ValueError(f'{model_path} not found. Please ensure it exists.')
        
    face_type = None
    
    model_dat = model_path / 'XSeg_data.dat'
    if model_dat.exists():
        dat = pickle.loads( model_dat.read_bytes() )
        dat_options = dat.get('options', None)
        if dat_options is not None:
            face_type = dat_options.get('face_type', None)
        
        
        
    if face_type is None:
        face_type = io.input_str ("XSeg model face type", 'same', ['h','mf','f','wf','head','same'], help_message="Specify face type of trained XSeg model. For example if XSeg model trained as WF, but faceset is HEAD, specify WF to apply xseg only on WF part of HEAD. Default is 'same'").lower()
        if face_type == 'same':
            face_type = None
    
    if face_type is not None:
        face_type = {'h'  : FaceType.HALF,
                     'mf' : FaceType.MID_FULL,
                     'f'  : FaceType.FULL,
                     'wf' : FaceType.WHOLE_FACE,
                     'head' : FaceType.HEAD}[face_type]
                     
    io.log_info(f'Applying trained XSeg model to {input_path.name}/ folder.')

    device_config = nn.DeviceConfig.ask_choose_device(choose_only_one=True)
    nn.initialize(device_config)
        
    
    
    xseg = XSegNet(name='XSeg', 
                    load_weights=True,
                    weights_file_root=model_path,
                    data_format=nn.data_format,
                    raise_on_no_model_files=True)
    xseg_res = xseg.get_resolution()
              
    images_paths = pathex.get_image_paths(input_path, return_Path_class=True)
    
    for filepath in io.progress_bar_generator(images_paths, "Processing"):
        dflimg = DFLIMG.load(filepath)
        if dflimg is None or not dflimg.has_data():
            io.log_info(f'{filepath} is not a DFLIMG')
            continue
        
        img = cv2_imread(filepath).astype(np.float32) / 255.0
        h,w,c = img.shape
        
        img_face_type = FaceType.fromString( dflimg.get_face_type() )
        if face_type is not None and img_face_type != face_type:
            lmrks = dflimg.get_source_landmarks()
            
            fmat = LandmarksProcessor.get_transform_mat(lmrks, w, face_type)
            imat = LandmarksProcessor.get_transform_mat(lmrks, w, img_face_type)
            
            g_p = LandmarksProcessor.transform_points (np.float32([(0,0),(w,0),(0,w) ]), fmat, True)
            g_p2 = LandmarksProcessor.transform_points (g_p, imat)
            
            mat = cv2.getAffineTransform( g_p2, np.float32([(0,0),(w,0),(0,w) ]) )
            
            img = cv2.warpAffine(img, mat, (w, w), cv2.INTER_LANCZOS4)
            img = cv2.resize(img, (xseg_res, xseg_res), interpolation=cv2.INTER_LANCZOS4)
        else:
            if w != xseg_res:
                img = cv2.resize( img, (xseg_res,xseg_res), interpolation=cv2.INTER_LANCZOS4 )    
                    
        if len(img.shape) == 2:
            img = img[...,None]            
    
        mask = xseg.extract(img)
        
        if face_type is not None and img_face_type != face_type:
            mask = cv2.resize(mask, (w, w), interpolation=cv2.INTER_LANCZOS4)
            mask = cv2.warpAffine( mask, mat, (w,w), np.zeros( (h,w,c), dtype=np.float), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4)
            mask = cv2.resize(mask, (xseg_res, xseg_res), interpolation=cv2.INTER_LANCZOS4)
        mask[mask < 0.5]=0
        mask[mask >= 0.5]=1    
        dflimg.set_xseg_mask(mask)
        dflimg.save()


        
def fetch_xseg(input_path):
    if not input_path.exists():
        raise ValueError(f'{input_path} not found. Please ensure it exists.')
    
    output_path = input_path.parent / (input_path.name + '_xseg')
    output_path.mkdir(exist_ok=True, parents=True)
    
    io.log_info(f'Copying faces containing XSeg polygons to {output_path.name}/ folder.')
    
    images_paths = pathex.get_image_paths(input_path, return_Path_class=True)
    
    
    files_copied = []
    for filepath in io.progress_bar_generator(images_paths, "Processing"):
        dflimg = DFLIMG.load(filepath)
        if dflimg is None or not dflimg.has_data():
            io.log_info(f'{filepath} is not a DFLIMG')
            continue
        
        ie_polys = dflimg.get_seg_ie_polys()

        if ie_polys.has_polys():
            files_copied.append(filepath)
            shutil.copy ( str(filepath), str(output_path / filepath.name) )
    
    io.log_info(f'Files copied: {len(files_copied)}')
    
    is_delete = io.input_bool (f"\r\nDelete original files?", True)
    if is_delete:
        for filepath in files_copied:
            Path(filepath).unlink()
            
    
def remove_xseg(input_path):
    if not input_path.exists():
        raise ValueError(f'{input_path} not found. Please ensure it exists.')
    
    io.log_info(f'Processing folder {input_path}')
    io.log_info('!!! WARNING : APPLIED XSEG MASKS WILL BE REMOVED FROM THE FRAMES !!!')
    io.log_info('!!! WARNING : APPLIED XSEG MASKS WILL BE REMOVED FROM THE FRAMES !!!')
    io.log_info('!!! WARNING : APPLIED XSEG MASKS WILL BE REMOVED FROM THE FRAMES !!!')
    io.input_str('Press enter to continue.')
                               
    images_paths = pathex.get_image_paths(input_path, return_Path_class=True)
    
    files_processed = 0
    for filepath in io.progress_bar_generator(images_paths, "Processing"):
        dflimg = DFLIMG.load(filepath)
        if dflimg is None or not dflimg.has_data():
            io.log_info(f'{filepath} is not a DFLIMG')
            continue
        
        if dflimg.has_xseg_mask():
            dflimg.set_xseg_mask(None)
            dflimg.save()
            files_processed += 1
    io.log_info(f'Files processed: {files_processed}')
    
def remove_xseg_labels(input_path):
    if not input_path.exists():
        raise ValueError(f'{input_path} not found. Please ensure it exists.')
    
    io.log_info(f'Processing folder {input_path}')
    io.log_info('!!! WARNING : LABELED XSEG POLYGONS WILL BE REMOVED FROM THE FRAMES !!!')
    io.log_info('!!! WARNING : LABELED XSEG POLYGONS WILL BE REMOVED FROM THE FRAMES !!!')
    io.log_info('!!! WARNING : LABELED XSEG POLYGONS WILL BE REMOVED FROM THE FRAMES !!!')
    io.input_str('Press enter to continue.')
    
    images_paths = pathex.get_image_paths(input_path, return_Path_class=True)
    
    files_processed = 0
    for filepath in io.progress_bar_generator(images_paths, "Processing"):
        dflimg = DFLIMG.load(filepath)
        if dflimg is None or not dflimg.has_data():
            io.log_info(f'{filepath} is not a DFLIMG')
            continue

        if dflimg.has_seg_ie_polys():
            dflimg.set_seg_ie_polys(None)
            dflimg.save()            
            files_processed += 1
            
    io.log_info(f'Files processed: {files_processed}')
mainscripts/XSegUtil.py
oo92/TorchFaceLab
import traceback
import json
import multiprocessing
import shutil
from pathlib import Path
import cv2
import numpy as np

from core import imagelib, pathex
from core.cv2ex import *
from core.interact import interact as io
from core.joblib import Subprocessor
from core.leras import nn
from DFLIMG import *
from facelib import FaceType, LandmarksProcessor
from . import Extractor, Sorter
from .Extractor import ExtractSubprocessor


def extract_vggface2_dataset(input_dir, device_args={} ):
    multi_gpu = device_args.get('multi_gpu', False)
    cpu_only = device_args.get('cpu_only', False)

    input_path = Path(input_dir)
    if not input_path.exists():
        raise ValueError('Input directory not found. Please ensure it exists.')

    bb_csv = input_path / 'loose_bb_train.csv'
    if not bb_csv.exists():
        raise ValueError('loose_bb_train.csv found. Please ensure it exists.')

    bb_lines = bb_csv.read_text().split('\n')
    bb_lines.pop(0)

    bb_dict = {}
    for line in bb_lines:
        name, l, t, w, h = line.split(',')
        name = name[1:-1]
        l, t, w, h = [ int(x) for x in (l, t, w, h) ]
        bb_dict[name] = (l,t,w, h)


    output_path = input_path.parent / (input_path.name + '_out')

    dir_names = pathex.get_all_dir_names(input_path)

    if not output_path.exists():
        output_path.mkdir(parents=True, exist_ok=True)

    data = []
    for dir_name in io.progress_bar_generator(dir_names, "Collecting"):
        cur_input_path = input_path / dir_name
        cur_output_path = output_path / dir_name

        if not cur_output_path.exists():
            cur_output_path.mkdir(parents=True, exist_ok=True)

        input_path_image_paths = pathex.get_image_paths(cur_input_path)

        for filename in input_path_image_paths:
            filename_path = Path(filename)

            name = filename_path.parent.name + '/' + filename_path.stem
            if name not in bb_dict:
                continue

            l,t,w,h = bb_dict[name]
            if min(w,h) < 128:
                continue

            data += [ ExtractSubprocessor.Data(filename=filename,rects=[ (l,t,l+w,t+h) ], landmarks_accurate=False, force_output_path=cur_output_path ) ]

    face_type = FaceType.fromString('full_face')

    io.log_info ('Performing 2nd pass...')
    data = ExtractSubprocessor (data, 'landmarks', 256, face_type, debug_dir=None, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False).run()

    io.log_info ('Performing 3rd pass...')
    ExtractSubprocessor (data, 'final', 256, face_type, debug_dir=None, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False, final_output_path=None).run()


"""
    import code
    code.interact(local=dict(globals(), **locals()))

    data_len = len(data)
    i = 0
    while i < data_len-1:
        i_name = Path(data[i].filename).parent.name

        sub_data = []

        for j in range (i, data_len):
            j_name = Path(data[j].filename).parent.name
            if i_name == j_name:
                sub_data += [ data[j] ]
            else:
                break
        i = j

        cur_output_path = output_path / i_name

        io.log_info (f"Processing: {str(cur_output_path)}, {i}/{data_len} ")

        if not cur_output_path.exists():
            cur_output_path.mkdir(parents=True, exist_ok=True)








    for dir_name in dir_names:

        cur_input_path = input_path / dir_name
        cur_output_path = output_path / dir_name

        input_path_image_paths = pathex.get_image_paths(cur_input_path)
        l = len(input_path_image_paths)
        #if l < 250 or l > 350:
        #    continue

        io.log_info (f"Processing: {str(cur_input_path)} ")

        if not cur_output_path.exists():
            cur_output_path.mkdir(parents=True, exist_ok=True)


        data = []
        for filename in input_path_image_paths:
            filename_path = Path(filename)

            name = filename_path.parent.name + '/' + filename_path.stem
            if name not in bb_dict:
                continue

            bb = bb_dict[name]
            l,t,w,h = bb
            if min(w,h) < 128:
                continue

            data += [ ExtractSubprocessor.Data(filename=filename,rects=[ (l,t,l+w,t+h) ], landmarks_accurate=False ) ]



        io.log_info ('Performing 2nd pass...')
        data = ExtractSubprocessor (data, 'landmarks', 256, face_type, debug_dir=None, multi_gpu=False, cpu_only=False, manual=False).run()

        io.log_info ('Performing 3rd pass...')
        data = ExtractSubprocessor (data, 'final', 256, face_type, debug_dir=None, multi_gpu=False, cpu_only=False, manual=False, final_output_path=cur_output_path).run()


        io.log_info (f"Sorting: {str(cur_output_path)} ")
        Sorter.main (input_path=str(cur_output_path), sort_by_method='hist')

        import code
        code.interact(local=dict(globals(), **locals()))

        #try:
        #    io.log_info (f"Removing: {str(cur_input_path)} ")
        #    shutil.rmtree(cur_input_path)
        #except:
        #    io.log_info (f"unable to remove: {str(cur_input_path)} ")




def extract_vggface2_dataset(input_dir, device_args={} ):
    multi_gpu = device_args.get('multi_gpu', False)
    cpu_only = device_args.get('cpu_only', False)

    input_path = Path(input_dir)
    if not input_path.exists():
        raise ValueError('Input directory not found. Please ensure it exists.')

    output_path = input_path.parent / (input_path.name + '_out')

    dir_names = pathex.get_all_dir_names(input_path)

    if not output_path.exists():
        output_path.mkdir(parents=True, exist_ok=True)



    for dir_name in dir_names:

        cur_input_path = input_path / dir_name
        cur_output_path = output_path / dir_name

        l = len(pathex.get_image_paths(cur_input_path))
        if l < 250 or l > 350:
            continue

        io.log_info (f"Processing: {str(cur_input_path)} ")

        if not cur_output_path.exists():
            cur_output_path.mkdir(parents=True, exist_ok=True)

        Extractor.main( str(cur_input_path),
              str(cur_output_path),
              detector='s3fd',
              image_size=256,
              face_type='full_face',
              max_faces_from_image=1,
              device_args=device_args )

        io.log_info (f"Sorting: {str(cur_input_path)} ")
        Sorter.main (input_path=str(cur_output_path), sort_by_method='hist')

        try:
            io.log_info (f"Removing: {str(cur_input_path)} ")
            shutil.rmtree(cur_input_path)
        except:
            io.log_info (f"unable to remove: {str(cur_input_path)} ")

"""

#unused in end user workflow
def dev_test_68(input_dir ):
    # process 68 landmarks dataset with .pts files
    input_path = Path(input_dir)
    if not input_path.exists():
        raise ValueError('input_dir not found. Please ensure it exists.')

    output_path = input_path.parent / (input_path.name+'_aligned')

    io.log_info(f'Output dir is % {output_path}')

    if output_path.exists():
        output_images_paths = pathex.get_image_paths(output_path)
        if len(output_images_paths) > 0:
            io.input_bool("WARNING !!! \n %s contains files! \n They will be deleted. \n Press enter to continue." % (str(output_path)), False )
            for filename in output_images_paths:
                Path(filename).unlink()
    else:
        output_path.mkdir(parents=True, exist_ok=True)

    images_paths = pathex.get_image_paths(input_path)

    for filepath in io.progress_bar_generator(images_paths, "Processing"):
        filepath = Path(filepath)


        pts_filepath = filepath.parent / (filepath.stem+'.pts')
        if pts_filepath.exists():
            pts = pts_filepath.read_text()
            pts_lines = pts.split('\n')

            lmrk_lines = None
            for pts_line in pts_lines:
                if pts_line == '{':
                    lmrk_lines = []
                elif pts_line == '}':
                    break
                else:
                    if lmrk_lines is not None:
                        lmrk_lines.append (pts_line)

            if lmrk_lines is not None and len(lmrk_lines) == 68:
                try:
                    lmrks = [ np.array ( lmrk_line.strip().split(' ') ).astype(np.float32).tolist() for lmrk_line in lmrk_lines]
                except Exception as e:
                    print(e)
                    print(filepath)
                    continue

                rect = LandmarksProcessor.get_rect_from_landmarks(lmrks)

                output_filepath = output_path / (filepath.stem+'.jpg')

                img = cv2_imread(filepath)
                img = imagelib.normalize_channels(img, 3)
                cv2_imwrite(output_filepath, img, [int(cv2.IMWRITE_JPEG_QUALITY), 95] )
                
                raise Exception("unimplemented")
                #DFLJPG.x(output_filepath, face_type=FaceType.toString(FaceType.MARK_ONLY),
                #                                landmarks=lmrks,
                #                                source_filename=filepath.name,
                #                                source_rect=rect,
                #                                source_landmarks=lmrks
                #                    )

    io.log_info("Done.")

#unused in end user workflow
def extract_umd_csv(input_file_csv,
                    face_type='full_face',
                    device_args={} ):

    #extract faces from umdfaces.io dataset csv file with pitch,yaw,roll info.
    multi_gpu = device_args.get('multi_gpu', False)
    cpu_only = device_args.get('cpu_only', False)
    face_type = FaceType.fromString(face_type)

    input_file_csv_path = Path(input_file_csv)
    if not input_file_csv_path.exists():
        raise ValueError('input_file_csv not found. Please ensure it exists.')

    input_file_csv_root_path = input_file_csv_path.parent
    output_path = input_file_csv_path.parent / ('aligned_' + input_file_csv_path.name)

    io.log_info("Output dir is %s." % (str(output_path)) )

    if output_path.exists():
        output_images_paths = pathex.get_image_paths(output_path)
        if len(output_images_paths) > 0:
            io.input_bool("WARNING !!! \n %s contains files! \n They will be deleted. \n Press enter to continue." % (str(output_path)), False )
            for filename in output_images_paths:
                Path(filename).unlink()
    else:
        output_path.mkdir(parents=True, exist_ok=True)

    try:
        with open( str(input_file_csv_path), 'r') as f:
            csv_file = f.read()
    except Exception as e:
        io.log_err("Unable to open or read file " + str(input_file_csv_path) + ": " + str(e) )
        return

    strings = csv_file.split('\n')
    keys = strings[0].split(',')
    keys_len = len(keys)
    csv_data = []
    for i in range(1, len(strings)):
        values = strings[i].split(',')
        if keys_len != len(values):
            io.log_err("Wrong string in csv file, skipping.")
            continue

        csv_data += [ { keys[n] : values[n] for n in range(keys_len) } ]

    data = []
    for d in csv_data:
        filename = input_file_csv_root_path / d['FILE']


        x,y,w,h = float(d['FACE_X']), float(d['FACE_Y']), float(d['FACE_WIDTH']), float(d['FACE_HEIGHT'])

        data += [ ExtractSubprocessor.Data(filename=filename, rects=[ [x,y,x+w,y+h] ]) ]

    images_found = len(data)
    faces_detected = 0
    if len(data) > 0:
        io.log_info ("Performing 2nd pass from csv file...")
        data = ExtractSubprocessor (data, 'landmarks', multi_gpu=multi_gpu, cpu_only=cpu_only).run()

        io.log_info ('Performing 3rd pass...')
        data = ExtractSubprocessor (data, 'final', face_type, None, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False, final_output_path=output_path).run()
        faces_detected += sum([d.faces_detected for d in data])


    io.log_info ('-------------------------')
    io.log_info ('Images found:        %d' % (images_found) )
    io.log_info ('Faces detected:      %d' % (faces_detected) )
    io.log_info ('-------------------------')


    
def dev_test1(input_dir):
    # LaPa dataset
    
    image_size = 1024
    face_type = FaceType.HEAD
    
    input_path = Path(input_dir)
    images_path = input_path / 'images'    
    if not images_path.exists:
        raise ValueError('LaPa dataset: images folder not found.')
    labels_path = input_path / 'labels'    
    if not labels_path.exists:
        raise ValueError('LaPa dataset: labels folder not found.')
    landmarks_path = input_path / 'landmarks'    
    if not landmarks_path.exists:
        raise ValueError('LaPa dataset: landmarks folder not found.')
    
    output_path = input_path / 'out'    
    if output_path.exists():
        output_images_paths = pathex.get_image_paths(output_path)
        if len(output_images_paths) != 0:
            io.input(f"\n WARNING !!! \n {output_path} contains files! \n They will be deleted. \n Press enter to continue.\n")
            for filename in output_images_paths:
                Path(filename).unlink()
    output_path.mkdir(parents=True, exist_ok=True)
    
    data = []
    
    img_paths = pathex.get_image_paths (images_path)
    for filename in img_paths:
        filepath = Path(filename)

        landmark_filepath = landmarks_path / (filepath.stem + '.txt')
        if not landmark_filepath.exists():
            raise ValueError(f'no landmarks for {filepath}')
        
        #img = cv2_imread(filepath)
        
        lm = landmark_filepath.read_text()
        lm = lm.split('\n')
        if int(lm[0]) != 106:
            raise ValueError(f'wrong landmarks format in {landmark_filepath}')
        
        lmrks = []
        for i in range(106):
            x,y = lm[i+1].split(' ')
            x,y = float(x), float(y)
            lmrks.append ( (x,y) )
            
        lmrks = np.array(lmrks)
        
        l,t = np.min(lmrks, 0)
        r,b = np.max(lmrks, 0)
        
        l,t,r,b = ( int(x) for x in (l,t,r,b) )
        
        #for x, y in lmrks:
        #    x,y = int(x), int(y)
        #    cv2.circle(img, (x, y), 1, (0,255,0) , 1, lineType=cv2.LINE_AA)   
         
        #imagelib.draw_rect(img, (l,t,r,b), (0,255,0) )
         
        
        data += [ ExtractSubprocessor.Data(filepath=filepath, rects=[ (l,t,r,b) ]) ]

        #cv2.imshow("", img) 
        #cv2.waitKey(0)
 
    if len(data) > 0:
        device_config = nn.DeviceConfig.BestGPU()
        
        io.log_info ("Performing 2nd pass...")
        data = ExtractSubprocessor (data, 'landmarks', image_size, 95, face_type,  device_config=device_config).run()
        io.log_info ("Performing 3rd pass...")
        data = ExtractSubprocessor (data, 'final', image_size, 95, face_type, final_output_path=output_path, device_config=device_config).run()


        for filename in pathex.get_image_paths (output_path):
            filepath = Path(filename)
            
            
            dflimg = DFLJPG.load(filepath)
            
            src_filename = dflimg.get_source_filename()
            image_to_face_mat = dflimg.get_image_to_face_mat()

            label_filepath = labels_path / ( Path(src_filename).stem + '.png')        
            if not label_filepath.exists():
                raise ValueError(f'{label_filepath} does not exist')
            
            mask = cv2_imread(label_filepath)        
            #mask[mask == 10] = 0 # remove hair
            mask[mask > 0] = 1
            mask = cv2.warpAffine(mask, image_to_face_mat, (image_size, image_size), cv2.INTER_LINEAR)
            mask = cv2.blur(mask, (3,3) )
            
            #cv2.imshow("", (mask*255).astype(np.uint8) ) 
            #cv2.waitKey(0)
            
            dflimg.set_xseg_mask(mask)
            dflimg.save()
        
    
    import code
    code.interact(local=dict(globals(), **locals()))
                    

def dev_resave_pngs(input_dir):
    input_path = Path(input_dir)
    if not input_path.exists():
        raise ValueError('input_dir not found. Please ensure it exists.')

    images_paths = pathex.get_image_paths(input_path, image_extensions=['.png'], subdirs=True, return_Path_class=True)

    for filepath in io.progress_bar_generator(images_paths,"Processing"):
        cv2_imwrite(filepath, cv2_imread(filepath))


def dev_segmented_trash(input_dir):
    input_path = Path(input_dir)
    if not input_path.exists():
        raise ValueError('input_dir not found. Please ensure it exists.')

    output_path = input_path.parent / (input_path.name+'_trash')
    output_path.mkdir(parents=True, exist_ok=True)

    images_paths = pathex.get_image_paths(input_path, return_Path_class=True)

    trash_paths = []
    for filepath in images_paths:
        json_file = filepath.parent / (filepath.stem +'.json')
        if not json_file.exists():
            trash_paths.append(filepath)

    for filepath in trash_paths:

        try:
            filepath.rename ( output_path / filepath.name )
        except:
            io.log_info ('fail to trashing %s' % (src.name) )



def dev_test(input_dir):
    """
    extract FaceSynthetics dataset https://github.com/microsoft/FaceSynthetics
    
    BACKGROUND = 0
    SKIN = 1
    NOSE = 2
    RIGHT_EYE = 3
    LEFT_EYE = 4
    RIGHT_BROW = 5
    LEFT_BROW = 6
    RIGHT_EAR = 7
    LEFT_EAR = 8
    MOUTH_INTERIOR = 9
    TOP_LIP = 10
    BOTTOM_LIP = 11
    NECK = 12
    HAIR = 13
    BEARD = 14
    CLOTHING = 15
    GLASSES = 16
    HEADWEAR = 17
    FACEWEAR = 18
    IGNORE = 255
    """
    
    
    image_size = 1024
    face_type = FaceType.WHOLE_FACE
    
    input_path = Path(input_dir)
    
    
    
    output_path = input_path.parent / f'{input_path.name}_out'    
    if output_path.exists():
        output_images_paths = pathex.get_image_paths(output_path)
        if len(output_images_paths) != 0:
            io.input(f"\n WARNING !!! \n {output_path} contains files! \n They will be deleted. \n Press enter to continue.\n")
            for filename in output_images_paths:
                Path(filename).unlink()
    output_path.mkdir(parents=True, exist_ok=True)
    
    data = []
    
    for filepath in io.progress_bar_generator(pathex.get_paths(input_path), "Processing"):
        if filepath.suffix == '.txt':
            
            image_filepath = filepath.parent / f'{filepath.name.split("_")[0]}.png'
            if not image_filepath.exists():
                print(f'{image_filepath} does not exist, skipping') 
                
            lmrks = []
            for lmrk_line in filepath.read_text().split('\n'):
                if len(lmrk_line) == 0:
                    continue
                    
                x, y = lmrk_line.split(' ')
                x, y = float(x), float(y)
                
                lmrks.append( (x,y) )
                
            lmrks = np.array(lmrks[:68], np.float32)
            rect = LandmarksProcessor.get_rect_from_landmarks(lmrks)
            data += [ ExtractSubprocessor.Data(filepath=image_filepath, rects=[rect], landmarks=[ lmrks ] ) ]

    if len(data) > 0:
        io.log_info ("Performing 3rd pass...")
        data = ExtractSubprocessor (data, 'final', image_size, 95, face_type, final_output_path=output_path, device_config=nn.DeviceConfig.CPU()).run()

        for filename in io.progress_bar_generator(pathex.get_image_paths (output_path), "Processing"):
            filepath = Path(filename)
            
            dflimg = DFLJPG.load(filepath)
            
            src_filename = dflimg.get_source_filename()
            image_to_face_mat = dflimg.get_image_to_face_mat()
            
            seg_filepath = input_path / ( Path(src_filename).stem + '_seg.png')        
            if not seg_filepath.exists():
                raise ValueError(f'{seg_filepath} does not exist')
            
            seg = cv2_imread(seg_filepath)     
            seg_inds = np.isin(seg, [1,2,3,4,5,6,9,10,11]) 
            seg[~seg_inds] = 0
            seg[seg_inds] = 1
            seg = seg.astype(np.float32)
            seg = cv2.warpAffine(seg, image_to_face_mat, (image_size, image_size), cv2.INTER_LANCZOS4)
            dflimg.set_xseg_mask(seg)
            dflimg.save()
            
mainscripts/dev_misc.py
oo92/TorchFaceLab
from pathlib import Path

class FrameInfo(object):
    def __init__(self, filepath=None, landmarks_list=None):
        self.filepath = filepath
        self.landmarks_list = landmarks_list or []
        self.motion_deg = 0
        self.motion_power = 0
merger/FrameInfo.py
oo92/TorchFaceLab
import multiprocessing
import os
import pickle
import sys
import traceback
from pathlib import Path

import numpy as np

from core import imagelib, pathex
from core.cv2ex import *
from core.interact import interact as io
from core.joblib import Subprocessor
from merger import MergeFaceAvatar, MergeMasked, MergerConfig

from .MergerScreen import Screen, ScreenManager

MERGER_DEBUG = False
class InteractiveMergerSubprocessor(Subprocessor):

    class Frame(object):
        def __init__(self, prev_temporal_frame_infos=None,
                           frame_info=None,
                           next_temporal_frame_infos=None):
            self.prev_temporal_frame_infos = prev_temporal_frame_infos
            self.frame_info = frame_info
            self.next_temporal_frame_infos = next_temporal_frame_infos
            self.output_filepath = None
            self.output_mask_filepath = None

            self.idx = None
            self.cfg = None
            self.is_done = False
            self.is_processing = False
            self.is_shown = False
            self.image = None

    class ProcessingFrame(object):
        def __init__(self, idx=None,
                           cfg=None,
                           prev_temporal_frame_infos=None,
                           frame_info=None,
                           next_temporal_frame_infos=None,
                           output_filepath=None,
                           output_mask_filepath=None,
                           need_return_image = False):
            self.idx = idx
            self.cfg = cfg
            self.prev_temporal_frame_infos = prev_temporal_frame_infos
            self.frame_info = frame_info
            self.next_temporal_frame_infos = next_temporal_frame_infos
            self.output_filepath = output_filepath
            self.output_mask_filepath = output_mask_filepath

            self.need_return_image = need_return_image
            if self.need_return_image:
                self.image = None

    class Cli(Subprocessor.Cli):

        #override
        def on_initialize(self, client_dict):
            self.log_info ('Running on %s.' % (client_dict['device_name']) )
            self.device_idx  = client_dict['device_idx']
            self.device_name = client_dict['device_name']
            self.predictor_func = client_dict['predictor_func']
            self.predictor_input_shape = client_dict['predictor_input_shape']
            self.face_enhancer_func = client_dict['face_enhancer_func']
            self.xseg_256_extract_func = client_dict['xseg_256_extract_func']


            #transfer and set stdin in order to work code.interact in debug subprocess
            stdin_fd         = client_dict['stdin_fd']
            if stdin_fd is not None:
                sys.stdin = os.fdopen(stdin_fd)

            return None

        #override
        def process_data(self, pf): #pf=ProcessingFrame
            cfg = pf.cfg.copy()

            frame_info = pf.frame_info
            filepath = frame_info.filepath

            if len(frame_info.landmarks_list) == 0:
                
                if cfg.mode == 'raw-predict':        
                    h,w,c = self.predictor_input_shape
                    img_bgr = np.zeros( (h,w,3), dtype=np.uint8)
                    img_mask = np.zeros( (h,w,1), dtype=np.uint8)               
                else:                
                    self.log_info (f'no faces found for {filepath.name}, copying without faces')
                    img_bgr = cv2_imread(filepath)
                    imagelib.normalize_channels(img_bgr, 3)                    
                    h,w,c = img_bgr.shape
                    img_mask = np.zeros( (h,w,1), dtype=img_bgr.dtype)
                    
                cv2_imwrite (pf.output_filepath, img_bgr)
                cv2_imwrite (pf.output_mask_filepath, img_mask)

                if pf.need_return_image:
                    pf.image = np.concatenate ([img_bgr, img_mask], axis=-1)

            else:
                if cfg.type == MergerConfig.TYPE_MASKED:
                    try:
                        final_img = MergeMasked (self.predictor_func, self.predictor_input_shape,
                                                 face_enhancer_func=self.face_enhancer_func,
                                                 xseg_256_extract_func=self.xseg_256_extract_func,
                                                 cfg=cfg,
                                                 frame_info=frame_info)
                    except Exception as e:
                        e_str = traceback.format_exc()
                        if 'MemoryError' in e_str:
                            raise Subprocessor.SilenceException
                        else:
                            raise Exception( f'Error while merging file [{filepath}]: {e_str}' )

                elif cfg.type == MergerConfig.TYPE_FACE_AVATAR:
                    final_img = MergeFaceAvatar (self.predictor_func, self.predictor_input_shape,
                                                   cfg, pf.prev_temporal_frame_infos,
                                                        pf.frame_info,
                                                        pf.next_temporal_frame_infos )

                cv2_imwrite (pf.output_filepath,      final_img[...,0:3] )
                cv2_imwrite (pf.output_mask_filepath, final_img[...,3:4] )

                if pf.need_return_image:
                    pf.image = final_img

            return pf

        #overridable
        def get_data_name (self, pf):
            #return string identificator of your data
            return pf.frame_info.filepath




    #override
    def __init__(self, is_interactive, merger_session_filepath, predictor_func, predictor_input_shape, face_enhancer_func, xseg_256_extract_func, merger_config, frames, frames_root_path, output_path, output_mask_path, model_iter, subprocess_count=4):
        if len (frames) == 0:
            raise ValueError ("len (frames) == 0")

        super().__init__('Merger', InteractiveMergerSubprocessor.Cli, io_loop_sleep_time=0.001)

        self.is_interactive = is_interactive
        self.merger_session_filepath = Path(merger_session_filepath)
        self.merger_config = merger_config

        self.predictor_func = predictor_func
        self.predictor_input_shape = predictor_input_shape

        self.face_enhancer_func = face_enhancer_func
        self.xseg_256_extract_func = xseg_256_extract_func

        self.frames_root_path = frames_root_path
        self.output_path = output_path
        self.output_mask_path = output_mask_path
        self.model_iter = model_iter

        self.prefetch_frame_count = self.process_count = subprocess_count

        session_data = None
        if self.is_interactive and self.merger_session_filepath.exists():
            io.input_skip_pending()
            if io.input_bool ("Use saved session?", True):
                try:
                    with open( str(self.merger_session_filepath), "rb") as f:
                        session_data = pickle.loads(f.read())

                except Exception as e:
                    pass

        rewind_to_frame_idx = None
        self.frames = frames
        self.frames_idxs = [ *range(len(self.frames)) ]
        self.frames_done_idxs = []

        if self.is_interactive and session_data is not None:
            # Loaded session data, check it
            s_frames = session_data.get('frames', None)
            s_frames_idxs = session_data.get('frames_idxs', None)
            s_frames_done_idxs = session_data.get('frames_done_idxs', None)
            s_model_iter = session_data.get('model_iter', None)

            frames_equal = (s_frames is not None) and \
                           (s_frames_idxs is not None) and \
                           (s_frames_done_idxs is not None) and \
                           (s_model_iter is not None) and \
                           (len(frames) == len(s_frames)) # frames count must match

            if frames_equal:
                for i in range(len(frames)):
                    frame = frames[i]
                    s_frame = s_frames[i]
                    # frames filenames must match
                    if frame.frame_info.filepath.name != s_frame.frame_info.filepath.name:
                        frames_equal = False
                    if not frames_equal:
                        break

            if frames_equal:
                io.log_info ('Using saved session from ' + '/'.join (self.merger_session_filepath.parts[-2:]) )

                for frame in s_frames:
                    if frame.cfg is not None:
                        # recreate MergerConfig class using constructor with get_config() as dict params
                        # so if any new param will be added, old merger session will work properly
                        frame.cfg = frame.cfg.__class__( **frame.cfg.get_config() )

                self.frames = s_frames
                self.frames_idxs = s_frames_idxs
                self.frames_done_idxs = s_frames_done_idxs

                if self.model_iter != s_model_iter:
                    # model was more trained, recompute all frames
                    rewind_to_frame_idx = -1
                    for frame in self.frames:
                        frame.is_done = False
                elif len(self.frames_idxs) == 0:
                    # all frames are done?
                    rewind_to_frame_idx = -1

                if len(self.frames_idxs) != 0:
                    cur_frame = self.frames[self.frames_idxs[0]]
                    cur_frame.is_shown = False

            if not frames_equal:
                session_data = None

        if session_data is None:
            for filename in pathex.get_image_paths(self.output_path): #remove all images in output_path
                Path(filename).unlink()

            for filename in pathex.get_image_paths(self.output_mask_path): #remove all images in output_mask_path
                Path(filename).unlink()


            frames[0].cfg = self.merger_config.copy()

        for i in range( len(self.frames) ):
            frame = self.frames[i]
            frame.idx = i
            frame.output_filepath      = self.output_path      / ( frame.frame_info.filepath.stem + '.png' )
            frame.output_mask_filepath = self.output_mask_path / ( frame.frame_info.filepath.stem + '.png' )

            if not frame.output_filepath.exists() or \
               not frame.output_mask_filepath.exists():
                # if some frame does not exist, recompute and rewind
                frame.is_done = False
                frame.is_shown = False

                if rewind_to_frame_idx is None:
                    rewind_to_frame_idx = i-1
                else:
                    rewind_to_frame_idx = min(rewind_to_frame_idx, i-1)

        if rewind_to_frame_idx is not None:
            while len(self.frames_done_idxs) > 0:
                if self.frames_done_idxs[-1] > rewind_to_frame_idx:
                    prev_frame = self.frames[self.frames_done_idxs.pop()]
                    self.frames_idxs.insert(0, prev_frame.idx)
                else:
                    break
    #override
    def process_info_generator(self):
        r = [0] if MERGER_DEBUG else range(self.process_count)

        for i in r:
            yield 'CPU%d' % (i), {}, {'device_idx': i,
                                      'device_name': 'CPU%d' % (i),
                                      'predictor_func': self.predictor_func,
                                      'predictor_input_shape' : self.predictor_input_shape,
                                      'face_enhancer_func': self.face_enhancer_func,
                                      'xseg_256_extract_func' : self.xseg_256_extract_func,
                                      'stdin_fd': sys.stdin.fileno() if MERGER_DEBUG else None
                                      }

    #overridable optional
    def on_clients_initialized(self):
        io.progress_bar ("Merging", len(self.frames_idxs)+len(self.frames_done_idxs), initial=len(self.frames_done_idxs) )

        self.process_remain_frames = not self.is_interactive
        self.is_interactive_quitting = not self.is_interactive

        if self.is_interactive:
            help_images = {
                    MergerConfig.TYPE_MASKED :      cv2_imread ( str(Path(__file__).parent / 'gfx' / 'help_merger_masked.jpg') ),
                    MergerConfig.TYPE_FACE_AVATAR : cv2_imread ( str(Path(__file__).parent / 'gfx' / 'help_merger_face_avatar.jpg') ),
                }

            self.main_screen = Screen(initial_scale_to_width=1368, image=None, waiting_icon=True)
            self.help_screen = Screen(initial_scale_to_height=768, image=help_images[self.merger_config.type], waiting_icon=False)
            self.screen_manager = ScreenManager( "Merger", [self.main_screen, self.help_screen], capture_keys=True )
            self.screen_manager.set_current (self.help_screen)
            self.screen_manager.show_current()

            self.masked_keys_funcs = {
                    '`' : lambda cfg,shift_pressed: cfg.set_mode(0),
                    '1' : lambda cfg,shift_pressed: cfg.set_mode(1),
                    '2' : lambda cfg,shift_pressed: cfg.set_mode(2),
                    '3' : lambda cfg,shift_pressed: cfg.set_mode(3),
                    '4' : lambda cfg,shift_pressed: cfg.set_mode(4),
                    '5' : lambda cfg,shift_pressed: cfg.set_mode(5),
                    '6' : lambda cfg,shift_pressed: cfg.set_mode(6),
                    'q' : lambda cfg,shift_pressed: cfg.add_hist_match_threshold(1 if not shift_pressed else 5),
                    'a' : lambda cfg,shift_pressed: cfg.add_hist_match_threshold(-1 if not shift_pressed else -5),
                    'w' : lambda cfg,shift_pressed: cfg.add_erode_mask_modifier(1 if not shift_pressed else 5),
                    's' : lambda cfg,shift_pressed: cfg.add_erode_mask_modifier(-1 if not shift_pressed else -5),
                    'e' : lambda cfg,shift_pressed: cfg.add_blur_mask_modifier(1 if not shift_pressed else 5),
                    'd' : lambda cfg,shift_pressed: cfg.add_blur_mask_modifier(-1 if not shift_pressed else -5),
                    'r' : lambda cfg,shift_pressed: cfg.add_motion_blur_power(1 if not shift_pressed else 5),
                    'f' : lambda cfg,shift_pressed: cfg.add_motion_blur_power(-1 if not shift_pressed else -5),
                    't' : lambda cfg,shift_pressed: cfg.add_super_resolution_power(1 if not shift_pressed else 5),
                    'g' : lambda cfg,shift_pressed: cfg.add_super_resolution_power(-1 if not shift_pressed else -5),
                    'y' : lambda cfg,shift_pressed: cfg.add_blursharpen_amount(1 if not shift_pressed else 5),
                    'h' : lambda cfg,shift_pressed: cfg.add_blursharpen_amount(-1 if not shift_pressed else -5),
                    'u' : lambda cfg,shift_pressed: cfg.add_output_face_scale(1 if not shift_pressed else 5),
                    'j' : lambda cfg,shift_pressed: cfg.add_output_face_scale(-1 if not shift_pressed else -5),
                    'i' : lambda cfg,shift_pressed: cfg.add_image_denoise_power(1 if not shift_pressed else 5),
                    'k' : lambda cfg,shift_pressed: cfg.add_image_denoise_power(-1 if not shift_pressed else -5),
                    'o' : lambda cfg,shift_pressed: cfg.add_bicubic_degrade_power(1 if not shift_pressed else 5),
                    'l' : lambda cfg,shift_pressed: cfg.add_bicubic_degrade_power(-1 if not shift_pressed else -5),
                    'p' : lambda cfg,shift_pressed: cfg.add_color_degrade_power(1 if not shift_pressed else 5),
                    ';' : lambda cfg,shift_pressed: cfg.add_color_degrade_power(-1),
                    ':' : lambda cfg,shift_pressed: cfg.add_color_degrade_power(-5),
                    'z' : lambda cfg,shift_pressed: cfg.toggle_masked_hist_match(),
                    'x' : lambda cfg,shift_pressed: cfg.toggle_mask_mode(),
                    'c' : lambda cfg,shift_pressed: cfg.toggle_color_transfer_mode(),
                    'n' : lambda cfg,shift_pressed: cfg.toggle_sharpen_mode(),
                    }
            self.masked_keys = list(self.masked_keys_funcs.keys())

    #overridable optional
    def on_clients_finalized(self):
        io.progress_bar_close()

        if self.is_interactive:
            self.screen_manager.finalize()

            for frame in self.frames:
                frame.output_filepath = None
                frame.output_mask_filepath = None
                frame.image = None

            session_data = {
                'frames': self.frames,
                'frames_idxs': self.frames_idxs,
                'frames_done_idxs': self.frames_done_idxs,
                'model_iter' : self.model_iter,
            }
            self.merger_session_filepath.write_bytes( pickle.dumps(session_data) )

            io.log_info ("Session is saved to " + '/'.join (self.merger_session_filepath.parts[-2:]) )

    #override
    def on_tick(self):
        io.process_messages()

        go_prev_frame = False
        go_first_frame = False
        go_prev_frame_overriding_cfg = False
        go_first_frame_overriding_cfg = False

        go_next_frame = self.process_remain_frames
        go_next_frame_overriding_cfg = False
        go_last_frame_overriding_cfg = False

        cur_frame = None
        if len(self.frames_idxs) != 0:
            cur_frame = self.frames[self.frames_idxs[0]]

        if self.is_interactive:

            screen_image = None if self.process_remain_frames else \
                                   self.main_screen.get_image()

            self.main_screen.set_waiting_icon( self.process_remain_frames or \
                                               self.is_interactive_quitting )

            if cur_frame is not None and not self.is_interactive_quitting:

                if not self.process_remain_frames:
                    if cur_frame.is_done:
                        if not cur_frame.is_shown:
                            if cur_frame.image is None:
                                image      = cv2_imread (cur_frame.output_filepath, verbose=False)
                                image_mask = cv2_imread (cur_frame.output_mask_filepath, verbose=False)
                                if image is None or image_mask is None:
                                    # unable to read? recompute then
                                    cur_frame.is_done = False
                                else:
                                    image = imagelib.normalize_channels(image, 3)
                                    image_mask = imagelib.normalize_channels(image_mask, 1)
                                    cur_frame.image = np.concatenate([image, image_mask], -1)

                            if cur_frame.is_done:
                                io.log_info (cur_frame.cfg.to_string( cur_frame.frame_info.filepath.name) )
                                cur_frame.is_shown = True
                                screen_image = cur_frame.image
                    else:
                        self.main_screen.set_waiting_icon(True)

            self.main_screen.set_image(screen_image)
            self.screen_manager.show_current()

            key_events = self.screen_manager.get_key_events()
            key, chr_key, ctrl_pressed, alt_pressed, shift_pressed = key_events[-1] if len(key_events) > 0 else (0,0,False,False,False)

            if key == 9: #tab
                self.screen_manager.switch_screens()
            else:
                if key == 27: #esc
                    self.is_interactive_quitting = True
                elif self.screen_manager.get_current() is self.main_screen:

                    if self.merger_config.type == MergerConfig.TYPE_MASKED and chr_key in self.masked_keys:
                        self.process_remain_frames = False

                        if cur_frame is not None:
                            cfg = cur_frame.cfg
                            prev_cfg = cfg.copy()

                            if cfg.type == MergerConfig.TYPE_MASKED:
                                self.masked_keys_funcs[chr_key](cfg, shift_pressed)

                            if prev_cfg != cfg:
                                io.log_info ( cfg.to_string(cur_frame.frame_info.filepath.name) )
                                cur_frame.is_done = False
                                cur_frame.is_shown = False
                    else:

                        if chr_key == ',' or chr_key == 'm':
                            self.process_remain_frames = False
                            go_prev_frame = True

                            if chr_key == ',':
                                if shift_pressed:
                                    go_first_frame = True

                            elif chr_key == 'm':
                                if not shift_pressed:
                                    go_prev_frame_overriding_cfg = True
                                else:
                                    go_first_frame_overriding_cfg = True

                        elif chr_key == '.' or chr_key == '/':
                            self.process_remain_frames = False
                            go_next_frame = True

                            if chr_key == '.':
                                if shift_pressed:
                                    self.process_remain_frames = not self.process_remain_frames

                            elif chr_key == '/':
                                if not shift_pressed:
                                    go_next_frame_overriding_cfg = True
                                else:
                                    go_last_frame_overriding_cfg = True

                        elif chr_key == '-':
                            self.screen_manager.get_current().diff_scale(-0.1)
                        elif chr_key == '=':
                            self.screen_manager.get_current().diff_scale(0.1)
                        elif chr_key == 'v':
                            self.screen_manager.get_current().toggle_show_checker_board()

        if go_prev_frame:
            if cur_frame is None or cur_frame.is_done:
                if cur_frame is not None:
                    cur_frame.image = None

                while True:
                    if len(self.frames_done_idxs) > 0:
                        prev_frame = self.frames[self.frames_done_idxs.pop()]
                        self.frames_idxs.insert(0, prev_frame.idx)
                        prev_frame.is_shown = False
                        io.progress_bar_inc(-1)

                        if cur_frame is not None and (go_prev_frame_overriding_cfg or go_first_frame_overriding_cfg):
                            if prev_frame.cfg != cur_frame.cfg:
                                prev_frame.cfg = cur_frame.cfg.copy()
                                prev_frame.is_done = False

                        cur_frame = prev_frame

                    if go_first_frame_overriding_cfg or go_first_frame:
                        if len(self.frames_done_idxs) > 0:
                            continue
                    break

        elif go_next_frame:
            if cur_frame is not None and cur_frame.is_done:
                cur_frame.image = None
                cur_frame.is_shown = True
                self.frames_done_idxs.append(cur_frame.idx)
                self.frames_idxs.pop(0)
                io.progress_bar_inc(1)

                f = self.frames

                if len(self.frames_idxs) != 0:
                    next_frame = f[ self.frames_idxs[0] ]
                    next_frame.is_shown = False

                    if go_next_frame_overriding_cfg or go_last_frame_overriding_cfg:

                        if go_next_frame_overriding_cfg:
                            to_frames = next_frame.idx+1
                        else:
                            to_frames = len(f)

                        for i in range( next_frame.idx, to_frames ):
                            f[i].cfg = None

                    for i in range( min(len(self.frames_idxs), self.prefetch_frame_count) ):
                        frame = f[ self.frames_idxs[i] ]
                        if frame.cfg is None:
                            if i == 0:
                                frame.cfg = cur_frame.cfg.copy()
                            else:
                                frame.cfg = f[ self.frames_idxs[i-1] ].cfg.copy()

                            frame.is_done = False #initiate solve again
                            frame.is_shown = False

            if len(self.frames_idxs) == 0:
                self.process_remain_frames = False

        return (self.is_interactive and self.is_interactive_quitting) or \
               (not self.is_interactive and self.process_remain_frames == False)


    #override
    def on_data_return (self, host_dict, pf):
        frame = self.frames[pf.idx]
        frame.is_done = False
        frame.is_processing = False

    #override
    def on_result (self, host_dict, pf_sent, pf_result):
        frame = self.frames[pf_result.idx]
        frame.is_processing = False
        if frame.cfg == pf_result.cfg:
            frame.is_done = True
            frame.image = pf_result.image

    #override
    def get_data(self, host_dict):
        if self.is_interactive and self.is_interactive_quitting:
            return None

        for i in range ( min(len(self.frames_idxs), self.prefetch_frame_count) ):
            frame = self.frames[ self.frames_idxs[i] ]

            if not frame.is_done and not frame.is_processing and frame.cfg is not None:
                frame.is_processing = True
                return InteractiveMergerSubprocessor.ProcessingFrame(idx=frame.idx,
                                                           cfg=frame.cfg.copy(),
                                                           prev_temporal_frame_infos=frame.prev_temporal_frame_infos,
                                                           frame_info=frame.frame_info,
                                                           next_temporal_frame_infos=frame.next_temporal_frame_infos,
                                                           output_filepath=frame.output_filepath,
                                                           output_mask_filepath=frame.output_mask_filepath,
                                                           need_return_image=True )

        return None

    #override
    def get_result(self):
        return 0
merger/InteractiveMergerSubprocessor.py
oo92/TorchFaceLab
import cv2
import numpy as np

from core import imagelib
from facelib import FaceType, LandmarksProcessor
from core.cv2ex import *

def process_frame_info(frame_info, inp_sh):
    img_uint8 = cv2_imread (frame_info.filename)
    img_uint8 = imagelib.normalize_channels (img_uint8, 3)
    img = img_uint8.astype(np.float32) / 255.0

    img_mat = LandmarksProcessor.get_transform_mat (frame_info.landmarks_list[0], inp_sh[0], face_type=FaceType.FULL_NO_ALIGN)
    img = cv2.warpAffine( img, img_mat, inp_sh[0:2], borderMode=cv2.BORDER_REPLICATE, flags=cv2.INTER_CUBIC )
    return img

def MergeFaceAvatar (predictor_func, predictor_input_shape, cfg, prev_temporal_frame_infos, frame_info, next_temporal_frame_infos):
    inp_sh = predictor_input_shape

    prev_imgs=[]
    next_imgs=[]
    for i in range(cfg.temporal_face_count):
        prev_imgs.append( process_frame_info(prev_temporal_frame_infos[i], inp_sh) )
        next_imgs.append( process_frame_info(next_temporal_frame_infos[i], inp_sh) )
    img = process_frame_info(frame_info, inp_sh)

    prd_f = predictor_func ( prev_imgs, img, next_imgs )

    #if cfg.super_resolution_mode != 0:
    #    prd_f = cfg.superres_func(cfg.super_resolution_mode, prd_f)

    if cfg.sharpen_mode != 0 and cfg.sharpen_amount != 0:
        prd_f = cfg.sharpen_func ( prd_f, cfg.sharpen_mode, 3, cfg.sharpen_amount)

    out_img = np.clip(prd_f, 0.0, 1.0)

    if cfg.add_source_image:
        out_img = np.concatenate ( [cv2.resize ( img, (prd_f.shape[1], prd_f.shape[0])  ),
                                    out_img], axis=1 )

    return (out_img*255).astype(np.uint8)

merger/MergeAvatar.py
oo92/TorchFaceLab
import sys
import traceback

import cv2
import numpy as np

from core import imagelib
from core.cv2ex import *
from core.interact import interact as io
from facelib import FaceType, LandmarksProcessor

is_windows = sys.platform[0:3] == 'win'
xseg_input_size = 256

def MergeMaskedFace (predictor_func, predictor_input_shape,
                     face_enhancer_func,
                     xseg_256_extract_func,
                     cfg, frame_info, img_bgr_uint8, img_bgr, img_face_landmarks):

    img_size = img_bgr.shape[1], img_bgr.shape[0]
    img_face_mask_a = LandmarksProcessor.get_image_hull_mask (img_bgr.shape, img_face_landmarks)

    input_size = predictor_input_shape[0]
    mask_subres_size = input_size*4
    output_size = input_size
    if cfg.super_resolution_power != 0:
        output_size *= 4

    face_mat        = LandmarksProcessor.get_transform_mat (img_face_landmarks, output_size, face_type=cfg.face_type)
    face_output_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, output_size, face_type=cfg.face_type, scale= 1.0 + 0.01*cfg.output_face_scale)

    if mask_subres_size == output_size:
        face_mask_output_mat = face_output_mat
    else:
        face_mask_output_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, mask_subres_size, face_type=cfg.face_type, scale= 1.0 + 0.01*cfg.output_face_scale)

    dst_face_bgr      = cv2.warpAffine( img_bgr        , face_mat, (output_size, output_size), flags=cv2.INTER_CUBIC )
    dst_face_bgr      = np.clip(dst_face_bgr, 0, 1)

    dst_face_mask_a_0 = cv2.warpAffine( img_face_mask_a, face_mat, (output_size, output_size), flags=cv2.INTER_CUBIC )
    dst_face_mask_a_0 = np.clip(dst_face_mask_a_0, 0, 1)

    predictor_input_bgr      = cv2.resize (dst_face_bgr, (input_size,input_size) )

    predicted = predictor_func (predictor_input_bgr)
    prd_face_bgr          = np.clip (predicted[0], 0, 1.0)
    prd_face_mask_a_0     = np.clip (predicted[1], 0, 1.0)
    prd_face_dst_mask_a_0 = np.clip (predicted[2], 0, 1.0)

    if cfg.super_resolution_power != 0:
        prd_face_bgr_enhanced = face_enhancer_func(prd_face_bgr, is_tanh=True, preserve_size=False)
        mod = cfg.super_resolution_power / 100.0
        prd_face_bgr = cv2.resize(prd_face_bgr, (output_size,output_size))*(1.0-mod) + prd_face_bgr_enhanced*mod
        prd_face_bgr = np.clip(prd_face_bgr, 0, 1)

    if cfg.super_resolution_power != 0:
        prd_face_mask_a_0     = cv2.resize (prd_face_mask_a_0,      (output_size, output_size), interpolation=cv2.INTER_CUBIC)
        prd_face_dst_mask_a_0 = cv2.resize (prd_face_dst_mask_a_0,  (output_size, output_size), interpolation=cv2.INTER_CUBIC)

    if cfg.mask_mode == 0: #full
        wrk_face_mask_a_0 = np.ones_like(dst_face_mask_a_0)
    elif cfg.mask_mode == 1: #dst
        wrk_face_mask_a_0 = cv2.resize (dst_face_mask_a_0, (output_size,output_size), interpolation=cv2.INTER_CUBIC)
    elif cfg.mask_mode == 2: #learned-prd
        wrk_face_mask_a_0 = prd_face_mask_a_0
    elif cfg.mask_mode == 3: #learned-dst
        wrk_face_mask_a_0 = prd_face_dst_mask_a_0
    elif cfg.mask_mode == 4: #learned-prd*learned-dst
        wrk_face_mask_a_0 = prd_face_mask_a_0*prd_face_dst_mask_a_0
    elif cfg.mask_mode == 5: #learned-prd+learned-dst
        wrk_face_mask_a_0 = np.clip( prd_face_mask_a_0+prd_face_dst_mask_a_0, 0, 1)
    elif cfg.mask_mode >= 6 and cfg.mask_mode <= 9:  #XSeg modes
        if cfg.mask_mode == 6 or cfg.mask_mode == 8 or cfg.mask_mode == 9:
            # obtain XSeg-prd
            prd_face_xseg_bgr = cv2.resize (prd_face_bgr, (xseg_input_size,)*2, interpolation=cv2.INTER_CUBIC)
            prd_face_xseg_mask = xseg_256_extract_func(prd_face_xseg_bgr)
            X_prd_face_mask_a_0 = cv2.resize ( prd_face_xseg_mask, (output_size, output_size), interpolation=cv2.INTER_CUBIC)

        if cfg.mask_mode >= 7 and cfg.mask_mode <= 9:
            # obtain XSeg-dst
            xseg_mat            = LandmarksProcessor.get_transform_mat (img_face_landmarks, xseg_input_size, face_type=cfg.face_type)
            dst_face_xseg_bgr   = cv2.warpAffine(img_bgr, xseg_mat, (xseg_input_size,)*2, flags=cv2.INTER_CUBIC )
            dst_face_xseg_mask  = xseg_256_extract_func(dst_face_xseg_bgr)
            X_dst_face_mask_a_0 = cv2.resize (dst_face_xseg_mask, (output_size,output_size), interpolation=cv2.INTER_CUBIC)

        if cfg.mask_mode == 6:   #'XSeg-prd'
            wrk_face_mask_a_0 = X_prd_face_mask_a_0
        elif cfg.mask_mode == 7: #'XSeg-dst'
            wrk_face_mask_a_0 = X_dst_face_mask_a_0
        elif cfg.mask_mode == 8: #'XSeg-prd*XSeg-dst'
            wrk_face_mask_a_0 = X_prd_face_mask_a_0 * X_dst_face_mask_a_0
        elif cfg.mask_mode == 9: #learned-prd*learned-dst*XSeg-prd*XSeg-dst
            wrk_face_mask_a_0 = prd_face_mask_a_0 * prd_face_dst_mask_a_0 * X_prd_face_mask_a_0 * X_dst_face_mask_a_0

    wrk_face_mask_a_0[ wrk_face_mask_a_0 < (1.0/255.0) ] = 0.0 # get rid of noise

    # resize to mask_subres_size
    if wrk_face_mask_a_0.shape[0] != mask_subres_size:
        wrk_face_mask_a_0 = cv2.resize (wrk_face_mask_a_0, (mask_subres_size, mask_subres_size), interpolation=cv2.INTER_CUBIC)

    # process mask in local predicted space
    if 'raw' not in cfg.mode:
        # add zero pad
        wrk_face_mask_a_0 = np.pad (wrk_face_mask_a_0, input_size)

        ero  = cfg.erode_mask_modifier
        blur = cfg.blur_mask_modifier

        if ero > 0:
            wrk_face_mask_a_0 = cv2.erode(wrk_face_mask_a_0, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(ero,ero)), iterations = 1 )
        elif ero < 0:
            wrk_face_mask_a_0 = cv2.dilate(wrk_face_mask_a_0, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(-ero,-ero)), iterations = 1 )

        # clip eroded/dilated mask in actual predict area
        # pad with half blur size in order to accuratelly fade to zero at the boundary
        clip_size = input_size + blur // 2

        wrk_face_mask_a_0[:clip_size,:] = 0
        wrk_face_mask_a_0[-clip_size:,:] = 0
        wrk_face_mask_a_0[:,:clip_size] = 0
        wrk_face_mask_a_0[:,-clip_size:] = 0

        if blur > 0:
            blur = blur + (1-blur % 2)
            wrk_face_mask_a_0 = cv2.GaussianBlur(wrk_face_mask_a_0, (blur, blur) , 0)

        wrk_face_mask_a_0 = wrk_face_mask_a_0[input_size:-input_size,input_size:-input_size]

        wrk_face_mask_a_0 = np.clip(wrk_face_mask_a_0, 0, 1)

    img_face_mask_a = cv2.warpAffine( wrk_face_mask_a_0, face_mask_output_mat, img_size, np.zeros(img_bgr.shape[0:2], dtype=np.float32), flags=cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC )[...,None]
    img_face_mask_a = np.clip (img_face_mask_a, 0.0, 1.0)
    img_face_mask_a [ img_face_mask_a < (1.0/255.0) ] = 0.0 # get rid of noise

    if wrk_face_mask_a_0.shape[0] != output_size:
        wrk_face_mask_a_0 = cv2.resize (wrk_face_mask_a_0, (output_size,output_size), interpolation=cv2.INTER_CUBIC)

    wrk_face_mask_a = wrk_face_mask_a_0[...,None]

    out_img = None
    out_merging_mask_a = None
    if cfg.mode == 'original':
        return img_bgr, img_face_mask_a

    elif 'raw' in cfg.mode:
        if cfg.mode == 'raw-rgb':
            out_img_face = cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.empty_like(img_bgr), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC)
            out_img_face_mask = cv2.warpAffine( np.ones_like(prd_face_bgr), face_output_mat, img_size, np.empty_like(img_bgr), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC)
            out_img = img_bgr*(1-out_img_face_mask) + out_img_face*out_img_face_mask
            out_merging_mask_a = img_face_mask_a
        elif cfg.mode == 'raw-predict':
            out_img = prd_face_bgr
            out_merging_mask_a = wrk_face_mask_a
        else:
            raise ValueError(f"undefined raw type {cfg.mode}")

        out_img = np.clip (out_img, 0.0, 1.0 )
    else:

        # Process if the mask meets minimum size
        maxregion = np.argwhere( img_face_mask_a >= 0.1 )
        if maxregion.size != 0:
            miny,minx = maxregion.min(axis=0)[:2]
            maxy,maxx = maxregion.max(axis=0)[:2]
            lenx = maxx - minx
            leny = maxy - miny
            if min(lenx,leny) >= 4:
                wrk_face_mask_area_a = wrk_face_mask_a.copy()
                wrk_face_mask_area_a[wrk_face_mask_area_a>0] = 1.0

                if 'seamless' not in cfg.mode and cfg.color_transfer_mode != 0:
                    if cfg.color_transfer_mode == 1: #rct
                        prd_face_bgr = imagelib.reinhard_color_transfer (prd_face_bgr, dst_face_bgr, target_mask=wrk_face_mask_area_a, source_mask=wrk_face_mask_area_a)
                    elif cfg.color_transfer_mode == 2: #lct
                        prd_face_bgr = imagelib.linear_color_transfer (prd_face_bgr, dst_face_bgr)
                    elif cfg.color_transfer_mode == 3: #mkl
                        prd_face_bgr = imagelib.color_transfer_mkl (prd_face_bgr, dst_face_bgr)
                    elif cfg.color_transfer_mode == 4: #mkl-m
                        prd_face_bgr = imagelib.color_transfer_mkl (prd_face_bgr*wrk_face_mask_area_a, dst_face_bgr*wrk_face_mask_area_a)
                    elif cfg.color_transfer_mode == 5: #idt
                        prd_face_bgr = imagelib.color_transfer_idt (prd_face_bgr, dst_face_bgr)
                    elif cfg.color_transfer_mode == 6: #idt-m
                        prd_face_bgr = imagelib.color_transfer_idt (prd_face_bgr*wrk_face_mask_area_a, dst_face_bgr*wrk_face_mask_area_a)
                    elif cfg.color_transfer_mode == 7: #sot-m
                        prd_face_bgr = imagelib.color_transfer_sot (prd_face_bgr*wrk_face_mask_area_a, dst_face_bgr*wrk_face_mask_area_a, steps=10, batch_size=30)
                        prd_face_bgr = np.clip (prd_face_bgr, 0.0, 1.0)
                    elif cfg.color_transfer_mode == 8: #mix-m
                        prd_face_bgr = imagelib.color_transfer_mix (prd_face_bgr*wrk_face_mask_area_a, dst_face_bgr*wrk_face_mask_area_a)

                if cfg.mode == 'hist-match':
                    hist_mask_a = np.ones ( prd_face_bgr.shape[:2] + (1,) , dtype=np.float32)

                    if cfg.masked_hist_match:
                        hist_mask_a *= wrk_face_mask_area_a

                    white =  (1.0-hist_mask_a)* np.ones ( prd_face_bgr.shape[:2] + (1,) , dtype=np.float32)

                    hist_match_1 = prd_face_bgr*hist_mask_a + white
                    hist_match_1[ hist_match_1 > 1.0 ] = 1.0

                    hist_match_2 = dst_face_bgr*hist_mask_a + white
                    hist_match_2[ hist_match_1 > 1.0 ] = 1.0

                    prd_face_bgr = imagelib.color_hist_match(hist_match_1, hist_match_2, cfg.hist_match_threshold ).astype(dtype=np.float32)

                if 'seamless' in cfg.mode:
                    #mask used for cv2.seamlessClone
                    img_face_seamless_mask_a = None
                    for i in range(1,10):
                        a = img_face_mask_a > i / 10.0
                        if len(np.argwhere(a)) == 0:
                            continue
                        img_face_seamless_mask_a = img_face_mask_a.copy()
                        img_face_seamless_mask_a[a] = 1.0
                        img_face_seamless_mask_a[img_face_seamless_mask_a <= i / 10.0] = 0.0
                        break

                out_img = cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.empty_like(img_bgr), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC )
                out_img = np.clip(out_img, 0.0, 1.0)

                if 'seamless' in cfg.mode:
                    try:
                        #calc same bounding rect and center point as in cv2.seamlessClone to prevent jittering (not flickering)
                        l,t,w,h = cv2.boundingRect( (img_face_seamless_mask_a*255).astype(np.uint8) )
                        s_maskx, s_masky = int(l+w/2), int(t+h/2)
                        out_img = cv2.seamlessClone( (out_img*255).astype(np.uint8), img_bgr_uint8, (img_face_seamless_mask_a*255).astype(np.uint8), (s_maskx,s_masky) , cv2.NORMAL_CLONE )
                        out_img = out_img.astype(dtype=np.float32) / 255.0
                    except Exception as e:
                        #seamlessClone may fail in some cases
                        e_str = traceback.format_exc()

                        if 'MemoryError' in e_str:
                            raise Exception("Seamless fail: " + e_str) #reraise MemoryError in order to reprocess this data by other processes
                        else:
                            print ("Seamless fail: " + e_str)

                cfg_mp = cfg.motion_blur_power / 100.0

                out_img = img_bgr*(1-img_face_mask_a) + (out_img*img_face_mask_a)

                if ('seamless' in cfg.mode and cfg.color_transfer_mode != 0) or \
                   cfg.mode == 'seamless-hist-match' or \
                   cfg_mp != 0 or \
                   cfg.blursharpen_amount != 0 or \
                   cfg.image_denoise_power != 0 or \
                   cfg.bicubic_degrade_power != 0:

                    out_face_bgr = cv2.warpAffine( out_img, face_mat, (output_size, output_size), flags=cv2.INTER_CUBIC )

                    if 'seamless' in cfg.mode and cfg.color_transfer_mode != 0:
                        if cfg.color_transfer_mode == 1:
                            out_face_bgr = imagelib.reinhard_color_transfer (out_face_bgr, dst_face_bgr, target_mask=wrk_face_mask_area_a, source_mask=wrk_face_mask_area_a)
                        elif cfg.color_transfer_mode == 2: #lct
                            out_face_bgr = imagelib.linear_color_transfer (out_face_bgr, dst_face_bgr)
                        elif cfg.color_transfer_mode == 3: #mkl
                            out_face_bgr = imagelib.color_transfer_mkl (out_face_bgr, dst_face_bgr)
                        elif cfg.color_transfer_mode == 4: #mkl-m
                            out_face_bgr = imagelib.color_transfer_mkl (out_face_bgr*wrk_face_mask_area_a, dst_face_bgr*wrk_face_mask_area_a)
                        elif cfg.color_transfer_mode == 5: #idt
                            out_face_bgr = imagelib.color_transfer_idt (out_face_bgr, dst_face_bgr)
                        elif cfg.color_transfer_mode == 6: #idt-m
                            out_face_bgr = imagelib.color_transfer_idt (out_face_bgr*wrk_face_mask_area_a, dst_face_bgr*wrk_face_mask_area_a)
                        elif cfg.color_transfer_mode == 7: #sot-m
                            out_face_bgr = imagelib.color_transfer_sot (out_face_bgr*wrk_face_mask_area_a, dst_face_bgr*wrk_face_mask_area_a, steps=10, batch_size=30)
                            out_face_bgr = np.clip (out_face_bgr, 0.0, 1.0)
                        elif cfg.color_transfer_mode == 8: #mix-m
                            out_face_bgr = imagelib.color_transfer_mix (out_face_bgr*wrk_face_mask_area_a, dst_face_bgr*wrk_face_mask_area_a)

                    if cfg.mode == 'seamless-hist-match':
                        out_face_bgr = imagelib.color_hist_match(out_face_bgr, dst_face_bgr, cfg.hist_match_threshold)

                    if cfg_mp != 0:
                        k_size = int(frame_info.motion_power*cfg_mp)
                        if k_size >= 1:
                            k_size = np.clip (k_size+1, 2, 50)
                            if cfg.super_resolution_power != 0:
                                k_size *= 2
                            out_face_bgr = imagelib.LinearMotionBlur (out_face_bgr, k_size , frame_info.motion_deg)

                    if cfg.blursharpen_amount != 0:
                        out_face_bgr = imagelib.blursharpen ( out_face_bgr, cfg.sharpen_mode, 3, cfg.blursharpen_amount)

                    if cfg.image_denoise_power != 0:
                        n = cfg.image_denoise_power
                        while n > 0:
                            img_bgr_denoised = cv2.medianBlur(img_bgr, 5)
                            if int(n / 100) != 0:
                                img_bgr = img_bgr_denoised
                            else:
                                pass_power = (n % 100) / 100.0
                                img_bgr = img_bgr*(1.0-pass_power)+img_bgr_denoised*pass_power
                            n = max(n-10,0)

                    if cfg.bicubic_degrade_power != 0:
                        p = 1.0 - cfg.bicubic_degrade_power / 101.0
                        img_bgr_downscaled = cv2.resize (img_bgr, ( int(img_size[0]*p), int(img_size[1]*p ) ), interpolation=cv2.INTER_CUBIC)
                        img_bgr = cv2.resize (img_bgr_downscaled, img_size, interpolation=cv2.INTER_CUBIC)

                    new_out = cv2.warpAffine( out_face_bgr, face_mat, img_size, np.empty_like(img_bgr), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC )

                    out_img =  np.clip( img_bgr*(1-img_face_mask_a) + (new_out*img_face_mask_a) , 0, 1.0 )

                if cfg.color_degrade_power != 0:
                    out_img_reduced = imagelib.reduce_colors(out_img, 256)
                    if cfg.color_degrade_power == 100:
                        out_img = out_img_reduced
                    else:
                        alpha = cfg.color_degrade_power / 100.0
                        out_img = (out_img*(1.0-alpha) + out_img_reduced*alpha)
        out_merging_mask_a = img_face_mask_a

    if out_img is None:
        out_img = img_bgr.copy()
        
    return out_img, out_merging_mask_a


def MergeMasked (predictor_func,
                 predictor_input_shape,
                 face_enhancer_func,
                 xseg_256_extract_func,
                 cfg,
                 frame_info):
    img_bgr_uint8 = cv2_imread(frame_info.filepath)
    img_bgr_uint8 = imagelib.normalize_channels (img_bgr_uint8, 3)
    img_bgr = img_bgr_uint8.astype(np.float32) / 255.0

    outs = []
    for face_num, img_landmarks in enumerate( frame_info.landmarks_list ):
        out_img, out_img_merging_mask = MergeMaskedFace (predictor_func, predictor_input_shape, face_enhancer_func, xseg_256_extract_func, cfg, frame_info, img_bgr_uint8, img_bgr, img_landmarks)
        outs += [ (out_img, out_img_merging_mask) ]

    #Combining multiple face outputs
    final_img = None
    final_mask = None
    for img, merging_mask in outs:
        h,w,c = img.shape

        if final_img is None:
            final_img = img
            final_mask = merging_mask
        else:
            final_img = final_img*(1-merging_mask) + img*merging_mask
            final_mask = np.clip (final_mask + merging_mask, 0, 1 )

    final_img = np.concatenate ( [final_img, final_mask], -1)

    return (final_img*255).astype(np.uint8)

merger/MergeMasked.py
oo92/TorchFaceLab
import numpy as np
import copy

from facelib import FaceType
from core.interact import interact as io


class MergerConfig(object):
    TYPE_NONE = 0
    TYPE_MASKED = 1
    TYPE_FACE_AVATAR = 2
    ####

    TYPE_IMAGE = 3
    TYPE_IMAGE_WITH_LANDMARKS = 4

    def __init__(self, type=0,
                       sharpen_mode=0,
                       blursharpen_amount=0,
                       **kwargs
                       ):
        self.type = type

        self.sharpen_dict = {0:"None", 1:'box', 2:'gaussian'}

        #default changeable params
        self.sharpen_mode = sharpen_mode
        self.blursharpen_amount = blursharpen_amount

    def copy(self):
        return copy.copy(self)

    #overridable
    def ask_settings(self):
        s = """Choose sharpen mode: \n"""
        for key in self.sharpen_dict.keys():
            s += f"""({key}) {self.sharpen_dict[key]}\n"""
        io.log_info(s)
        self.sharpen_mode = io.input_int ("", 0, valid_list=self.sharpen_dict.keys(), help_message="Enhance details by applying sharpen filter.")

        if self.sharpen_mode != 0:
            self.blursharpen_amount = np.clip ( io.input_int ("Choose blur/sharpen amount", 0, add_info="-100..100"), -100, 100 )

    def toggle_sharpen_mode(self):
        a = list( self.sharpen_dict.keys() )
        self.sharpen_mode = a[ (a.index(self.sharpen_mode)+1) % len(a) ]

    def add_blursharpen_amount(self, diff):
        self.blursharpen_amount = np.clip ( self.blursharpen_amount+diff, -100, 100)

    #overridable
    def get_config(self):
        d = self.__dict__.copy()
        d.pop('type')
        return d

    #overridable
    def __eq__(self, other):
        #check equality of changeable params

        if isinstance(other, MergerConfig):
            return self.sharpen_mode == other.sharpen_mode and \
                   self.blursharpen_amount == other.blursharpen_amount

        return False

    #overridable
    def to_string(self, filename):
        r = ""
        r += f"sharpen_mode : {self.sharpen_dict[self.sharpen_mode]}\n"
        r += f"blursharpen_amount : {self.blursharpen_amount}\n"
        return r

mode_dict = {0:'original',
             1:'overlay',
             2:'hist-match',
             3:'seamless',
             4:'seamless-hist-match',
             5:'raw-rgb',
             6:'raw-predict'}

mode_str_dict = { mode_dict[key] : key for key in mode_dict.keys() }

mask_mode_dict = {0:'full',
                  1:'dst',
                  2:'learned-prd',
                  3:'learned-dst',
                  4:'learned-prd*learned-dst',
                  5:'learned-prd+learned-dst',
                  6:'XSeg-prd',
                  7:'XSeg-dst',
                  8:'XSeg-prd*XSeg-dst',
                  9:'learned-prd*learned-dst*XSeg-prd*XSeg-dst'
                  }


ctm_dict = { 0: "None", 1:"rct", 2:"lct", 3:"mkl", 4:"mkl-m", 5:"idt", 6:"idt-m", 7:"sot-m", 8:"mix-m" }
ctm_str_dict = {None:0, "rct":1, "lct":2, "mkl":3, "mkl-m":4, "idt":5, "idt-m":6, "sot-m":7, "mix-m":8 }

class MergerConfigMasked(MergerConfig):

    def __init__(self, face_type=FaceType.FULL,
                       default_mode = 'overlay',
                       mode='overlay',
                       masked_hist_match=True,
                       hist_match_threshold = 238,
                       mask_mode = 4,
                       erode_mask_modifier = 0,
                       blur_mask_modifier = 0,
                       motion_blur_power = 0,
                       output_face_scale = 0,
                       super_resolution_power = 0,
                       color_transfer_mode = ctm_str_dict['rct'],
                       image_denoise_power = 0,
                       bicubic_degrade_power = 0,
                       color_degrade_power = 0,
                       **kwargs
                       ):

        super().__init__(type=MergerConfig.TYPE_MASKED, **kwargs)

        self.face_type = face_type
        if self.face_type not in [FaceType.HALF, FaceType.MID_FULL, FaceType.FULL, FaceType.WHOLE_FACE, FaceType.HEAD ]:
            raise ValueError("MergerConfigMasked does not support this type of face.")

        self.default_mode = default_mode

        #default changeable params
        if mode not in mode_str_dict:
            mode = mode_dict[1]

        self.mode = mode
        self.masked_hist_match = masked_hist_match
        self.hist_match_threshold = hist_match_threshold
        self.mask_mode = mask_mode
        self.erode_mask_modifier = erode_mask_modifier
        self.blur_mask_modifier = blur_mask_modifier
        self.motion_blur_power = motion_blur_power
        self.output_face_scale = output_face_scale
        self.super_resolution_power = super_resolution_power
        self.color_transfer_mode = color_transfer_mode
        self.image_denoise_power = image_denoise_power
        self.bicubic_degrade_power = bicubic_degrade_power
        self.color_degrade_power = color_degrade_power

    def copy(self):
        return copy.copy(self)

    def set_mode (self, mode):
        self.mode = mode_dict.get (mode, self.default_mode)

    def toggle_masked_hist_match(self):
        if self.mode == 'hist-match':
            self.masked_hist_match = not self.masked_hist_match

    def add_hist_match_threshold(self, diff):
        if self.mode == 'hist-match' or self.mode == 'seamless-hist-match':
            self.hist_match_threshold = np.clip ( self.hist_match_threshold+diff , 0, 255)

    def toggle_mask_mode(self):
        a = list( mask_mode_dict.keys() )
        self.mask_mode = a[ (a.index(self.mask_mode)+1) % len(a) ]

    def add_erode_mask_modifier(self, diff):
        self.erode_mask_modifier = np.clip ( self.erode_mask_modifier+diff , -400, 400)

    def add_blur_mask_modifier(self, diff):
        self.blur_mask_modifier = np.clip ( self.blur_mask_modifier+diff , 0, 400)

    def add_motion_blur_power(self, diff):
        self.motion_blur_power = np.clip ( self.motion_blur_power+diff, 0, 100)

    def add_output_face_scale(self, diff):
        self.output_face_scale = np.clip ( self.output_face_scale+diff , -50, 50)

    def toggle_color_transfer_mode(self):
        self.color_transfer_mode = (self.color_transfer_mode+1) % ( max(ctm_dict.keys())+1 )

    def add_super_resolution_power(self, diff):
        self.super_resolution_power = np.clip ( self.super_resolution_power+diff , 0, 100)

    def add_color_degrade_power(self, diff):
        self.color_degrade_power = np.clip ( self.color_degrade_power+diff , 0, 100)

    def add_image_denoise_power(self, diff):
        self.image_denoise_power = np.clip ( self.image_denoise_power+diff, 0, 500)

    def add_bicubic_degrade_power(self, diff):
        self.bicubic_degrade_power = np.clip ( self.bicubic_degrade_power+diff, 0, 100)

    def ask_settings(self):
        s = """Choose mode: \n"""
        for key in mode_dict.keys():
            s += f"""({key}) {mode_dict[key]}\n"""
        io.log_info(s)
        mode = io.input_int ("", mode_str_dict.get(self.default_mode, 1) )

        self.mode = mode_dict.get (mode, self.default_mode )

        if 'raw' not in self.mode:
            if self.mode == 'hist-match':
                self.masked_hist_match = io.input_bool("Masked hist match?", True)

            if self.mode == 'hist-match' or self.mode == 'seamless-hist-match':
                self.hist_match_threshold = np.clip ( io.input_int("Hist match threshold", 255, add_info="0..255"), 0, 255)

        s = """Choose mask mode: \n"""
        for key in mask_mode_dict.keys():
            s += f"""({key}) {mask_mode_dict[key]}\n"""
        io.log_info(s)
        self.mask_mode = io.input_int ("", 1, valid_list=mask_mode_dict.keys() )

        if 'raw' not in self.mode:
            self.erode_mask_modifier = np.clip ( io.input_int ("Choose erode mask modifier", 0, add_info="-400..400"), -400, 400)
            self.blur_mask_modifier =  np.clip ( io.input_int ("Choose blur mask modifier", 0, add_info="0..400"), 0, 400)
            self.motion_blur_power = np.clip ( io.input_int ("Choose motion blur power", 0, add_info="0..100"), 0, 100)

        self.output_face_scale = np.clip (io.input_int ("Choose output face scale modifier", 0, add_info="-50..50" ), -50, 50)

        if 'raw' not in self.mode:
            self.color_transfer_mode = io.input_str ( "Color transfer to predicted face", None, valid_list=list(ctm_str_dict.keys())[1:] )
            self.color_transfer_mode = ctm_str_dict[self.color_transfer_mode]

        super().ask_settings()

        self.super_resolution_power = np.clip ( io.input_int ("Choose super resolution power", 0, add_info="0..100", help_message="Enhance details by applying superresolution network."), 0, 100)

        if 'raw' not in self.mode:
            self.image_denoise_power = np.clip ( io.input_int ("Choose image degrade by denoise power", 0, add_info="0..500"), 0, 500)
            self.bicubic_degrade_power = np.clip ( io.input_int ("Choose image degrade by bicubic rescale power", 0, add_info="0..100"), 0, 100)
            self.color_degrade_power = np.clip (  io.input_int ("Degrade color power of final image", 0, add_info="0..100"), 0, 100)

        io.log_info ("")

    def __eq__(self, other):
        #check equality of changeable params

        if isinstance(other, MergerConfigMasked):
            return super().__eq__(other) and \
                   self.mode == other.mode and \
                   self.masked_hist_match == other.masked_hist_match and \
                   self.hist_match_threshold == other.hist_match_threshold and \
                   self.mask_mode == other.mask_mode and \
                   self.erode_mask_modifier == other.erode_mask_modifier and \
                   self.blur_mask_modifier == other.blur_mask_modifier and \
                   self.motion_blur_power == other.motion_blur_power and \
                   self.output_face_scale == other.output_face_scale and \
                   self.color_transfer_mode == other.color_transfer_mode and \
                   self.super_resolution_power == other.super_resolution_power and \
                   self.image_denoise_power == other.image_denoise_power and \
                   self.bicubic_degrade_power == other.bicubic_degrade_power and \
                   self.color_degrade_power == other.color_degrade_power

        return False

    def to_string(self, filename):
        r = (
            f"""MergerConfig {filename}:\n"""
            f"""Mode: {self.mode}\n"""
            )

        if self.mode == 'hist-match':
            r += f"""masked_hist_match: {self.masked_hist_match}\n"""

        if self.mode == 'hist-match' or self.mode == 'seamless-hist-match':
            r += f"""hist_match_threshold: {self.hist_match_threshold}\n"""

        r += f"""mask_mode: { mask_mode_dict[self.mask_mode] }\n"""

        if 'raw' not in self.mode:
            r += (f"""erode_mask_modifier: {self.erode_mask_modifier}\n"""
                  f"""blur_mask_modifier: {self.blur_mask_modifier}\n"""
                  f"""motion_blur_power: {self.motion_blur_power}\n""")

        r += f"""output_face_scale: {self.output_face_scale}\n"""

        if 'raw' not in self.mode:
            r += f"""color_transfer_mode: {ctm_dict[self.color_transfer_mode]}\n"""
            r += super().to_string(filename)

        r += f"""super_resolution_power: {self.super_resolution_power}\n"""

        if 'raw' not in self.mode:
            r += (f"""image_denoise_power: {self.image_denoise_power}\n"""
                  f"""bicubic_degrade_power: {self.bicubic_degrade_power}\n"""
                  f"""color_degrade_power: {self.color_degrade_power}\n""")

        r += "================"

        return r


class MergerConfigFaceAvatar(MergerConfig):

    def __init__(self, temporal_face_count=0,
                       add_source_image=False):
        super().__init__(type=MergerConfig.TYPE_FACE_AVATAR)
        self.temporal_face_count = temporal_face_count

        #changeable params
        self.add_source_image = add_source_image

    def copy(self):
        return copy.copy(self)

    #override
    def ask_settings(self):
        self.add_source_image = io.input_bool("Add source image?", False, help_message="Add source image for comparison.")
        super().ask_settings()

    def toggle_add_source_image(self):
        self.add_source_image = not self.add_source_image

    #override
    def __eq__(self, other):
        #check equality of changeable params

        if isinstance(other, MergerConfigFaceAvatar):
            return super().__eq__(other) and \
                   self.add_source_image == other.add_source_image

        return False

    #override
    def to_string(self, filename):
        return (f"MergerConfig {filename}:\n"
                f"add_source_image : {self.add_source_image}\n") + \
                super().to_string(filename) + "================"


merger/MergerConfig.py
oo92/TorchFaceLab
import math
from pathlib import Path

import numpy as np

from core import imagelib
from core.interact import interact as io
from core.cv2ex import *
from core import osex


class ScreenAssets(object):
    waiting_icon_image = cv2_imread ( str(Path(__file__).parent / 'gfx' / 'sand_clock_64.png') )

    @staticmethod
    def build_checkerboard_a( sh, size=5):
        h,w = sh[0], sh[1]
        tile = np.array([[0,1],[1,0]]).repeat(size, axis=0).repeat(size, axis=1)
        grid = np.tile(tile,(int(math.ceil((h+0.0)/(2*size))),int(math.ceil((w+0.0)/(2*size)))))
        return grid[:h,:w,None]

class Screen(object):
    def __init__(self, initial_scale_to_width=0, initial_scale_to_height=0, image=None, waiting_icon=False, **kwargs):
        self.initial_scale_to_width = initial_scale_to_width
        self.initial_scale_to_height = initial_scale_to_height
        self.image = image
        self.waiting_icon = waiting_icon

        self.state = -1
        self.scale = 1
        self.force_update = True
        self.is_first_appear = True
        self.show_checker_board = False

        self.last_screen_shape = (480,640,3)
        self.checkerboard_image = None
        self.set_image (image)
        self.scrn_manager = None

    def set_waiting_icon(self, b):
        self.waiting_icon = b

    def toggle_show_checker_board(self):
        self.show_checker_board = not self.show_checker_board
        self.force_update = True
    
    def get_image(self):
        return self.image
        
    def set_image(self, img):
        if not img is self.image:
            self.force_update = True

        self.image = img

        if self.image is not None:
            self.last_screen_shape = self.image.shape

            if self.initial_scale_to_width != 0:
                if self.last_screen_shape[1] > self.initial_scale_to_width:
                    self.scale = self.initial_scale_to_width / self.last_screen_shape[1]
                    self.force_update = True
                self.initial_scale_to_width = 0

            elif self.initial_scale_to_height != 0:
                if self.last_screen_shape[0] > self.initial_scale_to_height:
                    self.scale = self.initial_scale_to_height / self.last_screen_shape[0]
                    self.force_update = True
                self.initial_scale_to_height = 0


    def diff_scale(self, diff):
        self.scale = np.clip (self.scale + diff, 0.1, 4.0)
        self.force_update = True

    def show(self, force=False):
        new_state = 0 | self.waiting_icon

        if self.state != new_state or self.force_update or force:
            self.state = new_state
            self.force_update = False

            if self.image is None:
                screen = np.zeros ( self.last_screen_shape, dtype=np.uint8 )
            else:
                screen = self.image.copy()

            if self.waiting_icon:
                imagelib.overlay_alpha_image (screen, ScreenAssets.waiting_icon_image, (0,0) )

            h,w,c = screen.shape
            if self.scale != 1.0:
                screen = cv2.resize ( screen, ( int(w*self.scale), int(h*self.scale) ) )

            if c == 4:
                if not self.show_checker_board:
                    screen = screen[...,0:3]
                else:
                    if self.checkerboard_image is None or self.checkerboard_image.shape[0:2] != screen.shape[0:2]:
                        self.checkerboard_image = ScreenAssets.build_checkerboard_a(screen.shape)

                    screen = screen[...,0:3]*0.75 + 64*self.checkerboard_image*(1- (screen[...,3:4].astype(np.float32)/255.0) )
                    screen = screen.astype(np.uint8)

            io.show_image(self.scrn_manager.wnd_name, screen)

            if self.is_first_appear:
                self.is_first_appear = False
                #center window
                desktop_w, desktop_h = osex.get_screen_size()
                h,w,c = screen.shape
                cv2.moveWindow(self.scrn_manager.wnd_name, max(0,(desktop_w-w) // 2), max(0, (desktop_h-h) // 2) )

            io.process_messages(0.0001)

class ScreenManager(object):
    def __init__(self, window_name="ScreenManager", screens=None, capture_keys=False ):
        self.screens = screens or []
        self.current_screen_id = 0

        if self.screens is not None:
            for screen in self.screens:
                screen.scrn_manager = self

        self.wnd_name = window_name
        io.named_window(self.wnd_name)


        if capture_keys:
            io.capture_keys(self.wnd_name)

    def finalize(self):
        io.destroy_all_windows()

    def get_key_events(self):
        return io.get_key_events(self.wnd_name)

    def switch_screens(self):
        self.current_screen_id = (self.current_screen_id + 1) % len(self.screens)
        self.screens[self.current_screen_id].show(force=True)

    def show_current(self):
        self.screens[self.current_screen_id].show()

    def get_current(self):
        return self.screens[self.current_screen_id]

    def set_current(self, screen):
        self.current_screen_id = self.screens.index(screen)

merger/MergerScreen/MergerScreen.py
oo92/TorchFaceLab
from .MergerScreen import Screen, ScreenManager
merger/MergerScreen/__init__.py
oo92/TorchFaceLab
from .FrameInfo import FrameInfo
from .MergerConfig import MergerConfig, MergerConfigMasked, MergerConfigFaceAvatar
from .MergeMasked import MergeMasked
from .MergeAvatar import MergeFaceAvatar
from .InteractiveMergerSubprocessor import InteractiveMergerSubprocessor
merger/__init__.py
oo92/TorchFaceLab
import colorsys
import inspect
import json
import multiprocessing
import operator
import os
import pickle
import shutil
import tempfile
import time
from pathlib import Path

import cv2
import numpy as np

from core import imagelib, pathex
from core.cv2ex import *
from core.interact import interact as io
from core.leras import nn
from samplelib import SampleGeneratorBase


class ModelBase(object):
    def __init__(self, is_training=False,
                       is_exporting=False,
                       saved_models_path=None,
                       training_data_src_path=None,
                       training_data_dst_path=None,
                       pretraining_data_path=None,
                       pretrained_model_path=None,
                       no_preview=False,
                       force_model_name=None,
                       force_gpu_idxs=None,
                       cpu_only=False,
                       debug=False,
                       force_model_class_name=None,
                       silent_start=False,
                       **kwargs):
        self.is_training = is_training
        self.is_exporting = is_exporting
        self.saved_models_path = saved_models_path
        self.training_data_src_path = training_data_src_path
        self.training_data_dst_path = training_data_dst_path
        self.pretraining_data_path = pretraining_data_path
        self.pretrained_model_path = pretrained_model_path
        self.no_preview = no_preview
        self.debug = debug

        self.model_class_name = model_class_name = Path(inspect.getmodule(self).__file__).parent.name.rsplit("_", 1)[1]

        if force_model_class_name is None:
            if force_model_name is not None:
                self.model_name = force_model_name
            else:
                while True:
                    # gather all model dat files
                    saved_models_names = []
                    for filepath in pathex.get_file_paths(saved_models_path):
                        filepath_name = filepath.name
                        if filepath_name.endswith(f'{model_class_name}_data.dat'):
                            saved_models_names += [ (filepath_name.split('_')[0], os.path.getmtime(filepath)) ]

                    # sort by modified datetime
                    saved_models_names = sorted(saved_models_names, key=operator.itemgetter(1), reverse=True )
                    saved_models_names = [ x[0] for x in saved_models_names ]


                    if len(saved_models_names) != 0:
                        if silent_start:
                            self.model_name = saved_models_names[0]
                            io.log_info(f'Silent start: choosed model "{self.model_name}"')
                        else:
                            io.log_info ("Choose one of saved models, or enter a name to create a new model.")
                            io.log_info ("[r] : rename")
                            io.log_info ("[d] : delete")
                            io.log_info ("")
                            for i, model_name in enumerate(saved_models_names):
                                s = f"[{i}] : {model_name} "
                                if i == 0:
                                    s += "- latest"
                                io.log_info (s)

                            inp = io.input_str(f"", "0", show_default_value=False )
                            model_idx = -1
                            try:
                                model_idx = np.clip ( int(inp), 0, len(saved_models_names)-1 )
                            except:
                                pass

                            if model_idx == -1:
                                if len(inp) == 1:
                                    is_rename = inp[0] == 'r'
                                    is_delete = inp[0] == 'd'

                                    if is_rename or is_delete:
                                        if len(saved_models_names) != 0:

                                            if is_rename:
                                                name = io.input_str(f"Enter the name of the model you want to rename")
                                            elif is_delete:
                                                name = io.input_str(f"Enter the name of the model you want to delete")

                                            if name in saved_models_names:

                                                if is_rename:
                                                    new_model_name = io.input_str(f"Enter new name of the model")

                                                for filepath in pathex.get_paths(saved_models_path):
                                                    filepath_name = filepath.name

                                                    model_filename, remain_filename = filepath_name.split('_', 1)
                                                    if model_filename == name:

                                                        if is_rename:
                                                            new_filepath = filepath.parent / ( new_model_name + '_' + remain_filename )
                                                            filepath.rename (new_filepath)
                                                        elif is_delete:
                                                            filepath.unlink()
                                        continue

                                self.model_name = inp
                            else:
                                self.model_name = saved_models_names[model_idx]

                    else:
                        self.model_name = io.input_str(f"No saved models found. Enter a name of a new model", "new")
                        self.model_name = self.model_name.replace('_', ' ')
                    break


            self.model_name = self.model_name + '_' + self.model_class_name
        else:
            self.model_name = force_model_class_name

        self.iter = 0
        self.options = {}
        self.options_show_override = {}
        self.loss_history = []
        self.sample_for_preview = None
        self.choosed_gpu_indexes = None

        model_data = {}
        self.model_data_path = Path( self.get_strpath_storage_for_file('data.dat') )
        if self.model_data_path.exists():
            io.log_info (f"Loading {self.model_name} model...")
            model_data = pickle.loads ( self.model_data_path.read_bytes() )
            self.iter = model_data.get('iter',0)
            if self.iter != 0:
                self.options = model_data['options']
                self.loss_history = model_data.get('loss_history', [])
                self.sample_for_preview = model_data.get('sample_for_preview', None)
                self.choosed_gpu_indexes = model_data.get('choosed_gpu_indexes', None)

        if self.is_first_run():
            io.log_info ("\nModel first run.")

        if silent_start:
            self.device_config = nn.DeviceConfig.BestGPU()
            io.log_info (f"Silent start: choosed device {'CPU' if self.device_config.cpu_only else self.device_config.devices[0].name}")
        else:
            self.device_config = nn.DeviceConfig.GPUIndexes( force_gpu_idxs or nn.ask_choose_device_idxs(suggest_best_multi_gpu=True)) \
                                if not cpu_only else nn.DeviceConfig.CPU()

        nn.initialize(self.device_config)

        ####
        self.default_options_path = saved_models_path / f'{self.model_class_name}_default_options.dat'
        self.default_options = {}
        if self.default_options_path.exists():
            try:
                self.default_options = pickle.loads ( self.default_options_path.read_bytes() )
            except:
                pass

        self.choose_preview_history = False
        self.batch_size = self.load_or_def_option('batch_size', 1)
        #####

        io.input_skip_pending()
        self.on_initialize_options()

        if self.is_first_run():
            # save as default options only for first run model initialize
            self.default_options_path.write_bytes( pickle.dumps (self.options) )

        self.autobackup_hour = self.options.get('autobackup_hour', 0)
        self.write_preview_history = self.options.get('write_preview_history', False)
        self.target_iter = self.options.get('target_iter',0)
        self.random_flip = self.options.get('random_flip',True)
        self.random_src_flip = self.options.get('random_src_flip', False)
        self.random_dst_flip = self.options.get('random_dst_flip', True)
        
        self.on_initialize()
        self.options['batch_size'] = self.batch_size

        self.preview_history_writer = None
        if self.is_training:
            self.preview_history_path = self.saved_models_path / ( f'{self.get_model_name()}_history' )
            self.autobackups_path     = self.saved_models_path / ( f'{self.get_model_name()}_autobackups' )

            if self.write_preview_history or io.is_colab():
                if not self.preview_history_path.exists():
                    self.preview_history_path.mkdir(exist_ok=True)
                else:
                    if self.iter == 0:
                        for filename in pathex.get_image_paths(self.preview_history_path):
                            Path(filename).unlink()

            if self.generator_list is None:
                raise ValueError( 'You didnt set_training_data_generators()')
            else:
                for i, generator in enumerate(self.generator_list):
                    if not isinstance(generator, SampleGeneratorBase):
                        raise ValueError('training data generator is not subclass of SampleGeneratorBase')

            self.update_sample_for_preview(choose_preview_history=self.choose_preview_history)

            if self.autobackup_hour != 0:
                self.autobackup_start_time = time.time()

                if not self.autobackups_path.exists():
                    self.autobackups_path.mkdir(exist_ok=True)

        io.log_info( self.get_summary_text() )

    def update_sample_for_preview(self, choose_preview_history=False, force_new=False):
        if self.sample_for_preview is None or choose_preview_history or force_new:
            if choose_preview_history and io.is_support_windows():
                wnd_name = "[p] - next. [space] - switch preview type. [enter] - confirm."
                io.log_info (f"Choose image for the preview history. {wnd_name}")
                io.named_window(wnd_name)
                io.capture_keys(wnd_name)
                choosed = False
                preview_id_counter = 0
                while not choosed:
                    self.sample_for_preview = self.generate_next_samples()
                    previews = self.get_history_previews()

                    io.show_image( wnd_name, ( previews[preview_id_counter % len(previews) ][1] *255).astype(np.uint8) )

                    while True:
                        key_events = io.get_key_events(wnd_name)
                        key, chr_key, ctrl_pressed, alt_pressed, shift_pressed = key_events[-1] if len(key_events) > 0 else (0,0,False,False,False)
                        if key == ord('\n') or key == ord('\r'):
                            choosed = True
                            break
                        elif key == ord(' '):
                            preview_id_counter += 1
                            break
                        elif key == ord('p'):
                            break

                        try:
                            io.process_messages(0.1)
                        except KeyboardInterrupt:
                            choosed = True

                io.destroy_window(wnd_name)
            else:
                self.sample_for_preview = self.generate_next_samples()

        try:
            self.get_history_previews()
        except:
            self.sample_for_preview = self.generate_next_samples()

        self.last_sample = self.sample_for_preview

    def load_or_def_option(self, name, def_value):
        options_val = self.options.get(name, None)
        if options_val is not None:
            return options_val

        def_opt_val = self.default_options.get(name, None)
        if def_opt_val is not None:
            return def_opt_val

        return def_value

    def ask_override(self):
        return self.is_training and self.iter != 0 and io.input_in_time ("Press enter in 2 seconds to override model settings.", 5 if io.is_colab() else 2 )

    def ask_autobackup_hour(self, default_value=0):
        default_autobackup_hour = self.options['autobackup_hour'] = self.load_or_def_option('autobackup_hour', default_value)
        self.options['autobackup_hour'] = io.input_int(f"Autobackup every N hour", default_autobackup_hour, add_info="0..24", help_message="Autobackup model files with preview every N hour. Latest backup located in model/<>_autobackups/01")

    def ask_write_preview_history(self, default_value=False):
        default_write_preview_history = self.load_or_def_option('write_preview_history', default_value)
        self.options['write_preview_history'] = io.input_bool(f"Write preview history", default_write_preview_history, help_message="Preview history will be writed to <ModelName>_history folder.")

        if self.options['write_preview_history']:
            if io.is_support_windows():
                self.choose_preview_history = io.input_bool("Choose image for the preview history", False)
            elif io.is_colab():
                self.choose_preview_history = io.input_bool("Randomly choose new image for preview history", False, help_message="Preview image history will stay stuck with old faces if you reuse the same model on different celebs. Choose no unless you are changing src/dst to a new person")

    def ask_target_iter(self, default_value=0):
        default_target_iter = self.load_or_def_option('target_iter', default_value)
        self.options['target_iter'] = max(0, io.input_int("Target iteration", default_target_iter))

    def ask_random_flip(self):
        default_random_flip = self.load_or_def_option('random_flip', True)
        self.options['random_flip'] = io.input_bool("Flip faces randomly", default_random_flip, help_message="Predicted face will look more naturally without this option, but src faceset should cover all face directions as dst faceset.")
    
    def ask_random_src_flip(self):
        default_random_src_flip = self.load_or_def_option('random_src_flip', False)
        self.options['random_src_flip'] = io.input_bool("Flip SRC faces randomly", default_random_src_flip, help_message="Random horizontal flip SRC faceset. Covers more angles, but the face may look less naturally.")

    def ask_random_dst_flip(self):
        default_random_dst_flip = self.load_or_def_option('random_dst_flip', True)
        self.options['random_dst_flip'] = io.input_bool("Flip DST faces randomly", default_random_dst_flip, help_message="Random horizontal flip DST faceset. Makes generalization of src->dst better, if src random flip is not enabled.")

    def ask_batch_size(self, suggest_batch_size=None, range=None):
        default_batch_size = self.load_or_def_option('batch_size', suggest_batch_size or self.batch_size)

        batch_size = max(0, io.input_int("Batch_size", default_batch_size, valid_range=range, help_message="Larger batch size is better for NN's generalization, but it can cause Out of Memory error. Tune this value for your videocard manually."))

        if range is not None:
            batch_size = np.clip(batch_size, range[0], range[1])

        self.options['batch_size'] = self.batch_size = batch_size


    #overridable
    def on_initialize_options(self):
        pass

    #overridable
    def on_initialize(self):
        '''
        initialize your models

        store and retrieve your model options in self.options['']

        check example
        '''
        pass

    #overridable
    def onSave(self):
        #save your models here
        pass

    #overridable
    def onTrainOneIter(self, sample, generator_list):
        #train your models here

        #return array of losses
        return ( ('loss_src', 0), ('loss_dst', 0) )

    #overridable
    def onGetPreview(self, sample, for_history=False):
        #you can return multiple previews
        #return [ ('preview_name',preview_rgb), ... ]
        return []

    #overridable if you want model name differs from folder name
    def get_model_name(self):
        return self.model_name

    #overridable , return [ [model, filename],... ]  list
    def get_model_filename_list(self):
        return []

    #overridable
    def get_MergerConfig(self):
        #return predictor_func, predictor_input_shape, MergerConfig() for the model
        raise NotImplementedError

    def get_pretraining_data_path(self):
        return self.pretraining_data_path

    def get_target_iter(self):
        return self.target_iter

    def is_reached_iter_goal(self):
        return self.target_iter != 0 and self.iter >= self.target_iter

    def get_previews(self):
        return self.onGetPreview ( self.last_sample )

    def get_history_previews(self):
        return self.onGetPreview (self.sample_for_preview, for_history=True)

    def get_preview_history_writer(self):
        if self.preview_history_writer is None:
            self.preview_history_writer = PreviewHistoryWriter()
        return self.preview_history_writer

    def save(self):
        Path( self.get_summary_path() ).write_text( self.get_summary_text() )

        self.onSave()

        model_data = {
            'iter': self.iter,
            'options': self.options,
            'loss_history': self.loss_history,
            'sample_for_preview' : self.sample_for_preview,
            'choosed_gpu_indexes' : self.choosed_gpu_indexes,
        }
        pathex.write_bytes_safe (self.model_data_path, pickle.dumps(model_data) )

        if self.autobackup_hour != 0:
            diff_hour = int ( (time.time() - self.autobackup_start_time) // 3600 )

            if diff_hour > 0 and diff_hour % self.autobackup_hour == 0:
                self.autobackup_start_time += self.autobackup_hour*3600
                self.create_backup()

    def create_backup(self):
        io.log_info ("Creating backup...", end='\r')

        if not self.autobackups_path.exists():
            self.autobackups_path.mkdir(exist_ok=True)

        bckp_filename_list = [ self.get_strpath_storage_for_file(filename) for _, filename in self.get_model_filename_list() ]
        bckp_filename_list += [ str(self.get_summary_path()), str(self.model_data_path) ]

        for i in range(24,0,-1):
            idx_str = '%.2d' % i
            next_idx_str = '%.2d' % (i+1)

            idx_backup_path = self.autobackups_path / idx_str
            next_idx_packup_path = self.autobackups_path / next_idx_str

            if idx_backup_path.exists():
                if i == 24:
                    pathex.delete_all_files(idx_backup_path)
                else:
                    next_idx_packup_path.mkdir(exist_ok=True)
                    pathex.move_all_files (idx_backup_path, next_idx_packup_path)

            if i == 1:
                idx_backup_path.mkdir(exist_ok=True)
                for filename in bckp_filename_list:
                    shutil.copy ( str(filename), str(idx_backup_path / Path(filename).name) )

                previews = self.get_previews()
                plist = []
                for i in range(len(previews)):
                    name, bgr = previews[i]
                    plist += [ (bgr, idx_backup_path / ( ('preview_%s.jpg') % (name))  )  ]

                if len(plist) != 0:
                    self.get_preview_history_writer().post(plist, self.loss_history, self.iter)

    def debug_one_iter(self):
        images = []
        for generator in self.generator_list:
            for i,batch in enumerate(next(generator)):
                if len(batch.shape) == 4:
                    images.append( batch[0] )

        return imagelib.equalize_and_stack_square (images)

    def generate_next_samples(self):
        sample = []
        for generator in self.generator_list:
            if generator.is_initialized():
                sample.append ( generator.generate_next() )
            else:
                sample.append ( [] )
        self.last_sample = sample
        return sample

    #overridable
    def should_save_preview_history(self):
        return (not io.is_colab() and self.iter % 10 == 0) or (io.is_colab() and self.iter % 100 == 0)

    def train_one_iter(self):

        iter_time = time.time()
        losses = self.onTrainOneIter()
        iter_time = time.time() - iter_time

        self.loss_history.append ( [float(loss[1]) for loss in losses] )

        if self.should_save_preview_history():
            plist = []

            if io.is_colab():
                previews = self.get_previews()
                for i in range(len(previews)):
                    name, bgr = previews[i]
                    plist += [ (bgr, self.get_strpath_storage_for_file('preview_%s.jpg' % (name) ) ) ]

            if self.write_preview_history:
                previews = self.get_history_previews()
                for i in range(len(previews)):
                    name, bgr = previews[i]
                    path = self.preview_history_path / name
                    plist += [ ( bgr, str ( path / ( f'{self.iter:07d}.jpg') ) ) ]
                    if not io.is_colab():
                        plist += [ ( bgr, str ( path / ( '_last.jpg' ) )) ]

            if len(plist) != 0:
                self.get_preview_history_writer().post(plist, self.loss_history, self.iter)

        self.iter += 1

        return self.iter, iter_time

    def pass_one_iter(self):
        self.generate_next_samples()

    def finalize(self):
        nn.close_session()

    def is_first_run(self):
        return self.iter == 0

    def is_debug(self):
        return self.debug

    def set_batch_size(self, batch_size):
        self.batch_size = batch_size

    def get_batch_size(self):
        return self.batch_size

    def get_iter(self):
        return self.iter

    def set_iter(self, iter):
        self.iter = iter
        self.loss_history = self.loss_history[:iter]

    def get_loss_history(self):
        return self.loss_history

    def set_training_data_generators (self, generator_list):
        self.generator_list = generator_list

    def get_training_data_generators (self):
        return self.generator_list

    def get_model_root_path(self):
        return self.saved_models_path

    def get_strpath_storage_for_file(self, filename):
        return str( self.saved_models_path / ( self.get_model_name() + '_' + filename) )

    def get_summary_path(self):
        return self.get_strpath_storage_for_file('summary.txt')

    def get_summary_text(self):
        visible_options = self.options.copy()
        visible_options.update(self.options_show_override)
        
        ###Generate text summary of model hyperparameters
        #Find the longest key name and value string. Used as column widths.
        width_name = max([len(k) for k in visible_options.keys()] + [17]) + 1 # Single space buffer to left edge. Minimum of 17, the length of the longest static string used "Current iteration"
        width_value = max([len(str(x)) for x in visible_options.values()] + [len(str(self.get_iter())), len(self.get_model_name())]) + 1 # Single space buffer to right edge
        if len(self.device_config.devices) != 0: #Check length of GPU names
            width_value = max([len(device.name)+1 for device in self.device_config.devices] + [width_value])
        width_total = width_name + width_value + 2 #Plus 2 for ": "

        summary_text = []
        summary_text += [f'=={" Model Summary ":=^{width_total}}=='] # Model/status summary
        summary_text += [f'=={" "*width_total}==']
        summary_text += [f'=={"Model name": >{width_name}}: {self.get_model_name(): <{width_value}}=='] # Name
        summary_text += [f'=={" "*width_total}==']
        summary_text += [f'=={"Current iteration": >{width_name}}: {str(self.get_iter()): <{width_value}}=='] # Iter
        summary_text += [f'=={" "*width_total}==']

        summary_text += [f'=={" Model Options ":-^{width_total}}=='] # Model options
        summary_text += [f'=={" "*width_total}==']
        for key in visible_options.keys():
            summary_text += [f'=={key: >{width_name}}: {str(visible_options[key]): <{width_value}}=='] # visible_options key/value pairs
        summary_text += [f'=={" "*width_total}==']

        summary_text += [f'=={" Running On ":-^{width_total}}=='] # Training hardware info
        summary_text += [f'=={" "*width_total}==']
        if len(self.device_config.devices) == 0:
            summary_text += [f'=={"Using device": >{width_name}}: {"CPU": <{width_value}}=='] # cpu_only
        else:
            for device in self.device_config.devices:
                summary_text += [f'=={"Device index": >{width_name}}: {device.index: <{width_value}}=='] # GPU hardware device index
                summary_text += [f'=={"Name": >{width_name}}: {device.name: <{width_value}}=='] # GPU name
                vram_str = f'{device.total_mem_gb:.2f}GB' # GPU VRAM - Formated as #.## (or ##.##)
                summary_text += [f'=={"VRAM": >{width_name}}: {vram_str: <{width_value}}==']
        summary_text += [f'=={" "*width_total}==']
        summary_text += [f'=={"="*width_total}==']
        summary_text = "\n".join (summary_text)
        return summary_text

    @staticmethod
    def get_loss_history_preview(loss_history, iter, w, c):
        loss_history = np.array (loss_history.copy())

        lh_height = 100
        lh_img = np.ones ( (lh_height,w,c) ) * 0.1

        if len(loss_history) != 0:
            loss_count = len(loss_history[0])
            lh_len = len(loss_history)

            l_per_col = lh_len / w
            plist_max = [   [   max (0.0, loss_history[int(col*l_per_col)][p],
                                                *[  loss_history[i_ab][p]
                                                    for i_ab in range( int(col*l_per_col), int((col+1)*l_per_col) )
                                                ]
                                    )
                                for p in range(loss_count)
                            ]
                            for col in range(w)
                        ]

            plist_min = [   [   min (plist_max[col][p], loss_history[int(col*l_per_col)][p],
                                                *[  loss_history[i_ab][p]
                                                    for i_ab in range( int(col*l_per_col), int((col+1)*l_per_col) )
                                                ]
                                    )
                                for p in range(loss_count)
                            ]
                            for col in range(w)
                        ]

            plist_abs_max = np.mean(loss_history[ len(loss_history) // 5 : ]) * 2

            for col in range(0, w):
                for p in range(0,loss_count):
                    point_color = [1.0]*c
                    point_color[0:3] = colorsys.hsv_to_rgb ( p * (1.0/loss_count), 1.0, 1.0 )

                    ph_max = int ( (plist_max[col][p] / plist_abs_max) * (lh_height-1) )
                    ph_max = np.clip( ph_max, 0, lh_height-1 )

                    ph_min = int ( (plist_min[col][p] / plist_abs_max) * (lh_height-1) )
                    ph_min = np.clip( ph_min, 0, lh_height-1 )

                    for ph in range(ph_min, ph_max+1):
                        lh_img[ (lh_height-ph-1), col ] = point_color

        lh_lines = 5
        lh_line_height = (lh_height-1)/lh_lines
        for i in range(0,lh_lines+1):
            lh_img[ int(i*lh_line_height), : ] = (0.8,)*c

        last_line_t = int((lh_lines-1)*lh_line_height)
        last_line_b = int(lh_lines*lh_line_height)

        lh_text = 'Iter: %d' % (iter) if iter != 0 else ''

        lh_img[last_line_t:last_line_b, 0:w] += imagelib.get_text_image (  (last_line_b-last_line_t,w,c), lh_text, color=[0.8]*c )
        return lh_img

class PreviewHistoryWriter():
    def __init__(self):
        self.sq = multiprocessing.Queue()
        self.p = multiprocessing.Process(target=self.process, args=( self.sq, ))
        self.p.daemon = True
        self.p.start()

    def process(self, sq):
        while True:
            while not sq.empty():
                plist, loss_history, iter = sq.get()

                preview_lh_cache = {}
                for preview, filepath in plist:
                    filepath = Path(filepath)
                    i = (preview.shape[1], preview.shape[2])

                    preview_lh = preview_lh_cache.get(i, None)
                    if preview_lh is None:
                        preview_lh = ModelBase.get_loss_history_preview(loss_history, iter, preview.shape[1], preview.shape[2])
                        preview_lh_cache[i] = preview_lh

                    img = (np.concatenate ( [preview_lh, preview], axis=0 ) * 255).astype(np.uint8)

                    filepath.parent.mkdir(parents=True, exist_ok=True)
                    cv2_imwrite (filepath, img )

            time.sleep(0.01)

    def post(self, plist, loss_history, iter):
        self.sq.put ( (plist, loss_history, iter) )

    # disable pickling
    def __getstate__(self):
        return dict()
    def __setstate__(self, d):
        self.__dict__.update(d)

models/ModelBase.py
oo92/TorchFaceLab
import multiprocessing
import operator
from functools import partial

import numpy as np

from core import mathlib
from core.interact import interact as io
from core.leras import nn
from facelib import FaceType
from models import ModelBase
from samplelib import *
from core.cv2ex import *

class AMPModel(ModelBase):

    #override
    def on_initialize_options(self):
        default_resolution         = self.options['resolution']         = self.load_or_def_option('resolution', 224)
        default_face_type          = self.options['face_type']          = self.load_or_def_option('face_type', 'wf')
        default_models_opt_on_gpu  = self.options['models_opt_on_gpu']  = self.load_or_def_option('models_opt_on_gpu', True)

        default_ae_dims            = self.options['ae_dims']            = self.load_or_def_option('ae_dims', 256)
        default_inter_dims         = self.options['inter_dims']         = self.load_or_def_option('inter_dims', 1024)

        default_e_dims             = self.options['e_dims']             = self.load_or_def_option('e_dims', 64)
        default_d_dims             = self.options['d_dims']             = self.options.get('d_dims', None)
        default_d_mask_dims        = self.options['d_mask_dims']        = self.options.get('d_mask_dims', None)
        default_morph_factor       = self.options['morph_factor']       = self.options.get('morph_factor', 0.5)
        default_uniform_yaw        = self.options['uniform_yaw']        = self.load_or_def_option('uniform_yaw', False)
        default_blur_out_mask      = self.options['blur_out_mask']      = self.load_or_def_option('blur_out_mask', False)
        default_lr_dropout         = self.options['lr_dropout']         = self.load_or_def_option('lr_dropout', 'n')
        default_random_warp        = self.options['random_warp']        = self.load_or_def_option('random_warp', True)
        default_ct_mode            = self.options['ct_mode']            = self.load_or_def_option('ct_mode', 'none')
        default_clipgrad           = self.options['clipgrad']           = self.load_or_def_option('clipgrad', False)

        ask_override = self.ask_override()
        if self.is_first_run() or ask_override:
            self.ask_autobackup_hour()
            self.ask_write_preview_history()
            self.ask_target_iter()
            self.ask_random_src_flip()
            self.ask_random_dst_flip()
            self.ask_batch_size(8)

        if self.is_first_run():
            resolution = io.input_int("Resolution", default_resolution, add_info="64-640", help_message="More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 32 .")
            resolution = np.clip ( (resolution // 32) * 32, 64, 640)
            self.options['resolution'] = resolution
            self.options['face_type'] = io.input_str ("Face type", default_face_type, ['f','wf','head'], help_message="whole face / head").lower()


        default_d_dims             = self.options['d_dims']             = self.load_or_def_option('d_dims', 64)

        default_d_mask_dims        = default_d_dims // 3
        default_d_mask_dims        += default_d_mask_dims % 2
        default_d_mask_dims        = self.options['d_mask_dims']        = self.load_or_def_option('d_mask_dims', default_d_mask_dims)

        if self.is_first_run():
            self.options['ae_dims']    = np.clip ( io.input_int("AutoEncoder dimensions", default_ae_dims, add_info="32-1024", help_message="All face information will packed to AE dims. If amount of AE dims are not enough, then for example closed eyes will not be recognized. More dims are better, but require more VRAM. You can fine-tune model size to fit your GPU." ), 32, 1024 )
            self.options['inter_dims'] = np.clip ( io.input_int("Inter dimensions", default_inter_dims, add_info="32-2048", help_message="Should be equal or more than AutoEncoder dimensions. More dims are better, but require more VRAM. You can fine-tune model size to fit your GPU." ), 32, 2048 )

            e_dims = np.clip ( io.input_int("Encoder dimensions", default_e_dims, add_info="16-256", help_message="More dims help to recognize more facial features and achieve sharper result, but require more VRAM. You can fine-tune model size to fit your GPU." ), 16, 256 )
            self.options['e_dims'] = e_dims + e_dims % 2

            d_dims = np.clip ( io.input_int("Decoder dimensions", default_d_dims, add_info="16-256", help_message="More dims help to recognize more facial features and achieve sharper result, but require more VRAM. You can fine-tune model size to fit your GPU." ), 16, 256 )
            self.options['d_dims'] = d_dims + d_dims % 2

            d_mask_dims = np.clip ( io.input_int("Decoder mask dimensions", default_d_mask_dims, add_info="16-256", help_message="Typical mask dimensions = decoder dimensions / 3. If you manually cut out obstacles from the dst mask, you can increase this parameter to achieve better quality." ), 16, 256 )
            self.options['d_mask_dims'] = d_mask_dims + d_mask_dims % 2

            morph_factor = np.clip ( io.input_number ("Morph factor.", default_morph_factor, add_info="0.1 .. 0.5", help_message="Typical fine value is 0.5"), 0.1, 0.5 )
            self.options['morph_factor'] = morph_factor

        if self.is_first_run() or ask_override:
            self.options['uniform_yaw'] = io.input_bool ("Uniform yaw distribution of samples", default_uniform_yaw, help_message='Helps to fix blurry side faces due to small amount of them in the faceset.')
            self.options['blur_out_mask'] = io.input_bool ("Blur out mask", default_blur_out_mask, help_message='Blurs nearby area outside of applied face mask of training samples. The result is the background near the face is smoothed and less noticeable on swapped face. The exact xseg mask in src and dst faceset is required.')
            self.options['lr_dropout']  = io.input_str (f"Use learning rate dropout", default_lr_dropout, ['n','y','cpu'], help_message="When the face is trained enough, you can enable this option to get extra sharpness and reduce subpixel shake for less amount of iterations. Enabled it before `disable random warp` and before GAN. \nn - disabled.\ny - enabled\ncpu - enabled on CPU. This allows not to use extra VRAM, sacrificing 20% time of iteration.")

        default_gan_power          = self.options['gan_power']          = self.load_or_def_option('gan_power', 0.0)
        default_gan_patch_size     = self.options['gan_patch_size']     = self.load_or_def_option('gan_patch_size', self.options['resolution'] // 8)
        default_gan_dims           = self.options['gan_dims']           = self.load_or_def_option('gan_dims', 16)

        if self.is_first_run() or ask_override:
            self.options['models_opt_on_gpu'] = io.input_bool ("Place models and optimizer on GPU", default_models_opt_on_gpu, help_message="When you train on one GPU, by default model and optimizer weights are placed on GPU to accelerate the process. You can place they on CPU to free up extra VRAM, thus set bigger dimensions.")

            self.options['random_warp'] = io.input_bool ("Enable random warp of samples", default_random_warp, help_message="Random warp is required to generalize facial expressions of both faces. When the face is trained enough, you can disable it to get extra sharpness and reduce subpixel shake for less amount of iterations.")

            self.options['gan_power'] = np.clip ( io.input_number ("GAN power", default_gan_power, add_info="0.0 .. 5.0", help_message="Forces the neural network to learn small details of the face. Enable it only when the face is trained enough with random_warp(off), and don't disable. The higher the value, the higher the chances of artifacts. Typical fine value is 0.1"), 0.0, 5.0 )

            if self.options['gan_power'] != 0.0:
                gan_patch_size = np.clip ( io.input_int("GAN patch size", default_gan_patch_size, add_info="3-640", help_message="The higher patch size, the higher the quality, the more VRAM is required. You can get sharper edges even at the lowest setting. Typical fine value is resolution / 8." ), 3, 640 )
                self.options['gan_patch_size'] = gan_patch_size

                gan_dims = np.clip ( io.input_int("GAN dimensions", default_gan_dims, add_info="4-512", help_message="The dimensions of the GAN network. The higher dimensions, the more VRAM is required. You can get sharper edges even at the lowest setting. Typical fine value is 16." ), 4, 512 )
                self.options['gan_dims'] = gan_dims

            self.options['ct_mode'] = io.input_str (f"Color transfer for src faceset", default_ct_mode, ['none','rct','lct','mkl','idt','sot'], help_message="Change color distribution of src samples close to dst samples. If src faceset is deverse enough, then lct mode is fine in most cases.")
            self.options['clipgrad'] = io.input_bool ("Enable gradient clipping", default_clipgrad, help_message="Gradient clipping reduces chance of model collapse, sacrificing speed of training.")

        self.gan_model_changed = (default_gan_patch_size != self.options['gan_patch_size']) or (default_gan_dims != self.options['gan_dims'])

    #override
    def on_initialize(self):
        device_config = nn.getCurrentDeviceConfig()
        devices = device_config.devices
        self.model_data_format = "NCHW"
        nn.initialize(data_format=self.model_data_format)
        tf = nn.tf

        input_ch=3
        resolution  = self.resolution = self.options['resolution']
        e_dims      = self.options['e_dims']
        ae_dims     = self.options['ae_dims']
        inter_dims  = self.inter_dims = self.options['inter_dims']
        inter_res   = self.inter_res = resolution // 32
        d_dims      = self.options['d_dims']
        d_mask_dims = self.options['d_mask_dims']
        face_type   = self.face_type = {'f'    : FaceType.FULL,
                                        'wf'   : FaceType.WHOLE_FACE,
                                        'head' : FaceType.HEAD}[ self.options['face_type'] ]
        morph_factor = self.options['morph_factor']
        gan_power    = self.gan_power = self.options['gan_power']
        random_warp  = self.options['random_warp']

        blur_out_mask = self.options['blur_out_mask']

        ct_mode = self.options['ct_mode']
        if ct_mode == 'none':
            ct_mode = None

        use_fp16 = False
        if self.is_exporting:
            use_fp16 = io.input_bool ("Export quantized?", False, help_message='Makes the exported model faster. If you have problems, disable this option.')

        conv_dtype = tf.float16 if use_fp16 else tf.float32

        class Downscale(nn.ModelBase):
            def on_build(self, in_ch, out_ch, kernel_size=5 ):
                self.conv1 = nn.Conv2D( in_ch, out_ch, kernel_size=kernel_size, strides=2, padding='SAME', dtype=conv_dtype)

            def forward(self, x):
                return tf.nn.leaky_relu(self.conv1(x), 0.1)

        class Upscale(nn.ModelBase):
            def on_build(self, in_ch, out_ch, kernel_size=3 ):
                self.conv1 = nn.Conv2D(in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype)

            def forward(self, x):
                x = nn.depth_to_space(tf.nn.leaky_relu(self.conv1(x), 0.1), 2)
                return x

        class ResidualBlock(nn.ModelBase):
            def on_build(self, ch, kernel_size=3 ):
                self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype)
                self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype)

            def forward(self, inp):
                x = self.conv1(inp)
                x = tf.nn.leaky_relu(x, 0.2)
                x = self.conv2(x)
                x = tf.nn.leaky_relu(inp+x, 0.2)
                return x

        class Encoder(nn.ModelBase):
            def on_build(self):
                self.down1 = Downscale(input_ch, e_dims, kernel_size=5)
                self.res1 = ResidualBlock(e_dims)
                self.down2 = Downscale(e_dims, e_dims*2, kernel_size=5)
                self.down3 = Downscale(e_dims*2, e_dims*4, kernel_size=5)
                self.down4 = Downscale(e_dims*4, e_dims*8, kernel_size=5)
                self.down5 = Downscale(e_dims*8, e_dims*8, kernel_size=5)
                self.res5 = ResidualBlock(e_dims*8)
                self.dense1 = nn.Dense( (( resolution//(2**5) )**2) * e_dims*8, ae_dims )

            def forward(self, x):
                if use_fp16:
                    x = tf.cast(x, tf.float16)
                x = self.down1(x)
                x = self.res1(x)
                x = self.down2(x)
                x = self.down3(x)
                x = self.down4(x)
                x = self.down5(x)
                x = self.res5(x)
                if use_fp16:
                    x = tf.cast(x, tf.float32)
                x = nn.pixel_norm(nn.flatten(x), axes=-1)
                x = self.dense1(x)
                return x


        class Inter(nn.ModelBase):
            def on_build(self):
                self.dense2 = nn.Dense(ae_dims, inter_res * inter_res * inter_dims)

            def forward(self, inp):
                x = inp
                x = self.dense2(x)
                x = nn.reshape_4D (x, inter_res, inter_res, inter_dims)
                return x


        class Decoder(nn.ModelBase):
            def on_build(self ):
                self.upscale0 = Upscale(inter_dims, d_dims*8, kernel_size=3)
                self.upscale1 = Upscale(d_dims*8, d_dims*8, kernel_size=3)
                self.upscale2 = Upscale(d_dims*8, d_dims*4, kernel_size=3)
                self.upscale3 = Upscale(d_dims*4, d_dims*2, kernel_size=3)

                self.res0 = ResidualBlock(d_dims*8, kernel_size=3)
                self.res1 = ResidualBlock(d_dims*8, kernel_size=3)
                self.res2 = ResidualBlock(d_dims*4, kernel_size=3)
                self.res3 = ResidualBlock(d_dims*2, kernel_size=3)

                self.upscalem0 = Upscale(inter_dims, d_mask_dims*8, kernel_size=3)
                self.upscalem1 = Upscale(d_mask_dims*8, d_mask_dims*8, kernel_size=3)
                self.upscalem2 = Upscale(d_mask_dims*8, d_mask_dims*4, kernel_size=3)
                self.upscalem3 = Upscale(d_mask_dims*4, d_mask_dims*2, kernel_size=3)
                self.upscalem4 = Upscale(d_mask_dims*2, d_mask_dims*1, kernel_size=3)
                self.out_convm = nn.Conv2D( d_mask_dims*1, 1, kernel_size=1, padding='SAME', dtype=conv_dtype)

                self.out_conv  = nn.Conv2D( d_dims*2, 3, kernel_size=1, padding='SAME', dtype=conv_dtype)
                self.out_conv1 = nn.Conv2D( d_dims*2, 3, kernel_size=3, padding='SAME', dtype=conv_dtype)
                self.out_conv2 = nn.Conv2D( d_dims*2, 3, kernel_size=3, padding='SAME', dtype=conv_dtype)
                self.out_conv3 = nn.Conv2D( d_dims*2, 3, kernel_size=3, padding='SAME', dtype=conv_dtype)

            def forward(self, z):
                if use_fp16:
                    z = tf.cast(z, tf.float16)

                x = self.upscale0(z)
                x = self.res0(x)
                x = self.upscale1(x)
                x = self.res1(x)
                x = self.upscale2(x)
                x = self.res2(x)
                x = self.upscale3(x)
                x = self.res3(x)

                x = tf.nn.sigmoid( nn.depth_to_space(tf.concat( (self.out_conv(x),
                                                                 self.out_conv1(x),
                                                                 self.out_conv2(x),
                                                                 self.out_conv3(x)), nn.conv2d_ch_axis), 2) )
                m = self.upscalem0(z)
                m = self.upscalem1(m)
                m = self.upscalem2(m)
                m = self.upscalem3(m)
                m = self.upscalem4(m)
                m = tf.nn.sigmoid(self.out_convm(m))

                if use_fp16:
                    x = tf.cast(x, tf.float32)
                    m = tf.cast(m, tf.float32)
                return x, m

        models_opt_on_gpu = False if len(devices) == 0 else self.options['models_opt_on_gpu']
        models_opt_device = nn.tf_default_device_name if models_opt_on_gpu and self.is_training else '/CPU:0'
        optimizer_vars_on_cpu = models_opt_device=='/CPU:0'

        bgr_shape = self.bgr_shape = nn.get4Dshape(resolution,resolution,input_ch)
        mask_shape = nn.get4Dshape(resolution,resolution,1)
        self.model_filename_list = []

        with tf.device ('/CPU:0'):
            #Place holders on CPU
            self.warped_src = tf.placeholder (nn.floatx, bgr_shape, name='warped_src')
            self.warped_dst = tf.placeholder (nn.floatx, bgr_shape, name='warped_dst')

            self.target_src = tf.placeholder (nn.floatx, bgr_shape, name='target_src')
            self.target_dst = tf.placeholder (nn.floatx, bgr_shape, name='target_dst')

            self.target_srcm    = tf.placeholder (nn.floatx, mask_shape, name='target_srcm')
            self.target_srcm_em = tf.placeholder (nn.floatx, mask_shape, name='target_srcm_em')
            self.target_dstm    = tf.placeholder (nn.floatx, mask_shape, name='target_dstm')
            self.target_dstm_em = tf.placeholder (nn.floatx, mask_shape, name='target_dstm_em')

            self.morph_value_t = tf.placeholder (nn.floatx, (1,), name='morph_value_t')

        # Initializing model classes
        with tf.device (models_opt_device):
            self.encoder = Encoder(name='encoder')
            self.inter_src = Inter(name='inter_src')
            self.inter_dst = Inter(name='inter_dst')
            self.decoder = Decoder(name='decoder')

            self.model_filename_list += [   [self.encoder,  'encoder.npy'],
                                            [self.inter_src, 'inter_src.npy'],
                                            [self.inter_dst , 'inter_dst.npy'],
                                            [self.decoder , 'decoder.npy'] ]

            if self.is_training:
                # Initialize optimizers
                clipnorm = 1.0 if self.options['clipgrad'] else 0.0
                if self.options['lr_dropout'] in ['y','cpu']:
                    lr_cos = 500
                    lr_dropout = 0.3
                else:
                    lr_cos = 0
                    lr_dropout = 1.0
                self.G_weights = self.encoder.get_weights() + self.decoder.get_weights()

                self.src_dst_opt = nn.AdaBelief(lr=5e-5, lr_dropout=lr_dropout, lr_cos=lr_cos, clipnorm=clipnorm, name='src_dst_opt')
                self.src_dst_opt.initialize_variables (self.G_weights, vars_on_cpu=optimizer_vars_on_cpu)
                self.model_filename_list += [ (self.src_dst_opt, 'src_dst_opt.npy') ]

                if gan_power != 0:
                    self.GAN = nn.UNetPatchDiscriminator(patch_size=self.options['gan_patch_size'], in_ch=input_ch, base_ch=self.options['gan_dims'], name="GAN")
                    self.GAN_opt = nn.AdaBelief(lr=5e-5, lr_dropout=lr_dropout, lr_cos=lr_cos, clipnorm=clipnorm, name='GAN_opt')
                    self.GAN_opt.initialize_variables ( self.GAN.get_weights(), vars_on_cpu=optimizer_vars_on_cpu)
                    self.model_filename_list += [ [self.GAN, 'GAN.npy'],
                                                  [self.GAN_opt, 'GAN_opt.npy'] ]

        if self.is_training:
            # Adjust batch size for multiple GPU
            gpu_count = max(1, len(devices) )
            bs_per_gpu = max(1, self.get_batch_size() // gpu_count)
            self.set_batch_size( gpu_count*bs_per_gpu)

            # Compute losses per GPU
            gpu_pred_src_src_list = []
            gpu_pred_dst_dst_list = []
            gpu_pred_src_dst_list = []
            gpu_pred_src_srcm_list = []
            gpu_pred_dst_dstm_list = []
            gpu_pred_src_dstm_list = []

            gpu_src_losses = []
            gpu_dst_losses = []
            gpu_G_loss_gradients = []
            gpu_GAN_loss_gradients = []

            def DLossOnes(logits):
                return tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.ones_like(logits), logits=logits), axis=[1,2,3])

            def DLossZeros(logits):
                return tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.zeros_like(logits), logits=logits), axis=[1,2,3])

            for gpu_id in range(gpu_count):
                with tf.device( f'/{devices[gpu_id].tf_dev_type}:{gpu_id}' if len(devices) != 0 else f'/CPU:0' ):
                    with tf.device(f'/CPU:0'):
                        # slice on CPU, otherwise all batch data will be transfered to GPU first
                        batch_slice = slice( gpu_id*bs_per_gpu, (gpu_id+1)*bs_per_gpu )
                        gpu_warped_src      = self.warped_src [batch_slice,:,:,:]
                        gpu_warped_dst      = self.warped_dst [batch_slice,:,:,:]
                        gpu_target_src      = self.target_src [batch_slice,:,:,:]
                        gpu_target_dst      = self.target_dst [batch_slice,:,:,:]
                        gpu_target_srcm     = self.target_srcm[batch_slice,:,:,:]
                        gpu_target_srcm_em  = self.target_srcm_em[batch_slice,:,:,:]
                        gpu_target_dstm     = self.target_dstm[batch_slice,:,:,:]
                        gpu_target_dstm_em  = self.target_dstm_em[batch_slice,:,:,:]

                    # process model tensors
                    gpu_src_code = self.encoder (gpu_warped_src)
                    gpu_dst_code = self.encoder (gpu_warped_dst)

                    gpu_src_inter_src_code, gpu_src_inter_dst_code = self.inter_src (gpu_src_code), self.inter_dst (gpu_src_code)
                    gpu_dst_inter_src_code, gpu_dst_inter_dst_code = self.inter_src (gpu_dst_code), self.inter_dst (gpu_dst_code)

                    inter_dims_bin = int(inter_dims*morph_factor)
                    with tf.device(f'/CPU:0'):
                        inter_rnd_binomial = tf.stack([tf.random.shuffle(tf.concat([tf.tile(tf.constant([1], tf.float32), ( inter_dims_bin, )),
                                                                                    tf.tile(tf.constant([0], tf.float32), ( inter_dims-inter_dims_bin, ))], 0 )) for _ in range(bs_per_gpu)], 0)

                        inter_rnd_binomial = tf.stop_gradient(inter_rnd_binomial[...,None,None])

                    gpu_src_code = gpu_src_inter_src_code * inter_rnd_binomial + gpu_src_inter_dst_code * (1-inter_rnd_binomial)
                    gpu_dst_code = gpu_dst_inter_dst_code

                    inter_dims_slice = tf.cast(inter_dims*self.morph_value_t[0], tf.int32)
                    gpu_src_dst_code = tf.concat( (tf.slice(gpu_dst_inter_src_code, [0,0,0,0],   [-1, inter_dims_slice , inter_res, inter_res]),
                                                   tf.slice(gpu_dst_inter_dst_code, [0,inter_dims_slice,0,0], [-1,inter_dims-inter_dims_slice, inter_res,inter_res]) ), 1 )

                    gpu_pred_src_src, gpu_pred_src_srcm = self.decoder(gpu_src_code)
                    gpu_pred_dst_dst, gpu_pred_dst_dstm = self.decoder(gpu_dst_code)
                    gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder(gpu_src_dst_code)

                    gpu_pred_src_src_list.append(gpu_pred_src_src), gpu_pred_src_srcm_list.append(gpu_pred_src_srcm)
                    gpu_pred_dst_dst_list.append(gpu_pred_dst_dst), gpu_pred_dst_dstm_list.append(gpu_pred_dst_dstm)
                    gpu_pred_src_dst_list.append(gpu_pred_src_dst), gpu_pred_src_dstm_list.append(gpu_pred_src_dstm)

                    gpu_target_srcm_anti = 1-gpu_target_srcm
                    gpu_target_dstm_anti = 1-gpu_target_dstm

                    gpu_target_srcm_gblur = nn.gaussian_blur(gpu_target_srcm, resolution // 32)
                    gpu_target_dstm_gblur = nn.gaussian_blur(gpu_target_dstm, resolution // 32)

                    gpu_target_srcm_blur = tf.clip_by_value(gpu_target_srcm_gblur, 0, 0.5) * 2
                    gpu_target_dstm_blur = tf.clip_by_value(gpu_target_dstm_gblur, 0, 0.5) * 2
                    gpu_target_srcm_anti_blur = 1.0-gpu_target_srcm_blur
                    gpu_target_dstm_anti_blur = 1.0-gpu_target_dstm_blur

                    if blur_out_mask:
                        sigma = resolution / 128

                        x = nn.gaussian_blur(gpu_target_src*gpu_target_srcm_anti, sigma)
                        y = 1-nn.gaussian_blur(gpu_target_srcm, sigma)
                        y = tf.where(tf.equal(y, 0), tf.ones_like(y), y)
                        gpu_target_src = gpu_target_src*gpu_target_srcm + (x/y)*gpu_target_srcm_anti

                        x = nn.gaussian_blur(gpu_target_dst*gpu_target_dstm_anti, sigma)
                        y = 1-nn.gaussian_blur(gpu_target_dstm, sigma)
                        y = tf.where(tf.equal(y, 0), tf.ones_like(y), y)
                        gpu_target_dst = gpu_target_dst*gpu_target_dstm + (x/y)*gpu_target_dstm_anti

                    gpu_target_src_masked = gpu_target_src*gpu_target_srcm_blur
                    gpu_target_dst_masked = gpu_target_dst*gpu_target_dstm_blur
                    gpu_target_src_anti_masked = gpu_target_src*gpu_target_srcm_anti_blur
                    gpu_target_dst_anti_masked = gpu_target_dst*gpu_target_dstm_anti_blur

                    gpu_pred_src_src_masked = gpu_pred_src_src*gpu_target_srcm_blur
                    gpu_pred_dst_dst_masked = gpu_pred_dst_dst*gpu_target_dstm_blur
                    gpu_pred_src_src_anti_masked = gpu_pred_src_src*gpu_target_srcm_anti_blur
                    gpu_pred_dst_dst_anti_masked = gpu_pred_dst_dst*gpu_target_dstm_anti_blur

                    # Structural loss
                    gpu_src_loss =  tf.reduce_mean (5*nn.dssim(gpu_target_src_masked, gpu_pred_src_src_masked, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1])
                    gpu_src_loss += tf.reduce_mean (5*nn.dssim(gpu_target_src_masked, gpu_pred_src_src_masked, max_val=1.0, filter_size=int(resolution/23.2)), axis=[1])
                    gpu_dst_loss =  tf.reduce_mean (5*nn.dssim(gpu_target_dst_masked, gpu_pred_dst_dst_masked, max_val=1.0, filter_size=int(resolution/11.6) ), axis=[1])
                    gpu_dst_loss += tf.reduce_mean (5*nn.dssim(gpu_target_dst_masked, gpu_pred_dst_dst_masked, max_val=1.0, filter_size=int(resolution/23.2) ), axis=[1])

                    # Pixel loss
                    gpu_src_loss += tf.reduce_mean (10*tf.square(gpu_target_src_masked-gpu_pred_src_src_masked), axis=[1,2,3])
                    gpu_dst_loss += tf.reduce_mean (10*tf.square(gpu_target_dst_masked-gpu_pred_dst_dst_masked), axis=[1,2,3])

                    # Eyes+mouth prio loss
                    gpu_src_loss += tf.reduce_mean (300*tf.abs (gpu_target_src*gpu_target_srcm_em-gpu_pred_src_src*gpu_target_srcm_em), axis=[1,2,3])
                    gpu_dst_loss += tf.reduce_mean (300*tf.abs (gpu_target_dst*gpu_target_dstm_em-gpu_pred_dst_dst*gpu_target_dstm_em), axis=[1,2,3])

                    # Mask loss
                    gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] )
                    gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] )

                    gpu_src_losses += [gpu_src_loss]
                    gpu_dst_losses += [gpu_dst_loss]
                    gpu_G_loss = gpu_src_loss + gpu_dst_loss
                    # dst-dst background weak loss
                    gpu_G_loss += tf.reduce_mean(0.1*tf.square(gpu_pred_dst_dst_anti_masked-gpu_target_dst_anti_masked),axis=[1,2,3] )
                    gpu_G_loss += 0.000001*nn.total_variation_mse(gpu_pred_dst_dst_anti_masked)


                    if gan_power != 0:
                        gpu_pred_src_src_d, gpu_pred_src_src_d2 = self.GAN(gpu_pred_src_src_masked)
                        gpu_pred_dst_dst_d, gpu_pred_dst_dst_d2 = self.GAN(gpu_pred_dst_dst_masked)
                        gpu_target_src_d, gpu_target_src_d2 = self.GAN(gpu_target_src_masked)
                        gpu_target_dst_d, gpu_target_dst_d2 = self.GAN(gpu_target_dst_masked)

                        gpu_GAN_loss = (DLossOnes (gpu_target_src_d)   + DLossOnes (gpu_target_src_d2) + \
                                        DLossZeros(gpu_pred_src_src_d) + DLossZeros(gpu_pred_src_src_d2) + \
                                        DLossOnes (gpu_target_dst_d)   + DLossOnes (gpu_target_dst_d2) + \
                                        DLossZeros(gpu_pred_dst_dst_d) + DLossZeros(gpu_pred_dst_dst_d2)
                                        ) * (1.0 / 8)

                        gpu_GAN_loss_gradients += [ nn.gradients (gpu_GAN_loss, self.GAN.get_weights() ) ]

                        gpu_G_loss += (DLossOnes(gpu_pred_src_src_d) + DLossOnes(gpu_pred_src_src_d2) + \
                                       DLossOnes(gpu_pred_dst_dst_d) + DLossOnes(gpu_pred_dst_dst_d2)
                                      ) * gan_power

                        # Minimal src-src-bg rec with total_variation_mse to suppress random bright dots from gan
                        gpu_G_loss += 0.000001*nn.total_variation_mse(gpu_pred_src_src)
                        gpu_G_loss += 0.02*tf.reduce_mean(tf.square(gpu_pred_src_src_anti_masked-gpu_target_src_anti_masked),axis=[1,2,3] )

                    gpu_G_loss_gradients += [ nn.gradients ( gpu_G_loss, self.G_weights ) ]

            # Average losses and gradients, and create optimizer update ops
            with tf.device(f'/CPU:0'):
                pred_src_src  = nn.concat(gpu_pred_src_src_list, 0)
                pred_dst_dst  = nn.concat(gpu_pred_dst_dst_list, 0)
                pred_src_dst  = nn.concat(gpu_pred_src_dst_list, 0)
                pred_src_srcm = nn.concat(gpu_pred_src_srcm_list, 0)
                pred_dst_dstm = nn.concat(gpu_pred_dst_dstm_list, 0)
                pred_src_dstm = nn.concat(gpu_pred_src_dstm_list, 0)

            with tf.device (models_opt_device):
                src_loss = tf.concat(gpu_src_losses, 0)
                dst_loss = tf.concat(gpu_dst_losses, 0)
                train_op = self.src_dst_opt.get_update_op (nn.average_gv_list (gpu_G_loss_gradients))

                if gan_power != 0:
                    GAN_train_op = self.GAN_opt.get_update_op (nn.average_gv_list(gpu_GAN_loss_gradients) )

            # Initializing training and view functions
            def train(warped_src, target_src, target_srcm, target_srcm_em,  \
                              warped_dst, target_dst, target_dstm, target_dstm_em, ):
                s, d, _ = nn.tf_sess.run ([src_loss, dst_loss, train_op],
                                            feed_dict={self.warped_src :warped_src,
                                                       self.target_src :target_src,
                                                       self.target_srcm:target_srcm,
                                                       self.target_srcm_em:target_srcm_em,
                                                       self.warped_dst :warped_dst,
                                                       self.target_dst :target_dst,
                                                       self.target_dstm:target_dstm,
                                                       self.target_dstm_em:target_dstm_em,
                                                       })
                return s, d
            self.train = train

            if gan_power != 0:
                def GAN_train(warped_src, target_src, target_srcm, target_srcm_em,  \
                              warped_dst, target_dst, target_dstm, target_dstm_em, ):
                    nn.tf_sess.run ([GAN_train_op], feed_dict={self.warped_src :warped_src,
                                                               self.target_src :target_src,
                                                               self.target_srcm:target_srcm,
                                                               self.target_srcm_em:target_srcm_em,
                                                               self.warped_dst :warped_dst,
                                                               self.target_dst :target_dst,
                                                               self.target_dstm:target_dstm,
                                                               self.target_dstm_em:target_dstm_em})
                self.GAN_train = GAN_train

            def AE_view(warped_src, warped_dst, morph_value):
                return nn.tf_sess.run ( [pred_src_src, pred_dst_dst, pred_dst_dstm, pred_src_dst, pred_src_dstm],
                                            feed_dict={self.warped_src:warped_src, self.warped_dst:warped_dst, self.morph_value_t:[morph_value] })

            self.AE_view = AE_view
        else:
            #Initializing merge function
            with tf.device( nn.tf_default_device_name if len(devices) != 0 else f'/CPU:0'):
                gpu_dst_code = self.encoder (self.warped_dst)
                gpu_dst_inter_src_code = self.inter_src (gpu_dst_code)
                gpu_dst_inter_dst_code = self.inter_dst (gpu_dst_code)

                inter_dims_slice = tf.cast(inter_dims*self.morph_value_t[0], tf.int32)
                gpu_src_dst_code =  tf.concat( ( tf.slice(gpu_dst_inter_src_code, [0,0,0,0],   [-1, inter_dims_slice , inter_res, inter_res]),
                                                 tf.slice(gpu_dst_inter_dst_code, [0,inter_dims_slice,0,0], [-1,inter_dims-inter_dims_slice, inter_res,inter_res]) ), 1 )

                gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder(gpu_src_dst_code)
                _, gpu_pred_dst_dstm = self.decoder(gpu_dst_inter_dst_code)

            def AE_merge(warped_dst, morph_value):
                return nn.tf_sess.run ( [gpu_pred_src_dst, gpu_pred_dst_dstm, gpu_pred_src_dstm], feed_dict={self.warped_dst:warped_dst, self.morph_value_t:[morph_value] })

            self.AE_merge = AE_merge

        # Loading/initializing all models/optimizers weights
        for model, filename in io.progress_bar_generator(self.model_filename_list, "Initializing models"):
            do_init = self.is_first_run()
            if self.is_training and gan_power != 0 and model == self.GAN:
                if self.gan_model_changed:
                    do_init = True
            if not do_init:
                do_init = not model.load_weights( self.get_strpath_storage_for_file(filename) )
            if do_init:
                model.init_weights()
        ###############

        # initializing sample generators
        if self.is_training:
            training_data_src_path = self.training_data_src_path #if not self.pretrain else self.get_pretraining_data_path()
            training_data_dst_path = self.training_data_dst_path #if not self.pretrain else self.get_pretraining_data_path()

            random_ct_samples_path=training_data_dst_path if ct_mode is not None else None #and not self.pretrain

            cpu_count = multiprocessing.cpu_count()
            src_generators_count = cpu_count // 2
            dst_generators_count = cpu_count // 2
            if ct_mode is not None:
                src_generators_count = int(src_generators_count * 1.5)



            self.set_training_data_generators ([
                    SampleGeneratorFace(training_data_src_path, random_ct_samples_path=random_ct_samples_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
                        sample_process_options=SampleProcessor.Options(scale_range=[-0.15, 0.15], random_flip=self.random_src_flip),
                        output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':random_warp, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'ct_mode': ct_mode,                                         'face_type':face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                                {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':False      , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'ct_mode': ct_mode,                                         'face_type':face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                                {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False      , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G,   'face_mask_type' : SampleProcessor.FaceMaskType.FULL_FACE,  'face_type':face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                                {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False      , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G,   'face_mask_type' : SampleProcessor.FaceMaskType.EYES_MOUTH, 'face_type':face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                              ],
                        uniform_yaw_distribution=self.options['uniform_yaw'],# or self.pretrain,
                        generators_count=src_generators_count ),

                    SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
                        sample_process_options=SampleProcessor.Options(scale_range=[-0.15, 0.15], random_flip=self.random_dst_flip),
                        output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':random_warp, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR,                                                             'face_type':face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                                {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':False      , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR,                                                             'face_type':face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                                {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False      , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G,   'face_mask_type' : SampleProcessor.FaceMaskType.FULL_FACE,  'face_type':face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                                {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False      , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G,   'face_mask_type' : SampleProcessor.FaceMaskType.EYES_MOUTH, 'face_type':face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                              ],
                        uniform_yaw_distribution=self.options['uniform_yaw'],# or self.pretrain,
                        generators_count=dst_generators_count )
                             ])

    def export_dfm (self):
        output_path=self.get_strpath_storage_for_file('model.dfm')

        io.log_info(f'Dumping .dfm to {output_path}')

        tf = nn.tf
        with tf.device (nn.tf_default_device_name):
            warped_dst = tf.placeholder (nn.floatx, (None, self.resolution, self.resolution, 3), name='in_face')
            warped_dst = tf.transpose(warped_dst, (0,3,1,2))
            morph_value = tf.placeholder (nn.floatx, (1,), name='morph_value')

            gpu_dst_code = self.encoder (warped_dst)
            gpu_dst_inter_src_code = self.inter_src ( gpu_dst_code)
            gpu_dst_inter_dst_code = self.inter_dst ( gpu_dst_code)

            inter_dims_slice = tf.cast(self.inter_dims*morph_value[0], tf.int32)
            gpu_src_dst_code =  tf.concat( (tf.slice(gpu_dst_inter_src_code, [0,0,0,0],   [-1, inter_dims_slice , self.inter_res, self.inter_res]),
                                            tf.slice(gpu_dst_inter_dst_code, [0,inter_dims_slice,0,0], [-1,self.inter_dims-inter_dims_slice, self.inter_res,self.inter_res]) ), 1 )

            gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder(gpu_src_dst_code)
            _, gpu_pred_dst_dstm = self.decoder(gpu_dst_inter_dst_code)

            gpu_pred_src_dst = tf.transpose(gpu_pred_src_dst, (0,2,3,1))
            gpu_pred_dst_dstm = tf.transpose(gpu_pred_dst_dstm, (0,2,3,1))
            gpu_pred_src_dstm = tf.transpose(gpu_pred_src_dstm, (0,2,3,1))

        tf.identity(gpu_pred_dst_dstm, name='out_face_mask')
        tf.identity(gpu_pred_src_dst, name='out_celeb_face')
        tf.identity(gpu_pred_src_dstm, name='out_celeb_face_mask')

        output_graph_def = tf.graph_util.convert_variables_to_constants(
            nn.tf_sess,
            tf.get_default_graph().as_graph_def(),
            ['out_face_mask','out_celeb_face','out_celeb_face_mask']
        )

        import tf2onnx
        with tf.device("/CPU:0"):
            model_proto, _ = tf2onnx.convert._convert_common(
                output_graph_def,
                name='AMP',
                input_names=['in_face:0','morph_value:0'],
                output_names=['out_face_mask:0','out_celeb_face:0','out_celeb_face_mask:0'],
                opset=12,
                output_path=output_path)

    #override
    def get_model_filename_list(self):
        return self.model_filename_list

    #override
    def onSave(self):
        for model, filename in io.progress_bar_generator(self.get_model_filename_list(), "Saving", leave=False):
            model.save_weights ( self.get_strpath_storage_for_file(filename) )

    #override
    def should_save_preview_history(self):
        return (not io.is_colab() and self.iter % ( 10*(max(1,self.resolution // 64)) ) == 0) or \
               (io.is_colab() and self.iter % 100 == 0)

    #override
    def onTrainOneIter(self):
        bs = self.get_batch_size()

        ( (warped_src, target_src, target_srcm, target_srcm_em), \
          (warped_dst, target_dst, target_dstm, target_dstm_em) ) = self.generate_next_samples()

        src_loss, dst_loss = self.train (warped_src, target_src, target_srcm, target_srcm_em, warped_dst, target_dst, target_dstm, target_dstm_em)

        if self.gan_power != 0:
            self.GAN_train (warped_src, target_src, target_srcm, target_srcm_em, warped_dst, target_dst, target_dstm, target_dstm_em)

        return ( ('src_loss', np.mean(src_loss) ), ('dst_loss', np.mean(dst_loss) ), )

    #override
    def onGetPreview(self, samples, for_history=False):
        ( (warped_src, target_src, target_srcm, target_srcm_em),
          (warped_dst, target_dst, target_dstm, target_dstm_em) ) = samples

        S, D, SS, DD, DDM_000, _, _ = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in ([target_src,target_dst] + self.AE_view (target_src, target_dst, 0.0)  ) ]

        _, _, DDM_025, SD_025, SDM_025 = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in self.AE_view (target_src, target_dst, 0.25) ]
        _, _, DDM_050, SD_050, SDM_050 = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in self.AE_view (target_src, target_dst, 0.50) ]
        _, _, DDM_065, SD_065, SDM_065 = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in self.AE_view (target_src, target_dst, 0.65) ]
        _, _, DDM_075, SD_075, SDM_075 = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in self.AE_view (target_src, target_dst, 0.75) ]
        _, _, DDM_100, SD_100, SDM_100 = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in self.AE_view (target_src, target_dst, 1.00) ]

        (DDM_000,
         DDM_025, SDM_025,
         DDM_050, SDM_050,
         DDM_065, SDM_065,
         DDM_075, SDM_075,
         DDM_100, SDM_100) = [ np.repeat (x, (3,), -1) for x in (DDM_000,
                                                                 DDM_025, SDM_025,
                                                                 DDM_050, SDM_050,
                                                                 DDM_065, SDM_065,
                                                                 DDM_075, SDM_075,
                                                                 DDM_100, SDM_100) ]

        target_srcm, target_dstm = [ nn.to_data_format(x,"NHWC", self.model_data_format) for x in ([target_srcm, target_dstm] )]

        n_samples = min(4, self.get_batch_size(), 800 // self.resolution )

        result = []

        i = np.random.randint(n_samples) if not for_history else 0

        st =  [ np.concatenate ((S[i],  D[i],  DD[i]*DDM_000[i]), axis=1) ]
        st += [ np.concatenate ((SS[i], DD[i], SD_100[i] ), axis=1) ]

        result += [ ('AMP morph 1.0', np.concatenate (st, axis=0 )), ]

        st =  [ np.concatenate ((DD[i], SD_025[i],  SD_050[i]), axis=1) ]
        st += [ np.concatenate ((SD_065[i], SD_075[i], SD_100[i]), axis=1) ]
        result += [ ('AMP morph list', np.concatenate (st, axis=0 )), ]

        st =  [ np.concatenate ((DD[i], SD_025[i]*DDM_025[i]*SDM_025[i],  SD_050[i]*DDM_050[i]*SDM_050[i]), axis=1) ]
        st += [ np.concatenate ((SD_065[i]*DDM_065[i]*SDM_065[i], SD_075[i]*DDM_075[i]*SDM_075[i], SD_100[i]*DDM_100[i]*SDM_100[i]), axis=1) ]
        result += [ ('AMP morph list masked', np.concatenate (st, axis=0 )), ]

        return result

    def predictor_func (self, face, morph_value):
        face = nn.to_data_format(face[None,...], self.model_data_format, "NHWC")

        bgr, mask_dst_dstm, mask_src_dstm = [ nn.to_data_format(x,"NHWC", self.model_data_format).astype(np.float32) for x in self.AE_merge (face, morph_value) ]

        return bgr[0], mask_src_dstm[0][...,0], mask_dst_dstm[0][...,0]

    #override
    def get_MergerConfig(self):
        morph_factor = np.clip ( io.input_number ("Morph factor", 1.0, add_info="0.0 .. 1.0"), 0.0, 1.0 )

        def predictor_morph(face):
            return self.predictor_func(face, morph_factor)


        import merger
        return predictor_morph, (self.options['resolution'], self.options['resolution'], 3), merger.MergerConfigMasked(face_type=self.face_type, default_mode = 'overlay')

Model = AMPModel

models/Model_AMP/Model.py
oo92/TorchFaceLab
import multiprocessing
from functools import partial

import numpy as np

from core import mathlib
from core.interact import interact as io
from core.leras import nn
from facelib import FaceType
from models import ModelBase
from samplelib import *

class QModel(ModelBase):
    #override
    def on_initialize(self):
        device_config = nn.getCurrentDeviceConfig()
        devices = device_config.devices
        self.model_data_format = "NCHW" if len(devices) != 0 and not self.is_debug() else "NHWC"
        nn.initialize(data_format=self.model_data_format)
        tf = nn.tf

        resolution = self.resolution = 96
        self.face_type = FaceType.FULL
        ae_dims = 128
        e_dims = 64
        d_dims = 64
        d_mask_dims = 16
        self.pretrain = False
        self.pretrain_just_disabled = False

        masked_training = True

        models_opt_on_gpu = len(devices) >= 1 and all([dev.total_mem_gb >= 4 for dev in devices])
        models_opt_device = nn.tf_default_device_name if models_opt_on_gpu and self.is_training else '/CPU:0'
        optimizer_vars_on_cpu = models_opt_device=='/CPU:0'

        input_ch = 3
        bgr_shape = nn.get4Dshape(resolution,resolution,input_ch)
        mask_shape = nn.get4Dshape(resolution,resolution,1)

        self.model_filename_list = []
        
        model_archi = nn.DeepFakeArchi(resolution, opts='ud')

        with tf.device ('/CPU:0'):
            #Place holders on CPU
            self.warped_src = tf.placeholder (nn.floatx, bgr_shape)
            self.warped_dst = tf.placeholder (nn.floatx, bgr_shape)

            self.target_src = tf.placeholder (nn.floatx, bgr_shape)
            self.target_dst = tf.placeholder (nn.floatx, bgr_shape)

            self.target_srcm = tf.placeholder (nn.floatx, mask_shape)
            self.target_dstm = tf.placeholder (nn.floatx, mask_shape)

        # Initializing model classes
        with tf.device (models_opt_device):
            self.encoder = model_archi.Encoder(in_ch=input_ch, e_ch=e_dims, name='encoder')
            encoder_out_ch = self.encoder.get_out_ch()*self.encoder.get_out_res(resolution)**2

            self.inter = model_archi.Inter (in_ch=encoder_out_ch, ae_ch=ae_dims, ae_out_ch=ae_dims, name='inter')
            inter_out_ch = self.inter.get_out_ch()

            self.decoder_src = model_archi.Decoder(in_ch=inter_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, name='decoder_src')
            self.decoder_dst = model_archi.Decoder(in_ch=inter_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, name='decoder_dst')

            self.model_filename_list += [ [self.encoder,     'encoder.npy'    ],
                                          [self.inter,       'inter.npy'      ],
                                          [self.decoder_src, 'decoder_src.npy'],
                                          [self.decoder_dst, 'decoder_dst.npy']  ]

            if self.is_training:
                self.src_dst_trainable_weights = self.encoder.get_weights() + self.inter.get_weights() + self.decoder_src.get_weights() + self.decoder_dst.get_weights()

                # Initialize optimizers
                self.src_dst_opt = nn.RMSprop(lr=2e-4, lr_dropout=0.3, name='src_dst_opt')
                self.src_dst_opt.initialize_variables(self.src_dst_trainable_weights, vars_on_cpu=optimizer_vars_on_cpu )
                self.model_filename_list += [ (self.src_dst_opt, 'src_dst_opt.npy') ]

        if self.is_training:
            # Adjust batch size for multiple GPU
            gpu_count = max(1, len(devices) )
            bs_per_gpu = max(1, 4 // gpu_count)
            self.set_batch_size( gpu_count*bs_per_gpu)

            # Compute losses per GPU
            gpu_pred_src_src_list = []
            gpu_pred_dst_dst_list = []
            gpu_pred_src_dst_list = []
            gpu_pred_src_srcm_list = []
            gpu_pred_dst_dstm_list = []
            gpu_pred_src_dstm_list = []

            gpu_src_losses = []
            gpu_dst_losses = []
            gpu_src_dst_loss_gvs = []
            
            for gpu_id in range(gpu_count):
                with tf.device( f'/{devices[gpu_id].tf_dev_type}:{gpu_id}' if len(devices) != 0 else f'/CPU:0' ):
                    batch_slice = slice( gpu_id*bs_per_gpu, (gpu_id+1)*bs_per_gpu )
                    with tf.device(f'/CPU:0'):
                        # slice on CPU, otherwise all batch data will be transfered to GPU first
                        gpu_warped_src   = self.warped_src [batch_slice,:,:,:]
                        gpu_warped_dst   = self.warped_dst [batch_slice,:,:,:]
                        gpu_target_src   = self.target_src [batch_slice,:,:,:]
                        gpu_target_dst   = self.target_dst [batch_slice,:,:,:]
                        gpu_target_srcm  = self.target_srcm[batch_slice,:,:,:]
                        gpu_target_dstm  = self.target_dstm[batch_slice,:,:,:]

                    # process model tensors
                    gpu_src_code     = self.inter(self.encoder(gpu_warped_src))
                    gpu_dst_code     = self.inter(self.encoder(gpu_warped_dst))
                    gpu_pred_src_src, gpu_pred_src_srcm = self.decoder_src(gpu_src_code)
                    gpu_pred_dst_dst, gpu_pred_dst_dstm = self.decoder_dst(gpu_dst_code)
                    gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(gpu_dst_code)

                    gpu_pred_src_src_list.append(gpu_pred_src_src)
                    gpu_pred_dst_dst_list.append(gpu_pred_dst_dst)
                    gpu_pred_src_dst_list.append(gpu_pred_src_dst)

                    gpu_pred_src_srcm_list.append(gpu_pred_src_srcm)
                    gpu_pred_dst_dstm_list.append(gpu_pred_dst_dstm)
                    gpu_pred_src_dstm_list.append(gpu_pred_src_dstm)

                    gpu_target_srcm_blur = nn.gaussian_blur(gpu_target_srcm,  max(1, resolution // 32) )
                    gpu_target_dstm_blur = nn.gaussian_blur(gpu_target_dstm,  max(1, resolution // 32) )

                    gpu_target_dst_masked      = gpu_target_dst*gpu_target_dstm_blur
                    gpu_target_dst_anti_masked = gpu_target_dst*(1.0 - gpu_target_dstm_blur)

                    gpu_target_src_masked_opt  = gpu_target_src*gpu_target_srcm_blur if masked_training else gpu_target_src
                    gpu_target_dst_masked_opt = gpu_target_dst_masked if masked_training else gpu_target_dst

                    gpu_pred_src_src_masked_opt = gpu_pred_src_src*gpu_target_srcm_blur if masked_training else gpu_pred_src_src
                    gpu_pred_dst_dst_masked_opt = gpu_pred_dst_dst*gpu_target_dstm_blur if masked_training else gpu_pred_dst_dst

                    gpu_psd_target_dst_masked = gpu_pred_src_dst*gpu_target_dstm_blur
                    gpu_psd_target_dst_anti_masked = gpu_pred_src_dst*(1.0 - gpu_target_dstm_blur)

                    gpu_src_loss =  tf.reduce_mean ( 10*nn.dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1])
                    gpu_src_loss += tf.reduce_mean ( 10*tf.square ( gpu_target_src_masked_opt - gpu_pred_src_src_masked_opt ), axis=[1,2,3])
                    gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] )

                    gpu_dst_loss  = tf.reduce_mean ( 10*nn.dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/11.6) ), axis=[1])
                    gpu_dst_loss += tf.reduce_mean ( 10*tf.square(  gpu_target_dst_masked_opt- gpu_pred_dst_dst_masked_opt ), axis=[1,2,3])
                    gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] )

                    gpu_src_losses += [gpu_src_loss]
                    gpu_dst_losses += [gpu_dst_loss]

                    gpu_G_loss = gpu_src_loss + gpu_dst_loss
                    gpu_src_dst_loss_gvs += [ nn.gradients ( gpu_G_loss, self.src_dst_trainable_weights ) ]


            # Average losses and gradients, and create optimizer update ops
            with tf.device (models_opt_device):
                pred_src_src  = nn.concat(gpu_pred_src_src_list, 0)
                pred_dst_dst  = nn.concat(gpu_pred_dst_dst_list, 0)
                pred_src_dst  = nn.concat(gpu_pred_src_dst_list, 0)
                pred_src_srcm = nn.concat(gpu_pred_src_srcm_list, 0)
                pred_dst_dstm = nn.concat(gpu_pred_dst_dstm_list, 0)
                pred_src_dstm = nn.concat(gpu_pred_src_dstm_list, 0)

                src_loss = nn.average_tensor_list(gpu_src_losses)
                dst_loss = nn.average_tensor_list(gpu_dst_losses)
                src_dst_loss_gv = nn.average_gv_list (gpu_src_dst_loss_gvs)
                src_dst_loss_gv_op = self.src_dst_opt.get_update_op (src_dst_loss_gv)

            # Initializing training and view functions
            def src_dst_train(warped_src, target_src, target_srcm, \
                              warped_dst, target_dst, target_dstm):
                s, d, _ = nn.tf_sess.run ( [ src_loss, dst_loss, src_dst_loss_gv_op],
                                            feed_dict={self.warped_src :warped_src,
                                                       self.target_src :target_src,
                                                       self.target_srcm:target_srcm,
                                                       self.warped_dst :warped_dst,
                                                       self.target_dst :target_dst,
                                                       self.target_dstm:target_dstm,
                                                       })
                s = np.mean(s)
                d = np.mean(d)
                return s, d
            self.src_dst_train = src_dst_train

            def AE_view(warped_src, warped_dst):
                return nn.tf_sess.run ( [pred_src_src, pred_dst_dst, pred_dst_dstm, pred_src_dst, pred_src_dstm],
                                            feed_dict={self.warped_src:warped_src,
                                                    self.warped_dst:warped_dst})

            self.AE_view = AE_view
        else:
            # Initializing merge function
            with tf.device( nn.tf_default_device_name if len(devices) != 0 else f'/CPU:0'):
                gpu_dst_code     = self.inter(self.encoder(self.warped_dst))
                gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(gpu_dst_code)
                _, gpu_pred_dst_dstm = self.decoder_dst(gpu_dst_code)

            def AE_merge( warped_dst):

                return nn.tf_sess.run ( [gpu_pred_src_dst, gpu_pred_dst_dstm, gpu_pred_src_dstm], feed_dict={self.warped_dst:warped_dst})

            self.AE_merge = AE_merge

        # Loading/initializing all models/optimizers weights
        for model, filename in io.progress_bar_generator(self.model_filename_list, "Initializing models"):
            if self.pretrain_just_disabled:
                do_init = False
                if model == self.inter:
                    do_init = True
            else:
                do_init = self.is_first_run()

            if not do_init:
                do_init = not model.load_weights( self.get_strpath_storage_for_file(filename) )

            if do_init and self.pretrained_model_path is not None:
                pretrained_filepath = self.pretrained_model_path / filename
                if pretrained_filepath.exists():
                    do_init = not model.load_weights(pretrained_filepath)

            if do_init:
                model.init_weights()

        # initializing sample generators
        if self.is_training:
            training_data_src_path = self.training_data_src_path if not self.pretrain else self.get_pretraining_data_path()
            training_data_dst_path = self.training_data_dst_path if not self.pretrain else self.get_pretraining_data_path()

            cpu_count = min(multiprocessing.cpu_count(), 8)
            src_generators_count = cpu_count // 2
            dst_generators_count = cpu_count // 2

            self.set_training_data_generators ([
                    SampleGeneratorFace(training_data_src_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
                        sample_process_options=SampleProcessor.Options(random_flip=True if self.pretrain else False),
                        output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':True,  'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR,                                                           'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                                {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':False, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR,                                                           'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                                {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G,   'face_mask_type' : SampleProcessor.FaceMaskType.FULL_FACE, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}
                                              ],
                        generators_count=src_generators_count ),

                    SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
                        sample_process_options=SampleProcessor.Options(random_flip=True if self.pretrain else False),
                        output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':True,  'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR,                                                           'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                                {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':False, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR,                                                           'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                                {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G,   'face_mask_type' : SampleProcessor.FaceMaskType.FULL_FACE, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}
                                               ],
                        generators_count=dst_generators_count )
                             ])

            self.last_samples = None

    #override
    def get_model_filename_list(self):
        return self.model_filename_list

    #override
    def onSave(self):
        for model, filename in io.progress_bar_generator(self.get_model_filename_list(), "Saving", leave=False):
            model.save_weights ( self.get_strpath_storage_for_file(filename) )

    #override
    def onTrainOneIter(self):

        if self.get_iter() % 3 == 0 and self.last_samples is not None:
            ( (warped_src, target_src, target_srcm), \
              (warped_dst, target_dst, target_dstm) ) = self.last_samples
            warped_src = target_src
            warped_dst = target_dst
        else:
            samples = self.last_samples = self.generate_next_samples()
            ( (warped_src, target_src, target_srcm), \
              (warped_dst, target_dst, target_dstm) ) = samples

        src_loss, dst_loss = self.src_dst_train (warped_src, target_src, target_srcm,
                                                 warped_dst, target_dst, target_dstm)

        return ( ('src_loss', src_loss), ('dst_loss', dst_loss), )

    #override
    def onGetPreview(self, samples, for_history=False):
        ( (warped_src, target_src, target_srcm),
          (warped_dst, target_dst, target_dstm) ) = samples

        S, D, SS, DD, DDM, SD, SDM = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in ([target_src,target_dst] + self.AE_view (target_src, target_dst) ) ]
        DDM, SDM, = [ np.repeat (x, (3,), -1) for x in [DDM, SDM] ]

        target_srcm, target_dstm = [ nn.to_data_format(x,"NHWC", self.model_data_format) for x in ([target_srcm, target_dstm] )]

        n_samples = min(4, self.get_batch_size() )
        result = []
        st = []
        for i in range(n_samples):
            ar = S[i], SS[i], D[i], DD[i], SD[i]
            st.append ( np.concatenate ( ar, axis=1) )

        result += [ ('Quick96', np.concatenate (st, axis=0 )), ]

        st_m = []
        for i in range(n_samples):
            ar = S[i]*target_srcm[i], SS[i], D[i]*target_dstm[i], DD[i]*DDM[i], SD[i]*(DDM[i]*SDM[i])
            st_m.append ( np.concatenate ( ar, axis=1) )

        result += [ ('Quick96 masked', np.concatenate (st_m, axis=0 )), ]

        return result

    def predictor_func (self, face=None):
        face = nn.to_data_format(face[None,...], self.model_data_format, "NHWC")

        bgr, mask_dst_dstm, mask_src_dstm = [ nn.to_data_format(x, "NHWC", self.model_data_format).astype(np.float32) for x in self.AE_merge (face) ]
        return bgr[0], mask_src_dstm[0][...,0], mask_dst_dstm[0][...,0]

    #override
    def get_MergerConfig(self):
        import merger
        return self.predictor_func, (self.resolution, self.resolution, 3), merger.MergerConfigMasked(face_type=self.face_type,
                                     default_mode = 'overlay',
                                    )

Model = QModel

models/Model_Quick96/Model.py
oo92/TorchFaceLab
import multiprocessing
import operator
from functools import partial

import numpy as np

from core import mathlib
from core.interact import interact as io
from core.leras import nn
from facelib import FaceType
from models import ModelBase
from samplelib import *

class SAEHDModel(ModelBase):

    #override
    def on_initialize_options(self):
        device_config = nn.getCurrentDeviceConfig()

        lowest_vram = 2
        if len(device_config.devices) != 0:
            lowest_vram = device_config.devices.get_worst_device().total_mem_gb

        if lowest_vram >= 4:
            suggest_batch_size = 8
        else:
            suggest_batch_size = 4

        yn_str = {True:'y',False:'n'}
        min_res = 64
        max_res = 640

        #default_usefp16            = self.options['use_fp16']           = self.load_or_def_option('use_fp16', False)
        default_resolution         = self.options['resolution']         = self.load_or_def_option('resolution', 128)
        default_face_type          = self.options['face_type']          = self.load_or_def_option('face_type', 'f')
        default_models_opt_on_gpu  = self.options['models_opt_on_gpu']  = self.load_or_def_option('models_opt_on_gpu', True)

        default_archi              = self.options['archi']              = self.load_or_def_option('archi', 'liae-ud')

        default_ae_dims            = self.options['ae_dims']            = self.load_or_def_option('ae_dims', 256)
        default_e_dims             = self.options['e_dims']             = self.load_or_def_option('e_dims', 64)
        default_d_dims             = self.options['d_dims']             = self.options.get('d_dims', None)
        default_d_mask_dims        = self.options['d_mask_dims']        = self.options.get('d_mask_dims', None)
        default_masked_training    = self.options['masked_training']    = self.load_or_def_option('masked_training', True)
        default_eyes_mouth_prio    = self.options['eyes_mouth_prio']    = self.load_or_def_option('eyes_mouth_prio', False)
        default_uniform_yaw        = self.options['uniform_yaw']        = self.load_or_def_option('uniform_yaw', False)
        default_blur_out_mask      = self.options['blur_out_mask']      = self.load_or_def_option('blur_out_mask', False)

        default_adabelief          = self.options['adabelief']          = self.load_or_def_option('adabelief', True)

        lr_dropout = self.load_or_def_option('lr_dropout', 'n')
        lr_dropout = {True:'y', False:'n'}.get(lr_dropout, lr_dropout) #backward comp
        default_lr_dropout         = self.options['lr_dropout'] = lr_dropout

        default_random_warp        = self.options['random_warp']        = self.load_or_def_option('random_warp', True)
        default_random_hsv_power   = self.options['random_hsv_power']   = self.load_or_def_option('random_hsv_power', 0.0)
        default_true_face_power    = self.options['true_face_power']    = self.load_or_def_option('true_face_power', 0.0)
        default_face_style_power   = self.options['face_style_power']   = self.load_or_def_option('face_style_power', 0.0)
        default_bg_style_power     = self.options['bg_style_power']     = self.load_or_def_option('bg_style_power', 0.0)
        default_ct_mode            = self.options['ct_mode']            = self.load_or_def_option('ct_mode', 'none')
        default_clipgrad           = self.options['clipgrad']           = self.load_or_def_option('clipgrad', False)
        default_pretrain           = self.options['pretrain']           = self.load_or_def_option('pretrain', False)

        ask_override = self.ask_override()
        if self.is_first_run() or ask_override:
            self.ask_autobackup_hour()
            self.ask_write_preview_history()
            self.ask_target_iter()
            self.ask_random_src_flip()
            self.ask_random_dst_flip()
            self.ask_batch_size(suggest_batch_size)
            #self.options['use_fp16'] = io.input_bool ("Use fp16", default_usefp16, help_message='Increases training/inference speed, reduces model size. Model may crash. Enable it after 1-5k iters.')

        if self.is_first_run():
            resolution = io.input_int("Resolution", default_resolution, add_info="64-640", help_message="More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 16 and 32 for -d archi.")
            resolution = np.clip ( (resolution // 16) * 16, min_res, max_res)
            self.options['resolution'] = resolution



            self.options['face_type'] = io.input_str ("Face type", default_face_type, ['h','mf','f','wf','head'], help_message="Half / mid face / full face / whole face / head. Half face has better resolution, but covers less area of cheeks. Mid face is 30% wider than half face. 'Whole face' covers full area of face include forehead. 'head' covers full head, but requires XSeg for src and dst faceset.").lower()

            while True:
                archi = io.input_str ("AE architecture", default_archi, help_message=\
"""
'df' keeps more identity-preserved face.
'liae' can fix overly different face shapes.
'-u' increased likeness of the face.
'-d' (experimental) doubling the resolution using the same computation cost.
Examples: df, liae, df-d, df-ud, liae-ud, ...
""").lower()

                archi_split = archi.split('-')

                if len(archi_split) == 2:
                    archi_type, archi_opts = archi_split
                elif len(archi_split) == 1:
                    archi_type, archi_opts = archi_split[0], None
                else:
                    continue

                if archi_type not in ['df', 'liae']:
                    continue

                if archi_opts is not None:
                    if len(archi_opts) == 0:
                        continue
                    if len([ 1 for opt in archi_opts if opt not in ['u','d','t','c'] ]) != 0:
                        continue

                    if 'd' in archi_opts:
                        self.options['resolution'] = np.clip ( (self.options['resolution'] // 32) * 32, min_res, max_res)

                break
            self.options['archi'] = archi

        default_d_dims             = self.options['d_dims']             = self.load_or_def_option('d_dims', 64)

        default_d_mask_dims        = default_d_dims // 3
        default_d_mask_dims        += default_d_mask_dims % 2
        default_d_mask_dims        = self.options['d_mask_dims']        = self.load_or_def_option('d_mask_dims', default_d_mask_dims)

        if self.is_first_run():
            self.options['ae_dims'] = np.clip ( io.input_int("AutoEncoder dimensions", default_ae_dims, add_info="32-1024", help_message="All face information will packed to AE dims. If amount of AE dims are not enough, then for example closed eyes will not be recognized. More dims are better, but require more VRAM. You can fine-tune model size to fit your GPU." ), 32, 1024 )

            e_dims = np.clip ( io.input_int("Encoder dimensions", default_e_dims, add_info="16-256", help_message="More dims help to recognize more facial features and achieve sharper result, but require more VRAM. You can fine-tune model size to fit your GPU." ), 16, 256 )
            self.options['e_dims'] = e_dims + e_dims % 2

            d_dims = np.clip ( io.input_int("Decoder dimensions", default_d_dims, add_info="16-256", help_message="More dims help to recognize more facial features and achieve sharper result, but require more VRAM. You can fine-tune model size to fit your GPU." ), 16, 256 )
            self.options['d_dims'] = d_dims + d_dims % 2

            d_mask_dims = np.clip ( io.input_int("Decoder mask dimensions", default_d_mask_dims, add_info="16-256", help_message="Typical mask dimensions = decoder dimensions / 3. If you manually cut out obstacles from the dst mask, you can increase this parameter to achieve better quality." ), 16, 256 )
            self.options['d_mask_dims'] = d_mask_dims + d_mask_dims % 2

        if self.is_first_run() or ask_override:
            if self.options['face_type'] == 'wf' or self.options['face_type'] == 'head':
                self.options['masked_training']  = io.input_bool ("Masked training", default_masked_training, help_message="This option is available only for 'whole_face' or 'head' type. Masked training clips training area to full_face mask or XSeg mask, thus network will train the faces properly.")

            self.options['eyes_mouth_prio'] = io.input_bool ("Eyes and mouth priority", default_eyes_mouth_prio, help_message='Helps to fix eye problems during training like "alien eyes" and wrong eyes direction. Also makes the detail of the teeth higher.')
            self.options['uniform_yaw'] = io.input_bool ("Uniform yaw distribution of samples", default_uniform_yaw, help_message='Helps to fix blurry side faces due to small amount of them in the faceset.')
            self.options['blur_out_mask'] = io.input_bool ("Blur out mask", default_blur_out_mask, help_message='Blurs nearby area outside of applied face mask of training samples. The result is the background near the face is smoothed and less noticeable on swapped face. The exact xseg mask in src and dst faceset is required.')

        default_gan_power          = self.options['gan_power']          = self.load_or_def_option('gan_power', 0.0)
        default_gan_patch_size     = self.options['gan_patch_size']     = self.load_or_def_option('gan_patch_size', self.options['resolution'] // 8)
        default_gan_dims           = self.options['gan_dims']           = self.load_or_def_option('gan_dims', 16)

        if self.is_first_run() or ask_override:
            self.options['models_opt_on_gpu'] = io.input_bool ("Place models and optimizer on GPU", default_models_opt_on_gpu, help_message="When you train on one GPU, by default model and optimizer weights are placed on GPU to accelerate the process. You can place they on CPU to free up extra VRAM, thus set bigger dimensions.")

            self.options['adabelief'] = io.input_bool ("Use AdaBelief optimizer?", default_adabelief, help_message="Use AdaBelief optimizer. It requires more VRAM, but the accuracy and the generalization of the model is higher.")

            self.options['lr_dropout']  = io.input_str (f"Use learning rate dropout", default_lr_dropout, ['n','y','cpu'], help_message="When the face is trained enough, you can enable this option to get extra sharpness and reduce subpixel shake for less amount of iterations. Enabled it before `disable random warp` and before GAN. \nn - disabled.\ny - enabled\ncpu - enabled on CPU. This allows not to use extra VRAM, sacrificing 20% time of iteration.")

            self.options['random_warp'] = io.input_bool ("Enable random warp of samples", default_random_warp, help_message="Random warp is required to generalize facial expressions of both faces. When the face is trained enough, you can disable it to get extra sharpness and reduce subpixel shake for less amount of iterations.")

            self.options['random_hsv_power'] = np.clip ( io.input_number ("Random hue/saturation/light intensity", default_random_hsv_power, add_info="0.0 .. 0.3", help_message="Random hue/saturation/light intensity applied to the src face set only at the input of the neural network. Stabilizes color perturbations during face swapping. Reduces the quality of the color transfer by selecting the closest one in the src faceset. Thus the src faceset must be diverse enough. Typical fine value is 0.05"), 0.0, 0.3 )

            self.options['gan_power'] = np.clip ( io.input_number ("GAN power", default_gan_power, add_info="0.0 .. 5.0", help_message="Forces the neural network to learn small details of the face. Enable it only when the face is trained enough with lr_dropout(on) and random_warp(off), and don't disable. The higher the value, the higher the chances of artifacts. Typical fine value is 0.1"), 0.0, 5.0 )

            if self.options['gan_power'] != 0.0:
                gan_patch_size = np.clip ( io.input_int("GAN patch size", default_gan_patch_size, add_info="3-640", help_message="The higher patch size, the higher the quality, the more VRAM is required. You can get sharper edges even at the lowest setting. Typical fine value is resolution / 8." ), 3, 640 )
                self.options['gan_patch_size'] = gan_patch_size

                gan_dims = np.clip ( io.input_int("GAN dimensions", default_gan_dims, add_info="4-512", help_message="The dimensions of the GAN network. The higher dimensions, the more VRAM is required. You can get sharper edges even at the lowest setting. Typical fine value is 16." ), 4, 512 )
                self.options['gan_dims'] = gan_dims

            if 'df' in self.options['archi']:
                self.options['true_face_power'] = np.clip ( io.input_number ("'True face' power.", default_true_face_power, add_info="0.0000 .. 1.0", help_message="Experimental option. Discriminates result face to be more like src face. Higher value - stronger discrimination. Typical value is 0.01 . Comparison - https://i.imgur.com/czScS9q.png"), 0.0, 1.0 )
            else:
                self.options['true_face_power'] = 0.0

            self.options['face_style_power'] = np.clip ( io.input_number("Face style power", default_face_style_power, add_info="0.0..100.0", help_message="Learn the color of the predicted face to be the same as dst inside mask. If you want to use this option with 'whole_face' you have to use XSeg trained mask. Warning: Enable it only after 10k iters, when predicted face is clear enough to start learn style. Start from 0.001 value and check history changes. Enabling this option increases the chance of model collapse."), 0.0, 100.0 )
            self.options['bg_style_power'] = np.clip ( io.input_number("Background style power", default_bg_style_power, add_info="0.0..100.0", help_message="Learn the area outside mask of the predicted face to be the same as dst. If you want to use this option with 'whole_face' you have to use XSeg trained mask. For whole_face you have to use XSeg trained mask. This can make face more like dst. Enabling this option increases the chance of model collapse. Typical value is 2.0"), 0.0, 100.0 )

            self.options['ct_mode'] = io.input_str (f"Color transfer for src faceset", default_ct_mode, ['none','rct','lct','mkl','idt','sot'], help_message="Change color distribution of src samples close to dst samples. Try all modes to find the best.")
            self.options['clipgrad'] = io.input_bool ("Enable gradient clipping", default_clipgrad, help_message="Gradient clipping reduces chance of model collapse, sacrificing speed of training.")

            self.options['pretrain'] = io.input_bool ("Enable pretraining mode", default_pretrain, help_message="Pretrain the model with large amount of various faces. After that, model can be used to train the fakes more quickly. Forces random_warp=N, random_flips=Y, gan_power=0.0, lr_dropout=N, styles=0.0, uniform_yaw=Y")

        if self.options['pretrain'] and self.get_pretraining_data_path() is None:
            raise Exception("pretraining_data_path is not defined")

        self.gan_model_changed = (default_gan_patch_size != self.options['gan_patch_size']) or (default_gan_dims != self.options['gan_dims'])

        self.pretrain_just_disabled = (default_pretrain == True and self.options['pretrain'] == False)

    #override
    def on_initialize(self):
        device_config = nn.getCurrentDeviceConfig()
        devices = device_config.devices
        self.model_data_format = "NCHW" if len(devices) != 0 and not self.is_debug() else "NHWC"
        nn.initialize(data_format=self.model_data_format)
        tf = nn.tf

        self.resolution = resolution = self.options['resolution']
        self.face_type = {'h'  : FaceType.HALF,
                          'mf' : FaceType.MID_FULL,
                          'f'  : FaceType.FULL,
                          'wf' : FaceType.WHOLE_FACE,
                          'head' : FaceType.HEAD}[ self.options['face_type'] ]

        if 'eyes_prio' in self.options:
            self.options.pop('eyes_prio')

        eyes_mouth_prio = self.options['eyes_mouth_prio']

        archi_split = self.options['archi'].split('-')

        if len(archi_split) == 2:
            archi_type, archi_opts = archi_split
        elif len(archi_split) == 1:
            archi_type, archi_opts = archi_split[0], None

        self.archi_type = archi_type

        ae_dims = self.options['ae_dims']
        e_dims = self.options['e_dims']
        d_dims = self.options['d_dims']
        d_mask_dims = self.options['d_mask_dims']
        self.pretrain = self.options['pretrain']
        if self.pretrain_just_disabled:
            self.set_iter(0)

        adabelief = self.options['adabelief']

        use_fp16 = False
        if self.is_exporting:
            use_fp16 = io.input_bool ("Export quantized?", False, help_message='Makes the exported model faster. If you have problems, disable this option.')

        self.gan_power = gan_power = 0.0 if self.pretrain else self.options['gan_power']
        random_warp = False if self.pretrain else self.options['random_warp']
        random_src_flip = self.random_src_flip if not self.pretrain else True
        random_dst_flip = self.random_dst_flip if not self.pretrain else True
        random_hsv_power = self.options['random_hsv_power'] if not self.pretrain else 0.0
        blur_out_mask = self.options['blur_out_mask']

        if self.pretrain:
            self.options_show_override['lr_dropout'] = 'n'
            self.options_show_override['random_warp'] = False
            self.options_show_override['gan_power'] = 0.0
            self.options_show_override['random_hsv_power'] = 0.0
            self.options_show_override['face_style_power'] = 0.0
            self.options_show_override['bg_style_power'] = 0.0
            self.options_show_override['uniform_yaw'] = True

        masked_training = self.options['masked_training']
        ct_mode = self.options['ct_mode']
        if ct_mode == 'none':
            ct_mode = None


        models_opt_on_gpu = False if len(devices) == 0 else self.options['models_opt_on_gpu']
        models_opt_device = nn.tf_default_device_name if models_opt_on_gpu and self.is_training else '/CPU:0'
        optimizer_vars_on_cpu = models_opt_device=='/CPU:0'

        input_ch=3
        bgr_shape = self.bgr_shape = nn.get4Dshape(resolution,resolution,input_ch)
        mask_shape = nn.get4Dshape(resolution,resolution,1)
        self.model_filename_list = []

        with tf.device ('/CPU:0'):
            #Place holders on CPU
            self.warped_src = tf.placeholder (nn.floatx, bgr_shape, name='warped_src')
            self.warped_dst = tf.placeholder (nn.floatx, bgr_shape, name='warped_dst')

            self.target_src = tf.placeholder (nn.floatx, bgr_shape, name='target_src')
            self.target_dst = tf.placeholder (nn.floatx, bgr_shape, name='target_dst')

            self.target_srcm    = tf.placeholder (nn.floatx, mask_shape, name='target_srcm')
            self.target_srcm_em = tf.placeholder (nn.floatx, mask_shape, name='target_srcm_em')
            self.target_dstm    = tf.placeholder (nn.floatx, mask_shape, name='target_dstm')
            self.target_dstm_em = tf.placeholder (nn.floatx, mask_shape, name='target_dstm_em')

        # Initializing model classes
        model_archi = nn.DeepFakeArchi(resolution, use_fp16=use_fp16, opts=archi_opts)

        with tf.device (models_opt_device):
            if 'df' in archi_type:
                self.encoder = model_archi.Encoder(in_ch=input_ch, e_ch=e_dims, name='encoder')
                encoder_out_ch = self.encoder.get_out_ch()*self.encoder.get_out_res(resolution)**2

                self.inter = model_archi.Inter (in_ch=encoder_out_ch, ae_ch=ae_dims, ae_out_ch=ae_dims, name='inter')
                inter_out_ch = self.inter.get_out_ch()

                self.decoder_src = model_archi.Decoder(in_ch=inter_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, name='decoder_src')
                self.decoder_dst = model_archi.Decoder(in_ch=inter_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, name='decoder_dst')

                self.model_filename_list += [ [self.encoder,     'encoder.npy'    ],
                                              [self.inter,       'inter.npy'      ],
                                              [self.decoder_src, 'decoder_src.npy'],
                                              [self.decoder_dst, 'decoder_dst.npy']  ]

                if self.is_training:
                    if self.options['true_face_power'] != 0:
                        self.code_discriminator = nn.CodeDiscriminator(ae_dims, code_res=self.inter.get_out_res(), name='dis' )
                        self.model_filename_list += [ [self.code_discriminator, 'code_discriminator.npy'] ]

            elif 'liae' in archi_type:
                self.encoder = model_archi.Encoder(in_ch=input_ch, e_ch=e_dims, name='encoder')
                encoder_out_ch = self.encoder.get_out_ch()*self.encoder.get_out_res(resolution)**2

                self.inter_AB = model_archi.Inter(in_ch=encoder_out_ch, ae_ch=ae_dims, ae_out_ch=ae_dims*2, name='inter_AB')
                self.inter_B  = model_archi.Inter(in_ch=encoder_out_ch, ae_ch=ae_dims, ae_out_ch=ae_dims*2, name='inter_B')

                inter_out_ch = self.inter_AB.get_out_ch()
                inters_out_ch = inter_out_ch*2
                self.decoder = model_archi.Decoder(in_ch=inters_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, name='decoder')

                self.model_filename_list += [ [self.encoder,  'encoder.npy'],
                                              [self.inter_AB, 'inter_AB.npy'],
                                              [self.inter_B , 'inter_B.npy'],
                                              [self.decoder , 'decoder.npy'] ]

            if self.is_training:
                if gan_power != 0:
                    self.D_src = nn.UNetPatchDiscriminator(patch_size=self.options['gan_patch_size'], in_ch=input_ch, base_ch=self.options['gan_dims'], name="D_src")
                    self.model_filename_list += [ [self.D_src, 'GAN.npy'] ]

                # Initialize optimizers
                lr=5e-5
                if self.options['lr_dropout'] in ['y','cpu'] and not self.pretrain:
                    lr_cos = 500
                    lr_dropout = 0.3
                else:
                    lr_cos = 0
                    lr_dropout = 1.0
                OptimizerClass = nn.AdaBelief if adabelief else nn.RMSprop
                clipnorm = 1.0 if self.options['clipgrad'] else 0.0

                if 'df' in archi_type:
                    self.src_dst_saveable_weights = self.encoder.get_weights() + self.inter.get_weights() + self.decoder_src.get_weights() + self.decoder_dst.get_weights()
                    self.src_dst_trainable_weights = self.src_dst_saveable_weights
                elif 'liae' in archi_type:
                    self.src_dst_saveable_weights = self.encoder.get_weights() + self.inter_AB.get_weights() + self.inter_B.get_weights() + self.decoder.get_weights()
                    if random_warp:
                        self.src_dst_trainable_weights = self.src_dst_saveable_weights
                    else:
                        self.src_dst_trainable_weights = self.encoder.get_weights() + self.inter_B.get_weights() + self.decoder.get_weights()

                self.src_dst_opt = OptimizerClass(lr=lr, lr_dropout=lr_dropout, lr_cos=lr_cos, clipnorm=clipnorm, name='src_dst_opt')
                self.src_dst_opt.initialize_variables (self.src_dst_saveable_weights, vars_on_cpu=optimizer_vars_on_cpu, lr_dropout_on_cpu=self.options['lr_dropout']=='cpu')
                self.model_filename_list += [ (self.src_dst_opt, 'src_dst_opt.npy') ]

                if self.options['true_face_power'] != 0:
                    self.D_code_opt = OptimizerClass(lr=lr, lr_dropout=lr_dropout, lr_cos=lr_cos, clipnorm=clipnorm, name='D_code_opt')
                    self.D_code_opt.initialize_variables ( self.code_discriminator.get_weights(), vars_on_cpu=optimizer_vars_on_cpu, lr_dropout_on_cpu=self.options['lr_dropout']=='cpu')
                    self.model_filename_list += [ (self.D_code_opt, 'D_code_opt.npy') ]

                if gan_power != 0:
                    self.D_src_dst_opt = OptimizerClass(lr=lr, lr_dropout=lr_dropout, lr_cos=lr_cos, clipnorm=clipnorm, name='GAN_opt')
                    self.D_src_dst_opt.initialize_variables ( self.D_src.get_weights(), vars_on_cpu=optimizer_vars_on_cpu, lr_dropout_on_cpu=self.options['lr_dropout']=='cpu')#+self.D_src_x2.get_weights()
                    self.model_filename_list += [ (self.D_src_dst_opt, 'GAN_opt.npy') ]

        if self.is_training:
            # Adjust batch size for multiple GPU
            gpu_count = max(1, len(devices) )
            bs_per_gpu = max(1, self.get_batch_size() // gpu_count)
            self.set_batch_size( gpu_count*bs_per_gpu)

            # Compute losses per GPU
            gpu_pred_src_src_list = []
            gpu_pred_dst_dst_list = []
            gpu_pred_src_dst_list = []
            gpu_pred_src_srcm_list = []
            gpu_pred_dst_dstm_list = []
            gpu_pred_src_dstm_list = []

            gpu_src_losses = []
            gpu_dst_losses = []
            gpu_G_loss_gvs = []
            gpu_D_code_loss_gvs = []
            gpu_D_src_dst_loss_gvs = []

            for gpu_id in range(gpu_count):
                with tf.device( f'/{devices[gpu_id].tf_dev_type}:{gpu_id}' if len(devices) != 0 else f'/CPU:0' ):
                    with tf.device(f'/CPU:0'):
                        # slice on CPU, otherwise all batch data will be transfered to GPU first
                        batch_slice = slice( gpu_id*bs_per_gpu, (gpu_id+1)*bs_per_gpu )
                        gpu_warped_src      = self.warped_src [batch_slice,:,:,:]
                        gpu_warped_dst      = self.warped_dst [batch_slice,:,:,:]
                        gpu_target_src      = self.target_src [batch_slice,:,:,:]
                        gpu_target_dst      = self.target_dst [batch_slice,:,:,:]
                        gpu_target_srcm     = self.target_srcm[batch_slice,:,:,:]
                        gpu_target_srcm_em  = self.target_srcm_em[batch_slice,:,:,:]
                        gpu_target_dstm     = self.target_dstm[batch_slice,:,:,:]
                        gpu_target_dstm_em  = self.target_dstm_em[batch_slice,:,:,:]

                    gpu_target_srcm_anti = 1-gpu_target_srcm
                    gpu_target_dstm_anti = 1-gpu_target_dstm

                    if blur_out_mask:
                        sigma = resolution / 128

                        x = nn.gaussian_blur(gpu_target_src*gpu_target_srcm_anti, sigma)
                        y = 1-nn.gaussian_blur(gpu_target_srcm, sigma)
                        y = tf.where(tf.equal(y, 0), tf.ones_like(y), y)
                        gpu_target_src = gpu_target_src*gpu_target_srcm + (x/y)*gpu_target_srcm_anti

                        x = nn.gaussian_blur(gpu_target_dst*gpu_target_dstm_anti, sigma)
                        y = 1-nn.gaussian_blur(gpu_target_dstm, sigma)
                        y = tf.where(tf.equal(y, 0), tf.ones_like(y), y)
                        gpu_target_dst = gpu_target_dst*gpu_target_dstm + (x/y)*gpu_target_dstm_anti


                    # process model tensors
                    if 'df' in archi_type:
                        gpu_src_code     = self.inter(self.encoder(gpu_warped_src))
                        gpu_dst_code     = self.inter(self.encoder(gpu_warped_dst))
                        gpu_pred_src_src, gpu_pred_src_srcm = self.decoder_src(gpu_src_code)
                        gpu_pred_dst_dst, gpu_pred_dst_dstm = self.decoder_dst(gpu_dst_code)
                        gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(gpu_dst_code)
                        gpu_pred_src_dst_no_code_grad, _ = self.decoder_src(tf.stop_gradient(gpu_dst_code))

                    elif 'liae' in archi_type:
                        gpu_src_code = self.encoder (gpu_warped_src)
                        gpu_src_inter_AB_code = self.inter_AB (gpu_src_code)
                        gpu_src_code = tf.concat([gpu_src_inter_AB_code,gpu_src_inter_AB_code], nn.conv2d_ch_axis  )
                        gpu_dst_code = self.encoder (gpu_warped_dst)
                        gpu_dst_inter_B_code = self.inter_B (gpu_dst_code)
                        gpu_dst_inter_AB_code = self.inter_AB (gpu_dst_code)
                        gpu_dst_code = tf.concat([gpu_dst_inter_B_code,gpu_dst_inter_AB_code], nn.conv2d_ch_axis )
                        gpu_src_dst_code = tf.concat([gpu_dst_inter_AB_code,gpu_dst_inter_AB_code], nn.conv2d_ch_axis )

                        gpu_pred_src_src, gpu_pred_src_srcm = self.decoder(gpu_src_code)
                        gpu_pred_dst_dst, gpu_pred_dst_dstm = self.decoder(gpu_dst_code)
                        gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder(gpu_src_dst_code)
                        gpu_pred_src_dst_no_code_grad, _ = self.decoder(tf.stop_gradient(gpu_src_dst_code))

                    gpu_pred_src_src_list.append(gpu_pred_src_src)
                    gpu_pred_dst_dst_list.append(gpu_pred_dst_dst)
                    gpu_pred_src_dst_list.append(gpu_pred_src_dst)

                    gpu_pred_src_srcm_list.append(gpu_pred_src_srcm)
                    gpu_pred_dst_dstm_list.append(gpu_pred_dst_dstm)
                    gpu_pred_src_dstm_list.append(gpu_pred_src_dstm)

                    gpu_target_srcm_blur = nn.gaussian_blur(gpu_target_srcm,  max(1, resolution // 32) )
                    gpu_target_srcm_blur = tf.clip_by_value(gpu_target_srcm_blur, 0, 0.5) * 2
                    gpu_target_srcm_anti_blur = 1.0-gpu_target_srcm_blur

                    gpu_target_dstm_blur = nn.gaussian_blur(gpu_target_dstm,  max(1, resolution // 32) )
                    gpu_target_dstm_blur = tf.clip_by_value(gpu_target_dstm_blur, 0, 0.5) * 2

                    gpu_style_mask_blur = nn.gaussian_blur(gpu_pred_src_dstm*gpu_pred_dst_dstm,  max(1, resolution // 32) )
                    gpu_style_mask_blur = tf.stop_gradient(tf.clip_by_value(gpu_target_srcm_blur, 0, 1.0))
                    gpu_style_mask_anti_blur = 1.0 - gpu_style_mask_blur

                    gpu_target_dst_masked = gpu_target_dst*gpu_target_dstm_blur

                    gpu_target_src_anti_masked = gpu_target_src*gpu_target_srcm_anti_blur
                    gpu_pred_src_src_anti_masked = gpu_pred_src_src*gpu_target_srcm_anti_blur

                    gpu_target_src_masked_opt  = gpu_target_src*gpu_target_srcm_blur if masked_training else gpu_target_src
                    gpu_target_dst_masked_opt  = gpu_target_dst_masked if masked_training else gpu_target_dst
                    gpu_pred_src_src_masked_opt = gpu_pred_src_src*gpu_target_srcm_blur if masked_training else gpu_pred_src_src
                    gpu_pred_dst_dst_masked_opt = gpu_pred_dst_dst*gpu_target_dstm_blur if masked_training else gpu_pred_dst_dst

                    if resolution < 256:
                        gpu_src_loss =  tf.reduce_mean ( 10*nn.dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1])
                    else:
                        gpu_src_loss =  tf.reduce_mean ( 5*nn.dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1])
                        gpu_src_loss += tf.reduce_mean ( 5*nn.dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/23.2)), axis=[1])
                    gpu_src_loss += tf.reduce_mean ( 10*tf.square ( gpu_target_src_masked_opt - gpu_pred_src_src_masked_opt ), axis=[1,2,3])

                    if eyes_mouth_prio:
                        gpu_src_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_src*gpu_target_srcm_em - gpu_pred_src_src*gpu_target_srcm_em ), axis=[1,2,3])

                    gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] )

                    face_style_power = self.options['face_style_power'] / 100.0
                    if face_style_power != 0 and not self.pretrain:
                        gpu_src_loss += nn.style_loss(gpu_pred_src_dst_no_code_grad*tf.stop_gradient(gpu_pred_src_dstm), tf.stop_gradient(gpu_pred_dst_dst*gpu_pred_dst_dstm), gaussian_blur_radius=resolution//8, loss_weight=10000*face_style_power)

                    bg_style_power = self.options['bg_style_power'] / 100.0
                    if bg_style_power != 0 and not self.pretrain:
                        gpu_target_dst_style_anti_masked = gpu_target_dst*gpu_style_mask_anti_blur
                        gpu_psd_style_anti_masked = gpu_pred_src_dst*gpu_style_mask_anti_blur

                        gpu_src_loss += tf.reduce_mean( (10*bg_style_power)*nn.dssim( gpu_psd_style_anti_masked,  gpu_target_dst_style_anti_masked, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1])
                        gpu_src_loss += tf.reduce_mean( (10*bg_style_power)*tf.square(gpu_psd_style_anti_masked - gpu_target_dst_style_anti_masked), axis=[1,2,3] )

                    if resolution < 256:
                        gpu_dst_loss = tf.reduce_mean ( 10*nn.dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/11.6) ), axis=[1])
                    else:
                        gpu_dst_loss = tf.reduce_mean ( 5*nn.dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/11.6) ), axis=[1])
                        gpu_dst_loss += tf.reduce_mean ( 5*nn.dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/23.2) ), axis=[1])
                    gpu_dst_loss += tf.reduce_mean ( 10*tf.square(  gpu_target_dst_masked_opt- gpu_pred_dst_dst_masked_opt ), axis=[1,2,3])

                    if eyes_mouth_prio:
                        gpu_dst_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_dst*gpu_target_dstm_em - gpu_pred_dst_dst*gpu_target_dstm_em ), axis=[1,2,3])

                    gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] )

                    gpu_src_losses += [gpu_src_loss]
                    gpu_dst_losses += [gpu_dst_loss]

                    gpu_G_loss = gpu_src_loss + gpu_dst_loss

                    def DLoss(labels,logits):
                        return tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits(labels=labels, logits=logits), axis=[1,2,3])

                    if self.options['true_face_power'] != 0:
                        gpu_src_code_d = self.code_discriminator( gpu_src_code )
                        gpu_src_code_d_ones  = tf.ones_like (gpu_src_code_d)
                        gpu_src_code_d_zeros = tf.zeros_like(gpu_src_code_d)
                        gpu_dst_code_d = self.code_discriminator( gpu_dst_code )
                        gpu_dst_code_d_ones = tf.ones_like(gpu_dst_code_d)

                        gpu_G_loss += self.options['true_face_power']*DLoss(gpu_src_code_d_ones, gpu_src_code_d)

                        gpu_D_code_loss = (DLoss(gpu_dst_code_d_ones , gpu_dst_code_d) + \
                                           DLoss(gpu_src_code_d_zeros, gpu_src_code_d) ) * 0.5

                        gpu_D_code_loss_gvs += [ nn.gradients (gpu_D_code_loss, self.code_discriminator.get_weights() ) ]

                    if gan_power != 0:
                        gpu_pred_src_src_d, \
                        gpu_pred_src_src_d2           = self.D_src(gpu_pred_src_src_masked_opt)

                        gpu_pred_src_src_d_ones  = tf.ones_like (gpu_pred_src_src_d)
                        gpu_pred_src_src_d_zeros = tf.zeros_like(gpu_pred_src_src_d)

                        gpu_pred_src_src_d2_ones  = tf.ones_like (gpu_pred_src_src_d2)
                        gpu_pred_src_src_d2_zeros = tf.zeros_like(gpu_pred_src_src_d2)

                        gpu_target_src_d, \
                        gpu_target_src_d2            = self.D_src(gpu_target_src_masked_opt)

                        gpu_target_src_d_ones    = tf.ones_like(gpu_target_src_d)
                        gpu_target_src_d2_ones    = tf.ones_like(gpu_target_src_d2)

                        gpu_D_src_dst_loss = (DLoss(gpu_target_src_d_ones      , gpu_target_src_d) + \
                                              DLoss(gpu_pred_src_src_d_zeros   , gpu_pred_src_src_d) ) * 0.5 + \
                                             (DLoss(gpu_target_src_d2_ones      , gpu_target_src_d2) + \
                                              DLoss(gpu_pred_src_src_d2_zeros   , gpu_pred_src_src_d2) ) * 0.5

                        gpu_D_src_dst_loss_gvs += [ nn.gradients (gpu_D_src_dst_loss, self.D_src.get_weights() ) ]#+self.D_src_x2.get_weights()

                        gpu_G_loss += gan_power*(DLoss(gpu_pred_src_src_d_ones, gpu_pred_src_src_d)  + \
                                                 DLoss(gpu_pred_src_src_d2_ones, gpu_pred_src_src_d2))

                        if masked_training:
                            # Minimal src-src-bg rec with total_variation_mse to suppress random bright dots from gan
                            gpu_G_loss += 0.000001*nn.total_variation_mse(gpu_pred_src_src)
                            gpu_G_loss += 0.02*tf.reduce_mean(tf.square(gpu_pred_src_src_anti_masked-gpu_target_src_anti_masked),axis=[1,2,3] )

                    gpu_G_loss_gvs += [ nn.gradients ( gpu_G_loss, self.src_dst_trainable_weights )]




            # Average losses and gradients, and create optimizer update ops
            with tf.device(f'/CPU:0'):
                pred_src_src  = nn.concat(gpu_pred_src_src_list, 0)
                pred_dst_dst  = nn.concat(gpu_pred_dst_dst_list, 0)
                pred_src_dst  = nn.concat(gpu_pred_src_dst_list, 0)
                pred_src_srcm = nn.concat(gpu_pred_src_srcm_list, 0)
                pred_dst_dstm = nn.concat(gpu_pred_dst_dstm_list, 0)
                pred_src_dstm = nn.concat(gpu_pred_src_dstm_list, 0)

            with tf.device (models_opt_device):
                src_loss = tf.concat(gpu_src_losses, 0)
                dst_loss = tf.concat(gpu_dst_losses, 0)
                src_dst_loss_gv_op = self.src_dst_opt.get_update_op (nn.average_gv_list (gpu_G_loss_gvs))

                if self.options['true_face_power'] != 0:
                    D_loss_gv_op = self.D_code_opt.get_update_op (nn.average_gv_list(gpu_D_code_loss_gvs))

                if gan_power != 0:
                    src_D_src_dst_loss_gv_op = self.D_src_dst_opt.get_update_op (nn.average_gv_list(gpu_D_src_dst_loss_gvs) )


            # Initializing training and view functions
            def src_dst_train(warped_src, target_src, target_srcm, target_srcm_em,  \
                              warped_dst, target_dst, target_dstm, target_dstm_em, ):
                s, d = nn.tf_sess.run ( [ src_loss, dst_loss, src_dst_loss_gv_op],
                                            feed_dict={self.warped_src :warped_src,
                                                       self.target_src :target_src,
                                                       self.target_srcm:target_srcm,
                                                       self.target_srcm_em:target_srcm_em,
                                                       self.warped_dst :warped_dst,
                                                       self.target_dst :target_dst,
                                                       self.target_dstm:target_dstm,
                                                       self.target_dstm_em:target_dstm_em,
                                                       })[:2]
                return s, d
            self.src_dst_train = src_dst_train

            if self.options['true_face_power'] != 0:
                def D_train(warped_src, warped_dst):
                    nn.tf_sess.run ([D_loss_gv_op], feed_dict={self.warped_src: warped_src, self.warped_dst: warped_dst})
                self.D_train = D_train

            if gan_power != 0:
                def D_src_dst_train(warped_src, target_src, target_srcm, target_srcm_em,  \
                                    warped_dst, target_dst, target_dstm, target_dstm_em, ):
                    nn.tf_sess.run ([src_D_src_dst_loss_gv_op], feed_dict={self.warped_src :warped_src,
                                                                           self.target_src :target_src,
                                                                           self.target_srcm:target_srcm,
                                                                           self.target_srcm_em:target_srcm_em,
                                                                           self.warped_dst :warped_dst,
                                                                           self.target_dst :target_dst,
                                                                           self.target_dstm:target_dstm,
                                                                           self.target_dstm_em:target_dstm_em})
                self.D_src_dst_train = D_src_dst_train


            def AE_view(warped_src, warped_dst):
                return nn.tf_sess.run ( [pred_src_src, pred_dst_dst, pred_dst_dstm, pred_src_dst, pred_src_dstm],
                                            feed_dict={self.warped_src:warped_src,
                                                    self.warped_dst:warped_dst})
            self.AE_view = AE_view
        else:
            # Initializing merge function
            with tf.device( nn.tf_default_device_name if len(devices) != 0 else f'/CPU:0'):
                if 'df' in archi_type:
                    gpu_dst_code     = self.inter(self.encoder(self.warped_dst))
                    gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(gpu_dst_code)
                    _, gpu_pred_dst_dstm = self.decoder_dst(gpu_dst_code)

                elif 'liae' in archi_type:
                    gpu_dst_code = self.encoder (self.warped_dst)
                    gpu_dst_inter_B_code = self.inter_B (gpu_dst_code)
                    gpu_dst_inter_AB_code = self.inter_AB (gpu_dst_code)
                    gpu_dst_code = tf.concat([gpu_dst_inter_B_code,gpu_dst_inter_AB_code], nn.conv2d_ch_axis)
                    gpu_src_dst_code = tf.concat([gpu_dst_inter_AB_code,gpu_dst_inter_AB_code], nn.conv2d_ch_axis)

                    gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder(gpu_src_dst_code)
                    _, gpu_pred_dst_dstm = self.decoder(gpu_dst_code)


            def AE_merge( warped_dst):
                return nn.tf_sess.run ( [gpu_pred_src_dst, gpu_pred_dst_dstm, gpu_pred_src_dstm], feed_dict={self.warped_dst:warped_dst})

            self.AE_merge = AE_merge

        # Loading/initializing all models/optimizers weights
        for model, filename in io.progress_bar_generator(self.model_filename_list, "Initializing models"):
            if self.pretrain_just_disabled:
                do_init = False
                if 'df' in archi_type:
                    if model == self.inter:
                        do_init = True
                elif 'liae' in archi_type:
                    if model == self.inter_AB or model == self.inter_B:
                        do_init = True
            else:
                do_init = self.is_first_run()
                if self.is_training and gan_power != 0 and model == self.D_src:
                    if self.gan_model_changed:
                        do_init = True

            if not do_init:
                do_init = not model.load_weights( self.get_strpath_storage_for_file(filename) )

            if do_init:
                model.init_weights()


        ###############

        # initializing sample generators
        if self.is_training:
            training_data_src_path = self.training_data_src_path if not self.pretrain else self.get_pretraining_data_path()
            training_data_dst_path = self.training_data_dst_path if not self.pretrain else self.get_pretraining_data_path()

            random_ct_samples_path=training_data_dst_path if ct_mode is not None and not self.pretrain else None

            cpu_count = multiprocessing.cpu_count()
            src_generators_count = cpu_count // 2
            dst_generators_count = cpu_count // 2
            if ct_mode is not None:
                src_generators_count = int(src_generators_count * 1.5)

            self.set_training_data_generators ([
                    SampleGeneratorFace(training_data_src_path, random_ct_samples_path=random_ct_samples_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
                        sample_process_options=SampleProcessor.Options(scale_range=[-0.15, 0.15], random_flip=random_src_flip),
                        output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':random_warp, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'ct_mode': ct_mode,   'random_hsv_shift_amount' : random_hsv_power,                                        'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                                {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':False                      , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'ct_mode': ct_mode,                           'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                                {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False                      , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G,   'face_mask_type' : SampleProcessor.FaceMaskType.FULL_FACE, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                                {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False                      , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G,   'face_mask_type' : SampleProcessor.FaceMaskType.EYES_MOUTH, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                              ],
                        uniform_yaw_distribution=self.options['uniform_yaw'] or self.pretrain,
                        generators_count=src_generators_count ),

                    SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
                        sample_process_options=SampleProcessor.Options(scale_range=[-0.15, 0.15], random_flip=random_dst_flip),
                        output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':random_warp, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR,                                                                'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                                {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':False                      , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR,                                                'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                                {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False                      , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G,   'face_mask_type' : SampleProcessor.FaceMaskType.FULL_FACE, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                                {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False                      , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G,   'face_mask_type' : SampleProcessor.FaceMaskType.EYES_MOUTH, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                              ],
                        uniform_yaw_distribution=self.options['uniform_yaw'] or self.pretrain,
                        generators_count=dst_generators_count )
                             ])

            if self.pretrain_just_disabled:
                self.update_sample_for_preview(force_new=True)

    def export_dfm (self):
        output_path=self.get_strpath_storage_for_file('model.dfm')

        io.log_info(f'Dumping .dfm to {output_path}')

        tf = nn.tf
        nn.set_data_format('NCHW')

        with tf.device (nn.tf_default_device_name):
            warped_dst = tf.placeholder (nn.floatx, (None, self.resolution, self.resolution, 3), name='in_face')
            warped_dst = tf.transpose(warped_dst, (0,3,1,2))


            if 'df' in self.archi_type:
                gpu_dst_code     = self.inter(self.encoder(warped_dst))
                gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(gpu_dst_code)
                _, gpu_pred_dst_dstm = self.decoder_dst(gpu_dst_code)

            elif 'liae' in self.archi_type:
                gpu_dst_code = self.encoder (warped_dst)
                gpu_dst_inter_B_code = self.inter_B (gpu_dst_code)
                gpu_dst_inter_AB_code = self.inter_AB (gpu_dst_code)
                gpu_dst_code = tf.concat([gpu_dst_inter_B_code,gpu_dst_inter_AB_code], nn.conv2d_ch_axis)
                gpu_src_dst_code = tf.concat([gpu_dst_inter_AB_code,gpu_dst_inter_AB_code], nn.conv2d_ch_axis)

                gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder(gpu_src_dst_code)
                _, gpu_pred_dst_dstm = self.decoder(gpu_dst_code)

            gpu_pred_src_dst = tf.transpose(gpu_pred_src_dst, (0,2,3,1))
            gpu_pred_dst_dstm = tf.transpose(gpu_pred_dst_dstm, (0,2,3,1))
            gpu_pred_src_dstm = tf.transpose(gpu_pred_src_dstm, (0,2,3,1))

        tf.identity(gpu_pred_dst_dstm, name='out_face_mask')
        tf.identity(gpu_pred_src_dst, name='out_celeb_face')
        tf.identity(gpu_pred_src_dstm, name='out_celeb_face_mask')

        output_graph_def = tf.graph_util.convert_variables_to_constants(
            nn.tf_sess,
            tf.get_default_graph().as_graph_def(),
            ['out_face_mask','out_celeb_face','out_celeb_face_mask']
        )

        import tf2onnx
        with tf.device("/CPU:0"):
            model_proto, _ = tf2onnx.convert._convert_common(
                output_graph_def,
                name='SAEHD',
                input_names=['in_face:0'],
                output_names=['out_face_mask:0','out_celeb_face:0','out_celeb_face_mask:0'],
                opset=12,
                output_path=output_path)

    #override
    def get_model_filename_list(self):
        return self.model_filename_list

    #override
    def onSave(self):
        for model, filename in io.progress_bar_generator(self.get_model_filename_list(), "Saving", leave=False):
            model.save_weights ( self.get_strpath_storage_for_file(filename) )

    #override
    def should_save_preview_history(self):
        return (not io.is_colab() and self.iter % ( 10*(max(1,self.resolution // 64)) ) == 0) or \
               (io.is_colab() and self.iter % 100 == 0)

    #override
    def onTrainOneIter(self):
        if self.get_iter() == 0 and not self.pretrain and not self.pretrain_just_disabled:
            io.log_info('You are training the model from scratch. It is strongly recommended to use a pretrained model to speed up the training and improve the quality.\n')

        ( (warped_src, target_src, target_srcm, target_srcm_em), \
          (warped_dst, target_dst, target_dstm, target_dstm_em) ) = self.generate_next_samples()

        src_loss, dst_loss = self.src_dst_train (warped_src, target_src, target_srcm, target_srcm_em, warped_dst, target_dst, target_dstm, target_dstm_em)

        if self.options['true_face_power'] != 0 and not self.pretrain:
            self.D_train (warped_src, warped_dst)

        if self.gan_power != 0:
            self.D_src_dst_train (warped_src, target_src, target_srcm, target_srcm_em, warped_dst, target_dst, target_dstm, target_dstm_em)

        return ( ('src_loss', np.mean(src_loss) ), ('dst_loss', np.mean(dst_loss) ), )

    #override
    def onGetPreview(self, samples, for_history=False):
        ( (warped_src, target_src, target_srcm, target_srcm_em),
          (warped_dst, target_dst, target_dstm, target_dstm_em) ) = samples

        S, D, SS, DD, DDM, SD, SDM = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in ([target_src,target_dst] + self.AE_view (target_src, target_dst) ) ]
        DDM, SDM, = [ np.repeat (x, (3,), -1) for x in [DDM, SDM] ]

        target_srcm, target_dstm = [ nn.to_data_format(x,"NHWC", self.model_data_format) for x in ([target_srcm, target_dstm] )]

        n_samples = min(4, self.get_batch_size(), 800 // self.resolution )

        if self.resolution <= 256:
            result = []

            st = []
            for i in range(n_samples):
                ar = S[i], SS[i], D[i], DD[i], SD[i]
                st.append ( np.concatenate ( ar, axis=1) )
            result += [ ('SAEHD', np.concatenate (st, axis=0 )), ]


            st_m = []
            for i in range(n_samples):
                SD_mask = DDM[i]*SDM[i] if self.face_type < FaceType.HEAD else SDM[i]

                ar = S[i]*target_srcm[i], SS[i], D[i]*target_dstm[i], DD[i]*DDM[i], SD[i]*SD_mask
                st_m.append ( np.concatenate ( ar, axis=1) )

            result += [ ('SAEHD masked', np.concatenate (st_m, axis=0 )), ]
        else:
            result = []

            st = []
            for i in range(n_samples):
                ar = S[i], SS[i]
                st.append ( np.concatenate ( ar, axis=1) )
            result += [ ('SAEHD src-src', np.concatenate (st, axis=0 )), ]

            st = []
            for i in range(n_samples):
                ar = D[i], DD[i]
                st.append ( np.concatenate ( ar, axis=1) )
            result += [ ('SAEHD dst-dst', np.concatenate (st, axis=0 )), ]

            st = []
            for i in range(n_samples):
                ar = D[i], SD[i]
                st.append ( np.concatenate ( ar, axis=1) )
            result += [ ('SAEHD pred', np.concatenate (st, axis=0 )), ]


            st_m = []
            for i in range(n_samples):
                ar = S[i]*target_srcm[i], SS[i]
                st_m.append ( np.concatenate ( ar, axis=1) )
            result += [ ('SAEHD masked src-src', np.concatenate (st_m, axis=0 )), ]

            st_m = []
            for i in range(n_samples):
                ar = D[i]*target_dstm[i], DD[i]*DDM[i]
                st_m.append ( np.concatenate ( ar, axis=1) )
            result += [ ('SAEHD masked dst-dst', np.concatenate (st_m, axis=0 )), ]

            st_m = []
            for i in range(n_samples):
                SD_mask = DDM[i]*SDM[i] if self.face_type < FaceType.HEAD else SDM[i]
                ar = D[i]*target_dstm[i], SD[i]*SD_mask
                st_m.append ( np.concatenate ( ar, axis=1) )
            result += [ ('SAEHD masked pred', np.concatenate (st_m, axis=0 )), ]

        return result

    def predictor_func (self, face=None):
        face = nn.to_data_format(face[None,...], self.model_data_format, "NHWC")

        bgr, mask_dst_dstm, mask_src_dstm = [ nn.to_data_format(x,"NHWC", self.model_data_format).astype(np.float32) for x in self.AE_merge (face) ]

        return bgr[0], mask_src_dstm[0][...,0], mask_dst_dstm[0][...,0]

    #override
    def get_MergerConfig(self):
        import merger
        return self.predictor_func, (self.options['resolution'], self.options['resolution'], 3), merger.MergerConfigMasked(face_type=self.face_type, default_mode = 'overlay')

Model = SAEHDModel

models/Model_SAEHD/Model.py
oo92/TorchFaceLab
import multiprocessing
import operator
from functools import partial

import numpy as np

from core import mathlib
from core.interact import interact as io
from core.leras import nn
from facelib import FaceType, XSegNet
from models import ModelBase
from samplelib import *

class XSegModel(ModelBase):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, force_model_class_name='XSeg', **kwargs)

    #override
    def on_initialize_options(self):
        ask_override = self.ask_override()

        if not self.is_first_run() and ask_override:
            if io.input_bool(f"Restart training?", False, help_message="Reset model weights and start training from scratch."):
                self.set_iter(0)

        default_face_type          = self.options['face_type']          = self.load_or_def_option('face_type', 'wf')
        default_pretrain           = self.options['pretrain']           = self.load_or_def_option('pretrain', False)

        if self.is_first_run():
            self.options['face_type'] = io.input_str ("Face type", default_face_type, ['h','mf','f','wf','head'], help_message="Half / mid face / full face / whole face / head. Choose the same as your deepfake model.").lower()

        if self.is_first_run() or ask_override:
            self.ask_batch_size(4, range=[2,16])
            self.options['pretrain'] = io.input_bool ("Enable pretraining mode", default_pretrain)
        
        if not self.is_exporting and (self.options['pretrain'] and self.get_pretraining_data_path() is None):
            raise Exception("pretraining_data_path is not defined")
            
        self.pretrain_just_disabled = (default_pretrain == True and self.options['pretrain'] == False)
        
    #override
    def on_initialize(self):
        device_config = nn.getCurrentDeviceConfig()
        self.model_data_format = "NCHW" if self.is_exporting or (len(device_config.devices) != 0 and not self.is_debug()) else "NHWC"
        nn.initialize(data_format=self.model_data_format)
        tf = nn.tf

        device_config = nn.getCurrentDeviceConfig()
        devices = device_config.devices

        self.resolution = resolution = 256


        self.face_type = {'h'  : FaceType.HALF,
                          'mf' : FaceType.MID_FULL,
                          'f'  : FaceType.FULL,
                          'wf' : FaceType.WHOLE_FACE,
                          'head' : FaceType.HEAD}[ self.options['face_type'] ]
        
            
        place_model_on_cpu = len(devices) == 0
        models_opt_device = '/CPU:0' if place_model_on_cpu else nn.tf_default_device_name

        bgr_shape = nn.get4Dshape(resolution,resolution,3)
        mask_shape = nn.get4Dshape(resolution,resolution,1)

        # Initializing model classes
        self.model = XSegNet(name='XSeg',
                               resolution=resolution,
                               load_weights=not self.is_first_run(),
                               weights_file_root=self.get_model_root_path(),
                               training=True,
                               place_model_on_cpu=place_model_on_cpu,
                               optimizer=nn.RMSprop(lr=0.0001, lr_dropout=0.3, name='opt'),
                               data_format=nn.data_format)
        
        self.pretrain = self.options['pretrain']
        if self.pretrain_just_disabled:
            self.set_iter(0)
            
        if self.is_training:
            # Adjust batch size for multiple GPU
            gpu_count = max(1, len(devices) )
            bs_per_gpu = max(1, self.get_batch_size() // gpu_count)
            self.set_batch_size( gpu_count*bs_per_gpu)

            # Compute losses per GPU
            gpu_pred_list = []

            gpu_losses = []
            gpu_loss_gvs = []

            for gpu_id in range(gpu_count):
                with tf.device(f'/{devices[gpu_id].tf_dev_type}:{gpu_id}' if len(devices) != 0 else f'/CPU:0' ):
                    with tf.device(f'/CPU:0'):
                        # slice on CPU, otherwise all batch data will be transfered to GPU first
                        batch_slice = slice( gpu_id*bs_per_gpu, (gpu_id+1)*bs_per_gpu )
                        gpu_input_t       = self.model.input_t [batch_slice,:,:,:]
                        gpu_target_t      = self.model.target_t [batch_slice,:,:,:]

                    # process model tensors
                    gpu_pred_logits_t, gpu_pred_t = self.model.flow(gpu_input_t, pretrain=self.pretrain)
                    gpu_pred_list.append(gpu_pred_t)
                    
                    
                    if self.pretrain:
                        # Structural loss
                        gpu_loss =  tf.reduce_mean (5*nn.dssim(gpu_target_t, gpu_pred_t, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1])
                        gpu_loss += tf.reduce_mean (5*nn.dssim(gpu_target_t, gpu_pred_t, max_val=1.0, filter_size=int(resolution/23.2)), axis=[1])
                        # Pixel loss
                        gpu_loss += tf.reduce_mean (10*tf.square(gpu_target_t-gpu_pred_t), axis=[1,2,3])
                    else:
                        gpu_loss = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits(labels=gpu_target_t, logits=gpu_pred_logits_t), axis=[1,2,3])
                      
                    gpu_losses += [gpu_loss]

                    gpu_loss_gvs += [ nn.gradients ( gpu_loss, self.model.get_weights() ) ]


            # Average losses and gradients, and create optimizer update ops
            #with tf.device(f'/CPU:0'): # Temporary fix. Unknown bug with training freeze starts from 2.4.0, but 2.3.1 was ok
            with tf.device (models_opt_device):
                pred = tf.concat(gpu_pred_list, 0)
                loss = tf.concat(gpu_losses, 0)
                loss_gv_op = self.model.opt.get_update_op (nn.average_gv_list (gpu_loss_gvs))


            # Initializing training and view functions
            if self.pretrain:
                def train(input_np, target_np):
                    l, _ = nn.tf_sess.run ( [loss, loss_gv_op], feed_dict={self.model.input_t :input_np, self.model.target_t :target_np})
                    return l
            else:
                def train(input_np, target_np):
                    l, _ = nn.tf_sess.run ( [loss, loss_gv_op], feed_dict={self.model.input_t :input_np, self.model.target_t :target_np })
                    return l
            self.train = train

            def view(input_np):
                return nn.tf_sess.run ( [pred], feed_dict={self.model.input_t :input_np})
            self.view = view

            # initializing sample generators
            cpu_count = min(multiprocessing.cpu_count(), 8)
            src_dst_generators_count = cpu_count // 2
            src_generators_count = cpu_count // 2
            dst_generators_count = cpu_count // 2
            
            if self.pretrain:
                pretrain_gen = SampleGeneratorFace(self.get_pretraining_data_path(), debug=self.is_debug(), batch_size=self.get_batch_size(),
                                    sample_process_options=SampleProcessor.Options(random_flip=True),
                                    output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':True, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                                            {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':True, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G,   'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},                                                            
                                                          ],
                                    uniform_yaw_distribution=False,
                                    generators_count=cpu_count )
                self.set_training_data_generators ([pretrain_gen])
            else:   
                srcdst_generator = SampleGeneratorFaceXSeg([self.training_data_src_path, self.training_data_dst_path],
                                                            debug=self.is_debug(),
                                                            batch_size=self.get_batch_size(),
                                                            resolution=resolution,
                                                            face_type=self.face_type,
                                                            generators_count=src_dst_generators_count,
                                                            data_format=nn.data_format)

                src_generator = SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
                                                    sample_process_options=SampleProcessor.Options(random_flip=False),
                                                    output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,  'warp':False, 'transform':False, 'channel_type' : SampleProcessor.ChannelType.BGR, 'border_replicate':False, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                                                        ],
                                                    generators_count=src_generators_count,
                                                    raise_on_no_data=False )
                dst_generator = SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
                                                    sample_process_options=SampleProcessor.Options(random_flip=False),
                                                    output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,  'warp':False, 'transform':False, 'channel_type' : SampleProcessor.ChannelType.BGR, 'border_replicate':False, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                                                        ],
                                                    generators_count=dst_generators_count,
                                                    raise_on_no_data=False )

                self.set_training_data_generators ([srcdst_generator, src_generator, dst_generator])

    #override
    def get_model_filename_list(self):
        return self.model.model_filename_list

    #override
    def onSave(self):
        self.model.save_weights()

    #override
    def onTrainOneIter(self):
        image_np, target_np = self.generate_next_samples()[0]
        loss = self.train (image_np, target_np)
        
        return ( ('loss', np.mean(loss) ), )

    #override
    def onGetPreview(self, samples, for_history=False):
        n_samples = min(4, self.get_batch_size(), 800 // self.resolution )
        
        if self.pretrain:
            srcdst_samples, = samples       
            image_np, mask_np = srcdst_samples     
        else:
            srcdst_samples, src_samples, dst_samples = samples
            image_np, mask_np = srcdst_samples

        I, M, IM, = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in ([image_np,mask_np] + self.view (image_np) ) ]
        M, IM, = [ np.repeat (x, (3,), -1) for x in [M, IM] ]

        green_bg = np.tile( np.array([0,1,0], dtype=np.float32)[None,None,...], (self.resolution,self.resolution,1) )

        result = []
        st = []
        for i in range(n_samples):
            if self.pretrain:
                ar = I[i], IM[i]
            else:
                ar = I[i]*M[i]+0.5*I[i]*(1-M[i])+0.5*green_bg*(1-M[i]), IM[i], I[i]*IM[i]+0.5*I[i]*(1-IM[i]) + 0.5*green_bg*(1-IM[i])
            st.append ( np.concatenate ( ar, axis=1) )
        result += [ ('XSeg training faces', np.concatenate (st, axis=0 )), ]

        if not self.pretrain and len(src_samples) != 0:
            src_np, = src_samples


            D, DM, = [ np.clip(nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in ([src_np] + self.view (src_np) ) ]
            DM, = [ np.repeat (x, (3,), -1) for x in [DM] ]

            st = []
            for i in range(n_samples):
                ar = D[i], DM[i], D[i]*DM[i] + 0.5*D[i]*(1-DM[i]) + 0.5*green_bg*(1-DM[i])
                st.append ( np.concatenate ( ar, axis=1) )

            result += [ ('XSeg src faces', np.concatenate (st, axis=0 )), ]

        if not self.pretrain and len(dst_samples) != 0:
            dst_np, = dst_samples


            D, DM, = [ np.clip(nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in ([dst_np] + self.view (dst_np) ) ]
            DM, = [ np.repeat (x, (3,), -1) for x in [DM] ]

            st = []
            for i in range(n_samples):
                ar = D[i], DM[i], D[i]*DM[i]  + 0.5*D[i]*(1-DM[i]) + 0.5*green_bg*(1-DM[i])
                st.append ( np.concatenate ( ar, axis=1) )

            result += [ ('XSeg dst faces', np.concatenate (st, axis=0 )), ]

        return result
        
    def export_dfm (self):
        output_path = self.get_strpath_storage_for_file(f'model.onnx')
        io.log_info(f'Dumping .onnx to {output_path}')
        tf = nn.tf
        
        with tf.device (nn.tf_default_device_name):
            input_t = tf.placeholder (nn.floatx, (None, self.resolution, self.resolution, 3), name='in_face')
            input_t = tf.transpose(input_t, (0,3,1,2))
            _, pred_t = self.model.flow(input_t)
            pred_t = tf.transpose(pred_t, (0,2,3,1))
            
        tf.identity(pred_t, name='out_mask')
        
        output_graph_def = tf.graph_util.convert_variables_to_constants(
            nn.tf_sess, 
            tf.get_default_graph().as_graph_def(), 
            ['out_mask']
        ) 
        
        import tf2onnx
        with tf.device("/CPU:0"):
            model_proto, _ = tf2onnx.convert._convert_common(
                output_graph_def,
                name='XSeg',
                input_names=['in_face:0'],
                output_names=['out_mask:0'],
                opset=13,
                output_path=output_path)
                
Model = XSegModel
models/Model_XSeg/Model.py
oo92/TorchFaceLab
from .ModelBase import ModelBase

def import_model(model_class_name):
    module = __import__('Model_'+model_class_name, globals(), locals(), [], 1)
    return getattr(module, 'Model')

models/__init__.py
oo92/TorchFaceLab
import pickle
import shutil
import struct
from pathlib import Path

import samplelib.SampleLoader
from core.interact import interact as io
from samplelib import Sample
from core import pathex

packed_faceset_filename = 'faceset.pak'

class PackedFaceset():
    VERSION = 1

    @staticmethod
    def pack(samples_path):
        samples_dat_path = samples_path / packed_faceset_filename

        if samples_dat_path.exists():
            io.log_info(f"{samples_dat_path} : file already exists !")
            io.input("Press enter to continue and overwrite.")

        as_person_faceset = False
        dir_names = pathex.get_all_dir_names(samples_path)
        if len(dir_names) != 0:
            as_person_faceset = io.input_bool(f"{len(dir_names)} subdirectories found, process as person faceset?", True)

        if as_person_faceset:
            image_paths = []

            for dir_name in dir_names:
                image_paths += pathex.get_image_paths(samples_path / dir_name)
        else:
            image_paths = pathex.get_image_paths(samples_path)

        samples = samplelib.SampleLoader.load_face_samples(image_paths)
        samples_len = len(samples)

        samples_configs = []
        for sample in io.progress_bar_generator (samples, "Processing"):
            sample_filepath = Path(sample.filename)
            sample.filename = sample_filepath.name

            if as_person_faceset:
                sample.person_name = sample_filepath.parent.name
            samples_configs.append ( sample.get_config() )
        samples_bytes = pickle.dumps(samples_configs, 4)

        of = open(samples_dat_path, "wb")
        of.write ( struct.pack ("Q", PackedFaceset.VERSION ) )
        of.write ( struct.pack ("Q", len(samples_bytes) ) )
        of.write ( samples_bytes )

        del samples_bytes   #just free mem
        del samples_configs

        sample_data_table_offset = of.tell()
        of.write ( bytes( 8*(samples_len+1) ) ) #sample data offset table

        data_start_offset = of.tell()
        offsets = []

        for sample in io.progress_bar_generator(samples, "Packing"):
            try:
                if sample.person_name is not None:
                    sample_path = samples_path / sample.person_name / sample.filename
                else:
                    sample_path = samples_path / sample.filename


                with open(sample_path, "rb") as f:
                   b = f.read()

                offsets.append ( of.tell() - data_start_offset )
                of.write(b)
            except:
                raise Exception(f"error while processing sample {sample_path}")

        offsets.append ( of.tell() )

        of.seek(sample_data_table_offset, 0)
        for offset in offsets:
            of.write ( struct.pack("Q", offset) )
        of.seek(0,2)
        of.close()
        
        if io.input_bool(f"Delete original files?", True):
            for filename in io.progress_bar_generator(image_paths, "Deleting files"):
                Path(filename).unlink()

            if as_person_faceset:
                for dir_name in io.progress_bar_generator(dir_names, "Deleting dirs"):
                    dir_path = samples_path / dir_name
                    try:
                        shutil.rmtree(dir_path)
                    except:
                        io.log_info (f"unable to remove: {dir_path} ")

    @staticmethod
    def unpack(samples_path):
        samples_dat_path = samples_path / packed_faceset_filename
        if not samples_dat_path.exists():
            io.log_info(f"{samples_dat_path} : file not found.")
            return

        samples = PackedFaceset.load(samples_path)

        for sample in io.progress_bar_generator(samples, "Unpacking"):
            person_name = sample.person_name
            if person_name is not None:
                person_path = samples_path / person_name
                person_path.mkdir(parents=True, exist_ok=True)

                target_filepath = person_path / sample.filename
            else:
                target_filepath = samples_path / sample.filename

            with open(target_filepath, "wb") as f:
                f.write( sample.read_raw_file() )

        samples_dat_path.unlink()

    @staticmethod
    def path_contains(samples_path):
        samples_dat_path = samples_path / packed_faceset_filename
        return samples_dat_path.exists()
    
    @staticmethod
    def load(samples_path):
        samples_dat_path = samples_path / packed_faceset_filename
        if not samples_dat_path.exists():
            return None

        f = open(samples_dat_path, "rb")
        version, = struct.unpack("Q", f.read(8) )
        if version != PackedFaceset.VERSION:
            raise NotImplementedError

        sizeof_samples_bytes, = struct.unpack("Q", f.read(8) )

        samples_configs = pickle.loads ( f.read(sizeof_samples_bytes) )
        samples = []
        for sample_config in samples_configs:
            sample_config = pickle.loads(pickle.dumps (sample_config))
            samples.append ( Sample (**sample_config) )

        offsets = [ struct.unpack("Q", f.read(8) )[0] for _ in range(len(samples)+1) ]
        data_start_offset = f.tell()
        f.close()

        for i, sample in enumerate(samples):
            start_offset, end_offset = offsets[i], offsets[i+1]
            sample.set_filename_offset_size( str(samples_dat_path), data_start_offset+start_offset, end_offset-start_offset )

        return samples

samplelib/PackedFaceset.py
oo92/TorchFaceLab
from enum import IntEnum
from pathlib import Path

import cv2
import numpy as np

from core.cv2ex import *
from facelib import LandmarksProcessor
from core import imagelib
from core.imagelib import SegIEPolys

class SampleType(IntEnum):
    IMAGE = 0 #raw image

    FACE_BEGIN = 1
    FACE = 1                        #aligned face unsorted
    FACE_PERSON = 2                 #aligned face person
    FACE_TEMPORAL_SORTED = 3        #sorted by source filename
    FACE_END = 3

    QTY = 4

class Sample(object):
    __slots__ = ['sample_type',
                 'filename',
                 'face_type',
                 'shape',
                 'landmarks',
                 'seg_ie_polys',
                 'xseg_mask',
                 'xseg_mask_compressed',
                 'eyebrows_expand_mod',
                 'source_filename',
                 'person_name',
                 'pitch_yaw_roll',
                 '_filename_offset_size',
                ]

    def __init__(self, sample_type=None,
                       filename=None,
                       face_type=None,
                       shape=None,
                       landmarks=None,
                       seg_ie_polys=None,
                       xseg_mask=None,
                       xseg_mask_compressed=None,
                       eyebrows_expand_mod=None,
                       source_filename=None,
                       person_name=None,
                       pitch_yaw_roll=None,
                       **kwargs):

        self.sample_type = sample_type if sample_type is not None else SampleType.IMAGE
        self.filename = filename
        self.face_type = face_type
        self.shape = shape
        self.landmarks = np.array(landmarks) if landmarks is not None else None
        
        if isinstance(seg_ie_polys, SegIEPolys):
            self.seg_ie_polys = seg_ie_polys
        else:
            self.seg_ie_polys = SegIEPolys.load(seg_ie_polys)
        
        self.xseg_mask = xseg_mask
        self.xseg_mask_compressed = xseg_mask_compressed
        
        if self.xseg_mask_compressed is None and self.xseg_mask is not None:
            xseg_mask = np.clip( imagelib.normalize_channels(xseg_mask, 1)*255, 0, 255 ).astype(np.uint8)        
            ret, xseg_mask_compressed = cv2.imencode('.png', xseg_mask)
            if not ret:
                raise Exception("Sample(): unable to generate xseg_mask_compressed")
            self.xseg_mask_compressed = xseg_mask_compressed
            self.xseg_mask = None
 
        self.eyebrows_expand_mod = eyebrows_expand_mod if eyebrows_expand_mod is not None else 1.0
        self.source_filename = source_filename
        self.person_name = person_name
        self.pitch_yaw_roll = pitch_yaw_roll

        self._filename_offset_size = None

    def has_xseg_mask(self):
        return self.xseg_mask is not None or self.xseg_mask_compressed is not None
        
    def get_xseg_mask(self):
        if self.xseg_mask_compressed is not None:
            xseg_mask = cv2.imdecode(self.xseg_mask_compressed, cv2.IMREAD_UNCHANGED)
            if len(xseg_mask.shape) == 2:
                xseg_mask = xseg_mask[...,None]
            return xseg_mask.astype(np.float32) / 255.0
        return self.xseg_mask
        
    def get_pitch_yaw_roll(self):
        if self.pitch_yaw_roll is None:
            self.pitch_yaw_roll = LandmarksProcessor.estimate_pitch_yaw_roll(self.landmarks, size=self.shape[1])
        return self.pitch_yaw_roll

    def set_filename_offset_size(self, filename, offset, size):
        self._filename_offset_size = (filename, offset, size)

    def read_raw_file(self, filename=None):
        if self._filename_offset_size is not None:
            filename, offset, size = self._filename_offset_size
            with open(filename, "rb") as f:
                f.seek( offset, 0)
                return f.read (size)
        else:
            with open(filename, "rb") as f:
                return f.read()

    def load_bgr(self):
        img = cv2_imread (self.filename, loader_func=self.read_raw_file).astype(np.float32) / 255.0
        return img

    def get_config(self):
        return {'sample_type': self.sample_type,
                'filename': self.filename,
                'face_type': self.face_type,
                'shape': self.shape,
                'landmarks': self.landmarks.tolist(),
                'seg_ie_polys': self.seg_ie_polys.dump(),
                'xseg_mask' : self.xseg_mask,
                'xseg_mask_compressed' : self.xseg_mask_compressed,
                'eyebrows_expand_mod': self.eyebrows_expand_mod,
                'source_filename': self.source_filename,
                'person_name': self.person_name
               }

samplelib/Sample.py
oo92/TorchFaceLab
from pathlib import Path

'''
You can implement your own SampleGenerator
'''
class SampleGeneratorBase(object):


    def __init__ (self, debug=False, batch_size=1):
        self.debug = debug
        self.batch_size = 1 if self.debug else batch_size
        self.last_generation = None
        self.active = True

    def set_active(self, is_active):
        self.active = is_active

    def generate_next(self):
        if not self.active and self.last_generation is not None:
            return self.last_generation
        self.last_generation = next(self)
        return self.last_generation

    #overridable
    def __iter__(self):
        #implement your own iterator
        return self

    def __next__(self):
        #implement your own iterator
        return None
    
    #overridable
    def is_initialized(self):
        return True
samplelib/SampleGeneratorBase.py
oo92/TorchFaceLab
import multiprocessing
import time
import traceback

import cv2
import numpy as np

from core import mplib
from core.interact import interact as io
from core.joblib import SubprocessGenerator, ThisThreadGenerator
from facelib import LandmarksProcessor
from samplelib import (SampleGeneratorBase, SampleLoader, SampleProcessor,
                       SampleType)


'''
arg
output_sample_types = [
                        [SampleProcessor.TypeFlags, size, (optional) {} opts ] ,
                        ...
                      ]
'''
class SampleGeneratorFace(SampleGeneratorBase):
    def __init__ (self, samples_path, debug=False, batch_size=1,
                        random_ct_samples_path=None,
                        sample_process_options=SampleProcessor.Options(),
                        output_sample_types=[],
                        uniform_yaw_distribution=False,
                        generators_count=4,
                        raise_on_no_data=True,                        
                        **kwargs):

        super().__init__(debug, batch_size)
        self.initialized = False
        self.sample_process_options = sample_process_options
        self.output_sample_types = output_sample_types
        
        if self.debug:
            self.generators_count = 1
        else:
            self.generators_count = max(1, generators_count)

        samples = SampleLoader.load (SampleType.FACE, samples_path)
        self.samples_len = len(samples)
        
        if self.samples_len == 0:
            if raise_on_no_data:
                raise ValueError('No training data provided.')
            else:
                return
                
        if uniform_yaw_distribution:
            samples_pyr = [ ( idx, sample.get_pitch_yaw_roll() ) for idx, sample in enumerate(samples) ]
            
            grads = 128
            #instead of math.pi / 2, using -1.2,+1.2 because actually maximum yaw for 2DFAN landmarks are -1.2+1.2
            grads_space = np.linspace (-1.2, 1.2,grads)

            yaws_sample_list = [None]*grads
            for g in io.progress_bar_generator ( range(grads), "Sort by yaw"):
                yaw = grads_space[g]
                next_yaw = grads_space[g+1] if g < grads-1 else yaw

                yaw_samples = []
                for idx, pyr in samples_pyr:
                    s_yaw = -pyr[1]
                    if (g == 0          and s_yaw < next_yaw) or \
                    (g < grads-1     and s_yaw >= yaw and s_yaw < next_yaw) or \
                    (g == grads-1    and s_yaw >= yaw):
                        yaw_samples += [ idx ]
                if len(yaw_samples) > 0:
                    yaws_sample_list[g] = yaw_samples
            
            yaws_sample_list = [ y for y in yaws_sample_list if y is not None ]
            
            index_host = mplib.Index2DHost( yaws_sample_list )
        else:
            index_host = mplib.IndexHost(self.samples_len)

        if random_ct_samples_path is not None:
            ct_samples = SampleLoader.load (SampleType.FACE, random_ct_samples_path)
            ct_index_host = mplib.IndexHost( len(ct_samples) )
        else:
            ct_samples = None
            ct_index_host = None

        if self.debug:
            self.generators = [ThisThreadGenerator ( self.batch_func, (samples, index_host.create_cli(), ct_samples, ct_index_host.create_cli() if ct_index_host is not None else None) )]
        else:
            self.generators = [SubprocessGenerator ( self.batch_func, (samples, index_host.create_cli(), ct_samples, ct_index_host.create_cli() if ct_index_host is not None else None), start_now=False ) \
                               for i in range(self.generators_count) ]
                               
            SubprocessGenerator.start_in_parallel( self.generators )

        self.generator_counter = -1
        
        self.initialized = True
        
    #overridable
    def is_initialized(self):
        return self.initialized
        
    def __iter__(self):
        return self

    def __next__(self):
        if not self.initialized:
            return []
            
        self.generator_counter += 1
        generator = self.generators[self.generator_counter % len(self.generators) ]
        return next(generator)

    def batch_func(self, param ):
        samples, index_host, ct_samples, ct_index_host = param
 
        bs = self.batch_size
        while True:
            batches = None

            indexes = index_host.multi_get(bs)
            ct_indexes = ct_index_host.multi_get(bs) if ct_samples is not None else None

            t = time.time()
            for n_batch in range(bs):
                sample_idx = indexes[n_batch]
                sample = samples[sample_idx]

                ct_sample = None
                if ct_samples is not None:
                    ct_sample = ct_samples[ct_indexes[n_batch]]

                try:
                    x, = SampleProcessor.process ([sample], self.sample_process_options, self.output_sample_types, self.debug, ct_sample=ct_sample)
                except:
                    raise Exception ("Exception occured in sample %s. Error: %s" % (sample.filename, traceback.format_exc() ) )

                if batches is None:
                    batches = [ [] for _ in range(len(x)) ]

                for i in range(len(x)):
                    batches[i].append ( x[i] )

            yield [ np.array(batch) for batch in batches]

samplelib/SampleGeneratorFace.py
oo92/TorchFaceLab
import multiprocessing
import pickle
import time
import traceback
from enum import IntEnum

import cv2
import numpy as np

from core import imagelib, mplib, pathex
from core.cv2ex import *
from core.interact import interact as io
from core.joblib import SubprocessGenerator, ThisThreadGenerator
from facelib import LandmarksProcessor
from samplelib import SampleGeneratorBase


class MaskType(IntEnum):
    none   = 0,
    cloth  = 1,
    ear_r  = 2,
    eye_g  = 3,
    hair   = 4,
    hat    = 5,
    l_brow = 6,
    l_ear  = 7,
    l_eye  = 8,
    l_lip  = 9,
    mouth  = 10,
    neck   = 11,
    neck_l = 12,
    nose   = 13,
    r_brow = 14,
    r_ear  = 15,
    r_eye  = 16,
    skin   = 17,
    u_lip  = 18



MaskType_to_name = {
    int(MaskType.none  ) : 'none',
    int(MaskType.cloth ) : 'cloth',
    int(MaskType.ear_r ) : 'ear_r',
    int(MaskType.eye_g ) : 'eye_g',
    int(MaskType.hair  ) : 'hair',
    int(MaskType.hat   ) : 'hat',
    int(MaskType.l_brow) : 'l_brow',
    int(MaskType.l_ear ) : 'l_ear',
    int(MaskType.l_eye ) : 'l_eye',
    int(MaskType.l_lip ) : 'l_lip',
    int(MaskType.mouth ) : 'mouth',
    int(MaskType.neck  ) : 'neck',
    int(MaskType.neck_l) : 'neck_l',
    int(MaskType.nose  ) : 'nose',
    int(MaskType.r_brow) : 'r_brow',
    int(MaskType.r_ear ) : 'r_ear',
    int(MaskType.r_eye ) : 'r_eye',
    int(MaskType.skin  ) : 'skin',
    int(MaskType.u_lip ) : 'u_lip',
}

MaskType_from_name = { MaskType_to_name[k] : k for k in MaskType_to_name.keys() }

class SampleGeneratorFaceCelebAMaskHQ(SampleGeneratorBase):
    def __init__ (self, root_path, debug=False, batch_size=1, resolution=256,
                        generators_count=4, data_format="NHWC",
                        **kwargs):

        super().__init__(debug, batch_size)
        self.initialized = False

        dataset_path = root_path / 'CelebAMask-HQ'
        if not dataset_path.exists():
            raise ValueError(f'Unable to find {dataset_path}')

        images_path = dataset_path /'CelebA-HQ-img'
        if not images_path.exists():
            raise ValueError(f'Unable to find {images_path}')

        masks_path = dataset_path / 'CelebAMask-HQ-mask-anno'
        if not masks_path.exists():
            raise ValueError(f'Unable to find {masks_path}')


        if self.debug:
            self.generators_count = 1
        else:
            self.generators_count = max(1, generators_count)

        source_images_paths = pathex.get_image_paths(images_path, return_Path_class=True)
        source_images_paths_len = len(source_images_paths)
        mask_images_paths = pathex.get_image_paths(masks_path, subdirs=True, return_Path_class=True)

        if source_images_paths_len == 0 or len(mask_images_paths) == 0:
            raise ValueError('No training data provided.')

        mask_file_id_hash = {}

        for filepath in io.progress_bar_generator(mask_images_paths, "Loading"):
            stem = filepath.stem

            file_id, mask_type = stem.split('_', 1)
            file_id = int(file_id)

            if file_id not in mask_file_id_hash:
                mask_file_id_hash[file_id] = {}

            mask_file_id_hash[file_id][ MaskType_from_name[mask_type] ] = str(filepath.relative_to(masks_path))

        source_file_id_set = set()

        for filepath in source_images_paths:
            stem = filepath.stem

            file_id = int(stem)
            source_file_id_set.update ( {file_id} )

        for k in mask_file_id_hash.keys():
            if k not in source_file_id_set:
                io.log_err (f"Corrupted dataset: {k} not in {images_path}")



        if self.debug:
            self.generators = [ThisThreadGenerator ( self.batch_func, (images_path, masks_path, mask_file_id_hash, data_format) )]
        else:
            self.generators = [SubprocessGenerator ( self.batch_func, (images_path, masks_path, mask_file_id_hash, data_format), start_now=False ) \
                               for i in range(self.generators_count) ]

            SubprocessGenerator.start_in_parallel( self.generators )

        self.generator_counter = -1

        self.initialized = True

    #overridable
    def is_initialized(self):
        return self.initialized

    def __iter__(self):
        return self

    def __next__(self):
        self.generator_counter += 1
        generator = self.generators[self.generator_counter % len(self.generators) ]
        return next(generator)

    def batch_func(self, param ):
        images_path, masks_path, mask_file_id_hash, data_format = param

        file_ids = list(mask_file_id_hash.keys())

        shuffle_file_ids = []

        resolution = 256
        random_flip = True
        rotation_range=[-15,15]
        scale_range=[-0.10, 0.95]
        tx_range=[-0.3, 0.3]
        ty_range=[-0.3, 0.3]

        random_bilinear_resize = (25,75)
        motion_blur = (25, 5)
        gaussian_blur = (25, 5)

        bs = self.batch_size
        while True:
            batches = None

            n_batch = 0
            while n_batch < bs:
                try:
                    if len(shuffle_file_ids) == 0:
                        shuffle_file_ids = file_ids.copy()
                        np.random.shuffle(shuffle_file_ids)

                    file_id = shuffle_file_ids.pop()
                    masks = mask_file_id_hash[file_id]
                    image_path = images_path / f'{file_id}.jpg'

                    skin_path = masks.get(MaskType.skin, None)
                    hair_path = masks.get(MaskType.hair, None)
                    hat_path = masks.get(MaskType.hat, None)
                    #neck_path = masks.get(MaskType.neck, None)

                    img = cv2_imread(image_path).astype(np.float32) / 255.0
                    mask = cv2_imread(masks_path / skin_path)[...,0:1].astype(np.float32) / 255.0

                    if hair_path is not None:
                        hair_path = masks_path / hair_path
                        if hair_path.exists():
                            hair = cv2_imread(hair_path)[...,0:1].astype(np.float32) / 255.0
                            mask *= (1-hair)

                    if hat_path is not None:
                        hat_path = masks_path / hat_path
                        if hat_path.exists():
                            hat = cv2_imread(hat_path)[...,0:1].astype(np.float32) / 255.0
                            mask *= (1-hat)
                    
                    #if neck_path is not None:
                    #    neck_path = masks_path / neck_path
                    #    if neck_path.exists():
                    #        neck = cv2_imread(neck_path)[...,0:1].astype(np.float32) / 255.0
                    #        mask = np.clip(mask+neck, 0, 1)
                            
                    warp_params = imagelib.gen_warp_params(resolution, random_flip, rotation_range=rotation_range, scale_range=scale_range, tx_range=tx_range, ty_range=ty_range )
  
                    img = cv2.resize( img, (resolution,resolution), cv2.INTER_LANCZOS4 )
                    h, s, v = cv2.split(cv2.cvtColor(img, cv2.COLOR_BGR2HSV))
                    h = ( h + np.random.randint(360) ) % 360
                    s = np.clip ( s + np.random.random()-0.5, 0, 1 )
                    v = np.clip ( v + np.random.random()/2-0.25, 0, 1 )                    
                    img = np.clip( cv2.cvtColor(cv2.merge([h, s, v]), cv2.COLOR_HSV2BGR) , 0, 1 )
                            
                    if motion_blur is not None:
                        chance, mb_max_size = motion_blur
                        chance = np.clip(chance, 0, 100)

                        mblur_rnd_chance = np.random.randint(100)
                        mblur_rnd_kernel = np.random.randint(mb_max_size)+1
                        mblur_rnd_deg    = np.random.randint(360)

                        if mblur_rnd_chance < chance:
                            img = imagelib.LinearMotionBlur (img, mblur_rnd_kernel, mblur_rnd_deg )

                    img = imagelib.warp_by_params (warp_params, img,  can_warp=True, can_transform=True, can_flip=True, border_replicate=False, cv2_inter=cv2.INTER_LANCZOS4)
                    
                    if gaussian_blur is not None:
                        chance, kernel_max_size = gaussian_blur
                        chance = np.clip(chance, 0, 100)

                        gblur_rnd_chance = np.random.randint(100)
                        gblur_rnd_kernel = np.random.randint(kernel_max_size)*2+1

                        if gblur_rnd_chance < chance:
                            img = cv2.GaussianBlur(img, (gblur_rnd_kernel,) *2 , 0)
                            
                    if random_bilinear_resize is not None:
                        chance, max_size_per = random_bilinear_resize
                        chance = np.clip(chance, 0, 100)                        
                        pick_chance = np.random.randint(100)                        
                        resize_to = resolution - int( np.random.rand()* int(resolution*(max_size_per/100.0)) )                        
                        img = cv2.resize (img, (resize_to,resize_to), cv2.INTER_LINEAR )
                        img = cv2.resize (img, (resolution,resolution), cv2.INTER_LINEAR )
                        
                            
                    mask = cv2.resize( mask, (resolution,resolution), cv2.INTER_LANCZOS4 )[...,None]
                    mask = imagelib.warp_by_params (warp_params, mask, can_warp=True, can_transform=True, can_flip=True, border_replicate=False, cv2_inter=cv2.INTER_LANCZOS4)
                    mask[mask < 0.5] = 0.0
                    mask[mask >= 0.5] = 1.0
                    mask = np.clip(mask, 0, 1)

                    if data_format == "NCHW":
                        img = np.transpose(img, (2,0,1) )
                        mask = np.transpose(mask, (2,0,1) )
                        
                    if batches is None:
                        batches = [ [], [] ]
                    
                    batches[0].append ( img )
                    batches[1].append ( mask )

                    n_batch += 1
                except:
                    io.log_err ( traceback.format_exc() )

            yield [ np.array(batch) for batch in batches]

samplelib/SampleGeneratorFaceCelebAMaskHQ.py
oo92/TorchFaceLab
import copy
import multiprocessing
import traceback

import cv2
import numpy as np

from core import mplib
from core.joblib import SubprocessGenerator, ThisThreadGenerator
from facelib import LandmarksProcessor
from samplelib import (SampleGeneratorBase, SampleLoader, SampleProcessor,
                       SampleType)



class Index2DHost():
    """
    Provides random shuffled 2D indexes for multiprocesses
    """
    def __init__(self, indexes2D):
        self.sq = multiprocessing.Queue()
        self.cqs = []
        self.clis = []
        self.thread = threading.Thread(target=self.host_thread, args=(indexes2D,) )
        self.thread.daemon = True
        self.thread.start()

    def host_thread(self, indexes2D):
        indexes_counts_len = len(indexes2D)

        idxs = [*range(indexes_counts_len)]
        idxs_2D = [None]*indexes_counts_len
        shuffle_idxs = []
        shuffle_idxs_2D = [None]*indexes_counts_len
        for i in range(indexes_counts_len):
            idxs_2D[i] = indexes2D[i]
            shuffle_idxs_2D[i] = []

        sq = self.sq

        while True:
            while not sq.empty():
                obj = sq.get()
                cq_id, cmd = obj[0], obj[1]

                if cmd == 0: #get_1D
                    count = obj[2]

                    result = []
                    for i in range(count):
                        if len(shuffle_idxs) == 0:
                            shuffle_idxs = idxs.copy()
                            np.random.shuffle(shuffle_idxs)
                        result.append(shuffle_idxs.pop())
                    self.cqs[cq_id].put (result)
                elif cmd == 1: #get_2D
                    targ_idxs,count = obj[2], obj[3]
                    result = []

                    for targ_idx in targ_idxs:
                        sub_idxs = []
                        for i in range(count):
                            ar = shuffle_idxs_2D[targ_idx]
                            if len(ar) == 0:
                                ar = shuffle_idxs_2D[targ_idx] = idxs_2D[targ_idx].copy()
                                np.random.shuffle(ar)
                            sub_idxs.append(ar.pop())
                        result.append (sub_idxs)
                    self.cqs[cq_id].put (result)

            time.sleep(0.001)

    def create_cli(self):
        cq = multiprocessing.Queue()
        self.cqs.append ( cq )
        cq_id = len(self.cqs)-1
        return Index2DHost.Cli(self.sq, cq, cq_id)

    # disable pickling
    def __getstate__(self):
        return dict()
    def __setstate__(self, d):
        self.__dict__.update(d)

    class Cli():
        def __init__(self, sq, cq, cq_id):
            self.sq = sq
            self.cq = cq
            self.cq_id = cq_id

        def get_1D(self, count):
            self.sq.put ( (self.cq_id,0, count) )

            while True:
                if not self.cq.empty():
                    return self.cq.get()
                time.sleep(0.001)

        def get_2D(self, idxs, count):
            self.sq.put ( (self.cq_id,1,idxs,count) )

            while True:
                if not self.cq.empty():
                    return self.cq.get()
                time.sleep(0.001)
                
'''
arg
output_sample_types = [
                        [SampleProcessor.TypeFlags, size, (optional) {} opts ] ,
                        ...
                      ]
'''
class SampleGeneratorFacePerson(SampleGeneratorBase):
    def __init__ (self, samples_path, debug=False, batch_size=1,
                        sample_process_options=SampleProcessor.Options(),
                        output_sample_types=[],
                        person_id_mode=1,
                        **kwargs):

        super().__init__(debug, batch_size)
        self.sample_process_options = sample_process_options
        self.output_sample_types = output_sample_types
        self.person_id_mode = person_id_mode

        raise NotImplementedError("Currently SampleGeneratorFacePerson is not implemented.")

        samples_host = SampleLoader.mp_host (SampleType.FACE, samples_path)
        samples = samples_host.get_list()
        self.samples_len = len(samples)

        if self.samples_len == 0:
            raise ValueError('No training data provided.')

        unique_person_names = { sample.person_name for sample in samples }
        persons_name_idxs = { person_name : [] for person_name in unique_person_names }
        for i,sample in enumerate(samples):
            persons_name_idxs[sample.person_name].append (i)
        indexes2D = [ persons_name_idxs[person_name] for person_name in unique_person_names ]
        index2d_host = Index2DHost(indexes2D)

        if self.debug:
            self.generators_count = 1
            self.generators = [iter_utils.ThisThreadGenerator ( self.batch_func, (samples_host.create_cli(), index2d_host.create_cli(),) )]
        else:
            self.generators_count = np.clip(multiprocessing.cpu_count(), 2, 4)
            self.generators = [iter_utils.SubprocessGenerator ( self.batch_func, (samples_host.create_cli(), index2d_host.create_cli(),) ) for i in range(self.generators_count) ]

        self.generator_counter = -1

    def __iter__(self):
        return self

    def __next__(self):
        self.generator_counter += 1
        generator = self.generators[self.generator_counter % len(self.generators) ]
        return next(generator)

    def batch_func(self, param ):
        samples, index2d_host, = param
        bs = self.batch_size

        while True:
            person_idxs = index2d_host.get_1D(bs)
            samples_idxs = index2d_host.get_2D(person_idxs, 1)

            batches = None
            for n_batch in range(bs):
                person_id = person_idxs[n_batch]
                sample_idx = samples_idxs[n_batch][0]

                sample = samples[ sample_idx ]
                try:
                    x, = SampleProcessor.process ([sample], self.sample_process_options, self.output_sample_types, self.debug)
                except:
                    raise Exception ("Exception occured in sample %s. Error: %s" % (sample.filename, traceback.format_exc() ) )

                if batches is None:
                    batches = [ [] for _ in range(len(x)) ]

                    batches += [ [] ]
                    i_person_id = len(batches)-1

                for i in range(len(x)):
                    batches[i].append ( x[i] )

                batches[i_person_id].append ( np.array([person_id]) )

            yield [ np.array(batch) for batch in batches]

    @staticmethod
    def get_person_id_max_count(samples_path):
        return SampleLoader.get_person_id_max_count(samples_path)

"""
if self.person_id_mode==1:
            samples_len = len(samples)
            samples_idxs = [*range(samples_len)]
            shuffle_idxs = []
        elif self.person_id_mode==2:
            persons_count = len(samples)

            person_idxs = []
            for j in range(persons_count):
                for i in range(j+1,persons_count):
                    person_idxs += [ [i,j] ]

            shuffle_person_idxs = []

            samples_idxs = [None]*persons_count
            shuffle_idxs = [None]*persons_count

            for i in range(persons_count):
                samples_idxs[i] = [*range(len(samples[i]))]
                shuffle_idxs[i] = []
        elif self.person_id_mode==3:
            persons_count = len(samples)

            person_idxs = [ *range(persons_count) ]
            shuffle_person_idxs = []

            samples_idxs = [None]*persons_count
            shuffle_idxs = [None]*persons_count

            for i in range(persons_count):
                samples_idxs[i] = [*range(len(samples[i]))]
                shuffle_idxs[i] = []

if self.person_id_mode==2:
                if len(shuffle_person_idxs) == 0:
                    shuffle_person_idxs = person_idxs.copy()
                    np.random.shuffle(shuffle_person_idxs)
                person_ids = shuffle_person_idxs.pop()


            batches = None
            for n_batch in range(self.batch_size):

                if self.person_id_mode==1:
                    if len(shuffle_idxs) == 0:
                        shuffle_idxs = samples_idxs.copy()
                        np.random.shuffle(shuffle_idxs) ###

                    idx = shuffle_idxs.pop()
                    sample = samples[ idx ]

                    try:
                        x, = SampleProcessor.process ([sample], self.sample_process_options, self.output_sample_types, self.debug)
                    except:
                        raise Exception ("Exception occured in sample %s. Error: %s" % (sample.filename, traceback.format_exc() ) )

                    if type(x) != tuple and type(x) != list:
                        raise Exception('SampleProcessor.process returns NOT tuple/list')

                    if batches is None:
                        batches = [ [] for _ in range(len(x)) ]

                        batches += [ [] ]
                        i_person_id = len(batches)-1

                    for i in range(len(x)):
                        batches[i].append ( x[i] )

                    batches[i_person_id].append ( np.array([sample.person_id]) )


                elif self.person_id_mode==2:
                    person_id1, person_id2 = person_ids

                    if len(shuffle_idxs[person_id1]) == 0:
                        shuffle_idxs[person_id1] = samples_idxs[person_id1].copy()
                        np.random.shuffle(shuffle_idxs[person_id1])

                    idx = shuffle_idxs[person_id1].pop()
                    sample1 = samples[person_id1][idx]

                    if len(shuffle_idxs[person_id2]) == 0:
                        shuffle_idxs[person_id2] = samples_idxs[person_id2].copy()
                        np.random.shuffle(shuffle_idxs[person_id2])

                    idx = shuffle_idxs[person_id2].pop()
                    sample2 = samples[person_id2][idx]

                    if sample1 is not None and sample2 is not None:
                        try:
                            x1, = SampleProcessor.process ([sample1], self.sample_process_options, self.output_sample_types, self.debug)
                        except:
                            raise Exception ("Exception occured in sample %s. Error: %s" % (sample1.filename, traceback.format_exc() ) )

                        try:
                            x2, = SampleProcessor.process ([sample2], self.sample_process_options, self.output_sample_types, self.debug)
                        except:
                            raise Exception ("Exception occured in sample %s. Error: %s" % (sample2.filename, traceback.format_exc() ) )

                        x1_len = len(x1)
                        if batches is None:
                            batches = [ [] for _ in range(x1_len) ]
                            batches += [ [] ]
                            i_person_id1 = len(batches)-1

                            batches += [ [] for _ in range(len(x2)) ]
                            batches += [ [] ]
                            i_person_id2 = len(batches)-1

                        for i in range(x1_len):
                            batches[i].append ( x1[i] )

                        for i in range(len(x2)):
                            batches[x1_len+1+i].append ( x2[i] )

                        batches[i_person_id1].append ( np.array([sample1.person_id]) )

                        batches[i_person_id2].append ( np.array([sample2.person_id]) )

                elif self.person_id_mode==3:
                    if len(shuffle_person_idxs) == 0:
                        shuffle_person_idxs = person_idxs.copy()
                        np.random.shuffle(shuffle_person_idxs)
                    person_id = shuffle_person_idxs.pop()

                    if len(shuffle_idxs[person_id]) == 0:
                        shuffle_idxs[person_id] = samples_idxs[person_id].copy()
                        np.random.shuffle(shuffle_idxs[person_id])

                    idx = shuffle_idxs[person_id].pop()
                    sample1 = samples[person_id][idx]

                    if len(shuffle_idxs[person_id]) == 0:
                        shuffle_idxs[person_id] = samples_idxs[person_id].copy()
                        np.random.shuffle(shuffle_idxs[person_id])

                    idx = shuffle_idxs[person_id].pop()
                    sample2 = samples[person_id][idx]

                    if sample1 is not None and sample2 is not None:
                        try:
                            x1, = SampleProcessor.process ([sample1], self.sample_process_options, self.output_sample_types, self.debug)
                        except:
                            raise Exception ("Exception occured in sample %s. Error: %s" % (sample1.filename, traceback.format_exc() ) )

                        try:
                            x2, = SampleProcessor.process ([sample2], self.sample_process_options, self.output_sample_types, self.debug)
                        except:
                            raise Exception ("Exception occured in sample %s. Error: %s" % (sample2.filename, traceback.format_exc() ) )

                        x1_len = len(x1)
                        if batches is None:
                            batches = [ [] for _ in range(x1_len) ]
                            batches += [ [] ]
                            i_person_id1 = len(batches)-1

                            batches += [ [] for _ in range(len(x2)) ]
                            batches += [ [] ]
                            i_person_id2 = len(batches)-1

                        for i in range(x1_len):
                            batches[i].append ( x1[i] )

                        for i in range(len(x2)):
                            batches[x1_len+1+i].append ( x2[i] )

                        batches[i_person_id1].append ( np.array([sample1.person_id]) )

                        batches[i_person_id2].append ( np.array([sample2.person_id]) )
"""

samplelib/SampleGeneratorFacePerson.py
oo92/TorchFaceLab
import multiprocessing
import pickle
import time
import traceback

import cv2
import numpy as np

from core import mplib
from core.joblib import SubprocessGenerator, ThisThreadGenerator
from facelib import LandmarksProcessor
from samplelib import (SampleGeneratorBase, SampleLoader, SampleProcessor,
                       SampleType)


class SampleGeneratorFaceTemporal(SampleGeneratorBase):
    def __init__ (self, samples_path, debug, batch_size,
                        temporal_image_count=3,
                        sample_process_options=SampleProcessor.Options(),
                        output_sample_types=[],
                        generators_count=2,
                        **kwargs):
        super().__init__(debug, batch_size)

        self.temporal_image_count = temporal_image_count
        self.sample_process_options = sample_process_options
        self.output_sample_types = output_sample_types

        if self.debug:
            self.generators_count = 1
        else:
            self.generators_count = generators_count

        samples = SampleLoader.load (SampleType.FACE_TEMPORAL_SORTED, samples_path)
        samples_len = len(samples)
        if samples_len == 0:
            raise ValueError('No training data provided.')

        mult_max = 1
        l = samples_len - ( (self.temporal_image_count)*mult_max - (mult_max-1)  )
        index_host = mplib.IndexHost(l+1)

        pickled_samples = pickle.dumps(samples, 4)
        if self.debug:
            self.generators = [ThisThreadGenerator ( self.batch_func, (pickled_samples, index_host.create_cli(),) )]
        else:
            self.generators = [SubprocessGenerator ( self.batch_func, (pickled_samples, index_host.create_cli(),) ) for i in range(self.generators_count) ]

        self.generator_counter = -1

    def __iter__(self):
        return self

    def __next__(self):
        self.generator_counter += 1
        generator = self.generators[self.generator_counter % len(self.generators) ]
        return next(generator)

    def batch_func(self, param):
        mult_max = 1
        bs = self.batch_size
        pickled_samples, index_host = param
        samples = pickle.loads(pickled_samples)

        while True:
            batches = None

            indexes = index_host.multi_get(bs)

            for n_batch in range(self.batch_size):
                idx = indexes[n_batch]

                temporal_samples = []
                mult = np.random.randint(mult_max)+1
                for i in range( self.temporal_image_count ):
                    sample = samples[ idx+i*mult ]
                    try:
                        temporal_samples += SampleProcessor.process ([sample], self.sample_process_options, self.output_sample_types, self.debug)[0]
                    except:
                        raise Exception ("Exception occured in sample %s. Error: %s" % (sample.filename, traceback.format_exc() ) )

                if batches is None:
                    batches = [ [] for _ in range(len(temporal_samples)) ]

                for i in range(len(temporal_samples)):
                    batches[i].append ( temporal_samples[i] )

            yield [ np.array(batch) for batch in batches]

samplelib/SampleGeneratorFaceTemporal.py
oo92/TorchFaceLab
import multiprocessing
import pickle
import time
import traceback
from enum import IntEnum

import cv2
import numpy as np
from pathlib import Path
from core import imagelib, mplib, pathex
from core.imagelib import sd
from core.cv2ex import *
from core.interact import interact as io
from core.joblib import Subprocessor, SubprocessGenerator, ThisThreadGenerator
from facelib import LandmarksProcessor
from samplelib import (SampleGeneratorBase, SampleLoader, SampleProcessor, SampleType)

class SampleGeneratorFaceXSeg(SampleGeneratorBase):
    def __init__ (self, paths, debug=False, batch_size=1, resolution=256, face_type=None,
                        generators_count=4, data_format="NHWC",
                        **kwargs):

        super().__init__(debug, batch_size)
        self.initialized = False

        samples = sum([ SampleLoader.load (SampleType.FACE, path) for path in paths ]  )
        seg_sample_idxs = SegmentedSampleFilterSubprocessor(samples).run()

        if len(seg_sample_idxs) == 0:
            seg_sample_idxs = SegmentedSampleFilterSubprocessor(samples, count_xseg_mask=True).run()
            if len(seg_sample_idxs) == 0:
                raise Exception(f"No segmented faces found.")
            else:
                io.log_info(f"Using {len(seg_sample_idxs)} xseg labeled samples.")
        else:
            io.log_info(f"Using {len(seg_sample_idxs)} segmented samples.")

        if self.debug:
            self.generators_count = 1
        else:
            self.generators_count = max(1, generators_count)

        args = (samples, seg_sample_idxs, resolution, face_type, data_format)
        if self.debug:
            self.generators = [ThisThreadGenerator ( self.batch_func, args )]
        else:
            self.generators = [SubprocessGenerator ( self.batch_func, args, start_now=False ) for i in range(self.generators_count) ]

            SubprocessGenerator.start_in_parallel( self.generators )

        self.generator_counter = -1

        self.initialized = True

    #overridable
    def is_initialized(self):
        return self.initialized

    def __iter__(self):
        return self

    def __next__(self):
        self.generator_counter += 1
        generator = self.generators[self.generator_counter % len(self.generators) ]
        return next(generator)

    def batch_func(self, param ):
        samples, seg_sample_idxs, resolution, face_type, data_format = param

        shuffle_idxs = []
        bg_shuffle_idxs = []

        random_flip = True
        rotation_range=[-10,10]
        scale_range=[-0.05, 0.05]
        tx_range=[-0.05, 0.05]
        ty_range=[-0.05, 0.05]

        random_bilinear_resize_chance, random_bilinear_resize_max_size_per = 25,75
        sharpen_chance, sharpen_kernel_max_size = 25, 5
        motion_blur_chance, motion_blur_mb_max_size = 25, 5
        gaussian_blur_chance, gaussian_blur_kernel_max_size = 25, 5
        random_jpeg_compress_chance = 25

        def gen_img_mask(sample):
            img = sample.load_bgr()
            h,w,c = img.shape

            if sample.seg_ie_polys.has_polys():
                mask = np.zeros ((h,w,1), dtype=np.float32)
                sample.seg_ie_polys.overlay_mask(mask)
            elif sample.has_xseg_mask():
                mask = sample.get_xseg_mask()
                mask[mask < 0.5] = 0.0
                mask[mask >= 0.5] = 1.0
            else:
                raise Exception(f'no mask in sample {sample.filename}')

            if face_type == sample.face_type:
                if w != resolution:
                    img = cv2.resize( img, (resolution, resolution), interpolation=cv2.INTER_LANCZOS4 )
                    mask = cv2.resize( mask, (resolution, resolution), interpolation=cv2.INTER_LANCZOS4 )
            else:
                mat = LandmarksProcessor.get_transform_mat (sample.landmarks, resolution, face_type)
                img  = cv2.warpAffine( img,  mat, (resolution,resolution), borderMode=cv2.BORDER_CONSTANT, flags=cv2.INTER_LANCZOS4 )
                mask = cv2.warpAffine( mask, mat, (resolution,resolution), borderMode=cv2.BORDER_CONSTANT, flags=cv2.INTER_LANCZOS4 )

            if len(mask.shape) == 2:
                mask = mask[...,None]
            return img, mask

        bs = self.batch_size
        while True:
            batches = [ [], [] ]

            n_batch = 0
            while n_batch < bs:
                try:
                    if len(shuffle_idxs) == 0:
                        shuffle_idxs = seg_sample_idxs.copy()
                        np.random.shuffle(shuffle_idxs)
                    sample = samples[shuffle_idxs.pop()]
                    img, mask = gen_img_mask(sample)

                    if np.random.randint(2) == 0:
                        if len(bg_shuffle_idxs) == 0:
                            bg_shuffle_idxs = seg_sample_idxs.copy()
                            np.random.shuffle(bg_shuffle_idxs)
                        bg_sample = samples[bg_shuffle_idxs.pop()]

                        bg_img, bg_mask = gen_img_mask(bg_sample)

                        bg_wp   = imagelib.gen_warp_params(resolution, True, rotation_range=[-180,180], scale_range=[-0.10, 0.10], tx_range=[-0.10, 0.10], ty_range=[-0.10, 0.10] )
                        bg_img  = imagelib.warp_by_params (bg_wp, bg_img,  can_warp=False, can_transform=True, can_flip=True, border_replicate=True)
                        bg_mask = imagelib.warp_by_params (bg_wp, bg_mask, can_warp=False, can_transform=True, can_flip=True, border_replicate=False)
                        bg_img = bg_img*(1-bg_mask)
                        if np.random.randint(2) == 0:
                            bg_img = imagelib.apply_random_hsv_shift(bg_img)
                        else:
                            bg_img = imagelib.apply_random_rgb_levels(bg_img)

                        c_mask = 1.0 - (1-bg_mask) * (1-mask)
                        rnd = 0.15 + np.random.uniform()*0.85
                        img = img*(c_mask) + img*(1-c_mask)*rnd + bg_img*(1-c_mask)*(1-rnd)

                    warp_params = imagelib.gen_warp_params(resolution, random_flip, rotation_range=rotation_range, scale_range=scale_range, tx_range=tx_range, ty_range=ty_range )
                    img   = imagelib.warp_by_params (warp_params, img,  can_warp=True, can_transform=True, can_flip=True, border_replicate=True)
                    mask  = imagelib.warp_by_params (warp_params, mask, can_warp=True, can_transform=True, can_flip=True, border_replicate=False)

                    img = np.clip(img.astype(np.float32), 0, 1)
                    mask[mask < 0.5] = 0.0
                    mask[mask >= 0.5] = 1.0
                    mask = np.clip(mask, 0, 1)
                    
                    if np.random.randint(2) == 0:
                        # random face flare
                        krn = np.random.randint( resolution//4, resolution )
                        krn = krn - krn % 2 + 1
                        img = img + cv2.GaussianBlur(img*mask, (krn,krn), 0)

                    if np.random.randint(2) == 0:
                        # random bg flare
                        krn = np.random.randint( resolution//4, resolution )
                        krn = krn - krn % 2 + 1
                        img = img + cv2.GaussianBlur(img*(1-mask), (krn,krn), 0)

                    if np.random.randint(2) == 0:
                        img = imagelib.apply_random_hsv_shift(img, mask=sd.random_circle_faded ([resolution,resolution]))
                    else:
                        img = imagelib.apply_random_rgb_levels(img, mask=sd.random_circle_faded ([resolution,resolution]))
                        
                    if np.random.randint(2) == 0:
                        img = imagelib.apply_random_sharpen( img, sharpen_chance, sharpen_kernel_max_size, mask=sd.random_circle_faded ([resolution,resolution]))
                    else:
                        img = imagelib.apply_random_motion_blur( img, motion_blur_chance, motion_blur_mb_max_size, mask=sd.random_circle_faded ([resolution,resolution]))
                        img = imagelib.apply_random_gaussian_blur( img, gaussian_blur_chance, gaussian_blur_kernel_max_size, mask=sd.random_circle_faded ([resolution,resolution]))
                        
                    if np.random.randint(2) == 0:
                        img = imagelib.apply_random_nearest_resize( img, random_bilinear_resize_chance, random_bilinear_resize_max_size_per, mask=sd.random_circle_faded ([resolution,resolution]))
                    else:
                        img = imagelib.apply_random_bilinear_resize( img, random_bilinear_resize_chance, random_bilinear_resize_max_size_per, mask=sd.random_circle_faded ([resolution,resolution]))
                    img = np.clip(img, 0, 1)

                    img = imagelib.apply_random_jpeg_compress( img, random_jpeg_compress_chance, mask=sd.random_circle_faded ([resolution,resolution]))

                    if data_format == "NCHW":
                        img = np.transpose(img, (2,0,1) )
                        mask = np.transpose(mask, (2,0,1) )

                    batches[0].append ( img )
                    batches[1].append ( mask )

                    n_batch += 1
                except:
                    io.log_err ( traceback.format_exc() )

            yield [ np.array(batch) for batch in batches]

class SegmentedSampleFilterSubprocessor(Subprocessor):
    #override
    def __init__(self, samples, count_xseg_mask=False ):
        self.samples = samples
        self.samples_len = len(self.samples)
        self.count_xseg_mask = count_xseg_mask

        self.idxs = [*range(self.samples_len)]
        self.result = []
        super().__init__('SegmentedSampleFilterSubprocessor', SegmentedSampleFilterSubprocessor.Cli, 60)

    #override
    def process_info_generator(self):
        for i in range(multiprocessing.cpu_count()):
            yield 'CPU%d' % (i), {}, {'samples':self.samples, 'count_xseg_mask':self.count_xseg_mask}

    #override
    def on_clients_initialized(self):
        io.progress_bar ("Filtering", self.samples_len)

    #override
    def on_clients_finalized(self):
        io.progress_bar_close()

    #override
    def get_data(self, host_dict):
        if len (self.idxs) > 0:
            return self.idxs.pop(0)

        return None

    #override
    def on_data_return (self, host_dict, data):
        self.idxs.insert(0, data)

    #override
    def on_result (self, host_dict, data, result):
        idx, is_ok = result
        if is_ok:
            self.result.append(idx)
        io.progress_bar_inc(1)
    def get_result(self):
        return self.result

    class Cli(Subprocessor.Cli):
        #overridable optional
        def on_initialize(self, client_dict):
            self.samples = client_dict['samples']
            self.count_xseg_mask = client_dict['count_xseg_mask']

        def process_data(self, idx):
            if self.count_xseg_mask:
                return idx, self.samples[idx].has_xseg_mask()
            else:
                return idx, self.samples[idx].seg_ie_polys.get_pts_count() != 0

"""
  bg_path = None
        for path in paths:
            bg_path = Path(path) / 'backgrounds'
            if bg_path.exists():

                break
        if bg_path is None:
            io.log_info(f'Random backgrounds will not be used. Place no face jpg images to aligned\backgrounds folder. ')
            bg_pathes = None
        else:
            bg_pathes = pathex.get_image_paths(bg_path, image_extensions=['.jpg'], return_Path_class=True)
            io.log_info(f'Using {len(bg_pathes)} random backgrounds from {bg_path}')

if bg_pathes is not None:
            bg_path = bg_pathes[ np.random.randint(len(bg_pathes)) ]

            bg_img = cv2_imread(bg_path)
            if bg_img is not None:
                bg_img = bg_img.astype(np.float32) / 255.0
                bg_img = imagelib.normalize_channels(bg_img, 3)

                bg_img = imagelib.random_crop(bg_img, resolution, resolution)
                bg_img = cv2.resize(bg_img, (resolution, resolution), interpolation=cv2.INTER_LINEAR)

            if np.random.randint(2) == 0:
                bg_img = imagelib.apply_random_hsv_shift(bg_img)
            else:
                bg_img = imagelib.apply_random_rgb_levels(bg_img)

            bg_wp   = imagelib.gen_warp_params(resolution, True, rotation_range=[-180,180], scale_range=[0,0], tx_range=[0,0], ty_range=[0,0])
            bg_img  = imagelib.warp_by_params (bg_wp, bg_img,  can_warp=False, can_transform=True, can_flip=True, border_replicate=True)

            bg = img*(1-mask)
            fg = img*mask

            c_mask = sd.random_circle_faded ([resolution,resolution])
            bg = ( bg_img*c_mask + bg*(1-c_mask) )*(1-mask)

            img = fg+bg

        else:
"""
samplelib/SampleGeneratorFaceXSeg.py
oo92/TorchFaceLab
import traceback

import cv2
import numpy as np

from core.joblib import SubprocessGenerator, ThisThreadGenerator
from samplelib import (SampleGeneratorBase, SampleLoader, SampleProcessor,
                       SampleType)


class SampleGeneratorImage(SampleGeneratorBase):
    def __init__ (self, samples_path, debug, batch_size, sample_process_options=SampleProcessor.Options(), output_sample_types=[], raise_on_no_data=True, **kwargs):
        super().__init__(debug, batch_size)
        self.initialized = False
        self.sample_process_options = sample_process_options
        self.output_sample_types = output_sample_types

        samples = SampleLoader.load (SampleType.IMAGE, samples_path)
        
        if len(samples) == 0:
            if raise_on_no_data:
                raise ValueError('No training data provided.')
            return
        
        self.generators = [ThisThreadGenerator ( self.batch_func, samples )] if self.debug else \
                          [SubprocessGenerator ( self.batch_func, samples )]

        self.generator_counter = -1
        self.initialized = True
        
    def __iter__(self):
        return self

    def __next__(self):
        self.generator_counter += 1
        generator = self.generators[self.generator_counter % len(self.generators) ]
        return next(generator)

    def batch_func(self, samples):
        samples_len = len(samples)
        

        idxs = [ *range(samples_len) ]
        shuffle_idxs = []

        while True:

            batches = None
            for n_batch in range(self.batch_size):

                if len(shuffle_idxs) == 0:
                    shuffle_idxs = idxs.copy()
                    np.random.shuffle (shuffle_idxs)

                idx = shuffle_idxs.pop()
                sample = samples[idx]
                
                x, = SampleProcessor.process ([sample], self.sample_process_options, self.output_sample_types, self.debug)

                if batches is None:
                    batches = [ [] for _ in range(len(x)) ]

                for i in range(len(x)):
                    batches[i].append ( x[i] )

            yield [ np.array(batch) for batch in batches]

samplelib/SampleGeneratorImage.py
oo92/TorchFaceLab
import traceback

import cv2
import numpy as np

from core.joblib import SubprocessGenerator, ThisThreadGenerator
from samplelib import (SampleGeneratorBase, SampleLoader, SampleProcessor,
                       SampleType)


'''
output_sample_types = [
                        [SampleProcessor.TypeFlags, size, (optional)random_sub_size] ,
                        ...
                      ]
'''
class SampleGeneratorImageTemporal(SampleGeneratorBase):
    def __init__ (self, samples_path, debug, batch_size, temporal_image_count, sample_process_options=SampleProcessor.Options(), output_sample_types=[], **kwargs):
        super().__init__(debug, batch_size)

        self.temporal_image_count = temporal_image_count
        self.sample_process_options = sample_process_options
        self.output_sample_types = output_sample_types

        self.samples = SampleLoader.load (SampleType.IMAGE, samples_path)

        self.generator_samples = [ self.samples ]
        self.generators = [iter_utils.ThisThreadGenerator ( self.batch_func, 0 )] if self.debug else \
                          [iter_utils.SubprocessGenerator ( self.batch_func, 0 )]

        self.generator_counter = -1

    def __iter__(self):
        return self

    def __next__(self):
        self.generator_counter += 1
        generator = self.generators[self.generator_counter % len(self.generators) ]
        return next(generator)

    def batch_func(self, generator_id):
        samples = self.generator_samples[generator_id]
        samples_len = len(samples)
        if samples_len == 0:
            raise ValueError('No training data provided.')

        mult_max = 4
        samples_sub_len = samples_len - ( (self.temporal_image_count)*mult_max - (mult_max-1)  )

        if samples_sub_len <= 0:
            raise ValueError('Not enough samples to fit temporal line.')

        shuffle_idxs = []

        while True:

            batches = None
            for n_batch in range(self.batch_size):

                if len(shuffle_idxs) == 0:
                    shuffle_idxs = [ *range(samples_sub_len) ]
                    np.random.shuffle (shuffle_idxs)

                idx = shuffle_idxs.pop()

                temporal_samples = []
                mult = np.random.randint(mult_max)+1
                for i in range( self.temporal_image_count ):
                    sample = samples[ idx+i*mult ]
                    try:
                        temporal_samples += SampleProcessor.process ([sample], self.sample_process_options, self.output_sample_types, self.debug)[0]
                    except:
                        raise Exception ("Exception occured in sample %s. Error: %s" % (sample.filename, traceback.format_exc() ) )

                if batches is None:
                    batches = [ [] for _ in range(len(temporal_samples)) ]

                for i in range(len(temporal_samples)):
                    batches[i].append ( temporal_samples[i] )

            yield [ np.array(batch) for batch in batches]

samplelib/SampleGeneratorImageTemporal.py
oo92/TorchFaceLab
import multiprocessing
import operator
import pickle
import traceback
from pathlib import Path

import samplelib.PackedFaceset
from core import pathex
from core.mplib import MPSharedList
from core.interact import interact as io
from core.joblib import Subprocessor
from DFLIMG import *
from facelib import FaceType, LandmarksProcessor

from .Sample import Sample, SampleType


class SampleLoader:
    samples_cache = dict()
    @staticmethod
    def get_person_id_max_count(samples_path):
        samples = None
        try:
            samples = samplelib.PackedFaceset.load(samples_path)
        except:
            io.log_err(f"Error occured while loading samplelib.PackedFaceset.load {str(samples_path)}, {traceback.format_exc()}")

        if samples is None:
            raise ValueError("packed faceset not found.")
        persons_name_idxs = {}
        for sample in samples:
            persons_name_idxs[sample.person_name] = 0
        return len(list(persons_name_idxs.keys()))

    @staticmethod
    def load(sample_type, samples_path, subdirs=False):
        """
        Return MPSharedList of samples
        """
        samples_cache = SampleLoader.samples_cache

        if str(samples_path) not in samples_cache.keys():
            samples_cache[str(samples_path)] = [None]*SampleType.QTY

        samples = samples_cache[str(samples_path)]

        if            sample_type == SampleType.IMAGE:
            if  samples[sample_type] is None:
                samples[sample_type] = [ Sample(filename=filename) for filename in io.progress_bar_generator( pathex.get_image_paths(samples_path, subdirs=subdirs), "Loading") ]

        elif          sample_type == SampleType.FACE:
            if  samples[sample_type] is None:
                try:
                    result = samplelib.PackedFaceset.load(samples_path)
                except:
                    io.log_err(f"Error occured while loading samplelib.PackedFaceset.load {str(samples_dat_path)}, {traceback.format_exc()}")

                if result is not None:
                    io.log_info (f"Loaded {len(result)} packed faces from {samples_path}")

                if result is None:
                    result = SampleLoader.load_face_samples( pathex.get_image_paths(samples_path, subdirs=subdirs) )

                samples[sample_type] = MPSharedList(result)
        elif          sample_type == SampleType.FACE_TEMPORAL_SORTED:
                result = SampleLoader.load (SampleType.FACE, samples_path)
                result = SampleLoader.upgradeToFaceTemporalSortedSamples(result)
                samples[sample_type] = MPSharedList(result)

        return samples[sample_type]

    @staticmethod
    def load_face_samples ( image_paths):
        result = FaceSamplesLoaderSubprocessor(image_paths).run()
        sample_list = []

        for filename, data in result:
            if data is None:
                continue
            ( face_type,
              shape,
              landmarks,
              seg_ie_polys,
              xseg_mask_compressed,
              eyebrows_expand_mod,
              source_filename ) = data
              
            sample_list.append( Sample(filename=filename,
                                        sample_type=SampleType.FACE,
                                        face_type=FaceType.fromString (face_type),
                                        shape=shape,
                                        landmarks=landmarks,
                                        seg_ie_polys=seg_ie_polys,
                                        xseg_mask_compressed=xseg_mask_compressed,
                                        eyebrows_expand_mod=eyebrows_expand_mod,
                                        source_filename=source_filename,
                                    ))
        return sample_list

    @staticmethod
    def upgradeToFaceTemporalSortedSamples( samples ):
        new_s = [ (s, s.source_filename) for s in samples]
        new_s = sorted(new_s, key=operator.itemgetter(1))

        return [ s[0] for s in new_s]


class FaceSamplesLoaderSubprocessor(Subprocessor):
    #override
    def __init__(self, image_paths ):
        self.image_paths = image_paths
        self.image_paths_len = len(image_paths)
        self.idxs = [*range(self.image_paths_len)]
        self.result = [None]*self.image_paths_len
        super().__init__('FaceSamplesLoader', FaceSamplesLoaderSubprocessor.Cli, 60)

    #override
    def on_clients_initialized(self):
        io.progress_bar ("Loading samples", len (self.image_paths))

    #override
    def on_clients_finalized(self):
        io.progress_bar_close()

    #override
    def process_info_generator(self):
        for i in range(min(multiprocessing.cpu_count(), 8) ):
            yield 'CPU%d' % (i), {}, {}

    #override
    def get_data(self, host_dict):
        if len (self.idxs) > 0:
            idx = self.idxs.pop(0)
            return idx, self.image_paths[idx]

        return None

    #override
    def on_data_return (self, host_dict, data):
        self.idxs.insert(0, data[0])

    #override
    def on_result (self, host_dict, data, result):
        idx, dflimg = result
        self.result[idx] = (self.image_paths[idx], dflimg)
        io.progress_bar_inc(1)

    #override
    def get_result(self):
        return self.result

    class Cli(Subprocessor.Cli):
        #override
        def process_data(self, data):
            idx, filename = data
            dflimg = DFLIMG.load (Path(filename))

            if dflimg is None or not dflimg.has_data():
                self.log_err (f"FaceSamplesLoader: {filename} is not a dfl image file.")
                data = None
            else:
                data = (dflimg.get_face_type(),
                        dflimg.get_shape(),
                        dflimg.get_landmarks(),
                        dflimg.get_seg_ie_polys(),
                        dflimg.get_xseg_mask_compressed(),
                        dflimg.get_eyebrows_expand_mod(),
                        dflimg.get_source_filename() )

            return idx, data

        #override
        def get_data_name (self, data):
            #return string identificator of your data
            return data[1]

samplelib/SampleLoader.py
oo92/TorchFaceLab
import collections
import math
from enum import IntEnum

import cv2
import numpy as np

from core import imagelib
from core.cv2ex import *
from core.imagelib import sd
from facelib import FaceType, LandmarksProcessor


class SampleProcessor(object):
    class SampleType(IntEnum):
        NONE = 0
        IMAGE = 1
        FACE_IMAGE = 2
        FACE_MASK  = 3
        LANDMARKS_ARRAY            = 4
        PITCH_YAW_ROLL             = 5
        PITCH_YAW_ROLL_SIGMOID     = 6

    class ChannelType(IntEnum):
        NONE = 0
        BGR                   = 1  #BGR
        G                     = 2  #Grayscale
        GGG                   = 3  #3xGrayscale

    class FaceMaskType(IntEnum):
        NONE          = 0
        FULL_FACE      = 1  # mask all hull as grayscale
        EYES           = 2  # mask eyes hull as grayscale
        EYES_MOUTH     = 3  # eyes and mouse

    class Options(object):
        def __init__(self, random_flip = True, rotation_range=[-10,10], scale_range=[-0.05, 0.05], tx_range=[-0.05, 0.05], ty_range=[-0.05, 0.05] ):
            self.random_flip = random_flip
            self.rotation_range = rotation_range
            self.scale_range = scale_range
            self.tx_range = tx_range
            self.ty_range = ty_range

    @staticmethod
    def process (samples, sample_process_options, output_sample_types, debug, ct_sample=None):
        SPST = SampleProcessor.SampleType
        SPCT = SampleProcessor.ChannelType
        SPFMT = SampleProcessor.FaceMaskType

        
        outputs = []
        for sample in samples:
            sample_rnd_seed = np.random.randint(0x80000000)
            
            sample_face_type = sample.face_type
            sample_bgr = sample.load_bgr()
            sample_landmarks = sample.landmarks
            ct_sample_bgr = None
            h,w,c = sample_bgr.shape
            
            def get_full_face_mask():   
                xseg_mask = sample.get_xseg_mask()                                     
                if xseg_mask is not None:           
                    if xseg_mask.shape[0] != h or xseg_mask.shape[1] != w:
                        xseg_mask = cv2.resize(xseg_mask, (w,h), interpolation=cv2.INTER_CUBIC)                    
                        xseg_mask = imagelib.normalize_channels(xseg_mask, 1)
                    return np.clip(xseg_mask, 0, 1)
                else:
                    full_face_mask = LandmarksProcessor.get_image_hull_mask (sample_bgr.shape, sample_landmarks, eyebrows_expand_mod=sample.eyebrows_expand_mod )
                    return np.clip(full_face_mask, 0, 1)
                
            def get_eyes_mask():
                eyes_mask = LandmarksProcessor.get_image_eye_mask (sample_bgr.shape, sample_landmarks)
                return np.clip(eyes_mask, 0, 1)
            
            def get_eyes_mouth_mask():                
                eyes_mask = LandmarksProcessor.get_image_eye_mask (sample_bgr.shape, sample_landmarks)
                mouth_mask = LandmarksProcessor.get_image_mouth_mask (sample_bgr.shape, sample_landmarks)
                mask = eyes_mask + mouth_mask
                return np.clip(mask, 0, 1)
                
            is_face_sample = sample_landmarks is not None

            if debug and is_face_sample:
                LandmarksProcessor.draw_landmarks (sample_bgr, sample_landmarks, (0, 1, 0))

            outputs_sample = []
            for opts in output_sample_types:
                resolution     = opts.get('resolution', 0)
                sample_type    = opts.get('sample_type', SPST.NONE)
                channel_type   = opts.get('channel_type', SPCT.NONE)                
                nearest_resize_to = opts.get('nearest_resize_to', None)
                warp           = opts.get('warp', False)
                transform      = opts.get('transform', False)
                random_hsv_shift_amount = opts.get('random_hsv_shift_amount', 0)
                normalize_tanh = opts.get('normalize_tanh', False)
                ct_mode        = opts.get('ct_mode', None)
                data_format    = opts.get('data_format', 'NHWC')
                
                rnd_seed_shift      = opts.get('rnd_seed_shift', 0)
                warp_rnd_seed_shift = opts.get('warp_rnd_seed_shift', rnd_seed_shift)
                
                rnd_state      = np.random.RandomState (sample_rnd_seed+rnd_seed_shift)
                warp_rnd_state = np.random.RandomState (sample_rnd_seed+warp_rnd_seed_shift)
                
                warp_params = imagelib.gen_warp_params(resolution, 
                                                       sample_process_options.random_flip, 
                                                       rotation_range=sample_process_options.rotation_range, 
                                                       scale_range=sample_process_options.scale_range, 
                                                       tx_range=sample_process_options.tx_range, 
                                                       ty_range=sample_process_options.ty_range, 
                                                       rnd_state=rnd_state,
                                                       warp_rnd_state=warp_rnd_state,
                                                       )
                
                if sample_type == SPST.FACE_MASK or sample_type == SPST.IMAGE: 
                    border_replicate = False
                elif sample_type == SPST.FACE_IMAGE:
                    border_replicate = True
                    
                    
                border_replicate = opts.get('border_replicate', border_replicate)
                borderMode = cv2.BORDER_REPLICATE if border_replicate else cv2.BORDER_CONSTANT
                
                
                if sample_type == SPST.FACE_IMAGE or sample_type == SPST.FACE_MASK:
                    if not is_face_sample:    
                        raise ValueError("face_samples should be provided for sample_type FACE_*")

                if sample_type == SPST.FACE_IMAGE or sample_type == SPST.FACE_MASK:
                    face_type      = opts.get('face_type', None)
                    face_mask_type = opts.get('face_mask_type', SPFMT.NONE)
                
                    if face_type is None:
                        raise ValueError("face_type must be defined for face samples")

                    if sample_type == SPST.FACE_MASK: 
                        if face_mask_type == SPFMT.FULL_FACE:
                            img = get_full_face_mask()
                        elif face_mask_type == SPFMT.EYES:
                            img = get_eyes_mask()
                        elif face_mask_type == SPFMT.EYES_MOUTH:
                            mask = get_full_face_mask().copy()
                            mask[mask != 0.0] = 1.0                            
                            img = get_eyes_mouth_mask()*mask
                        else:
                            img = np.zeros ( sample_bgr.shape[0:2]+(1,), dtype=np.float32)

                        if sample_face_type == FaceType.MARK_ONLY:
                            raise NotImplementedError()
                            mat  = LandmarksProcessor.get_transform_mat (sample_landmarks, warp_resolution, face_type)
                            img = cv2.warpAffine( img, mat, (warp_resolution, warp_resolution), flags=cv2.INTER_LINEAR )
                            
                            img = imagelib.warp_by_params (warp_params, img, warp, transform, can_flip=True, border_replicate=border_replicate, cv2_inter=cv2.INTER_LINEAR)
                            img = cv2.resize( img, (resolution,resolution), interpolation=cv2.INTER_LINEAR )
                        else:
                            if face_type != sample_face_type:
                                mat = LandmarksProcessor.get_transform_mat (sample_landmarks, resolution, face_type)                            
                                img = cv2.warpAffine( img, mat, (resolution,resolution), borderMode=borderMode, flags=cv2.INTER_LINEAR )
                            else:
                                if w != resolution:
                                    img = cv2.resize( img, (resolution, resolution), interpolation=cv2.INTER_LINEAR )
                                
                            img = imagelib.warp_by_params (warp_params, img, warp, transform, can_flip=True, border_replicate=border_replicate, cv2_inter=cv2.INTER_LINEAR)

                        if face_mask_type == SPFMT.EYES_MOUTH:
                            div = img.max()
                            if div != 0.0:
                                img = img / div # normalize to 1.0 after warp
                            
                        if len(img.shape) == 2:
                            img = img[...,None]
                            
                        if channel_type == SPCT.G:
                            out_sample = img.astype(np.float32)
                        else:
                            raise ValueError("only channel_type.G supported for the mask")

                    elif sample_type == SPST.FACE_IMAGE:
                        img = sample_bgr                      
                            
                        if face_type != sample_face_type:
                            mat = LandmarksProcessor.get_transform_mat (sample_landmarks, resolution, face_type)
                            img = cv2.warpAffine( img, mat, (resolution,resolution), borderMode=borderMode, flags=cv2.INTER_CUBIC )
                        else:
                            if w != resolution:
                                img = cv2.resize( img, (resolution, resolution), interpolation=cv2.INTER_CUBIC )
                                
                        # Apply random color transfer                        
                        if ct_mode is not None and ct_sample is not None:
                            if ct_sample_bgr is None:
                               ct_sample_bgr = ct_sample.load_bgr()
                            img = imagelib.color_transfer (ct_mode, img, cv2.resize( ct_sample_bgr, (resolution,resolution), interpolation=cv2.INTER_LINEAR ) )
                        
                        if random_hsv_shift_amount != 0:
                            a = random_hsv_shift_amount
                            h_amount = max(1, int(360*a*0.5))
                            img_h, img_s, img_v = cv2.split(cv2.cvtColor(img, cv2.COLOR_BGR2HSV))
                            img_h = (img_h + rnd_state.randint(-h_amount, h_amount+1) ) % 360
                            img_s = np.clip (img_s + (rnd_state.random()-0.5)*a, 0, 1 )
                            img_v = np.clip (img_v + (rnd_state.random()-0.5)*a, 0, 1 )
                            img = np.clip( cv2.cvtColor(cv2.merge([img_h, img_s, img_v]), cv2.COLOR_HSV2BGR) , 0, 1 )

                        img  = imagelib.warp_by_params (warp_params, img,  warp, transform, can_flip=True, border_replicate=border_replicate)
  
                        img = np.clip(img.astype(np.float32), 0, 1)

                        # Transform from BGR to desired channel_type
                        if channel_type == SPCT.BGR:
                            out_sample = img
                        elif channel_type == SPCT.G:
                            out_sample = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)[...,None]
                        elif channel_type == SPCT.GGG:
                            out_sample = np.repeat ( np.expand_dims(cv2.cvtColor(img, cv2.COLOR_BGR2GRAY),-1), (3,), -1)

                    # Final transformations
                    if nearest_resize_to is not None:
                        out_sample = cv2_resize(out_sample, (nearest_resize_to,nearest_resize_to), interpolation=cv2.INTER_NEAREST)
                        
                    if not debug:
                        if normalize_tanh:
                            out_sample = np.clip (out_sample * 2.0 - 1.0, -1.0, 1.0)
                    if data_format == "NCHW":
                        out_sample = np.transpose(out_sample, (2,0,1) )
                elif sample_type == SPST.IMAGE:
                    img = sample_bgr      
                    img  = imagelib.warp_by_params (warp_params, img,  warp, transform, can_flip=True, border_replicate=True)
                    img  = cv2.resize( img,  (resolution, resolution), interpolation=cv2.INTER_CUBIC )
                    out_sample = img
                    
                    if data_format == "NCHW":
                        out_sample = np.transpose(out_sample, (2,0,1) )
                    
                    
                elif sample_type == SPST.LANDMARKS_ARRAY:
                    l = sample_landmarks
                    l = np.concatenate ( [ np.expand_dims(l[:,0] / w,-1), np.expand_dims(l[:,1] / h,-1) ], -1 )
                    l = np.clip(l, 0.0, 1.0)
                    out_sample = l
                elif sample_type == SPST.PITCH_YAW_ROLL or sample_type == SPST.PITCH_YAW_ROLL_SIGMOID:
                    pitch,yaw,roll = sample.get_pitch_yaw_roll()
                    if warp_params['flip']:
                        yaw = -yaw

                    if sample_type == SPST.PITCH_YAW_ROLL_SIGMOID:
                        pitch = np.clip( (pitch / math.pi) / 2.0 + 0.5, 0, 1)
                        yaw   = np.clip( (yaw / math.pi) / 2.0 + 0.5, 0, 1)
                        roll  = np.clip( (roll / math.pi) / 2.0 + 0.5, 0, 1)

                    out_sample = (pitch, yaw)
                else:
                    raise ValueError ('expected sample_type')

                outputs_sample.append ( out_sample )
            outputs += [outputs_sample]

        return outputs


samplelib/SampleProcessor.py
oo92/TorchFaceLab
from .Sample import Sample
from .Sample import SampleType
from .SampleLoader import SampleLoader
from .SampleProcessor import SampleProcessor
from .SampleGeneratorBase import SampleGeneratorBase
from .SampleGeneratorFace import SampleGeneratorFace
from .SampleGeneratorFacePerson import SampleGeneratorFacePerson
from .SampleGeneratorFaceTemporal import SampleGeneratorFaceTemporal
from .SampleGeneratorImage import SampleGeneratorImage
from .SampleGeneratorImageTemporal import SampleGeneratorImageTemporal
from .SampleGeneratorFaceCelebAMaskHQ import SampleGeneratorFaceCelebAMaskHQ
from .SampleGeneratorFaceXSeg import SampleGeneratorFaceXSeg
from .PackedFaceset import PackedFaceset
samplelib/__init__.py
oo92/TorchFaceLab