face_swap.h

#ifndef FACE_SWAP_FACE_SWAP_H
#define FACE_SWAP_FACE_SWAP_H

#include "face_swap/cnn_3dmm_expr.h"
#include "face_swap/basel_3dmm.h"
#include "face_swap/face_renderer.h"

// OpenCV
#include <opencv2/core.hpp>

// sfl
#include <sfl/sequence_face_landmarks.h>
#include <sfl/utilities.h>

// face_seg
#include <face_seg/face_seg.h>

namespace face_swap
{
    class FaceSwap
    {
    public:
        FaceSwap(const std::string& landmarks_path, const std::string& model_3dmm_h5_path,
            const std::string& model_3dmm_dat_path, const std::string& reg_model_path,
            const std::string& reg_deploy_path, const std::string& reg_mean_path,
            bool generic = false, bool with_expr = true, bool with_gpu = true, 
            int gpu_device_id = 0);

        void setSegmentationModel(const std::string& seg_model_path, 
            const std::string& seg_deploy_path);

        void clearSegmentationModel();

        bool isSegmentationModelInit();

        bool setSource(const cv::Mat& img, const cv::Mat& seg = cv::Mat());

        bool setTarget(const cv::Mat& img, const cv::Mat& seg = cv::Mat());

        cv::Mat swap();

        const Mesh& getSourceMesh() const;

        const Mesh& getTargetMesh() const;

    private:

        bool preprocessImages(const cv::Mat& img, const cv::Mat& seg,
            std::vector<cv::Point>& landmarks, std::vector<cv::Point>& cropped_landmarks,
            cv::Mat& cropped_img, cv::Mat& cropped_seg, cv::Rect& bbox = cv::Rect());

        void generateTexture(const Mesh& mesh, const cv::Mat& img, const cv::Mat& seg,
            const cv::Mat& vecR, const cv::Mat& vecT, const cv::Mat& K,
            cv::Mat& tex, cv::Mat& uv);

        cv::Mat generateTextureCoordinates(const Mesh& mesh, const cv::Size& img_size,
            const cv::Mat& vecR, const cv::Mat& vecT, const cv::Mat& K);

        cv::Mat blend(const cv::Mat& src, const cv::Mat& dst,
            const cv::Mat& dst_seg = cv::Mat());

    private:
        std::shared_ptr<sfl::SequenceFaceLandmarks> m_sfl;
        std::unique_ptr<CNN3DMMExpr> m_cnn_3dmm_expr;
        std::unique_ptr<Basel3DMM> m_basel_3dmm;
        std::unique_ptr<FaceRenderer> m_face_renderer;
        std::unique_ptr<face_seg::FaceSeg> m_face_seg;

        bool m_with_gpu;
        int m_gpu_device_id;

        Mesh m_src_mesh, m_dst_mesh;
        cv::Mat m_vecR, m_vecT, m_K;
        cv::Mat m_tex, m_uv;
        cv::Mat m_tgt_cropped_img, m_tgt_cropped_seg;
        cv::Mat m_target_img, m_target_seg;
        cv::Rect m_target_bbox;

        cv::Mat m_source_img;
        cv::Mat m_src_cropped_img, m_src_cropped_seg;
        cv::Mat m_tgt_rendered_img;
        std::vector<cv::Point> m_src_cropped_landmarks;
        std::vector<cv::Point> m_tgt_cropped_landmarks;
        cv::Mat m_src_vecR, m_src_vecT, m_src_K;
        std::vector<cv::Point> m_src_landmarks, m_tgt_landmarks;

    public:
        cv::Mat debugSource();
        cv::Mat debugTarget();
        cv::Mat debug();
        cv::Mat debugSourceMesh();
        cv::Mat debugTargetMesh();
        cv::Mat debugMesh(const cv::Mat& img, const cv::Mat& seg, 
            const cv::Mat& uv, const Mesh& mesh,
            const cv::Mat& vecR, const cv::Mat& vecT, const cv::Mat& K);
        cv::Mat debugSourceLandmarks();
        cv::Mat debugTargetLandmarks();
        cv::Mat debugRender();
    };
}   // namespace face_swap

#endif // FACE_SWAP_FACE_SWAP_H