An easy use face swap tool for images and tools only depend on onnxruntime.
from .face_det import FaceAnalysis
from .face_swap import InSwapper
from .face_enhance import GFPGAN
from .super_resolution import BSRGAN
from .pose import PoseEstimator
from .face_core import FaceSwapper
from .pose_core import PoseSwapper

dofaker/__init__.py
justld/dofaker
import os

import numpy as np
import cv2
from moviepy.editor import VideoFileClip

from .face_det import FaceAnalysis
from .super_resolution import BSRGAN
from dofaker.face_swap import get_swapper_model
from dofaker.face_enhance import GFPGAN


class FaceSwapper:

    def __init__(self,
                 face_det_model='buffalo_l',
                 face_swap_model='inswapper',
                 image_sr_model='bsrgan',
                 face_enhance_model='gfpgan',
                 face_det_model_dir='weights/models',
                 face_swap_model_dir='weights/models',
                 image_sr_model_dir='weights/models',
                 face_enhance_model_dir='weights/models',
                 face_sim_thre=0.5,
                 log_iters=10,
                 use_enhancer=True,
                 use_sr=True,
                 scale=1):
        self.face_sim_thre = face_sim_thre
        self.log_iters = log_iters

        self.det_model = FaceAnalysis(name=face_det_model,
                                      root=face_det_model_dir)
        self.det_model.prepare(ctx_id=1, det_size=(640, 640))

        self.swapper_model = get_swapper_model(name=face_swap_model,
                                               root=face_swap_model_dir)
        if use_enhancer:
            self.face_enhance = GFPGAN(name=face_enhance_model,
                                       root=face_enhance_model_dir)
        else:
            self.face_enhance = None

        if use_sr:
            self.sr = BSRGAN(name=image_sr_model,
                             root=image_sr_model_dir,
                             scale=scale)
            self.scale = scale
        else:
            self.sr = None
            self.scale = scale

    def run(self,
            input_path: str,
            dst_face_paths,
            src_face_paths,
            output_dir='output'):
        if isinstance(dst_face_paths, str):
            dst_face_paths = [dst_face_paths]
        if isinstance(src_face_paths, str):
            src_face_paths = [src_face_paths]
        if input_path.lower().endswith(('jpg', 'jpeg', 'webp', 'png', 'bmp')):
            return self.swap_image(input_path, dst_face_paths, src_face_paths,
                                   output_dir)
        else:
            return self.swap_video(input_path, dst_face_paths, src_face_paths,
                                   output_dir)

    def swap_video(self,
                   input_video_path,
                   dst_face_paths,
                   src_face_paths,
                   output_dir='output'):
        assert os.path.exists(
            input_video_path), 'The input video path {} not exist.'
        os.makedirs(output_dir, exist_ok=True)
        src_faces = self.get_faces(src_face_paths)
        if dst_face_paths is not None:
            dst_faces = self.get_faces(dst_face_paths)
            dst_face_embeddings = self.get_faces_embeddings(dst_faces)
            assert len(dst_faces) == len(
                src_faces
            ), 'The detected faces in source images not equal target image faces.'

        video = cv2.VideoCapture(input_video_path)
        fps = video.get(cv2.CAP_PROP_FPS)
        total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
        width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
        frame_size = (width, height)
        print('video fps: {}, total_frames: {}, width: {}, height: {}'.format(
            fps, total_frames, width, height))

        video_name = os.path.basename(input_video_path).split('.')[0]
        four_cc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
        temp_video_path = os.path.join(output_dir,
                                       'temp_{}.mp4'.format(video_name))
        save_video_path = os.path.join(output_dir, '{}.mp4'.format(video_name))
        output_video = cv2.VideoWriter(
            temp_video_path, four_cc, fps,
            (int(frame_size[0] * self.scale), int(frame_size[1] * self.scale)))

        i = 0
        while video.isOpened():
            ret, frame = video.read()
            if ret:
                if dst_face_paths is not None:
                    swapped_image = self.swap_faces(frame,
                                                    dst_face_embeddings,
                                                    src_faces=src_faces)
                else:
                    swapped_image = self.swap_all_faces(frame,
                                                        src_faces=src_faces)
                i += 1
                if i % self.log_iters == 0:
                    print('processing {}/{}'.format(i, total_frames))
                output_video.write(swapped_image)
            else:
                break

        video.release()
        output_video.release()
        self.add_audio_to_video(input_video_path, temp_video_path,
                                save_video_path)
        os.remove(temp_video_path)
        return save_video_path

    def swap_image(self,
                   image_path,
                   dst_face_paths,
                   src_face_paths,
                   output_dir='output'):
        os.makedirs(output_dir, exist_ok=True)
        src_faces = self.get_faces(src_face_paths)
        if dst_face_paths is not None:
            dst_faces = self.get_faces(dst_face_paths)
            dst_face_embeddings = self.get_faces_embeddings(dst_faces)
            assert len(dst_faces) == len(
                src_faces
            ), 'The detected faces in source images not equal target image faces.'

        image = cv2.imread(image_path)
        if dst_face_paths is not None:
            swapped_image = self.swap_faces(image,
                                            dst_face_embeddings,
                                            src_faces=src_faces)
        else:
            swapped_image = self.swap_all_faces(image, src_faces=src_faces)
        base_name = os.path.basename(image_path)
        save_path = os.path.join(output_dir, base_name)
        cv2.imwrite(save_path, swapped_image)
        return save_path

    def add_audio_to_video(self, src_video_path, target_video_path,
                           save_video_path):
        audio = VideoFileClip(src_video_path).audio
        target_video = VideoFileClip(target_video_path)
        target_video = target_video.set_audio(audio)
        target_video.write_videofile(save_video_path)
        return target_video_path

    def get_faces(self, image_paths):
        if isinstance(image_paths, str):
            image_paths = [image_paths]
        faces = []
        for image_path in image_paths:
            image = cv2.imread(image_path)
            assert image is not None, "the source image is None, please check your image {} format.".format(
                image_path)
            img_faces = self.det_model.get(image, max_num=1)
            assert len(
                img_faces
            ) == 1, 'The detected face in image {} must be 1, but got {}, please ensure your image including one face.'.format(
                image_path, len(img_faces))
            faces += img_faces
        return faces

    def swap_faces(self, image, dst_face_embeddings: np.ndarray,
                   src_faces: list) -> np.ndarray:
        res = image.copy()
        image_faces = self.det_model.get(image)
        if len(image_faces) == 0:
            return res
        image_face_embeddings = self.get_faces_embeddings(image_faces)
        sim = np.dot(dst_face_embeddings, image_face_embeddings.T)

        for i in range(dst_face_embeddings.shape[0]):
            index = np.where(sim[i] > self.face_sim_thre)[0].tolist()
            for idx in index:
                res = self.swapper_model.get(res,
                                             image_faces[idx],
                                             src_faces[i],
                                             paste_back=True)
                if self.face_enhance is not None:
                    res = self.face_enhance.get(res,
                                                image_faces[idx],
                                                paste_back=True)

        if self.sr is not None:
            res = self.sr.get(res, image_format='bgr')
        return res

    def swap_all_faces(self, image, src_faces: list) -> np.ndarray:
        assert len(
            src_faces
        ) == 1, 'If replace all faces in source, the number of src face should be 1, but got {}.'.format(
            len(src_faces))
        res = image.copy()
        image_faces = self.det_model.get(image)
        if len(image_faces) == 0:
            return res
        for image_face in image_faces:
            res = self.swapper_model.get(res,
                                         image_face,
                                         src_faces[0],
                                         paste_back=True)
            if self.face_enhance is not None:
                res = self.face_enhance.get(res, image_face, paste_back=True)
        if self.sr is not None:
            res = self.sr.get(res, image_format='bgr')
        return res

    def get_faces_embeddings(self, faces):
        feats = []
        for face in faces:
            feats.append(face.normed_embedding)
        if len(feats) == 1:
            feats = np.array(feats, dtype=np.float32).reshape(1, -1)
        else:
            feats = np.array(feats, dtype=np.float32)
        return feats

dofaker/face_core.py
justld/dofaker
from .face_analysis import FaceAnalysis

dofaker/face_det/__init__.py
justld/dofaker
'''
The following code references:: https://github.com/deepinsight/insightface
'''

import glob
import os.path as osp

import onnxruntime

from insightface import model_zoo
from insightface.utils import ensure_available
from insightface.app.common import Face

from dofaker.utils import download_file, get_model_url

__all__ = ['FaceAnalysis']


class FaceAnalysis:

    def __init__(self,
                 name='buffalo_l',
                 root='weights',
                 allowed_modules=None,
                 **kwargs):
        self.model_dir, _ = download_file(get_model_url(name),
                                          save_dir=root,
                                          overwrite=False)
        onnxruntime.set_default_logger_severity(3)

        self.models = {}
        print(self.model_dir)
        onnx_files = glob.glob(osp.join(self.model_dir, '*.onnx'))
        onnx_files = sorted(onnx_files)
        for onnx_file in onnx_files:
            model = model_zoo.get_model(onnx_file, **kwargs)
            if model is None:
                print('model not recognized:', onnx_file)
            elif allowed_modules is not None and model.taskname not in allowed_modules:
                print('model ignore:', onnx_file, model.taskname)
                del model
            elif model.taskname not in self.models and (allowed_modules is None
                                                        or model.taskname
                                                        in allowed_modules):
                print('find model:', onnx_file, model.taskname,
                      model.input_shape, model.input_mean, model.input_std)
                self.models[model.taskname] = model
            else:
                print('duplicated model task type, ignore:', onnx_file,
                      model.taskname)
                del model
        assert 'detection' in self.models
        self.det_model = self.models['detection']

    def prepare(self, ctx_id, det_thresh=0.5, det_size=(640, 640)):
        self.det_thresh = det_thresh
        assert det_size is not None, "det_size can't be None."
        self.det_size = det_size
        for taskname, model in self.models.items():
            if taskname == 'detection':
                model.prepare(ctx_id,
                              input_size=det_size,
                              det_thresh=det_thresh)
            else:
                model.prepare(ctx_id)

    def get(self, img, max_num=0):
        bboxes, kpss = self.det_model.detect(img,
                                             max_num=max_num,
                                             metric='default')
        if bboxes.shape[0] == 0:
            return []
        ret = []
        for i in range(bboxes.shape[0]):
            bbox = bboxes[i, 0:4]
            det_score = bboxes[i, 4]
            kps = None
            if kpss is not None:
                kps = kpss[i]
            face = Face(bbox=bbox, kps=kps, det_score=det_score)
            for taskname, model in self.models.items():
                if taskname == 'detection':
                    continue
                model.get(img, face)
            ret.append(face)
        return ret

    def draw_on(self, img, faces):
        import cv2
        dimg = img.copy()
        for i in range(len(faces)):
            face = faces[i]
            box = face.bbox.astype('int')
            color = (0, 0, 255)
            cv2.rectangle(dimg, (box[0], box[1]), (box[2], box[3]), color, 2)
            if face.kps is not None:
                kps = face.kps.astype('int')
                for l in range(kps.shape[0]):
                    color = (0, 0, 255)
                    if l == 0 or l == 3:
                        color = (0, 255, 0)
                    cv2.circle(dimg, (kps[l][0], kps[l][1]), 1, color, 2)
            if face.gender is not None and face.age is not None:
                cv2.putText(dimg, '%s,%d' % (face.sex, face.age),
                            (box[0] - 1, box[1] - 4), cv2.FONT_HERSHEY_COMPLEX,
                            0.7, (0, 255, 0), 1)
        return dimg

dofaker/face_det/face_analysis.py
justld/dofaker
from .gfpgan import GFPGAN

dofaker/face_enhance/__init__.py
justld/dofaker
import numpy as np

import cv2

from insightface.utils import face_align
from insightface import model_zoo
from dofaker.utils import download_file, get_model_url


class GFPGAN:

    def __init__(self, name='gfpgan', root='weights/models') -> None:
        _, model_file = download_file(get_model_url(name),
                                      save_dir=root,
                                      overwrite=False)
        providers = model_zoo.model_zoo.get_default_providers()
        self.session = model_zoo.model_zoo.PickableInferenceSession(
            model_file, providers=providers)

        self.input_mean = 127.5
        self.input_std = 127.5
        inputs = self.session.get_inputs()
        self.input_names = []
        for inp in inputs:
            self.input_names.append(inp.name)
        outputs = self.session.get_outputs()
        output_names = []
        for out in outputs:
            output_names.append(out.name)
        self.output_names = output_names
        assert len(
            self.output_names
        ) == 1, "The output number of GFPGAN model should be 1, but got {}, please check your model.".format(
            len(self.output_names))
        output_shape = outputs[0].shape
        input_cfg = inputs[0]
        input_shape = input_cfg.shape
        self.input_shape = input_shape
        print('face_enhance-shape:', self.input_shape)
        self.input_size = tuple(input_shape[2:4][::-1])

    def forward(self, image, image_format='bgr'):
        if isinstance(image, str):
            image = cv2.imread(image, 1)
        elif isinstance(image, np.ndarray):
            if image_format == 'bgr':
                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            elif image_format == 'rgb':
                pass
            else:
                raise UserWarning(
                    "gfpgan not support image format {}".format(image_format))
        else:
            raise UserWarning(
                "gfpgan input must be str or np.ndarray, but got {}.".format(
                    type(image)))
        img = (image - self.input_mean) / self.input_std
        pred = self.session.run(self.output_names,
                                {self.input_names[0]: img})[0]
        return pred

    def _get(self, img, image_format='bgr'):
        if image_format.lower() == 'bgr':
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        elif image_format.lower() == 'rgb':
            pass
        else:
            raise UserWarning(
                "gfpgan not support image format {}".format(image_format))
        h, w, c = img.shape
        img = cv2.resize(img, (self.input_shape[-1], self.input_shape[-2]))
        blob = cv2.dnn.blobFromImage(
            img,
            1.0 / self.input_std,
            self.input_size,
            (self.input_mean, self.input_mean, self.input_mean),
            swapRB=False)
        pred = self.session.run(self.output_names,
                                {self.input_names[0]: blob})[0]
        image_aug = pred.transpose((0, 2, 3, 1))[0]
        rgb_aug = np.clip(self.input_std * image_aug + self.input_mean, 0,
                          255).astype(np.uint8)
        rgb_aug = cv2.resize(rgb_aug, (w, h))
        bgr_image = rgb_aug[:, :, ::-1]
        return bgr_image

    def get(self, img, target_face, paste_back=True, image_format='bgr'):
        aimg, M = face_align.norm_crop2(img, target_face.kps,
                                        self.input_size[0])
        bgr_fake = self._get(aimg, image_format='bgr')
        if not paste_back:
            return bgr_fake, M
        else:
            target_img = img
            fake_diff = bgr_fake.astype(np.float32) - aimg.astype(np.float32)
            fake_diff = np.abs(fake_diff).mean(axis=2)
            fake_diff[:2, :] = 0
            fake_diff[-2:, :] = 0
            fake_diff[:, :2] = 0
            fake_diff[:, -2:] = 0
            IM = cv2.invertAffineTransform(M)
            img_white = np.full((aimg.shape[0], aimg.shape[1]),
                                255,
                                dtype=np.float32)
            bgr_fake = cv2.warpAffine(
                bgr_fake,
                IM, (target_img.shape[1], target_img.shape[0]),
                borderValue=0.0)
            img_white = cv2.warpAffine(
                img_white,
                IM, (target_img.shape[1], target_img.shape[0]),
                borderValue=0.0)
            fake_diff = cv2.warpAffine(
                fake_diff,
                IM, (target_img.shape[1], target_img.shape[0]),
                borderValue=0.0)
            img_white[img_white > 20] = 255
            fthresh = 10
            fake_diff[fake_diff < fthresh] = 0
            fake_diff[fake_diff >= fthresh] = 255
            img_mask = img_white
            mask_h_inds, mask_w_inds = np.where(img_mask == 255)
            mask_h = np.max(mask_h_inds) - np.min(mask_h_inds)
            mask_w = np.max(mask_w_inds) - np.min(mask_w_inds)
            mask_size = int(np.sqrt(mask_h * mask_w))
            k = max(mask_size // 10, 10)
            #k = max(mask_size//20, 6)
            #k = 6
            kernel = np.ones((k, k), np.uint8)
            img_mask = cv2.erode(img_mask, kernel, iterations=1)
            kernel = np.ones((2, 2), np.uint8)
            fake_diff = cv2.dilate(fake_diff, kernel, iterations=1)
            k = max(mask_size // 20, 5)
            kernel_size = (k, k)
            blur_size = tuple(2 * i + 1 for i in kernel_size)
            img_mask = cv2.GaussianBlur(img_mask, blur_size, 0)
            k = 5
            kernel_size = (k, k)
            blur_size = tuple(2 * i + 1 for i in kernel_size)
            fake_diff = cv2.GaussianBlur(fake_diff, blur_size, 0)
            img_mask /= 255
            fake_diff /= 255
            img_mask = np.reshape(img_mask,
                                  [img_mask.shape[0], img_mask.shape[1], 1])
            fake_merged = img_mask * bgr_fake + (
                1 - img_mask) * target_img.astype(np.float32)
            fake_merged = fake_merged.astype(np.uint8)
            return fake_merged

dofaker/face_enhance/gfpgan.py
justld/dofaker
from .inswapper import InSwapper


def get_swapper_model(name='', root=None, **kwargs):
    if name.lower() == 'inswapper':
        return InSwapper(name=name, root=root, **kwargs)
    else:
        raise UserWarning('The swapper model {} not support.'.format(name))

dofaker/face_swap/__init__.py
justld/dofaker
class BaseSwapper:

    def forward(self, img, latent, *args, **kwargs):
        raise NotImplementedError

    def get(self,
            img,
            target_face,
            source_face,
            paste_back=True,
            *args,
            **kwargs):
        raise NotImplementedError

dofaker/face_swap/base_swapper.py
justld/dofaker
import numpy as np
import cv2
import onnx
from onnx import numpy_helper

from insightface import model_zoo
from insightface.utils import face_align
from .base_swapper import BaseSwapper

from dofaker.utils import download_file, get_model_url


class InSwapper(BaseSwapper):

    def __init__(self, name='inswapper', root='weights/models'):
        _, model_file = download_file(get_model_url(name),
                                      save_dir=root,
                                      overwrite=False)
        providers = model_zoo.model_zoo.get_default_providers()
        self.session = model_zoo.model_zoo.PickableInferenceSession(
            model_file, providers=providers)

        model = onnx.load(model_file)
        graph = model.graph
        self.emap = numpy_helper.to_array(graph.initializer[-1])
        self.input_mean = 0.0
        self.input_std = 255.0

        inputs = self.session.get_inputs()
        self.input_names = []
        for inp in inputs:
            self.input_names.append(inp.name)
        outputs = self.session.get_outputs()
        output_names = []
        for out in outputs:
            output_names.append(out.name)
        self.output_names = output_names
        assert len(
            self.output_names
        ) == 1, "The output number of inswapper model should be 1, but got {}, please check your model.".format(
            len(self.output_names))
        output_shape = outputs[0].shape
        input_cfg = inputs[0]
        input_shape = input_cfg.shape
        self.input_shape = input_shape
        print('inswapper-shape:', self.input_shape)
        self.input_size = tuple(input_shape[2:4][::-1])

    def forward(self, img, latent):
        img = (img - self.input_mean) / self.input_std
        pred = self.session.run(self.output_names, {
            self.input_names[0]: img,
            self.input_names[1]: latent
        })[0]
        return pred

    def get(self, img, target_face, source_face, paste_back=True):
        aimg, M = face_align.norm_crop2(img, target_face.kps,
                                        self.input_size[0])
        blob = cv2.dnn.blobFromImage(
            aimg,
            1.0 / self.input_std,
            self.input_size,
            (self.input_mean, self.input_mean, self.input_mean),
            swapRB=True)
        latent = source_face.normed_embedding.reshape((1, -1))
        latent = np.dot(latent, self.emap)
        latent /= np.linalg.norm(latent)
        pred = self.session.run(self.output_names, {
            self.input_names[0]: blob,
            self.input_names[1]: latent
        })[0]
        img_fake = pred.transpose((0, 2, 3, 1))[0]
        bgr_fake = np.clip(255 * img_fake, 0, 255).astype(np.uint8)[:, :, ::-1]
        if not paste_back:
            return bgr_fake, M
        else:
            target_img = img
            fake_diff = bgr_fake.astype(np.float32) - aimg.astype(np.float32)
            fake_diff = np.abs(fake_diff).mean(axis=2)
            fake_diff[:2, :] = 0
            fake_diff[-2:, :] = 0
            fake_diff[:, :2] = 0
            fake_diff[:, -2:] = 0
            IM = cv2.invertAffineTransform(M)
            img_white = np.full((aimg.shape[0], aimg.shape[1]),
                                255,
                                dtype=np.float32)
            bgr_fake = cv2.warpAffine(
                bgr_fake,
                IM, (target_img.shape[1], target_img.shape[0]),
                borderValue=0.0)
            img_white = cv2.warpAffine(
                img_white,
                IM, (target_img.shape[1], target_img.shape[0]),
                borderValue=0.0)
            fake_diff = cv2.warpAffine(
                fake_diff,
                IM, (target_img.shape[1], target_img.shape[0]),
                borderValue=0.0)
            img_white[img_white > 20] = 255
            fthresh = 10
            fake_diff[fake_diff < fthresh] = 0
            fake_diff[fake_diff >= fthresh] = 255
            img_mask = img_white
            mask_h_inds, mask_w_inds = np.where(img_mask == 255)
            mask_h = np.max(mask_h_inds) - np.min(mask_h_inds)
            mask_w = np.max(mask_w_inds) - np.min(mask_w_inds)
            mask_size = int(np.sqrt(mask_h * mask_w))
            k = max(mask_size // 10, 10)
            #k = max(mask_size//20, 6)
            #k = 6
            kernel = np.ones((k, k), np.uint8)
            img_mask = cv2.erode(img_mask, kernel, iterations=1)
            kernel = np.ones((2, 2), np.uint8)
            fake_diff = cv2.dilate(fake_diff, kernel, iterations=1)
            k = max(mask_size // 20, 5)
            #k = 3
            #k = 3
            kernel_size = (k, k)
            blur_size = tuple(2 * i + 1 for i in kernel_size)
            img_mask = cv2.GaussianBlur(img_mask, blur_size, 0)
            k = 5
            kernel_size = (k, k)
            blur_size = tuple(2 * i + 1 for i in kernel_size)
            fake_diff = cv2.GaussianBlur(fake_diff, blur_size, 0)
            img_mask /= 255
            fake_diff /= 255
            #img_mask = fake_diff
            img_mask = np.reshape(img_mask,
                                  [img_mask.shape[0], img_mask.shape[1], 1])
            fake_merged = img_mask * bgr_fake + (
                1 - img_mask) * target_img.astype(np.float32)
            fake_merged = fake_merged.astype(np.uint8)
            return fake_merged

dofaker/face_swap/inswapper.py
justld/dofaker
from .pose_estimator import PoseEstimator
from .pose_transfer import PoseTransfer

dofaker/pose/__init__.py
justld/dofaker
import numpy as np

import cv2
from scipy.ndimage.filters import gaussian_filter

from .pose_utils import _get_keypoints, _pad_image
from insightface import model_zoo
from dofaker.utils import download_file, get_model_url


class PoseEstimator:

    def __init__(self, name='openpose_body', root='weights/models'):
        _, model_file = download_file(get_model_url(name),
                                      save_dir=root,
                                      overwrite=False)
        providers = model_zoo.model_zoo.get_default_providers()
        self.session = model_zoo.model_zoo.PickableInferenceSession(
            model_file, providers=providers)

        self.input_mean = 127.5
        self.input_std = 255.0
        inputs = self.session.get_inputs()
        self.input_names = []
        for inp in inputs:
            self.input_names.append(inp.name)
        outputs = self.session.get_outputs()
        output_names = []
        for out in outputs:
            output_names.append(out.name)
        self.output_names = output_names
        assert len(
            self.output_names
        ) == 2, "The output number of PoseEstimator model should be 2, but got {}, please check your model.".format(
            len(self.output_names))
        output_shape = outputs[0].shape
        input_cfg = inputs[0]
        input_shape = input_cfg.shape
        self.input_shape = input_shape
        print('pose estimator shape:', self.input_shape)

    def forward(self, image, image_format='rgb'):
        if isinstance(image, str):
            image = cv2.imread(image, 1)
            image_format = 'bgr'
        elif isinstance(image, np.ndarray):
            if image_format == 'bgr':
                pass
            elif image_format == 'rgb':
                image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
                image_format = 'bgr'
            else:
                raise UserWarning(
                    "PoseEstimator not support image format {}".format(
                        image_format))
        else:
            raise UserWarning(
                "PoseEstimator input must be str or np.ndarray, but got {}.".
                format(type(image)))

        scales = [0.5]
        stride = 8
        bboxsize = 368
        padvalue = 128
        thresh_1 = 0.1
        thresh_2 = 0.05

        multipliers = [scale * bboxsize / image.shape[0] for scale in scales]
        heatmap_avg = np.zeros((image.shape[0], image.shape[1], 19))
        paf_avg = np.zeros((image.shape[0], image.shape[1], 38))

        for scale in multipliers:
            image_scaled = cv2.resize(image, (0, 0),
                                      fx=scale,
                                      fy=scale,
                                      interpolation=cv2.INTER_CUBIC)
            image_padded, pads = _pad_image(image_scaled, stride, padvalue)

            image_tensor = np.expand_dims(np.transpose(image_padded, (2, 0, 1)),
                                          0)
            blob = (np.float32(image_tensor) - self.input_mean) / self.input_std

            pred = self.session.run(self.output_names,
                                    {self.input_names[0]: blob})
            Mconv7_stage6_L1, Mconv7_stage6_L2 = pred[0], pred[1]

            heatmap = np.transpose(np.squeeze(Mconv7_stage6_L2), (1, 2, 0))
            heatmap = cv2.resize(heatmap, (0, 0),
                                 fx=stride,
                                 fy=stride,
                                 interpolation=cv2.INTER_CUBIC)
            heatmap = heatmap[:image_padded.shape[0] -
                              pads[3], :image_padded.shape[1] - pads[2], :]
            heatmap = cv2.resize(heatmap, (image.shape[1], image.shape[0]),
                                 interpolation=cv2.INTER_CUBIC)

            paf = np.transpose(np.squeeze(Mconv7_stage6_L1), (1, 2, 0))
            paf = cv2.resize(paf, (0, 0),
                             fx=stride,
                             fy=stride,
                             interpolation=cv2.INTER_CUBIC)
            paf = paf[:image_padded.shape[0] - pads[3], :image_padded.shape[1] -
                      pads[2], :]
            paf = cv2.resize(paf, (image.shape[1], image.shape[0]),
                             interpolation=cv2.INTER_CUBIC)

            heatmap_avg += (heatmap / len(multipliers))
            paf_avg += (paf / len(multipliers))

        all_peaks = []
        num_peaks = 0

        for part in range(18):
            map_orig = heatmap_avg[:, :, part]
            map_filt = gaussian_filter(map_orig, sigma=3)

            map_L = np.zeros_like(map_filt)
            map_T = np.zeros_like(map_filt)
            map_R = np.zeros_like(map_filt)
            map_B = np.zeros_like(map_filt)
            map_L[1:, :] = map_filt[:-1, :]
            map_T[:, 1:] = map_filt[:, :-1]
            map_R[:-1, :] = map_filt[1:, :]
            map_B[:, :-1] = map_filt[:, 1:]

            peaks_binary = np.logical_and.reduce(
                (map_filt >= map_L, map_filt >= map_T, map_filt
                 >= map_R, map_filt >= map_B, map_filt > thresh_1))
            peaks = list(
                zip(np.nonzero(peaks_binary)[1],
                    np.nonzero(peaks_binary)[0]))
            peaks_ids = range(num_peaks, num_peaks + len(peaks))
            peaks_with_scores = [
                peak + (map_orig[peak[1], peak[0]], ) for peak in peaks
            ]
            peaks_with_scores_and_ids = [peaks_with_scores[i] + (peaks_ids[i],) \
                                         for i in range(len(peaks_ids))]
            all_peaks.append(peaks_with_scores_and_ids)
            num_peaks += len(peaks)

        map_idx = [[31, 32], [39, 40], [33, 34], [35, 36], [41, 42], [43, 44],
                   [19, 20], [21, 22], [23, 24], [25, 26], [27, 28], [29, 30],
                   [47, 48], [49, 50], [53, 54], [51, 52], [55, 56], [37, 38],
                   [45, 46]]
        limbseq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9],
                   [9, 10], [10, 11], [2, 12], [12, 13], [13, 14], [2, 1],
                   [1, 15], [15, 17], [1, 16], [16, 18], [3, 17], [6, 18]]

        all_connections = []
        spl_k = []
        mid_n = 10

        for k in range(len(map_idx)):
            score_mid = paf_avg[:, :, [x - 19 for x in map_idx[k]]]
            candidate_A = all_peaks[limbseq[k][0] - 1]
            candidate_B = all_peaks[limbseq[k][1] - 1]
            n_A = len(candidate_A)
            n_B = len(candidate_B)
            index_A, index_B = limbseq[k]
            if n_A != 0 and n_B != 0:
                connection_candidates = []
                for i in range(n_A):
                    for j in range(n_B):
                        v = np.subtract(candidate_B[j][:2], candidate_A[i][:2])
                        n = np.sqrt(v[0] * v[0] + v[1] * v[1])
                        v = np.divide(v, n)

                        ab = list(
                            zip(
                                np.linspace(candidate_A[i][0],
                                            candidate_B[j][0],
                                            num=mid_n),
                                np.linspace(candidate_A[i][1],
                                            candidate_B[j][1],
                                            num=mid_n)))
                        vx = np.array([
                            score_mid[int(round(ab[x][1])),
                                      int(round(ab[x][0])), 0]
                            for x in range(len(ab))
                        ])
                        vy = np.array([
                            score_mid[int(round(ab[x][1])),
                                      int(round(ab[x][0])), 1]
                            for x in range(len(ab))
                        ])
                        score_midpoints = np.multiply(vx, v[0]) + np.multiply(
                            vy, v[1])
                        score_with_dist_prior = sum(
                            score_midpoints) / len(score_midpoints) + min(
                                0.5 * image.shape[0] / n - 1, 0)
                        criterion_1 = len(
                            np.nonzero(score_midpoints > thresh_2)
                            [0]) > 0.8 * len(score_midpoints)
                        criterion_2 = score_with_dist_prior > 0
                        if criterion_1 and criterion_2:
                            connection_candidate = [
                                i, j, score_with_dist_prior,
                                score_with_dist_prior + candidate_A[i][2] +
                                candidate_B[j][2]
                            ]
                            connection_candidates.append(connection_candidate)
                connection_candidates = sorted(connection_candidates,
                                               key=lambda x: x[2],
                                               reverse=True)
                connection = np.zeros((0, 5))
                for candidate in connection_candidates:
                    i, j, s = candidate[0:3]
                    if i not in connection[:, 3] and j not in connection[:, 4]:
                        connection = np.vstack([
                            connection,
                            [candidate_A[i][3], candidate_B[j][3], s, i, j]
                        ])
                        if len(connection) >= min(n_A, n_B):
                            break
                all_connections.append(connection)
            else:
                spl_k.append(k)
                all_connections.append([])

        candidate = np.array(
            [item for sublist in all_peaks for item in sublist])
        subset = np.ones((0, 20)) * -1

        for k in range(len(map_idx)):
            if k not in spl_k:
                part_As = all_connections[k][:, 0]
                part_Bs = all_connections[k][:, 1]
                index_A, index_B = np.array(limbseq[k]) - 1
                for i in range(len(all_connections[k])):
                    found = 0
                    subset_idx = [-1, -1]
                    for j in range(len(subset)):
                        if subset[j][index_A] == part_As[i] or subset[j][
                                index_B] == part_Bs[i]:
                            subset_idx[found] = j
                            found += 1
                    if found == 1:
                        j = subset_idx[0]
                        if subset[j][index_B] != part_Bs[i]:
                            subset[j][index_B] = part_Bs[i]
                            subset[j][-1] += 1
                            subset[j][-2] += candidate[
                                part_Bs[i].astype(int),
                                2] + all_connections[k][i][2]
                    elif found == 2:
                        j1, j2 = subset_idx
                        membership = ((subset[j1] >= 0).astype(int) +
                                      (subset[j2] >= 0).astype(int))[:-2]
                        if len(np.nonzero(membership == 2)[0]) == 0:
                            subset[j1][:-2] += (subset[j2][:-2] + 1)
                            subset[j1][-2:] += subset[j2][-2:]
                            subset[j1][-2] += all_connections[k][i][2]
                            subset = np.delete(subset, j2, 0)
                        else:
                            subset[j1][index_B] = part_Bs[i]
                            subset[j1][-1] += 1
                            subset[j1][-2] += candidate[
                                part_Bs[i].astype(int),
                                2] + all_connections[k][i][2]
                    elif not found and k < 17:
                        row = np.ones(20) * -1
                        row[index_A] = part_As[i]
                        row[index_B] = part_Bs[i]
                        row[-1] = 2
                        row[-2] = sum(
                            candidate[all_connections[k][i, :2].astype(int),
                                      2]) + all_connections[k][i][2]
                        subset = np.vstack([subset, row])

        del_idx = []

        for i in range(len(subset)):
            if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:
                del_idx.append(i)
        subset = np.delete(subset, del_idx, axis=0)

        return _get_keypoints(candidate, subset)

    def get(self, image, image_format='rgb'):
        return self.forward(image, image_format)

dofaker/pose/pose_estimator.py
justld/dofaker
import cv2
import numpy as np
from scipy.ndimage.filters import gaussian_filter

from .pose_utils import _get_keypoints, _pad_image
from insightface import model_zoo
from dofaker.utils import download_file, get_model_url
from dofaker.transforms import center_crop, pad


class PoseTransfer:

    def __init__(self,
                 name='pose_transfer',
                 root='weights/models',
                 pose_estimator=None):
        assert pose_estimator is not None, "The pose_estimator of PoseTransfer shouldn't be None"
        self.pose_estimator = pose_estimator
        _, model_file = download_file(get_model_url(name),
                                      save_dir=root,
                                      overwrite=False)
        providers = model_zoo.model_zoo.get_default_providers()
        self.session = model_zoo.model_zoo.PickableInferenceSession(
            model_file, providers=providers)

        self.input_mean = 127.5
        self.input_std = 127.5
        inputs = self.session.get_inputs()
        self.input_names = []
        for inp in inputs:
            self.input_names.append(inp.name)
        outputs = self.session.get_outputs()
        output_names = []
        for out in outputs:
            output_names.append(out.name)
        self.output_names = output_names
        assert len(
            self.output_names
        ) == 1, "The output number of PoseTransfer model should be 1, but got {}, please check your model.".format(
            len(self.output_names))
        output_shape = outputs[0].shape
        input_cfg = inputs[0]
        input_shape = input_cfg.shape
        self.input_shape = input_shape
        print('pose transfer shape:', self.input_shape)

    def forward(self, source_image, target_image, image_format='rgb'):
        h, w, c = source_image.shape
        if image_format == 'rgb':
            pass
        elif image_format == 'bgr':
            source_image = cv2.cvtColor(source_image, cv2.COLOR_BGR2RGB)
            target_image = cv2.cvtColor(target_image, cv2.COLOR_BGR2RGB)
            image_format = 'rgb'
        else:
            raise UserWarning(
                "PoseTransfer not support image format {}".format(image_format))
        imgA = self._resize_and_pad_image(source_image)
        kptA = self._estimate_keypoints(imgA, image_format=image_format)
        mapA = self._keypoints2heatmaps(kptA)

        imgB = self._resize_and_pad_image(target_image)
        kptB = self._estimate_keypoints(imgB)
        mapB = self._keypoints2heatmaps(kptB)

        imgA_t = (imgA.astype('float32') - self.input_mean) / self.input_std
        imgA_t = imgA_t.transpose([2, 0, 1])[None, ...]
        mapA_t = mapA.transpose([2, 0, 1])[None, ...]
        mapB_t = mapB.transpose([2, 0, 1])[None, ...]
        mapAB_t = np.concatenate((mapA_t, mapB_t), axis=1)
        pred = self.session.run(self.output_names, {
            self.input_names[0]: imgA_t,
            self.input_names[1]: mapAB_t
        })[0]
        target_image = pred.transpose((0, 2, 3, 1))[0]
        bgr_target_image = np.clip(
            self.input_std * target_image + self.input_mean, 0,
            255).astype(np.uint8)[:, :, ::-1]
        crop_size = (256,
                     min((256 * target_image.shape[1] // target_image.shape[0]),
                         176))
        bgr_image = center_crop(bgr_target_image, crop_size)
        bgr_image = cv2.resize(bgr_image, (w, h), interpolation=cv2.INTER_CUBIC)
        return bgr_image

    def get(self, source_image, target_image, image_format='rgb'):
        return self.forward(source_image, target_image, image_format)

    def _resize_and_pad_image(self, image: np.ndarray, size=256):
        w = size * image.shape[1] // image.shape[0]
        w_box = min(w, size * 11 // 16)
        image = cv2.resize(image, (w, size), interpolation=cv2.INTER_CUBIC)
        image = center_crop(image, (size, w_box))
        image = pad(image,
                    size - w_box,
                    size - w_box,
                    size - w_box,
                    size - w_box,
                    fill=255)
        image = center_crop(image, (size, size))
        return image

    def _estimate_keypoints(self, image: np.ndarray, image_format='rgb'):
        keypoints = self.pose_estimator.get(image, image_format)
        keypoints = keypoints[0] if len(keypoints) > 0 else np.zeros(
            (18, 3), dtype=np.int32)
        keypoints[np.where(keypoints[:, 2] == 0), :2] = -1
        keypoints = keypoints[:, :2]
        return keypoints

    def _keypoints2heatmaps(self, keypoints, size=256):
        heatmaps = np.zeros((size, size, keypoints.shape[0]), dtype=np.float32)
        for k in range(keypoints.shape[0]):
            x, y = keypoints[k]
            if x == -1 or y == -1:
                continue
            heatmaps[y, x, k] = 1.0
        return heatmaps

dofaker/pose/pose_transfer.py
justld/dofaker
import numpy as np


def _pad_image(image, stride=1, padvalue=0):
    assert len(image.shape) == 2 or len(image.shape) == 3
    h, w = image.shape[:2]
    pads = [None] * 4
    pads[0] = 0  # left
    pads[1] = 0  # top
    pads[2] = 0 if (w % stride == 0) else stride - (w % stride)  # right
    pads[3] = 0 if (h % stride == 0) else stride - (h % stride)  # bottom
    num_channels = 1 if len(image.shape) == 2 else image.shape[2]
    image_padded = np.ones(
        (h + pads[3], w + pads[2], num_channels), dtype=np.uint8) * padvalue
    image_padded = np.squeeze(image_padded)
    image_padded[:h, :w] = image
    return image_padded, pads


def _get_keypoints(candidates, subsets):
    k = subsets.shape[0]
    keypoints = np.zeros((k, 18, 3), dtype=np.int32)
    for i in range(k):
        for j in range(18):
            index = np.int32(subsets[i][j])
            if index != -1:
                x, y = np.int32(candidates[index][:2])
                keypoints[i][j] = (x, y, 1)
    return keypoints

dofaker/pose/pose_utils.py
justld/dofaker
import os
import cv2

import numpy as np
from moviepy.editor import VideoFileClip

from .pose import PoseEstimator, PoseTransfer
from .face_enhance import GFPGAN
from .super_resolution import BSRGAN
from .face_det import FaceAnalysis


class PoseSwapper:

    def __init__(self,
                 pose_estimator_name='openpose_body',
                 pose_estimator_model_dir='weights/models',
                 pose_transfer_name='pose_transfer',
                 pose_transfer_model_dir='weights/models',
                 image_sr_model='bsrgan',
                 image_sr_model_dir='weights/models',
                 face_enhance_name='gfpgan',
                 face_enhance_model_dir='weights/models/',
                 face_det_model='buffalo_l',
                 face_det_model_dir='weights/models',
                 log_iters=10,
                 use_enhancer=True,
                 use_sr=True,
                 scale=1):
        pose_estimator = PoseEstimator(name=pose_estimator_name,
                                       root=pose_estimator_model_dir)
        self.pose_transfer = PoseTransfer(name=pose_transfer_name,
                                          root=pose_transfer_model_dir,
                                          pose_estimator=pose_estimator)

        if use_enhancer:
            self.det_model = FaceAnalysis(name=face_det_model,
                                          root=face_det_model_dir)
            self.det_model.prepare(ctx_id=1, det_size=(640, 640))
            self.face_enhance = GFPGAN(name=face_enhance_name,
                                       root=face_enhance_model_dir)
        self.use_enhancer = use_enhancer

        if use_sr:
            self.sr_model = BSRGAN(name=image_sr_model,
                                   root=image_sr_model_dir,
                                   scale=scale)
            self.scale = scale
        else:
            self.scale = 1
        self.use_sr = use_sr
        self.log_iters = log_iters

    def run(self, input_path, target_path, output_dir='output'):
        assert os.path.exists(
            input_path), "The input path {} not exists.".format(input_path)
        assert os.path.exists(
            target_path), "The target path {} not exists.".format(target_path)
        os.makedirs(output_dir, exist_ok=True)
        assert input_path.lower().endswith(
            ('jpg', 'jpeg', 'webp', 'png', 'bmp')
        ), "pose swapper input must be image endswith ('jpg', 'jpeg', 'webp', 'png', 'bmp'), but got {}.".format(
            input_path)
        if target_path.lower().endswith(('jpg', 'jpeg', 'webp', 'png', 'bmp')):
            return self.transfer_image(input_path, target_path, output_dir)
        else:
            return self.transfer_video(input_path, target_path, output_dir)

    def transfer_image(self, input_path, target_path, output_dir):
        source = cv2.imread(input_path)
        target = cv2.imread(target_path)
        transferred_image = self.transfer_pose(source,
                                               target,
                                               image_format='bgr')
        base_name = os.path.basename(input_path)
        output_path = os.path.join(output_dir, base_name)
        cv2.imwrite(output_path, transferred_image)
        return output_path

    def transfer_video(self, input_path, target_path, output_dir):
        source = cv2.imread(input_path)
        video = cv2.VideoCapture(target_path)
        fps = video.get(cv2.CAP_PROP_FPS)
        total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
        height, width, _ = source.shape
        frame_size = (width, height)
        print('video fps: {}, total_frames: {}, width: {}, height: {}'.format(
            fps, total_frames, width, height))

        video_name = os.path.basename(input_path).split('.')[0]
        four_cc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
        temp_video_path = os.path.join(output_dir,
                                       'temp_{}.mp4'.format(video_name))
        save_video_path = os.path.join(output_dir, '{}.mp4'.format(video_name))
        output_video = cv2.VideoWriter(
            temp_video_path, four_cc, fps,
            (int(frame_size[0] * self.scale), int(frame_size[1] * self.scale)))
        i = 0
        while video.isOpened():
            ret, frame = video.read()
            if ret:
                transferred_image = self.transfer_pose(source,
                                                       frame,
                                                       image_format='bgr')
                i += 1
                if i % self.log_iters == 0:
                    print('processing {}/{}'.format(i, total_frames))
                output_video.write(transferred_image)
            else:
                break

        video.release()
        output_video.release()
        print(temp_video_path)
        self.add_audio_to_video(target_path, temp_video_path, save_video_path)
        os.remove(temp_video_path)
        return save_video_path

    def transfer_pose(self, source, target, image_format='bgr'):
        transferred_image = self.pose_transfer.get(source,
                                                   target,
                                                   image_format=image_format)
        if self.use_enhancer:
            faces = self.det_model.get(transferred_image, max_num=1)
            for face in faces:
                transferred_image = self.face_enhance.get(
                    transferred_image,
                    face,
                    paste_back=True,
                    image_format=image_format)

        if self.use_sr:
            transferred_image = self.sr_model.get(transferred_image,
                                                  image_format=image_format)
        return transferred_image

    def add_audio_to_video(self, src_video_path, target_video_path,
                           save_video_path):
        audio = VideoFileClip(src_video_path).audio
        target_video = VideoFileClip(target_video_path)
        target_video = target_video.set_audio(audio)
        target_video.write_videofile(save_video_path)
        return target_video_path

dofaker/pose_core.py
justld/dofaker
from .bsrgan import BSRGAN

dofaker/super_resolution/__init__.py
justld/dofaker
import numpy as np

import cv2

from insightface import model_zoo
from dofaker.utils import download_file, get_model_url


class BSRGAN:

    def __init__(self, name='bsrgan', root='weights/models', scale=1) -> None:
        _, model_file = download_file(get_model_url(name),
                                      save_dir=root,
                                      overwrite=False)
        self.scale = scale
        providers = model_zoo.model_zoo.get_default_providers()
        self.session = model_zoo.model_zoo.PickableInferenceSession(
            model_file, providers=providers)

        self.input_mean = 0.0
        self.input_std = 255.0
        inputs = self.session.get_inputs()
        self.input_names = []
        for inp in inputs:
            self.input_names.append(inp.name)
        outputs = self.session.get_outputs()
        output_names = []
        for out in outputs:
            output_names.append(out.name)
        self.output_names = output_names
        assert len(
            self.output_names
        ) == 1, "The output number of BSRGAN model should be 1, but got {}, please check your model.".format(
            len(self.output_names))
        output_shape = outputs[0].shape
        input_cfg = inputs[0]
        input_shape = input_cfg.shape
        self.input_shape = input_shape
        print('image super resolution shape:', self.input_shape)

    def forward(self, image, image_format='bgr'):
        if isinstance(image, str):
            image = cv2.imread(image, 1)
            image_format = 'bgr'
        elif isinstance(image, np.ndarray):
            if image_format == 'bgr':
                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            elif image_format == 'rgb':
                pass
            else:
                raise UserWarning(
                    "BSRGAN not support image format {}".format(image_format))
        else:
            raise UserWarning(
                "BSRGAN input must be str or np.ndarray, but got {}.".format(
                    type(image)))
        img = (image - self.input_mean) / self.input_std
        pred = self.session.run(self.output_names,
                                {self.input_names[0]: img})[0]
        return pred

    def get(self, img, image_format='bgr'):
        if image_format.lower() == 'bgr':
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        elif image_format.lower() == 'rgb':
            pass
        else:
            raise UserWarning(
                "gfpgan not support image format {}".format(image_format))
        h, w, c = img.shape
        blob = cv2.dnn.blobFromImage(
            img,
            1.0 / self.input_std, (w, h),
            (self.input_mean, self.input_mean, self.input_mean),
            swapRB=False)
        pred = self.session.run(self.output_names,
                                {self.input_names[0]: blob})[0]
        image_aug = pred.transpose((0, 2, 3, 1))[0]
        rgb_aug = np.clip(self.input_std * image_aug + self.input_mean, 0,
                          255).astype(np.uint8)
        rgb_aug = cv2.resize(rgb_aug,
                             (int(w * self.scale), int(h * self.scale)))
        bgr_aug = rgb_aug[:, :, ::-1]
        return bgr_aug

dofaker/super_resolution/bsrgan.py
justld/dofaker
from .functional import center_crop, pad

dofaker/transforms/__init__.py
justld/dofaker
import numbers
import numpy as np

import cv2


def center_crop(image: np.ndarray, output_size):
    if isinstance(output_size, numbers.Number):
        output_size = (int(output_size), int(output_size))
    elif isinstance(output_size, (tuple, list)) and len(output_size) == 1:
        output_size = (output_size[0], output_size[0])

    image_height, image_width, c = image.shape
    crop_height, crop_width = output_size

    if crop_width > image_width or crop_height > image_height:
        padding_ltrb = [
            (crop_width - image_width) // 2 if crop_width > image_width else 0,
            (crop_height - image_height) //
            2 if crop_height > image_height else 0,
            (crop_width - image_width + 1) //
            2 if crop_width > image_width else 0,
            (crop_height - image_height + 1) //
            2 if crop_height > image_height else 0,
        ]
        image = cv2.copyMakeBorder(image,
                                   padding_ltrb[1],
                                   padding_ltrb[3],
                                   padding_ltrb[0],
                                   padding_ltrb[2],
                                   cv2.BORDER_CONSTANT,
                                   value=(0, 0, 0))
        image_height, image_width, c = image.shape
        if crop_width == image_width and crop_height == image_height:
            return image

    crop_top = int(round((image_height - crop_height) / 2.0))
    crop_left = int(round((image_width - crop_width) / 2.0))
    return image[crop_top:crop_top + crop_height,
                 crop_left:crop_left + crop_width]


def pad(image,
        left,
        top,
        right,
        bottom,
        fill: int = 0,
        padding_mode: str = "constant"):
    if padding_mode == 'constant':
        return cv2.copyMakeBorder(image,
                                  top,
                                  bottom,
                                  left,
                                  right,
                                  cv2.BORDER_CONSTANT,
                                  value=(fill, fill, fill))
    else:
        raise UserWarning('padding mode {} not supported.'.format(padding_mode))

dofaker/transforms/functional.py
justld/dofaker
from .download import download_file
from .weights_urls import get_model_url

dofaker/utils/__init__.py
justld/dofaker
import os
import requests
import zipfile

from tqdm import tqdm


def download_file(url: str, save_dir='./', overwrite=False, unzip=True):
    os.makedirs(save_dir, exist_ok=True)
    file_name = url.split('/')[-1]
    file_path = os.path.join(save_dir, file_name)

    if os.path.exists(file_path) and not overwrite:
        pass
    else:
        print('Downloading file {} from {}...'.format(file_path, url))

        r = requests.get(url, stream=True)
        print(r.status_code)
        if r.status_code != 200:
            raise RuntimeError('Failed downloading url {}!'.format(url))
        total_length = r.headers.get('content-length')
        with open(file_path, 'wb') as f:
            if total_length is None:  # no content length header
                for chunk in r.iter_content(chunk_size=1024):
                    if chunk:  # filter out keep-alive new chunks
                        f.write(chunk)
            else:
                total_length = int(total_length)
                print('file length: ', int(total_length / 1024. + 0.5))
                for chunk in tqdm(r.iter_content(chunk_size=1024),
                                  total=int(total_length / 1024. + 0.5),
                                  unit='KB',
                                  unit_scale=False,
                                  dynamic_ncols=True):
                    f.write(chunk)
    if unzip and file_path.endswith('.zip'):
        save_dir = file_path.split('.')[0]
        if os.path.isdir(save_dir) and os.path.exists(save_dir):
            pass
        else:
            with zipfile.ZipFile(file_path, 'r') as zip_ref:
                zip_ref.extractall(save_dir)

    return save_dir, file_path

dofaker/utils/download.py
justld/dofaker
WEIGHT_URLS = {
    'buffalo_l':
    'https://github.com/justld/dofaker/releases/download/v0.1/buffalo_l.zip',
    'buffalo_s':
    'https://github.com/justld/dofaker/releases/download/v0.1/buffalo_s.zip',
    'buffalo_sc':
    'https://github.com/justld/dofaker/releases/download/v0.1/buffalo_sc.zip',
    'inswapper':
    'https://github.com/justld/dofaker/releases/download/v0.1/inswapper_128.onnx',
    'gfpgan':
    'https://github.com/justld/dofaker/releases/download/v0.1/GFPGANv1.3.onnx',
    'bsrgan':
    'https://github.com/justld/dofaker/releases/download/v0.1/bsrgan_4.onnx',
    'openpose_body':
    'https://github.com/justld/dofaker/releases/download/v0.1/openpose_body.onnx',
    'pose_transfer':
    'https://github.com/justld/dofaker/releases/download/v0.1/pose_transfer.onnx',
}


def get_model_url(model_name):
    return WEIGHT_URLS[model_name]

dofaker/utils/weights_urls.py
justld/dofaker
import argparse
from dofaker import FaceSwapper


def parse_args():
    parser = argparse.ArgumentParser(description='running face swap')
    parser.add_argument('--source',
                        help='select an image or video to be swapped',
                        dest='source',
                        required=True)
    parser.add_argument('--dst_face_paths',
                        help='select images in source to be swapped',
                        dest='dst_face_paths',
                        nargs='+',
                        default=None)
    parser.add_argument(
        '--src_face_paths',
        help='select images to replace dst_faces in source image or video.',
        dest='src_face_paths',
        nargs='+',
        required=True)
    parser.add_argument('--output_dir',
                        help='output directory',
                        dest='output_dir',
                        default='output')
    parser.add_argument('--det_model_name',
                        help='detection model name for insightface',
                        dest='det_model_name',
                        default='buffalo_l')
    parser.add_argument('--det_model_dir',
                        help='detection model dir for insightface',
                        dest='det_model_dir',
                        default='weights/models')
    parser.add_argument('--swap_model_name',
                        help='swap model name',
                        dest='swap_model_name',
                        default='inswapper')
    parser.add_argument('--image_sr_model',
                        help='image super resolution model',
                        dest='image_sr_model',
                        default='bsrgan')
    parser.add_argument('--face_swap_model_dir',
                        help='swap model path',
                        dest='face_swap_model_dir',
                        default='weights/models')
    parser.add_argument('--image_sr_model_dir',
                        help='image super resolution model dir',
                        dest='image_sr_model_dir',
                        default='weights/models')
    parser.add_argument('--face_enhance_name',
                        help='face enhance model',
                        dest='face_enhance_name',
                        default='gfpgan')
    parser.add_argument('--face_enhance_model_dir',
                        help='face enhance model dir',
                        dest='face_enhance_model_dir',
                        default='weights/models')
    parser.add_argument('--face_sim_thre',
                        help='similarity of face embedding threshold',
                        dest='face_sim_thre',
                        default=0.5)
    parser.add_argument('--log_iters',
                        help='print log intervals',
                        dest='log_iters',
                        default=10,
                        type=int)
    parser.add_argument('--use_enhancer',
                        help='whether use face enhance model',
                        dest='use_enhancer',
                        action='store_true')
    parser.add_argument('--use_sr',
                        help='whether use image super resolution model',
                        dest='use_sr',
                        action='store_true')
    parser.add_argument('--sr_scale',
                        help='image super resolution scale',
                        dest='sr_scale',
                        default=1,
                        type=float)
    return parser.parse_args()


if __name__ == '__main__':
    args = parse_args()
    faker = FaceSwapper(
        face_det_model=args.det_model_name,
        face_det_model_dir=args.det_model_dir,
        face_swap_model=args.swap_model_name,
        face_swap_model_dir=args.face_swap_model_dir,
        image_sr_model=args.image_sr_model,
        image_sr_model_dir=args.image_sr_model_dir,
        face_enhance_model=args.face_enhance_name,
        face_enhance_model_dir=args.face_enhance_model_dir,
        face_sim_thre=args.face_sim_thre,
        log_iters=args.log_iters,
        use_enhancer=args.use_enhancer,
        use_sr=args.use_sr,
        scale=args.sr_scale,
    )

    faker.run(
        input_path=args.source,
        dst_face_paths=args.dst_face_paths,
        src_face_paths=args.src_face_paths,
        output_dir=args.output_dir,
    )

run_faceswapper.py
justld/dofaker
import argparse
from dofaker import PoseSwapper


def parse_args():
    parser = argparse.ArgumentParser(description='running face swap')
    parser.add_argument('--source',
                        help='select an image or video to be swapped',
                        dest='source',
                        required=True)
    parser.add_argument('--target',
                        help='the target pose image',
                        dest='target',
                        required=True)
    parser.add_argument('--output_dir',
                        help='output directory',
                        dest='output_dir',
                        default='output')
    parser.add_argument('--pose_estimator_name',
                        help='pose estimator name',
                        dest='pose_estimator_name',
                        default='openpose_body')
    parser.add_argument('--pose_estimator_model_dir',
                        help='pose estimator model dir',
                        dest='pose_estimator_model_dir',
                        default='weights/models')
    parser.add_argument('--pose_transfer_name',
                        help='pose transfer name',
                        dest='pose_transfer_name',
                        default='pose_transfer')
    parser.add_argument('--pose_transfer_model_dir',
                        help='pose transfer model dir',
                        dest='pose_transfer_model_dir',
                        default='weights/models')
    parser.add_argument('--det_model_name',
                        help='detection model name for insightface',
                        dest='det_model_name',
                        default='buffalo_l')
    parser.add_argument('--det_model_dir',
                        help='detection model dir for insightface',
                        dest='det_model_dir',
                        default='weights/models')
    parser.add_argument('--image_sr_model',
                        help='image super resolution model',
                        dest='image_sr_model',
                        default='bsrgan')
    parser.add_argument('--image_sr_model_dir',
                        help='image super resolution model dir',
                        dest='image_sr_model_dir',
                        default='weights/models')
    parser.add_argument('--face_enhance_name',
                        help='face enhance model',
                        dest='face_enhance_name',
                        default='gfpgan')
    parser.add_argument('--face_enhance_model_dir',
                        help='face enhance model dir',
                        dest='face_enhance_model_dir',
                        default='weights/models')
    parser.add_argument('--log_iters',
                        help='print log intervals',
                        dest='log_iters',
                        default=10,
                        type=int)
    parser.add_argument('--use_enhancer',
                        help='whether use face enhance model',
                        dest='use_enhancer',
                        action='store_true')
    parser.add_argument('--use_sr',
                        help='whether use image super resolution model',
                        dest='use_sr',
                        action='store_true')
    parser.add_argument('--sr_scale',
                        help='image super resolution scale',
                        dest='sr_scale',
                        default=1,
                        type=float)
    return parser.parse_args()


if __name__ == '__main__':
    args = parse_args()
    faker = PoseSwapper(
        pose_estimator_name=args.pose_estimator_name,
        pose_estimator_model_dir=args.pose_estimator_model_dir,
        pose_transfer_name=args.pose_transfer_name,
        pose_transfer_model_dir=args.pose_transfer_model_dir,
        face_det_model=args.det_model_name,
        face_det_model_dir=args.det_model_dir,
        image_sr_model=args.image_sr_model,
        image_sr_model_dir=args.image_sr_model_dir,
        face_enhance_name=args.face_enhance_name,
        face_enhance_model_dir=args.face_enhance_model_dir,
        log_iters=args.log_iters,
        use_enhancer=args.use_enhancer,
        use_sr=args.use_sr,
        scale=args.sr_scale,
    )

    faker.run(
        input_path=args.source,
        target_path=args.target,
        output_dir=args.output_dir,
    )

run_posetransfer.py
justld/dofaker
from setuptools import setup, find_packages

with open('requirements.txt') as file:
    REQUIRED_PACKAGES = file.read()

setup(name='dofaker',
      version='0.1',
      keywords=('face swap'),
      description='A simple face swap tool',
      url='https://github.com/justld/dofaker',
      author='justld',
      author_email='1207540056@qq.com',
      packages=find_packages(),
      include_package_data=True,
      platforms='any',
      install_requires=REQUIRED_PACKAGES,
      scripts=[],
      license='GPL 3.0',
      entry_points={'console_scripts': [
          'dofaker = web_ui:main',
      ]})

setup.py
justld/dofaker
import argparse

import gradio as gr
from dofaker import FaceSwapper, PoseSwapper


def parse_args():
    parser = argparse.ArgumentParser(description='running face swap')
    parser.add_argument(
        '--inbrowser',
        help=
        'whether to automatically launch the interface in a new tab on the default browser.',
        dest='inbrowser',
        default=True)
    parser.add_argument(
        '--server_port',
        help=
        'will start gradio app on this port (if available). Can be set by environment variable GRADIO_SERVER_PORT. If None, will search for an available port starting at 7860.',
        dest='server_port',
        type=int,
        default=None)
    return parser.parse_args()


def swap_face(input_path, dst_path, src_path, use_enhancer, use_sr, scale,
              face_sim_thre):
    faker = FaceSwapper(use_enhancer=use_enhancer,
                        use_sr=use_sr,
                        scale=scale,
                        face_sim_thre=face_sim_thre)
    output_path = faker.run(input_path, dst_path, src_path)
    return output_path


def swap_pose(input_path, target_path, use_enhancer, use_sr, scale):
    faker = PoseSwapper(use_enhancer=use_enhancer, use_sr=use_sr, scale=scale)
    output_path = faker.run(input_path, target_path)
    return output_path


def main():
    args = parse_args()

    with gr.Blocks(title='DoFaker') as web_ui:
        gr.Markdown('DoFaker: Face Swap and pose swap web ui')
        with gr.Tab('FaceSwapper'):
            gr.Markdown('DoFaker: Face Swap Web UI')
            with gr.Tab('Image'):
                with gr.Row():
                    with gr.Column():
                        gr.Markdown('The source image to be swapped')
                        image_input = gr.Image(type='filepath')
                        with gr.Row():
                            with gr.Column():
                                gr.Markdown(
                                    'target face included in source image')
                                dst_face_image = gr.Image(type='filepath')
                            with gr.Column():
                                gr.Markdown(
                                    'source face to replace target face')
                                src_face_image = gr.Image(type='filepath')

                    with gr.Column():
                        output_image = gr.Image(type='filepath')
                        use_enhancer = gr.Checkbox(
                            label="face enhance",
                            info="Whether use face enhance model.")
                        with gr.Row():
                            use_sr = gr.Checkbox(
                                label="super resolution",
                                info="Whether use image resolution model.")
                            scale = gr.Number(
                                value=1, label='image super resolution scale')
                        with gr.Row():
                            face_sim_thre = gr.Number(
                                value=0.6,
                                label='face similarity threshold',
                                minimum=0.0,
                                maximum=1.0)
                        convert_button = gr.Button('Swap')
                        convert_button.click(fn=swap_face,
                                             inputs=[
                                                 image_input, dst_face_image,
                                                 src_face_image, use_enhancer,
                                                 use_sr, scale, face_sim_thre
                                             ],
                                             outputs=[output_image],
                                             api_name='image swap')

            with gr.Tab('Video'):
                with gr.Row():
                    with gr.Column():
                        gr.Markdown('The source video to be swapped')
                        video_input = gr.Video(type='filepath')
                        with gr.Row():
                            with gr.Column():
                                gr.Markdown(
                                    'target face included in source image')
                                dst_face_image = gr.Image(type='filepath')
                            with gr.Column():
                                gr.Markdown(
                                    'source face to replace target face')
                                src_face_image = gr.Image(type='filepath')

                    with gr.Column():
                        output_video = gr.Video(type='filepath')
                        use_enhancer = gr.Checkbox(
                            label="face enhance",
                            info="Whether use face enhance model.")
                        with gr.Row():
                            use_sr = gr.Checkbox(
                                label="super resolution",
                                info="Whether use image resolution model.")
                            scale = gr.Number(
                                value=1, label='image super resolution scale')
                        with gr.Row():
                            face_sim_thre = gr.Number(
                                value=0.6,
                                label='face similarity threshold',
                                minimum=0.0,
                                maximum=1.0)
                        convert_button = gr.Button('Swap')
                        convert_button.click(fn=swap_face,
                                             inputs=[
                                                 video_input, dst_face_image,
                                                 src_face_image, use_enhancer,
                                                 use_sr, scale, face_sim_thre
                                             ],
                                             outputs=[output_video],
                                             api_name='video swap')

        with gr.Tab('PoseSwapper'):
            gr.Markdown('DoFaker: Pose Swap Web UI')
            with gr.Tab('Image'):
                with gr.Row():
                    with gr.Column():
                        gr.Markdown('The source image to be swapped')
                        image_input = gr.Image(type='filepath')
                        gr.Markdown('The target image with pose')
                        target = gr.Image(type='filepath')

                    with gr.Column():
                        output_image = gr.Image(type='filepath')
                        use_enhancer = gr.Checkbox(
                            label="face enhance",
                            info="Whether use face enhance model.")
                        with gr.Row():
                            use_sr = gr.Checkbox(
                                label="super resolution",
                                info="Whether use image resolution model.")
                            scale = gr.Number(
                                value=1, label='image super resolution scale')
                        convert_button = gr.Button('Swap')
                        convert_button.click(fn=swap_pose,
                                             inputs=[
                                                 image_input, target,
                                                 use_enhancer, use_sr, scale
                                             ],
                                             outputs=[output_image],
                                             api_name='image swap')

            # with gr.Tab('Video'):
            #     with gr.Row():
            #         with gr.Column():
            #             gr.Markdown('The source video to be swapped')
            #             video_input = gr.Image(type='filepath')
            #             gr.Markdown('The target image with pose')
            #             target = gr.Video(type='filepath')

            #         with gr.Column():
            #             output_video = gr.Video(type='filepath')
            #             use_enhancer = gr.Checkbox(
            #                 label="face enhance",
            #                 info="Whether use face enhance model.")
            #             with gr.Row():
            #                 use_sr = gr.Checkbox(
            #                     label="super resolution",
            #                     info="Whether use image resolution model.")
            #                 scale = gr.Number(value=1,
            #                                 label='image super resolution scale')
            #             convert_button = gr.Button('Swap')
            #             convert_button.click(fn=swap_pose,
            #                                 inputs=[
            #                                     video_input, target, use_enhancer,
            #                                     use_sr, scale
            #                                 ],
            #                                 outputs=[output_video],
            #                                 api_name='video swap')

    web_ui.launch(inbrowser=args.inbrowser, server_port=args.server_port)


if __name__ == '__main__':
    main()

web_ui.py
justld/dofaker