inference_sequence.cpp

//
// Created by haoyuefan on 2023/11/02.
//

#include <memory>
#include <chrono>
#include "utils.h"
#include "light_glue.h"
#include "super_point.h"

// Moved here from utils.h -- these are only used in this file, and utils.h is included by many
// other translation units (transitively, via read_config.h/super_point.h) that never call them,
// which triggered -Wunused-function everywhere except here.
static bool GetFileNames(const std::string &path, std::vector<std::string> &filenames) {
    DIR *pDir;
    struct dirent *ptr;
    if (!(pDir = opendir(path.c_str()))) {
        std::cerr << "Current folder doesn't exist!" << std::endl;
        return false;
    }
    while ((ptr = readdir(pDir)) != nullptr) {
        if (strcmp(ptr->d_name, ".") != 0 && strcmp(ptr->d_name, "..") != 0) {
            filenames.push_back(path + "/" + ptr->d_name);
        }
    }
    closedir(pDir);
    std::sort(filenames.begin(), filenames.end());
    return true;
}

static void VisualizeMatching(const cv::Mat &image0, const std::vector<cv::KeyPoint> &keypoints0, const cv::Mat &image1,
                              const std::vector<cv::KeyPoint> &keypoints1,
                              const std::vector<cv::DMatch> &superglue_matches, cv::Mat &output_image, double cost_time = -1) {
    if(image0.size != image1.size) return;
    cv::drawMatches(image0, keypoints0, image1, keypoints1, superglue_matches, output_image, cv::Scalar(0, 255, 0), cv::Scalar(0, 0, 255));
    double sc = std::min(image0.rows / 640., 2.0);
    int ht = int(30 * sc);
    std::string title_str = "SuperPoint LightGlue TensorRT";
    cv::putText(output_image, title_str, cv::Point(int(8*sc), ht), cv::FONT_HERSHEY_DUPLEX,1.0*sc, cv::Scalar(0, 0, 0), 2, cv::LINE_AA);
    cv::putText(output_image, title_str, cv::Point(int(8*sc), ht), cv::FONT_HERSHEY_DUPLEX,1.0*sc, cv::Scalar(255, 255, 255), 1, cv::LINE_AA);
    std::string feature_points_str = "Keypoints: " + std::to_string(keypoints0.size()) + ":" + std::to_string(keypoints1.size());
    cv::putText(output_image, feature_points_str, cv::Point(int(8*sc), ht*2), cv::FONT_HERSHEY_DUPLEX,1.0*sc, cv::Scalar(0, 0, 0), 2, cv::LINE_AA);
    cv::putText(output_image, feature_points_str, cv::Point(int(8*sc), ht*2), cv::FONT_HERSHEY_DUPLEX,1.0*sc, cv::Scalar(255, 255, 255), 1, cv::LINE_AA);
    std::string match_points_str = "Matches: " + std::to_string(superglue_matches.size());
    cv::putText(output_image, match_points_str, cv::Point(int(8*sc), ht*3), cv::FONT_HERSHEY_DUPLEX,1.0*sc, cv::Scalar(0, 0, 0), 2, cv::LINE_AA);
    cv::putText(output_image, match_points_str, cv::Point(int(8*sc), ht*3), cv::FONT_HERSHEY_DUPLEX,1.0*sc, cv::Scalar(255, 255, 255), 1, cv::LINE_AA);
    if(cost_time != -1) {
        std::string time_str = "FPS: " + std::to_string(1000 / cost_time);
        cv::putText(output_image, time_str, cv::Point(int(8 * sc), ht * 4), cv::FONT_HERSHEY_DUPLEX, 1.0 * sc,
                    cv::Scalar(0, 0, 0), 2, cv::LINE_AA);
        cv::putText(output_image, time_str, cv::Point(int(8 * sc), ht * 4), cv::FONT_HERSHEY_DUPLEX, 1.0 * sc,
                    cv::Scalar(255, 255, 255), 1, cv::LINE_AA);
    }
}

int main(int argc, char** argv){
//    if (argc != 5) {
//        std::cerr << "./superpoint_lightglue_sequence config_path model_dir image_folder_absolutely_path output_folder_path" << std::endl;
//        return 0;
//    }

//    std::string config_path = argv[1];
//    std::string model_dir = argv[2];
//    std::string image_path = argv[3];
//    std::string output_path = argv[4];
    std::string config_path = "../config/config.yaml";
    std::string model_dir = "../weights/";
    std::string image_path = "../image/freiburg_sequence/";
    std::string output_path = "../image/freiburg_sequence/match_images/";
    std::vector<std::string> image_names;
    GetFileNames(image_path, image_names);
    Configs configs(config_path, model_dir);
    int width = configs.superpoint_lightglue_config.image_width;
    int height = configs.superpoint_lightglue_config.image_height;

    std::cout << "Building Inference Engine......" << std::endl;
    auto superpoint = std::make_shared<SuperPoint>(configs.superpoint_config);
    if (!superpoint->build()) {
        std::cerr << "Error in SuperPoint building engine. Please check your onnx model path." << std::endl;
        return 0;
    }
    auto superpoint_lightglue = std::make_shared<SuperPointLightGlue>(configs.superpoint_lightglue_config);
    if (!superpoint_lightglue->build()) {
        std::cerr << "Error in SuperGlue building engine. Please check your onnx model path." << std::endl;
        return 0;
    }
    std::cout << "SuperPoint and SuperGlue inference engine build success." << std::endl;

    Eigen::Matrix<double, 258, Eigen::Dynamic> feature_points0;
    Eigen::Matrix<double, 1, Eigen::Dynamic> feature_scores0;
    cv::Mat image0 = cv::imread(image_names[0], cv::IMREAD_GRAYSCALE);
    if(image0.empty()) {
        std::cerr << "First image in the image folder is empty." << std::endl;
        return 0;
    }
    cv::resize(image0, image0, cv::Size(width, height));
    std::cout << "First image size: " << image0.cols << "x" << image0.rows << std::endl;
    if (!superpoint->infer(image0, feature_points0, feature_scores0)) {
        std::cerr << "Failed when extracting features from first image." << std::endl;
        return 0;
    }
    std::vector<cv::DMatch> init_matches;
    superpoint_lightglue->matching_points(feature_points0, feature_points0, init_matches);
    std::string mkdir_cmd = "mkdir -p " + output_path;
    system(mkdir_cmd.c_str());

    for (int index = 1; index < image_names.size(); ++index) {
        Eigen::Matrix<double, 258, Eigen::Dynamic> feature_points1;
        Eigen::Matrix<double, 1, Eigen::Dynamic> feature_scores1;
        std::vector<cv::DMatch> superglue_matches;
        cv::Mat image1 = cv::imread(image_names[index], cv::IMREAD_GRAYSCALE);
        if(image1.empty()) continue;
        cv::resize(image1, image1, cv::Size(width, height));

        std::cout << "Second image size: " << image1.cols << "x" << image1.rows << std::endl;
        auto start = std::chrono::high_resolution_clock::now();
        if (!superpoint->infer(image1, feature_points1, feature_scores1)) {
            std::cerr << "Failed when extracting features from second image." << std::endl;
            return 0;
        }
        superpoint_lightglue->matching_points(feature_points0, feature_points1, superglue_matches);
        auto end = std::chrono::high_resolution_clock::now();
        auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
        cv::Mat match_image;
        std::vector<cv::KeyPoint> keypoints0, keypoints1;
        for (size_t i = 0; i < feature_points0.cols(); ++i) {
            double score = feature_scores0(0, i);
            double x = feature_points0(0, i);
            double y = feature_points0(1, i);
            keypoints0.emplace_back(x, y, 8, -1, score);
        }
        for (size_t i = 0; i < feature_points1.cols(); ++i) {
            double score = feature_scores1(0, i);
            double x = feature_points1(0, i);
            double y = feature_points1(1, i);
            keypoints1.emplace_back(x, y, 8, -1, score);
        }
        VisualizeMatching(image0, keypoints0, image1, keypoints1, superglue_matches, match_image, duration.count());
        cv::imwrite(output_path + "/" + std::to_string(index) + ".png", match_image);
    }
  
  return 0;
}