swap_mukham.py

import cv2
import numpy as np

import default_paths as dp
from utils.device import get_device_and_provider
from utils.face_alignment import get_cropped_head
from utils.image import paste_to_whole, mix_two_image

from face_swapper import Inswapper
from face_parsing import FaceParser
from face_upscaler import get_available_upscalers_names, cv2_upscalers, load_face_upscaler
from face_analyser import AnalyseFace, single_face_detect_conditions, face_detect_conditions, get_single_face, is_similar_face

from nsfw_checker import NSFWChecker

get_device_name = lambda x: x.lower().replace("executionprovider", "")

class SwapMukham:
    def __init__(self, device='cpu'):
        self.load_nsfw_detector(device=device)
        self.load_face_swapper(device=device)
        self.load_face_analyser(device=device)
        # self.load_face_parser(device=device)
        # self.load_face_upscaler(device=device)

        self.face_parser = None
        self.face_upscaler = None
        self.face_upscaler_name = ""

    def set_values(self, args):
        self.age = args.get('age', 0)
        self.detect_condition = args.get('detect_condition', "left most")
        self.similarity = args.get('similarity', 0.6)
        self.swap_condition = args.get('swap_condition', 'left most')
        self.face_scale = args.get('face_scale', 1.0)
        self.num_of_pass = args.get('num_of_pass', 1)
        self.mask_crop_values = args.get('mask_crop_values', (0,0,0,0))
        self.mask_erode_amount = args.get('mask_erode_amount', 0.1)
        self.mask_blur_amount = args.get('mask_blur_amount', 0.1)
        self.use_laplacian_blending = args.get('use_laplacian_blending', False)
        self.use_face_parsing = args.get('use_face_parsing', False)
        self.face_parse_regions = args.get('face_parse_regions', [1,2,3,4,5,10,11,12,13])
        self.face_upscaler_opacity = args.get('face_upscaler_opacity', 1.)
        self.parse_from_target = args.get('parse_from_target', False)
        self.averaging_method = args.get('averaging_method', 'mean')

        self.analyser.detection_threshold = args.get('face_detection_threshold', 0.5)
        self.analyser.detection_size = args.get('face_detection_size', (640, 640))
        self.analyser.detect_condition = args.get('face_detection_condition', 'best detection')

    def load_nsfw_detector(self, device='cpu'):
        device, provider, options = get_device_and_provider(device=device)
        self.nsfw_detector = NSFWChecker(model_path=dp.OPEN_NSFW_PATH, provider=provider, session_options=options)
        _device = get_device_name(self.nsfw_detector.session.get_providers()[0])
        print(f"[{_device}] NSFW detector model loaded.")

    def load_face_swapper(self, device='cpu'):
        device, provider, options = get_device_and_provider(device=device)
        self.swapper = Inswapper(model_file=dp.INSWAPPER_PATH, provider=provider, session_options=options)
        _device = get_device_name(self.swapper.session.get_providers()[0])
        print(f"[{_device}] Face swapper model loaded.")

    def load_face_analyser(self, device='cpu'):
        device, provider, options = get_device_and_provider(device=device)
        self.analyser = AnalyseFace(provider=provider, session_options=options)
        _device_d = get_device_name(self.analyser.detector.session.get_providers()[0])
        print(f"[{_device_d}] Face detection model loaded.")
        _device_r = get_device_name(self.analyser.recognizer.session.get_providers()[0])
        print(f"[{_device_r}] Face recognition model loaded.")
        _device_g = get_device_name(self.analyser.gender_age.session.get_providers()[0])
        print(f"[{_device_g}] Gender & Age detection model loaded.")

    def load_face_parser(self, device='cpu'):
        device, provider, options = get_device_and_provider(device=device)
        self.face_parser = FaceParser(model_path=dp.FACE_PARSER_PATH, provider=provider, session_options=options)
        _device = get_device_name(self.face_parser.session.get_providers()[0])
        print(f"[{_device}] Face parsing model loaded.")

    def load_face_upscaler(self, name, device='cpu'):
        device, provider, options = get_device_and_provider(device=device)
        if name in get_available_upscalers_names():
            self.face_upscaler = load_face_upscaler(name=name, provider=provider, session_options=options)
            self.face_upscaler_name = name
            _device = get_device_name(self.face_upscaler[0].session.get_providers()[0])
            print(f"[{_device}] Face upscaler model ({name}) loaded.")
        else:
            self.face_upscaler_name = ""
            self.face_upscaler = None

    def collect_heads(self, frame):
        faces = self.analyser.get_faces(frame, skip_task=['embedding', 'gender_age'])
        return [get_cropped_head(frame, face.kps) for face in faces if face["det_score"] > 0.5]

    def analyse_source_faces(self, source_specific):
        analysed_source_specific = []
        for i, (source, specific) in enumerate(source_specific):
            if source is not None:
                analysed_source = self.analyser.get_averaged_face(source, method=self.averaging_method)
                if specific is not None:
                    analysed_specific = self.analyser.get_face(specific)
                else:
                    analysed_specific = None
                analysed_source_specific.append((analysed_source, analysed_specific))
        self.analysed_source_specific = analysed_source_specific

    def process_frame(self, data):
        frame, custom_mask = data

        if len(frame.shape) == 2 or (len(frame.shape) == 3 and frame.shape[2] == 1):
            frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)

        alpha = None
        if frame.shape[2] == 4:
            alpha = frame[:, :, 3]
            frame = frame[:, :, :3]

        _frame = frame.copy()
        condition = self.swap_condition

        skip_task = []
        if condition != "specific face":
            skip_task.append('embedding')
        if condition not in ['age less than', 'age greater than', 'all male', 'all female']:
            skip_task.append('gender_age')

        analysed_target_faces = self.analyser.get_faces(frame, scale=self.face_scale, skip_task=skip_task)

        for analysed_target in analysed_target_faces:
            if (condition == "all face" or
                (condition == "age less than" and analysed_target["age"] <= self.age) or
                (condition == "age greater than" and analysed_target["age"] > self.age) or
                (condition == "all male" and analysed_target["gender"] == 1) or
                (condition == "all female" and analysed_target["gender"] == 0)):

                trg_face = analysed_target
                src_face = self.analysed_source_specific[0][0]
                _frame = self.swap_face(_frame, trg_face, src_face)

            elif condition == "specific face":
                for analysed_source, analysed_specific in self.analysed_source_specific:
                    if is_similar_face(analysed_specific, analysed_target, threshold=self.similarity):
                        trg_face = analysed_target
                        src_face = analysed_source
                        _frame = self.swap_face(_frame, trg_face, src_face)

        if condition in single_face_detect_conditions and len(analysed_target_faces) > 0:
            analysed_target = get_single_face(analysed_target_faces, method=condition)
            trg_face = analysed_target
            src_face = self.analysed_source_specific[0][0]
            _frame = self.swap_face(_frame, trg_face, src_face)

        if custom_mask is not None:
            _mask = cv2.resize(custom_mask, _frame.shape[:2][::-1])
            _frame = _mask * frame.astype('float32') + (1 - _mask) * _frame.astype('float32')
            _frame = _frame.clip(0,255).astype('uint8')

        if alpha is not None:
            _frame = np.dstack((_frame, alpha))

        return _frame

    def swap_face(self, frame, trg_face, src_face):
        target_face, generated_face, matrix = self.swapper.forward(frame, trg_face, src_face, n_pass=self.num_of_pass)
        upscaled_face, matrix = self.upscale_face(generated_face, matrix)
        if self.parse_from_target:
            mask = self.face_parsed_mask(target_face)
        else:
            mask = self.face_parsed_mask(upscaled_face)
        result = paste_to_whole(
            upscaled_face,
            frame,
            matrix,
            mask=mask,
            crop_mask=self.mask_crop_values,
            blur_amount=self.mask_blur_amount,
            erode_amount = self.mask_erode_amount
        )
        return result

    def upscale_face(self, face, matrix):
        face_size = face.shape[0]
        _face = cv2.resize(face, (512,512))
        if self.face_upscaler is not None:
            model, runner = self.face_upscaler
            face = runner(face, model)
        upscaled_face = cv2.resize(face, (512,512))
        upscaled_face = mix_two_image(_face, upscaled_face, self.face_upscaler_opacity)
        return upscaled_face, matrix * (512/face_size)

    def face_parsed_mask(self, face):
        if self.face_parser is not None and self.use_face_parsing:
            mask = self.face_parser.parse(face, regions=self.face_parse_regions)
        else:
            mask = None
        return mask