DeepFaceLab/mainscripts/MaskEditorTool.py

import os
import sys
import time
import traceback
from pathlib import Path

import cv2
import numpy as np
import numpy.linalg as npl

import imagelib
from DFLIMG import *
from facelib import LandmarksProcessor
from imagelib import IEPolys
from interact import interact as io
from utils import Path_utils
from utils.cv2_utils import *


class MaskEditor:
    STATE_NONE=0
    STATE_MASKING=1

    def __init__(self, img, prev_images, next_images, mask=None, ie_polys=None, get_status_lines_func=None):
        self.img = imagelib.normalize_channels (img,3)
        h, w, c = img.shape

        if h != w and w != 256:
            #to support any square res, scale img,mask and ie_polys to 256, then scale ie_polys back on .get_ie_polys()
            raise Exception ("MaskEditor does not support image size != 256x256")

        ph, pw = h // 4, w // 4 #pad wh

        self.prev_images = prev_images
        self.next_images = next_images

        if mask is not None:
            self.mask = imagelib.normalize_channels (mask,3)
        else:
            self.mask = np.zeros ( (h,w,3) )
        self.get_status_lines_func = get_status_lines_func

        self.state_prop = self.STATE_NONE

        self.w, self.h = w, h
        self.pw, self.ph = pw, ph
        self.pwh = np.array([self.pw, self.ph])
        self.pwh2 = np.array([self.pw*2, self.ph*2])
        self.sw, self.sh = w+pw*2, h+ph*2
        self.prwh = 64 #preview wh

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

        self.polys_mask = None
        self.preview_images = None

        self.mouse_x = self.mouse_y = 9999
        self.screen_status_block = None
        self.screen_status_block_dirty = True
        self.screen_changed = True

    def set_state(self, state):
        self.state = state

    @property
    def state(self):
        return self.state_prop

    @state.setter
    def state(self, value):
        self.state_prop = value
        if value == self.STATE_MASKING:
            self.ie_polys.dirty = True

    def get_mask(self):
        if self.ie_polys.switch_dirty():
            self.screen_status_block_dirty = True
            self.ie_mask = img = self.mask.copy()

            self.ie_polys.overlay_mask(img)

            return img
        return self.ie_mask

    def get_screen_overlay(self):
        img = np.zeros ( (self.sh, self.sw, 3) )

        if self.state == self.STATE_MASKING:
            mouse_xy = self.mouse_xy.copy() + self.pwh
            l = self.ie_polys.n_list()
            if l.n > 0:
                p = l.cur_point().copy() + self.pwh
                color = (0,1,0) if l.type == 1 else (0,0,1)
                cv2.line(img, tuple(p), tuple(mouse_xy), color )

        return img

    def undo_to_begin_point(self):
        while not self.undo_point():
            pass

    def undo_point(self):
        self.screen_changed = True
        if self.state == self.STATE_NONE:
            if self.ie_polys.n > 0:
                self.state = self.STATE_MASKING

        if self.state == self.STATE_MASKING:
            if self.ie_polys.n_list().n_dec() == 0 and \
               self.ie_polys.n_dec() == 0:
                self.state = self.STATE_NONE
            else:
                return False

        return True

    def redo_to_end_point(self):
        while not self.redo_point():
            pass

    def redo_point(self):
        self.screen_changed = True
        if self.state == self.STATE_NONE:
            if self.ie_polys.n_max > 0:
                self.state = self.STATE_MASKING
                if self.ie_polys.n == 0:
                    self.ie_polys.n_inc()

        if self.state == self.STATE_MASKING:
            while True:
                l = self.ie_polys.n_list()
                if l.n_inc() == l.n_max:
                    if self.ie_polys.n == self.ie_polys.n_max:
                        break
                    self.ie_polys.n_inc()
                else:
                    return False

        return True

    def combine_screens(self, screens):

        screens_len = len(screens)

        new_screens = []
        for screen, padded_overlay in screens:
            screen_img = np.zeros( (self.sh, self.sw, 3), dtype=np.float32 )

            screen = imagelib.normalize_channels (screen, 3)
            h,w,c = screen.shape

            screen_img[self.ph:-self.ph, self.pw:-self.pw, :] = screen

            if padded_overlay is not None:
                screen_img = screen_img + padded_overlay

            screen_img = np.clip(screen_img*255, 0, 255).astype(np.uint8)
            new_screens.append(screen_img)

        return np.concatenate (new_screens, axis=1)

    def get_screen_status_block(self, w, c):
        if self.screen_status_block_dirty:
            self.screen_status_block_dirty = False
            lines = [
                    'Polys current/max = %d/%d' % (self.ie_polys.n, self.ie_polys.n_max),
                    ]
            if self.get_status_lines_func is not None:
                lines += self.get_status_lines_func()

            lines_count = len(lines)


            h_line = 21
            h = lines_count * h_line
            img = np.ones ( (h,w,c) ) * 0.1

            for i in range(lines_count):
                img[ i*h_line:(i+1)*h_line, 0:w] += \
                    imagelib.get_text_image (  (h_line,w,c), lines[i], color=[0.8]*c )

            self.screen_status_block = np.clip(img*255, 0, 255).astype(np.uint8)

        return self.screen_status_block

    def set_screen_status_block_dirty(self):
        self.screen_status_block_dirty = True

    def set_screen_changed(self):
        self.screen_changed = True

    def switch_screen_changed(self):
        result = self.screen_changed
        self.screen_changed = False
        return result

    def make_screen(self):
        screen_overlay = self.get_screen_overlay()
        final_mask = self.get_mask()

        masked_img = self.img*final_mask*0.5 + self.img*(1-final_mask)

        pink = np.full ( (self.h, self.w, 3), (1,0,1) )
        pink_masked_img = self.img*final_mask + pink*(1-final_mask)




        screens = [ (self.img, screen_overlay),
                    (masked_img, screen_overlay),
                    (pink_masked_img, screen_overlay),
                    ]
        screens = self.combine_screens(screens)

        if self.preview_images is None:
            sh,sw,sc = screens.shape

            prh, prw = self.prwh, self.prwh

            total_w = sum ([ img.shape[1] for (t,img) in self.prev_images ]) + \
                      sum ([ img.shape[1] for (t,img) in self.next_images ])

            total_images_len = len(self.prev_images) + len(self.next_images)

            max_hor_images_count = sw // prw
            max_side_images_count = (max_hor_images_count - 1) // 2

            prev_images = self.prev_images[-max_side_images_count:]
            next_images = self.next_images[:max_side_images_count]

            border = 2

            max_wh_bordered = (prw-border*2, prh-border*2)

            prev_images = [ (t, cv2.resize( imagelib.normalize_channels(img, 3), max_wh_bordered )) for t,img in prev_images ]
            next_images = [ (t, cv2.resize( imagelib.normalize_channels(img, 3), max_wh_bordered )) for t,img in next_images ]

            for images in [prev_images, next_images]:
                for i, (t, img) in enumerate(images):
                    new_img = np.zeros ( (prh,prw, sc) )
                    new_img[border:-border,border:-border] = img

                    if t == 2:
                        cv2.line (new_img, (       prw//2, int(prh//1.5) ), (int(prw/1.5), prh      ) , (0,1,0), thickness=2 )
                        cv2.line (new_img, ( int(prw/1.5), prh           ), (         prw, prh // 2 ) , (0,1,0), thickness=2 )
                    elif t == 1:
                        cv2.line (new_img, ( prw//2, prh//2 ), ( prw, prh      ) , (0,0,1), thickness=2 )
                        cv2.line (new_img, ( prw//2, prh    ), ( prw, prh // 2 ) , (0,0,1), thickness=2 )

                    images[i] = new_img


            preview_images = []
            if len(prev_images) > 0:
                preview_images += [ np.concatenate (prev_images, axis=1) ]

            img = np.full ( (prh,prw, sc), (0,0,1), dtype=np.float )
            img[border:-border,border:-border] = cv2.resize( self.img, max_wh_bordered )

            preview_images += [ img ]

            if len(next_images) > 0:
                preview_images += [ np.concatenate (next_images, axis=1) ]

            preview_images = np.concatenate ( preview_images, axis=1  )

            left_pad = sw // 2 - len(prev_images) * prw - prw // 2
            right_pad = sw // 2 - len(next_images) * prw - prw // 2

            preview_images = np.concatenate ([np.zeros ( (preview_images.shape[0], left_pad, preview_images.shape[2]) ),
                                              preview_images,
                                              np.zeros ( (preview_images.shape[0], right_pad, preview_images.shape[2]) )
                                             ], axis=1)
            self.preview_images = np.clip(preview_images * 255, 0, 255 ).astype(np.uint8)

        status_img = self.get_screen_status_block( screens.shape[1], screens.shape[2] )

        result = np.concatenate ( [self.preview_images, screens, status_img], axis=0  )

        return result

    def mask_finish(self, n_clip=True):
        if self.state == self.STATE_MASKING:
            self.screen_changed = True
            if self.ie_polys.n_list().n <= 2:
                self.ie_polys.n_dec()
            self.state = self.STATE_NONE
            if n_clip:
                self.ie_polys.n_clip()

    def set_mouse_pos(self,x,y):
        if self.preview_images is not None:
            y -= self.preview_images.shape[0]

        mouse_x = x % (self.sw) - self.pw
        mouse_y = y % (self.sh) - self.ph



        if mouse_x != self.mouse_x or mouse_y != self.mouse_y:
            self.mouse_xy = np.array( [mouse_x, mouse_y] )
            self.mouse_x, self.mouse_y = self.mouse_xy
            self.screen_changed = True

    def mask_point(self, type):
        self.screen_changed = True
        if self.state == self.STATE_MASKING and \
           self.ie_polys.n_list().type != type:
            self.mask_finish()

        elif self.state == self.STATE_NONE:
            self.state = self.STATE_MASKING
            self.ie_polys.add(type)

        if self.state == self.STATE_MASKING:
            self.ie_polys.n_list().add (self.mouse_x, self.mouse_y)

    def get_ie_polys(self):
        return self.ie_polys
    
    def set_ie_polys(self, saved_ie_polys):
        self.state = self.STATE_NONE
        self.ie_polys = saved_ie_polys
        self.redo_to_end_point()
        self.mask_finish()
        
        
def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None, no_default_mask=False):
    input_path = Path(input_dir)

    confirmed_path = Path(confirmed_dir)
    skipped_path = Path(skipped_dir)

    if not input_path.exists():
        raise ValueError('Input directory not found. Please ensure it exists.')

    if not confirmed_path.exists():
        confirmed_path.mkdir(parents=True)

    if not skipped_path.exists():
        skipped_path.mkdir(parents=True)
        
    if not no_default_mask:
        eyebrows_expand_mod = np.clip ( io.input_int ("Default eyebrows expand modifier? (0..400, skip:100) : ", 100), 0, 400 ) / 100.0
    else:
        eyebrows_expand_mod = None

    wnd_name = "MaskEditor tool"
    io.named_window (wnd_name)
    io.capture_mouse(wnd_name)
    io.capture_keys(wnd_name)

    cached_images = {}

    image_paths = [ Path(x) for x in Path_utils.get_image_paths(input_path)]
    done_paths = []
    done_images_types = {}
    image_paths_total = len(image_paths)
    saved_ie_polys = IEPolys()
    zoom_factor = 1.0
    preview_images_count = 9
    target_wh = 256

    do_prev_count = 0
    do_save_move_count = 0
    do_save_count = 0
    do_skip_move_count = 0
    do_skip_count = 0
    
    def jobs_count():
        return do_prev_count + do_save_move_count + do_save_count + do_skip_move_count + do_skip_count

    is_exit = False
    while not is_exit:

        if len(image_paths) > 0:
            filepath = image_paths.pop(0)
        else:
            filepath = None

        next_image_paths = image_paths[0:preview_images_count]
        next_image_paths_names = [ path.name for path in next_image_paths ]
        prev_image_paths = done_paths[-preview_images_count:]
        prev_image_paths_names = [ path.name for path in prev_image_paths ]

        for key in list( cached_images.keys() ):
            if key not in prev_image_paths_names and \
               key not in next_image_paths_names:
                cached_images.pop(key)

        for paths in [prev_image_paths, next_image_paths]:
            for path in paths:
                if path.name not in cached_images:
                    cached_images[path.name] = cv2_imread(str(path)) / 255.0

        if filepath is not None:
            dflimg = DFLIMG.load (filepath)

            if dflimg is None:
                io.log_err ("%s is not a dfl image file" % (filepath.name) )
                continue
            else:
                lmrks = dflimg.get_landmarks()
                ie_polys = IEPolys.load(dflimg.get_ie_polys())
                fanseg_mask = dflimg.get_fanseg_mask()

                if filepath.name in cached_images:
                    img = cached_images[filepath.name]
                else:
                    img = cached_images[filepath.name] = cv2_imread(str(filepath)) / 255.0

                if fanseg_mask is not None:
                    mask = fanseg_mask
                else:
                    if no_default_mask:
                        mask = np.zeros ( (target_wh,target_wh,3) )
                    else:
                        mask = LandmarksProcessor.get_image_hull_mask( img.shape, lmrks, eyebrows_expand_mod=eyebrows_expand_mod)
        else:
            img = np.zeros ( (target_wh,target_wh,3) )
            mask = np.ones ( (target_wh,target_wh,3) )
            ie_polys = None

        def get_status_lines_func():
            return ['Progress: %d / %d . Current file: %s' % (len(done_paths), image_paths_total, str(filepath.name) if filepath is not None else "end" ),
                    '[Left mouse button] - mark include mask.',
                    '[Right mouse button] - mark exclude mask.',
                    '[Middle mouse button] - finish current poly.',
                    '[Mouse wheel] - undo/redo poly or point. [+ctrl] - undo to begin/redo to end',
                    '[r] - applies edits made to last saved image.',
                    '[q] - prev image. [w] - skip and move to %s. [e] - save and move to %s. ' % (skipped_path.name, confirmed_path.name),
                    '[z] - prev image. [x] - skip. [c] - save. ',
                    'hold [shift] - speed up the frame counter by 10.',
                    '[-/+] - window zoom [esc] - quit',
                    ]

        try:
            ed = MaskEditor(img,
                            [ (done_images_types[name], cached_images[name]) for name in prev_image_paths_names ],
                            [ (0, cached_images[name]) for name in next_image_paths_names ],
                            mask, ie_polys, get_status_lines_func)
        except Exception as e:
            print(e)
            continue

        next = False
        while not next:
            io.process_messages(0.005)

            if jobs_count() == 0:
                for (x,y,ev,flags) in io.get_mouse_events(wnd_name):
                    x, y = int (x / zoom_factor), int(y / zoom_factor)
                    ed.set_mouse_pos(x, y)
                    if filepath is not None:
                        if ev == io.EVENT_LBUTTONDOWN:
                            ed.mask_point(1)
                        elif ev == io.EVENT_RBUTTONDOWN:
                            ed.mask_point(0)
                        elif ev == io.EVENT_MBUTTONDOWN:
                            ed.mask_finish()
                        elif ev == io.EVENT_MOUSEWHEEL:
                            if flags & 0x80000000 != 0:
                                if flags & 0x8 != 0:
                                    ed.undo_to_begin_point()
                                else:
                                    ed.undo_point()
                            else:
                                if flags & 0x8 != 0:
                                    ed.redo_to_end_point()
                                else:
                                    ed.redo_point()

                for key, chr_key, ctrl_pressed, alt_pressed, shift_pressed in io.get_key_events(wnd_name):
                    if chr_key == 'q' or chr_key == 'z':
                        do_prev_count = 1 if not shift_pressed else 10
                    elif chr_key == '-':
                        zoom_factor = np.clip (zoom_factor-0.1, 0.1, 4.0)
                        ed.set_screen_changed()
                    elif chr_key == '+':
                        zoom_factor = np.clip (zoom_factor+0.1, 0.1, 4.0)
                        ed.set_screen_changed()
                    elif key == 27: #esc
                        is_exit = True
                        next = True
                        break
                    elif filepath is not None:
                        if chr_key == 'e':
                            saved_ie_polys = ed.ie_polys
                            do_save_move_count = 1 if not shift_pressed else 10
                        elif chr_key == 'c':
                            saved_ie_polys = ed.ie_polys
                            do_save_count = 1 if not shift_pressed else 10
                        elif chr_key == 'w':
                            do_skip_move_count = 1 if not shift_pressed else 10
                        elif chr_key == 'x':
                            do_skip_count = 1 if not shift_pressed else 10
                        elif chr_key == 'r' and saved_ie_polys != None:
                            ed.set_ie_polys(saved_ie_polys)

            if do_prev_count > 0:
                do_prev_count -= 1
                if len(done_paths) > 0:
                    if filepath is not None:
                        image_paths.insert(0, filepath)

                    filepath = done_paths.pop(-1)
                    done_images_types[filepath.name] = 0

                    if filepath.parent != input_path:
                        new_filename_path = input_path / filepath.name
                        filepath.rename ( new_filename_path )
                        image_paths.insert(0, new_filename_path)
                    else:
                        image_paths.insert(0, filepath)

                    next = True
            elif filepath is not None:
                if do_save_move_count > 0:
                    do_save_move_count -= 1

                    ed.mask_finish()
                    dflimg.embed_and_set (str(filepath), ie_polys=ed.get_ie_polys(), eyebrows_expand_mod=eyebrows_expand_mod )

                    done_paths += [ confirmed_path / filepath.name ]
                    done_images_types[filepath.name] = 2
                    filepath.rename(done_paths[-1])

                    next = True
                elif do_save_count > 0:
                    do_save_count -= 1

                    ed.mask_finish()
                    dflimg.embed_and_set (str(filepath), ie_polys=ed.get_ie_polys(), eyebrows_expand_mod=eyebrows_expand_mod )

                    done_paths += [ filepath ]
                    done_images_types[filepath.name] = 2

                    next = True
                elif do_skip_move_count > 0:
                    do_skip_move_count -= 1

                    done_paths += [ skipped_path / filepath.name ]
                    done_images_types[filepath.name] = 1
                    filepath.rename(done_paths[-1])

                    next = True
                elif do_skip_count > 0:
                    do_skip_count -= 1

                    done_paths += [ filepath ]
                    done_images_types[filepath.name] = 1

                    next = True
            else:
                do_save_move_count = do_save_count = do_skip_move_count = do_skip_count = 0

            if jobs_count() == 0:
                if ed.switch_screen_changed():
                    screen = ed.make_screen()
                    if zoom_factor != 1.0:
                        h,w,c = screen.shape
                        screen = cv2.resize ( screen, ( int(w*zoom_factor), int(h*zoom_factor) ) )
                    io.show_image (wnd_name, screen )


        io.process_messages(0.005)

    io.destroy_all_windows()