task8.cpp

#include "task8.h"
#include "task3.h"
#include "task6.h"
#include "config.h"

#include <iostream>
#include <iomanip>
#include <sstream>
#include <vector>
#include <string>
#include <cmath>
#include <algorithm>

#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif

// ─────────────────────────────────────────────────────────────────────────────
// Helpers
// ─────────────────────────────────────────────────────────────────────────────

static std::string stemOf(const std::string& p)
{
    auto sl = p.find_last_of("/\\");
    std::string f = (sl == std::string::npos) ? p : p.substr(sl+1);
    auto dt = f.find_last_of('.');
    return (dt == std::string::npos) ? f : f.substr(0, dt);
}

static std::string fmt3(const cv::Mat& v, int dec=2)
{
    std::ostringstream ss;
    ss << std::fixed << std::setprecision(dec)
       << "(" << v.at<double>(0)
       << ", " << v.at<double>(1)
       << ", " << v.at<double>(2) << ")";
    return ss.str();
}

static cv::VideoCapture openCameraSafe(int argc, char* argv[])
{
    cv::VideoCapture cap;

    if (argc > 3) {
        std::string src(argv[3]);
        bool isNum = !src.empty() &&
                     std::all_of(src.begin(), src.end(), ::isdigit);
        if (isNum) cap.open(std::stoi(src));
        else       cap.open(src);
        if (cap.isOpened()) return cap;
        std::cerr << "[Task8] Could not open '" << src << "'\n";
    }

    for (int idx : {0, 1, 2, 4}) {
        cap.open(idx);
        if (cap.isOpened()) {
            std::cout << "[Task8] Opened webcam at index " << idx << "\n";
            return cap;
        }
        cap.release();
    }
    return cap;
}

// ─────────────────────────────────────────────────────────────────────────────
// loadCustomTarget
// ─────────────────────────────────────────────────────────────────────────────

CustomTarget loadCustomTarget(const std::string& imagePath,
                              float physW, float physH)
{
    CustomTarget t;
    t.refImage = cv::imread(imagePath);
    if (t.refImage.empty()) {
        std::cerr << "[Task8] Cannot read reference image: " << imagePath << "\n";
        return t;
    }
    t.physW = physW;
    t.physH = physH;
    t.name  = stemOf(imagePath);

    cv::Mat gray;
    cv::cvtColor(t.refImage, gray, cv::COLOR_BGR2GRAY);

    cv::Ptr<cv::CLAHE> clahe = cv::createCLAHE(3.0, cv::Size(8,8));
    clahe->apply(gray, gray);

    auto orb = cv::ORB::create(4000, 1.2f, 8, 31, 0, 2,
                               cv::ORB::HARRIS_SCORE, 31, 20);
    orb->detectAndCompute(gray, cv::noArray(), t.refKps, t.refDesc);
    t.orb = orb;

    std::cout << "[Task8] Reference '" << t.name << "': "
              << t.refKps.size() << " ORB keypoints\n"
              << "        Image: " << t.refImage.cols << "x" << t.refImage.rows
              << "  World: " << physW << " x " << physH << " units\n";

    if ((int)t.refKps.size() < 50)
        std::cout << "[Task8] WARNING: Only " << t.refKps.size()
                  << " keypoints — use a more textured target.\n";
    else if ((int)t.refKps.size() < 200)
        std::cout << "[Task8] NOTE: " << t.refKps.size()
                  << " keypoints — tracking may be unstable.\n";

    return t;
}

// ─────────────────────────────────────────────────────────────────────────────
// detectCustomTarget
// ─────────────────────────────────────────────────────────────────────────────

bool detectCustomTarget(const cv::Mat& frame,
                        const cv::Mat& grayRaw,
                        const CustomTarget& target,
                        const cv::Mat& K,
                        const cv::Mat& D,
                        cv::Mat& rvec,
                        cv::Mat& tvec,
                        std::vector<cv::Point2f>& trackedCorners,
                        int* outMatchCount)
{
    trackedCorners.clear();
    if (outMatchCount) *outMatchCount = 0;
    if (target.refDesc.empty() || target.orb.empty()) return false;

    cv::Mat gray;
    cv::Ptr<cv::CLAHE> clahe = cv::createCLAHE(3.0, cv::Size(8,8));
    clahe->apply(grayRaw, gray);

    std::vector<cv::KeyPoint> frameKps;
    cv::Mat frameDesc;
    target.orb->detectAndCompute(gray, cv::noArray(), frameKps, frameDesc);
    if (frameDesc.empty() || (int)frameKps.size() < 8) return false;

    cv::BFMatcher matcher(cv::NORM_HAMMING);
    std::vector<std::vector<cv::DMatch>> knnMatches;
    matcher.knnMatch(target.refDesc, frameDesc, knnMatches, 2);

    const float RATIO = 0.75f;
    std::vector<cv::DMatch> good;
    for (const auto& pair : knnMatches)
        if (pair.size() == 2 && pair[0].distance < RATIO * pair[1].distance)
            good.push_back(pair[0]);

    if (outMatchCount) *outMatchCount = static_cast<int>(good.size());
    if ((int)good.size() < 8) return false;

    std::vector<cv::Point2f> refPts, framePts;
    refPts .reserve(good.size());
    framePts.reserve(good.size());
    for (const auto& m : good) {
        refPts  .push_back(target.refKps[m.queryIdx].pt);
        framePts.push_back(frameKps     [m.trainIdx].pt);
    }

    cv::Mat mask;
    cv::Mat H = cv::findHomography(refPts, framePts, cv::RANSAC, 3.0, mask);
    if (H.empty()) return false;

    int inliers = cv::countNonZero(mask);
    if (outMatchCount) *outMatchCount = inliers;
    if (inliers < 8) return false;

    float rW = static_cast<float>(target.refImage.cols);
    float rH = static_cast<float>(target.refImage.rows);
    std::vector<cv::Point2f> refCorners = {
        {0.f, 0.f}, {rW, 0.f}, {rW, rH}, {0.f, rH}
    };
    cv::perspectiveTransform(refCorners, trackedCorners, H);

    double area      = cv::contourArea(trackedCorners);
    double frameArea = static_cast<double>(frame.cols * frame.rows);
    if (area < frameArea * 0.003) return false;
    if (area > frameArea * 0.98)  return false;

    std::vector<cv::Point2f> hull;
    cv::convexHull(trackedCorners, hull);
    if ((int)hull.size() != 4) return false;

    cv::Rect frameBounds(-50, -50, frame.cols+100, frame.rows+100);
    for (const auto& pt : trackedCorners)
        if (!frameBounds.contains(cv::Point(static_cast<int>(pt.x),
                                            static_cast<int>(pt.y))))
            return false;

    float W = target.physW, Ht = target.physH;
    std::vector<cv::Point3f> worldPts = {
        {0.f,   0.f,  0.f},
        {W,     0.f,  0.f},
        {W,    -Ht,   0.f},
        {0.f,  -Ht,   0.f}
    };
    return cv::solvePnP(worldPts, trackedCorners, K, D,
                        rvec, tvec, false, cv::SOLVEPNP_ITERATIVE);
}

// ─────────────────────────────────────────────────────────────────────────────
// drawTargetOverlay
// ─────────────────────────────────────────────────────────────────────────────

void drawTargetOverlay(cv::Mat& frame,
                       const std::vector<cv::Point2f>& corners,
                       const cv::Mat& K,
                       const cv::Mat& D,
                       const cv::Mat& rvec,
                       const cv::Mat& tvec)
{
    if ((int)corners.size() != 4) return;

    for (int i = 0; i < 4; ++i)
        cv::line(frame, corners[i], corners[(i+1)%4],
                 cv::Scalar(0, 220, 255), 2, cv::LINE_AA);

    const std::string labels[4] = {"TL","TR","BR","BL"};
    const cv::Scalar  colors[4] = {
        {0,255,255}, {0,255,0}, {0,120,255}, {255,80,80}
    };
    for (int i = 0; i < 4; ++i) {
        cv::circle(frame, corners[i], 7, colors[i], cv::FILLED, cv::LINE_AA);
        cv::circle(frame, corners[i], 7, cv::Scalar(0,0,0), 1, cv::LINE_AA);
        cv::Point2f lp = corners[i] + cv::Point2f(10.f, -8.f);
        cv::putText(frame, labels[i], lp, cv::FONT_HERSHEY_SIMPLEX,
                    0.60, cv::Scalar(0,0,0), 3, cv::LINE_AA);
        cv::putText(frame, labels[i], lp, cv::FONT_HERSHEY_SIMPLEX,
                    0.60, colors[i], 1, cv::LINE_AA);
    }

    drawAxes(frame, K, D, rvec, tvec, 0.3f);
}

// ─────────────────────────────────────────────────────────────────────────────
// Info panel
// ─────────────────────────────────────────────────────────────────────────────

static void drawInfoPanel(cv::Mat& frame,
                          const std::string& name, bool detected,
                          int matchCount, int kpCount,
                          const cv::Mat& rvec, const cv::Mat& tvec)
{
    const int PX=10, PY=10, PW=250, LH=18, PAD=7;
    std::vector<std::string> lines;
    lines.push_back(detected ? "TARGET FOUND" : "SEARCHING...");
    lines.push_back("Ref: " + name + " (" + std::to_string(kpCount) + " kp)");
    lines.push_back("Inliers: " + std::to_string(matchCount));
    if (detected) {
        lines.push_back("R: " + fmt3(rvec));
        lines.push_back("T: " + fmt3(tvec));
    } else {
        lines.push_back(kpCount < 100
            ? "Low keypoints - use textured target"
            : matchCount < 8
                ? "Not enough matches - show more of target"
                : "Adjust angle / lighting");
    }

    int PH = static_cast<int>(lines.size())*LH + PAD*2;
    cv::Rect safe = cv::Rect(PX,PY,PW,PH) & cv::Rect(0,0,frame.cols,frame.rows);
    if (safe.area() > 0) {
        cv::Mat roi=frame(safe), bg(safe.height,safe.width,CV_8UC3,cv::Scalar(18,18,26));
        cv::addWeighted(roi,0.25,bg,0.75,0,roi);
    }
    cv::rectangle(frame, {PX,PY}, {PX+PW,PY+PH},
                  detected ? cv::Scalar(60,160,60) : cv::Scalar(60,60,160), 1);

    for (int i=0; i<(int)lines.size(); ++i) {
        cv::Scalar col = (i==0)
            ? (detected ? cv::Scalar(80,255,80) : cv::Scalar(100,100,255))
            : (i==(int)lines.size()-1 && !detected)
                ? cv::Scalar(80,140,255)
                : cv::Scalar(180,190,210);
        cv::putText(frame, lines[i],
                    {PX+PAD, PY+PAD+(i+1)*LH-3},
                    cv::FONT_HERSHEY_SIMPLEX, 0.40, col, 1, cv::LINE_AA);
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// runTask8
// ─────────────────────────────────────────────────────────────────────────────

void runTask8(int argc, char* argv[])
{
    if (argc < 3) {
        std::cout << "=== Task 8: Custom Target AR ===\n"
                  << "Usage:\n"
                  << "  ./ar_project 8 reference.jpg\n"
                  << "  ./ar_project 8 reference.jpg 1          # webcam 1\n"
                  << "  ./ar_project 8 reference.jpg video.mp4  # video file\n\n"
                  << "  reference.jpg — photo of your flat target\n"
                  << "                  (book cover, poster, magazine, box face)\n"
                  << "  Requires: calibration.yml\n";
        return;
    }

    std::string refPath(argv[2]);

    cv::Mat K, D;
    if (!loadCalibration("calibration.yml", K, D)) return;

    CustomTarget target = loadCustomTarget(refPath, 1.0f, -1.0f);
    if (target.refImage.empty()) return;

    target.physH = target.physW *
                   static_cast<float>(target.refImage.rows) /
                   static_cast<float>(target.refImage.cols);

    if ((int)target.refKps.size() < 20) {
        std::cerr << "[Task8] Reference image has too few keypoints to track.\n"
                  << "        Please use a more textured image.\n";
        return;
    }

    cv::VideoCapture cap = openCameraSafe(argc, argv);
    if (!cap.isOpened()) {
        std::cerr << "[Task8] Cannot open any video source.\n"
                  << "        Try: ./ar_project 8 ref.jpg 1\n";
        return;
    }

    // ── Build virtual objects ─────────────────────────────────────────────────
    // Eiffel Tower — positioned on the left third of the target
    float tcx = target.physW * 0.50f;
    float tcy = -target.physH * 0.50f;
    std::vector<Line3D> tower = buildEiffelTower(tcx, tcy);

    // // UFO — hovering above the right third of the target
    float ucx = target.physW * 0.50f;
    float ucy = -target.physH * 0.50f;
    std::vector<Line3D> ufo = buildUFO(ucx, ucy, 0.25f);

    std::cout << "\n=== Task 8: Custom Target AR ===\n"
              << "  Target  : " << target.name
              << " (" << target.refKps.size() << " keypoints)\n"
              << "  World   : " << target.physW << " x " << target.physH << " units\n"
              << "  Objects : UFO \n\n"
              << "  s    — save snapshot\n"
              << "  q/ESC — quit\n\n";

    const std::string WIN = "Task 8 - Custom Target AR";
    cv::namedWindow(WIN, cv::WINDOW_AUTOSIZE);

    {
        cv::Mat thumb;
        int tw = 200;
        cv::resize(target.refImage, thumb,
                   cv::Size(tw, tw * target.refImage.rows / target.refImage.cols));
        cv::putText(thumb, "Reference", {4,14},
                    cv::FONT_HERSHEY_SIMPLEX, 0.42,
                    cv::Scalar(0,220,255), 1, cv::LINE_AA);
        cv::imshow("Task 8 - Reference", thumb);
        cv::moveWindow("Task 8 - Reference", 10, 10);
    }

    cv::Mat frame, gray;
    int savedCount = 0;

    while (true) {
        cap >> frame;
        if (frame.empty()) {
            cap.set(cv::CAP_PROP_POS_FRAMES, 0);
            cap >> frame;
            if (frame.empty()) break;
        }

        cv::cvtColor(frame, gray, cv::COLOR_BGR2GRAY);

        cv::Mat rvec, tvec;
        std::vector<cv::Point2f> trackedCorners;
        int matchCount = 0;
        bool found = detectCustomTarget(frame, gray, target, K, D,
                                        rvec, tvec, trackedCorners,
                                        &matchCount);

        cv::Mat display = frame.clone();

        if (found) {
            drawTargetOverlay(display, trackedCorners, K, D, rvec, tvec);
            // drawVirtualObject(display, tower, K, D, rvec, tvec);
            drawVirtualObject(display, ufo,   K, D, rvec, tvec);
        }

        drawInfoPanel(display, target.name, found, matchCount,
                      static_cast<int>(target.refKps.size()),
                      rvec, tvec);

        drawStatusBar(display,
                      "Task 8 | " + target.name +
                      (found ? " - TRACKED  inliers=" + std::to_string(matchCount)
                             : " - SEARCHING..."),
                      found);

        cv::imshow(WIN, display);

        int key = cv::waitKey(1) & 0xFF;
        if (key == 'q' || key == 27) break;
        if (key == 's') {
            std::string fn = "task8_" + target.name + "_" +
                             std::to_string(++savedCount) + ".png";
            cv::imwrite(fn, display);
            std::cout << "[Task8][SAVED] " << fn << "\n";
        }
    }

    cap.release();
    cv::destroyAllWindows();
}