functions.cpp

#include "functions.h"
#include <fstream>
#include <filesystem>

namespace fs = std::filesystem;

// --- Feature Extraction Functions ---
void make_sift(const cv::Mat& img, std::vector<cv::KeyPoint>& keypoints, cv::Mat& descriptors) {
    auto detector = cv::SIFT::create();
    detector->detectAndCompute(img, cv::noArray(), keypoints, descriptors);
}

void make_surf(const cv::Mat& img, std::vector<cv::KeyPoint>& keypoints, cv::Mat& descriptors) {
    auto detector = cv::xfeatures2d::SURF::create();
    detector->detectAndCompute(img, cv::noArray(), keypoints, descriptors);
}

void make_orb(const cv::Mat& img, 
    std::vector<cv::KeyPoint>& keypoints, 
    cv::Mat& descriptors,
    int nfeatures,
    float scaleFactor,
    int nlevels,
    int edgeThreshold,
    int fastThreshold) {

    auto detector = cv::ORB::create(
    nfeatures, 
    scaleFactor, 
    nlevels, 
    edgeThreshold, 
    0, 
    2, 
    cv::ORB::HARRIS_SCORE, 
    31, 
    fastThreshold
    );
    detector->detectAndCompute(img, cv::noArray(), keypoints, descriptors);
}

double process_all_orb(const std::vector<std::string>& object_folders,
    int nfeatures, float scaleFactor, int nlevels, int edgeThreshold,
    int patchSize, int fastThreshold, int wta_k, int score_type){

        auto detector = cv::ORB::create(
            nfeatures,
            scaleFactor,
            nlevels,
            edgeThreshold,
            0,             // firstLevel (fixed)
            wta_k,
            score_type,
            patchSize,
            fastThreshold
        );
        
    }



// --- Matching Function ---
int match_descriptors(const cv::Mat& desc1, const cv::Mat& desc2, const std::string& algo, 
                      std::vector<cv::DMatch>& good_matches, double ratio_thresh) {
    if (desc1.empty() || desc2.empty()) return 0;

    cv::Ptr<cv::DescriptorMatcher> matcher;
    if (algo == "ORB") matcher = cv::BFMatcher::create(cv::NORM_HAMMING);
    else matcher = cv::BFMatcher::create(cv::NORM_L2);

    std::vector<std::vector<cv::DMatch>> knn_matches;
    matcher->knnMatch(desc1, desc2, knn_matches, 2);

    good_matches.clear();
    for (auto& m : knn_matches) {
        if (m.size() == 2 && m[0].distance < ratio_thresh * m[1].distance) {
            good_matches.push_back(m[0]);
        }
    }
    return static_cast<int>(good_matches.size());
}

// --- Bounding Box from Matched Points ---
cv::Rect getBoundingBoxFromMatchedPoints(const std::vector<cv::KeyPoint>& keypoints_test, const std::vector<cv::DMatch>& matches) {
    if (matches.empty()) return cv::Rect(0, 0, 1, 1);

    float min_x = keypoints_test[matches[0].queryIdx].pt.x;
    float max_x = keypoints_test[matches[0].queryIdx].pt.x;
    float min_y = keypoints_test[matches[0].queryIdx].pt.y;
    float max_y = keypoints_test[matches[0].queryIdx].pt.y;

    for (const auto& m : matches) {
        min_x = std::min(min_x, keypoints_test[m.queryIdx].pt.x);
        max_x = std::max(max_x, keypoints_test[m.queryIdx].pt.x);
        min_y = std::min(min_y, keypoints_test[m.queryIdx].pt.y);
        max_y = std::max(max_y, keypoints_test[m.queryIdx].pt.y);
    }

    if (min_x == max_x) max_x += 1;
    if (min_y == max_y) max_y += 1;

    return cv::Rect(cv::Point(min_x, min_y), cv::Point(max_x, max_y));
}

// --- Save Bounding Box ---
void save_bbox_txt(const std::string& file_path, const std::string& object_name, const cv::Rect& bbox) {
    std::ofstream file(file_path);
    if (file.is_open()) {
        file << object_name << " "
             << bbox.x << " " << bbox.y << " "
             << bbox.x + bbox.width << " " << bbox.y + bbox.height << std::endl;
        file.close();
    }
}

// --- Load Ground Truth ---
bool load_ground_truth_bbox(const std::string& label_path, cv::Rect& gt_bbox, std::string& object_name) {
    std::ifstream file(label_path);
    if (!file.is_open()) return false;

    int x1, y1, x2, y2;
    file >> object_name >> x1 >> y1 >> x2 >> y2;
    gt_bbox = cv::Rect(cv::Point(x1, y1), cv::Point(x2, y2));
    return true;
}

// --- IoU Calculation ---
double calculateIoU(const cv::Rect& box1, const cv::Rect& box2) {
    int intersectionArea = (box1 & box2).area();
    int unionArea = box1.area() + box2.area() - intersectionArea;
    if (unionArea <= 0) return 0.0;
    return static_cast<double>(intersectionArea) / unionArea;
}

// --- Draw Bounding Boxes ---
void draw_bboxes(cv::Mat& img, const cv::Rect& calc_box, const cv::Rect& gt_box) {
    cv::rectangle(img, calc_box, cv::Scalar(0, 255, 0), 2); // Green predicted
    cv::rectangle(img, gt_box, cv::Scalar(0, 0, 255), 2);   // Red ground truth
}

// --- Main Process ---
double process_all_orb(const std::vector<std::string>& object_folders,
    int nfeatures, float scaleFactor, int nlevels, int edgeThreshold,
    int patchSize, int fastThreshold, int wta_k, int score_type) {

    // Overloaded version: use full control of ORB parameters
    std::map<std::string, std::vector<std::pair<std::vector<cv::KeyPoint>, cv::Mat>>> models_per_object;

    for (const auto& object_folder : object_folders) {
        std::string models_path = object_folder + "/models";
        std::string object_name = fs::path(object_folder).filename().string();

        for (const auto& entry : fs::directory_iterator(models_path)) {
            if (entry.is_regular_file()) {
                cv::Mat img = cv::imread(entry.path().string(), cv::IMREAD_GRAYSCALE);
                if (img.empty()) continue;

                std::vector<cv::KeyPoint> keypoints;
                cv::Mat descriptors;
                auto detector = cv::ORB::create(nfeatures, scaleFactor, nlevels,
                                                edgeThreshold, 0, wta_k, score_type,
                                                patchSize, fastThreshold);
                detector->detectAndCompute(img, cv::noArray(), keypoints, descriptors);
                models_per_object[object_name].push_back({keypoints, descriptors});
            }
        }
    }

    int correct = 0;
    int total = 0;
    double total_miou = 0.0;

    for (const auto& test_object_folder : object_folders) {
        std::string test_images_path = test_object_folder + "/test_images";
        std::string labels_path = test_object_folder + "/labels";

        for (const auto& entry : fs::directory_iterator(test_images_path)) {
            if (!entry.is_regular_file()) continue;

            std::string img_name = entry.path().filename().string();
            std::string label_file = labels_path + "/" + img_name.substr(0, img_name.find("-color")) + "-box.txt";

            cv::Rect gt_bbox;
            std::string gt_object_name;
            if (!load_ground_truth_bbox(label_file, gt_bbox, gt_object_name)) {
                continue;
            }

            cv::Mat img_color = cv::imread(entry.path().string());
            cv::Mat img_gray;
            cv::cvtColor(img_color, img_gray, cv::COLOR_BGR2GRAY);

            std::vector<cv::KeyPoint> keypoints_test;
            cv::Mat descriptors_test;
            auto detector = cv::ORB::create(nfeatures, scaleFactor, nlevels,
                                            edgeThreshold, 0, wta_k, score_type,
                                            patchSize, fastThreshold);
            detector->detectAndCompute(img_gray, cv::noArray(), keypoints_test, descriptors_test);

            std::string best_object = "unknown";
            std::vector<cv::DMatch> best_all_matches;
            int best_total_matches = 0;

            for (auto& [object_name, model_list] : models_per_object) {
                std::vector<cv::DMatch> total_matches_for_object;
                for (auto& [model_keypoints, model_descriptors] : model_list) {
                    std::vector<cv::DMatch> matches;
                    match_descriptors(descriptors_test, model_descriptors, "ORB", matches);
                    total_matches_for_object.insert(total_matches_for_object.end(), matches.begin(), matches.end());
                }

                if (total_matches_for_object.size() > best_total_matches) {
                    best_total_matches = total_matches_for_object.size();
                    best_all_matches = total_matches_for_object;
                    best_object = object_name;
                }
            }

            cv::Rect calc_bbox = getBoundingBoxFromMatchedPoints(keypoints_test, best_all_matches);
            double iou = calculateIoU(gt_bbox, calc_bbox);
            total_miou += iou;
            if (iou > 0.5) correct++;
            total++;
        }
    }

    return (total > 0) ? (total_miou / total) : 0.0;
}