DeepFaceLab/mainscripts/FacesetResizer.py

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)