#include <iostream>
|
#include <memory>
|
#include <string>
|
#include <vector>
|
#include <algorithm>
|
#include "../nodes/vp_file_src_node.h"
|
#include "../nodes/infers/vp_yunet_face_detector_node.h"
|
#include "../nodes/infers/vp_face_swap_node.h"
|
#include "../nodes/osd/vp_face_osd_node.h"
|
#include "../nodes/vp_screen_des_node.h"
|
#include "../nodes/vp_file_des_node.h"
|
#include "../utils/analysis_board/vp_analysis_board.h"
|
#include <experimental/filesystem> // C++17 必须
|
|
/*
|
* ## face_swap_sample ##
|
* swap face for any video/images, no training need before running.
|
*/
|
|
std::string get_filename_without_extension(const std::string& path) {
|
std::experimental::filesystem::path fs_path(path);
|
return fs_path.stem().string(); // stem() 获取不带扩展名的文件名
|
}
|
// 解析命令行参数的函数
|
bool has_argument(const std::vector<std::string>& args, const std::string& option) {
|
return std::find(args.begin(), args.end(), option) != args.end();
|
}
|
|
void ensure_directory_exists(const std::string& save_dir) {
|
std::experimental::filesystem::path dir_path(save_dir);
|
|
// 检查目录是否存在
|
if (!std::experimental::filesystem::exists(dir_path)) {
|
try {
|
// 创建目录
|
std::experimental::filesystem::create_directories(dir_path); // 创建父目录及子目录
|
std::cout << "Directory created: " << save_dir << std::endl;
|
} catch (const std::experimental::filesystem::filesystem_error& e) {
|
std::cerr << "Error creating directory: " << e.what() << std::endl;
|
}
|
} else {
|
std::cout << "Directory already exists: " << save_dir << std::endl;
|
}
|
}
|
|
int main(int argc, char* argv[]) {
|
VP_SET_LOG_LEVEL(vp_utils::vp_log_level::INFO);
|
VP_SET_LOG_INCLUDE_CODE_LOCATION(false);
|
VP_SET_LOG_INCLUDE_THREAD_ID(false);
|
VP_LOGGER_INIT();
|
|
// 默认的命令行参数
|
std::string video_path = "./vp_data/test_video/face.mp4";
|
std::string image_path = "./vp_data/inswapper/data/mans1.jpeg";
|
bool show_screen = false;
|
bool show_board = false;
|
|
float resize_ratio = 0.5;
|
|
// 将命令行参数存储为 vector
|
std::vector<std::string> args(argv + 1, argv + argc);
|
|
// 如果提供了第一个和第二个参数，则覆盖默认的路径
|
if (args.size() >= 2) {
|
video_path = args[0];
|
image_path = args[1];
|
if (args.size() >= 3){
|
resize_ratio =std::stof(args[2]);
|
}
|
} else {
|
std::cout << "Usage: " << argv[0] << " <video_path> <image_path> <resize_ration> [--show_screen] [--show_board]\n";
|
std::cout << "Example: " << argv[0] << " ./vp_data/test_video/face.mp4 ./vp_data/inswapper/data/mans1.jpeg 0.5 --show_screen --show_board\n";
|
return 1;
|
}
|
std::string video_filename = get_filename_without_extension(video_path);
|
std::string image_filename = get_filename_without_extension(image_path);
|
|
ensure_directory_exists(video_filename);
|
|
|
// 检测可选参数
|
show_screen = has_argument(args, "--show_screen");
|
show_board = has_argument(args, "--show_board");
|
|
std::cout << "Video path: " << video_path << "\n";
|
std::cout << "Image path: " << image_path << "\n";
|
std::cout << "Video Resize ratio: " << resize_ratio << "\n";
|
std::cout << "Show screen: " << (show_screen ? "true" : "false") << "\n";
|
std::cout << "Show board: " << (show_board ? "true" : "false") << "\n";
|
|
// 创建节点
|
//auto file_src_0 = std::make_shared<vp_nodes::vp_file_src_node>("file_src_0", 0, video_path, 1.0, true, "avdec_h264", 4);
|
auto file_src_0 = std::make_shared<vp_nodes::vp_file_src_node>("file_src_0", 0, video_path, resize_ratio, false, "avdec_h264", 4);
|
|
auto yunet_face_detector = std::make_shared<vp_nodes::vp_yunet_face_detector_node>("yunet_face_detector", "./vp_data/models/face/face_detection_yunet_2022mar.onnx");
|
auto face_swap = std::make_shared<vp_nodes::vp_face_swap_node>(
|
"face_swap",
|
"./vp_data/models/face/face_detection_yunet_2022mar.onnx",
|
"./vp_data/models/face/swap/w600k_r50.onnx",
|
"./vp_data/models/face/swap/emap.txt",
|
"./vp_data/models/face/swap/inswapper_128.onnx",
|
image_path
|
);
|
|
|
auto file_des_node_0 = std::make_shared<vp_nodes::vp_file_des_node>("file_des_0", 0, video_filename, image_filename, 1);
|
|
// 构建管道
|
yunet_face_detector->attach_to({file_src_0});
|
face_swap->attach_to({yunet_face_detector});
|
|
if (show_screen) {
|
auto screen_des_0_ori = std::make_shared<vp_nodes::vp_screen_des_node>("screen_des_0_ori", 0, false);
|
auto screen_des_0_osd = std::make_shared<vp_nodes::vp_screen_des_node>("screen_des_0_osd", 0);
|
screen_des_0_ori->attach_to({face_swap});
|
screen_des_0_osd->attach_to({face_swap});
|
}
|
|
file_des_node_0->attach_to({face_swap});
|
|
file_src_0->start();
|
|
// 如果用户要求显示 board，则显示
|
if (show_board) {
|
vp_utils::vp_analysis_board board({file_src_0});
|
board.display(1, false);
|
}
|
|
std::string wait;
|
std::getline(std::cin, wait);
|
file_src_0->detach_recursively();
|
|
return 0;
|
}
|
