Files
VideoPipe/nodes/infers/vp_face_swap_node.h
2026-06-03 12:43:14 +08:00

65 lines
2.8 KiB
C++

#pragma once
#include "../vp_primary_infer_node.h"
namespace vp_nodes {
// face swap node
// used to swap faces in videos/images using a specific face
class vp_face_swap_node: public vp_primary_infer_node
{
private:
// inner temporary use
struct face_box {
int x;
int y;
int width;
int height;
float score;
std::vector<std::pair<int, int>> kps;
};
/* onnx network using opencv::dnn as backend */
cv::dnn::Net face_extract_net;
cv::dnn::Net face_encoding_net;
cv::dnn::Net face_swap_net;
cv::Mat read_emap_mat_from_txt(const std::string& emap_file_for_embeddings, int rows = 512, int cols = 512);
cv::Mat process_embeddings_using_emap(const cv::Mat& source_face_normed_embedding, const cv::Mat& emap);
cv::Mat get_similarity_transform_matrix(float src[5][2]);
cv::Mat align_crop(cv::Mat& _src_img, std::vector<std::pair<int, int>>& kps, cv::Mat& _aligned_img);
void extract_faces(const cv::Mat& frame, std::vector<face_box>& faces);
void init_source_face_embeddings(std::string& swap_source_image, int swap_source_face_index, std::string& emap_file_for_embeddings);
void swap(cv::Mat& aligned_face, cv::Mat& swapped_face);
void paste_back(cv::Mat& bg, cv::Mat& swapped_face, const cv::Mat& transform_matrix);
cv::Mat source_face_embeddings;
bool act_as_primary_detector = false;
bool swap_on_osd = true;
int min_face_w_h = 50;
/* used for extract faces */
std::vector<cv::Rect2f> priors;
void generatePriors(int inputW, int inputH);
protected:
// we need a totally new logic for the whole infer combinations
// no separate step pre-defined needed in base class
virtual void run_infer_combinations(const std::vector<std::shared_ptr<vp_objects::vp_frame_meta>>& frame_meta_with_batch) override;
// override pure virtual method, for compile pass
virtual void postprocess(const std::vector<cv::Mat>& raw_outputs, const std::vector<std::shared_ptr<vp_objects::vp_frame_meta>>& frame_meta_with_batch) override;
public:
vp_face_swap_node(std::string node_name,
std::string yunet_face_detect_model,
std::string buffalo_l_face_encoding_model,
std::string emap_file_for_embeddings,
std::string insightface_swap_model,
std::string swap_source_image,
int swap_source_face_index = 0,
int min_face_w_h = 50,
bool swap_on_osd = true,
bool act_as_primary_detector = false);
~vp_face_swap_node();
};
}