benchmark_multiresolution.cpp

//
// Multi-resolution benchmark for SuperPoint + LightGlue TensorRT
// Tests 320x240, 640x480, 160x120, 80x60 and prints a detailed timing table.
//

#include <chrono>
#include <memory>
#include <numeric>
#include <algorithm>
#include <iomanip>
#include <iostream>
#include <sstream>
#include <vector>
#include <cmath>

#include "light_glue.h"
#include "super_point.h"
#include "utils.h"

// ── Timing helpers ────────────────────────────────────────────────────────────

using Clock     = std::chrono::high_resolution_clock;
using TimePoint = std::chrono::time_point<Clock>;
using Ms        = std::chrono::duration<double, std::milli>;

static double elapsed_ms(TimePoint a, TimePoint b) {
    return std::chrono::duration_cast<Ms>(b - a).count();
}

struct Stats {
    double mean{}, min_v{}, max_v{}, stddev{}, p95{}, p99{};
    int count{};
};

static Stats compute_stats(std::vector<double> v) {
    if (v.empty()) return {};
    Stats s;
    s.count = (int)v.size();
    s.mean  = std::accumulate(v.begin(), v.end(), 0.0) / s.count;
    s.min_v = *std::min_element(v.begin(), v.end());
    s.max_v = *std::max_element(v.begin(), v.end());
    double sq = 0;
    for (double x : v) sq += (x - s.mean) * (x - s.mean);
    s.stddev = std::sqrt(sq / s.count);
    std::sort(v.begin(), v.end());
    s.p95 = v[std::max(0, (int)(s.count * 0.95) - 1)];
    s.p99 = v[std::max(0, (int)(s.count * 0.99) - 1)];
    return s;
}

// ── Pretty printing ───────────────────────────────────────────────────────────

static void sep(char c = '-', int w = 100) { std::cout << std::string(w, c) << "\n"; }

static void print_stats_block(const std::string& label, const Stats& s) {
    std::cout << "    " << std::left << std::setw(26) << label
              << "mean=" << std::fixed << std::setprecision(2) << std::setw(8) << s.mean << " ms"
              << "  min=" << std::setw(8) << s.min_v
              << "  max=" << std::setw(8) << s.max_v
              << "  std=" << std::setw(7) << s.stddev
              << "  p95=" << std::setw(8) << s.p95
              << "  p99=" << std::setw(8) << s.p99
              << "\n";
}

static void print_summary_header() {
    sep();
    std::cout << std::left
              << std::setw(14) << "Resolution"
              << std::setw(10) << "W x H"
              << std::setw(8)  << "Kpts0"
              << std::setw(8)  << "Kpts1"
              << std::setw(9)  << "Matches"
              << std::setw(7)  << "Rate%"
              << std::setw(13) << "SP0 mean(ms)"
              << std::setw(13) << "SP1 mean(ms)"
              << std::setw(13) << "LG  mean(ms)"
              << std::setw(13) << "Total  (ms)"
              << std::setw(11) << "Pairs/sec"
              << "\n";
    sep();
}

// ── Per-resolution run ────────────────────────────────────────────────────────

struct RunResult {
    std::string label;
    int width{}, height{};
    int kpts0{}, kpts1{}, matches{};
    Stats sp0, sp1, lg, total;
};

static RunResult profile_resolution(
        const std::shared_ptr<SuperPoint>&          superpoint,
        const std::shared_ptr<SuperPointLightGlue>& lightglue,
        const cv::Mat& image0_orig,
        const cv::Mat& image1_orig,
        int target_w, int target_h,
        const std::string& label,
        int warmup  = 5,
        int repeats = 100)
{
    cv::Mat img0, img1;
    cv::resize(image0_orig, img0, cv::Size(target_w, target_h));
    cv::resize(image1_orig, img1, cv::Size(target_w, target_h));

    Eigen::Matrix<double, 258, Eigen::Dynamic> fp0, fp1;
    Eigen::Matrix<double, 1,   Eigen::Dynamic> fs0, fs1;
    std::vector<cv::DMatch> matches;

    std::vector<double> t_sp0, t_sp1, t_lg, t_total;

    // Warmup
    for (int i = 0; i < warmup; ++i) {
        superpoint->infer(img0, fp0, fs0);
        superpoint->infer(img1, fp1, fs1);
        lightglue->matching_points(fp0, fp1, matches);
    }

    // Timed runs
    for (int i = 0; i < repeats; ++i) {
        auto t0 = Clock::now();
        superpoint->infer(img0, fp0, fs0);
        auto t1 = Clock::now();
        superpoint->infer(img1, fp1, fs1);
        auto t2 = Clock::now();
        lightglue->matching_points(fp0, fp1, matches);
        auto t3 = Clock::now();

        t_sp0.push_back(elapsed_ms(t0, t1));
        t_sp1.push_back(elapsed_ms(t1, t2));
        t_lg.push_back(elapsed_ms(t2, t3));
        t_total.push_back(elapsed_ms(t0, t3));
    }

    RunResult r;
    r.label   = label;
    r.width   = target_w;
    r.height  = target_h;
    r.kpts0   = (int)fp0.cols();
    r.kpts1   = (int)fp1.cols();
    r.matches = (int)matches.size();
    r.sp0     = compute_stats(t_sp0);
    r.sp1     = compute_stats(t_sp1);
    r.lg      = compute_stats(t_lg);
    r.total   = compute_stats(t_total);
    return r;
}

// ── Save match image ──────────────────────────────────────────────────────────

static void save_match_image(
        const std::shared_ptr<SuperPoint>&          superpoint,
        const std::shared_ptr<SuperPointLightGlue>& lightglue,
        const cv::Mat& image0_orig,
        const cv::Mat& image1_orig,
        const RunResult& r)
{
    cv::Mat img0, img1;
    cv::resize(image0_orig, img0, cv::Size(r.width, r.height));
    cv::resize(image1_orig, img1, cv::Size(r.width, r.height));

    Eigen::Matrix<double, 258, Eigen::Dynamic> fp0, fp1;
    Eigen::Matrix<double, 1,   Eigen::Dynamic> fs0, fs1;
    superpoint->infer(img0, fp0, fs0);
    superpoint->infer(img1, fp1, fs1);

    std::vector<cv::DMatch> matches;
    lightglue->matching_points(fp0, fp1, matches);

    std::vector<cv::KeyPoint> kp0, kp1;
    for (int i = 0; i < (int)fp0.cols(); ++i)
        kp0.emplace_back((float)fp0(0,i), (float)fp0(1,i), 8, -1, (float)fs0(0,i));
    for (int i = 0; i < (int)fp1.cols(); ++i)
        kp1.emplace_back((float)fp1(0,i), (float)fp1(1,i), 8, -1, (float)fs1(0,i));

    cv::Mat out;
    cv::drawMatches(img0, kp0, img1, kp1, matches, out,
                    cv::Scalar(0,255,0), cv::Scalar(0,0,255));

    std::string text = r.label
        + "  kpts=" + std::to_string(r.kpts0) + "/" + std::to_string(r.kpts1)
        + "  matches=" + std::to_string((int)matches.size())
        + "  " + std::to_string((int)(1000.0 / r.total.mean)) + " pairs/s";
    cv::putText(out, text, cv::Point(8, 22),
                cv::FONT_HERSHEY_DUPLEX, 0.55, cv::Scalar(0,0,0),     2, cv::LINE_AA);
    cv::putText(out, text, cv::Point(8, 22),
                cv::FONT_HERSHEY_DUPLEX, 0.55, cv::Scalar(255,255,255),1, cv::LINE_AA);

    std::string fname = r.label;
    for (char& c : fname) if (c == ' ' || c == '(' || c == ')' || c == '.') c = '_';
    fname = "match_" + fname + ".png";
    cv::imwrite(fname, out);
    std::cout << "  Saved: " << fname << "\n";
}

// ── Main ──────────────────────────────────────────────────────────────────────

int main(int argc, char** argv) {
    if (argc != 5) {
        std::cerr << "Usage: ./benchmark_multiresolution config_path model_dir image0 image1\n";
        return 1;
    }

    const std::string config_path = argv[1];
    const std::string model_dir   = argv[2];
    const std::string image0_path = argv[3];
    const std::string image1_path = argv[4];

    const int WARMUP  = 5;
    const int REPEATS = 100;

    // ── Load images ──────────────────────────────────────────────────────
    cv::Mat image0_orig = cv::imread(image0_path, cv::IMREAD_GRAYSCALE);
    cv::Mat image1_orig = cv::imread(image1_path, cv::IMREAD_GRAYSCALE);
    if (image0_orig.empty() || image1_orig.empty()) {
        std::cerr << "Failed to load images.\n";
        return 1;
    }
    int orig_w = image0_orig.cols;
    int orig_h = image0_orig.rows;
    std::cout << "Original image size: " << orig_w << " x " << orig_h << "\n";

    // ── Build engines ────────────────────────────────────────────────────
    Configs configs(config_path, model_dir);
    std::cout << "Building inference engines...\n";

    auto superpoint = std::make_shared<SuperPoint>(configs.superpoint_config);
    if (!superpoint->build()) {
        std::cerr << "Failed to build SuperPoint engine.\n";
        return 1;
    }
    auto lightglue = std::make_shared<SuperPointLightGlue>(configs.superpoint_lightglue_config);
    if (!lightglue->build()) {
        std::cerr << "Failed to build LightGlue engine.\n";
        return 1;
    }
    std::cout << "Engines ready.\n\n";

    // ── Define resolutions ───────────────────────────────────────────────
    struct Scale { std::string label; int w, h; };
    std::vector<Scale> scales = {
        { "160x120", 160,  120 },
        { "320x240", 320,  240 },
        { "640x480", 640,  480 },
    };

    // Also add original resolution if not already in list
    bool orig_covered = false;
    for (auto& s : scales)
        if (s.w == orig_w && s.h == orig_h) { orig_covered = true; break; }
    if (!orig_covered)
        scales.push_back({"original (" + std::to_string(orig_w) + "x" + std::to_string(orig_h) + ")",
                          orig_w, orig_h});

    // ── Run profiling ────────────────────────────────────────────────────
    std::vector<RunResult> results;
    for (auto& sc : scales) {
        std::cout << "Profiling " << sc.label << " (" << sc.w << "x" << sc.h << ")...\n";
        results.push_back(profile_resolution(
            superpoint, lightglue,
            image0_orig, image1_orig,
            sc.w, sc.h, sc.label,
            WARMUP, REPEATS));
    }

    // ── Summary table ────────────────────────────────────────────────────
    std::cout << "\n\n";
    sep('=', 100);
    std::cout << "  BENCHMARK SUMMARY — SuperPoint + LightGlue TensorRT\n";
    std::cout << "  Images : " << image0_path << "  |  " << image1_path << "\n";
    std::cout << "  Warmup : " << WARMUP << "  |  Timed runs: " << REPEATS << "\n";
    sep('=', 100);
    std::cout << "\n";

    print_summary_header();
    for (auto& r : results) {
        double rate = r.kpts0 > 0 ? 100.0 * r.matches / r.kpts0 : 0.0;
        std::cout << std::left
                  << std::setw(14) << r.label
                  << std::setw(10) << (std::to_string(r.width) + "x" + std::to_string(r.height))
                  << std::setw(8)  << r.kpts0
                  << std::setw(8)  << r.kpts1
                  << std::setw(9)  << r.matches
                  << std::setw(7)  << std::fixed << std::setprecision(1) << rate
                  << std::setw(13) << std::setprecision(2) << r.sp0.mean
                  << std::setw(13) << r.sp1.mean
                  << std::setw(13) << r.lg.mean
                  << std::setw(13) << r.total.mean
                  << std::setw(11) << std::setprecision(1) << (1000.0 / r.total.mean)
                  << "\n";
    }
    sep();

    // ── Detailed breakdown ───────────────────────────────────────────────
    std::cout << "\n  DETAILED BREAKDOWN\n";
    for (auto& r : results) {
        std::cout << "\n";
        sep();
        std::cout << "  [ " << r.label << "  —  " << r.width << " x " << r.height << " ]\n";
        std::cout << "  Keypoints : img0=" << r.kpts0
                  << "  img1=" << r.kpts1
                  << "  matches=" << r.matches
                  << "  match_rate=" << std::fixed << std::setprecision(1)
                  << (r.kpts0 > 0 ? 100.0 * r.matches / r.kpts0 : 0.0) << "%\n";
        sep('-', 100);
        print_stats_block("SuperPoint img0 (ms)", r.sp0);
        print_stats_block("SuperPoint img1 (ms)", r.sp1);
        print_stats_block("LightGlue match (ms)", r.lg);
        sep('-', 100);
        print_stats_block("Total pipeline  (ms)", r.total);
        std::cout << "  Throughput : "
                  << std::fixed << std::setprecision(1) << (1000.0 / r.total.mean) << " pairs/sec"
                  << "   " << std::setprecision(1) << (2000.0 / r.total.mean) << " images/sec\n";
    }
    sep();

    // ── Save match images ────────────────────────────────────────────────
    std::cout << "\n  Saving match images...\n";
    for (auto& r : results)
        save_match_image(superpoint, lightglue, image0_orig, image1_orig, r);

    std::cout << "\nDone.\n";
    return 0;
}