task3.cpp

// Sakshi B.
#include "task3.h"
#include "task2.h"
#include "task1.h"
#include "config.h"
#include <iostream>
#include <sstream>
#include <iomanip>
#include <cmath>

// ── calibrateFromImages ───────────────────────────────────────────────────────

double calibrateFromImages(int argc, char* argv[],
                            cv::Mat& cameraMatrix,
                            cv::Mat& distCoeffs,
                            const std::string& outputFile)
{
    // Types exactly as required by the assignment
    std::vector<cv::Vec3f>                point_set;   // one frame's 3-D points
    std::vector<std::vector<cv::Vec3f>>   point_list;  // accumulated across frames
    std::vector<std::vector<cv::Point2f>> corner_list; // accumulated 2-D corners

    // Pre-build point_set once — same for every frame
    point_set = buildPointSet();

    cv::Size imageSize;

    // ── Step 1: load any calib_N.png files already saved by Task 2 ───────────
    // Scan for calib_1.png, calib_2.png, ... until one is missing
    {
        int idx = 1;
        while (true) {
            std::string fname = "calib_" + std::to_string(idx) + ".png";
            cv::Mat img = cv::imread(fname);
            if (img.empty()) break;   // no more files

            imageSize = img.size();
            cv::Mat grayImg;
            cv::cvtColor(img, grayImg, cv::COLOR_BGR2GRAY);

            std::vector<cv::Point2f> corners;
            if (detectCorners(grayImg, corners)) {
                corner_list.push_back(corners);
                point_list.push_back(point_set);
                std::cout << "[Task3] Loaded " << fname << " ("
                          << corners.size() << " corners)\n";
            } else {
                std::cout << "[Task3] WARNING: no corners found in " << fname << ", skipping.\n";
            }
            ++idx;
        }
        if (!corner_list.empty())
            std::cout << "[Task3] Pre-loaded " << corner_list.size()
                      << " frame(s) from disk.\n\n";
    }

    // ── Step 2: open live camera so user can add more frames or calibrate ─────
    cv::VideoCapture cap = openVideoSource(argc, argv);
    if (!cap.isOpened()) {
        std::cerr << "[Task3][ERROR] Cannot open video source.\n";
        return -1.0;
    }

    std::cout << "=== Task 3: Camera Calibration ===\n"
              << "  Pre-loaded: " << corner_list.size() << " frame(s) from calib_N.png\n"
              << "  's'     — save another frame from live camera\n"
              << "  'c'     — calibrate now (need >= 5 frames total)\n"
              << "  'q'/ESC — quit without calibrating\n\n";

    cv::Mat frame, gray;
    cv::Mat      lastGoodFrame;
    std::vector<cv::Point2f> lastCorners;
    bool haveGoodCorners = false;
    int savedCount = static_cast<int>(corner_list.size());  // continue numbering from Task 2

    while (true) {
        cap >> frame;
        if (frame.empty()) { std::cout << "[Task3] End of stream.\n"; break; }

        imageSize = frame.size();

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

        std::vector<cv::Point2f> corners;
        bool found = detectCorners(gray, corners);

        if (found) {
            lastCorners     = corners;
            lastGoodFrame   = frame.clone();
            haveGoodCorners = true;
        }

        cv::Mat display = frame.clone();
        drawCornersOverlay(display, found ? corners : lastCorners,
                           found || haveGoodCorners);

        std::string info = "Saved: " + std::to_string(savedCount) +
                           "   's'=save  'c'=calibrate";
        cv::putText(display, info, cv::Point(10, 60),
                    cv::FONT_HERSHEY_SIMPLEX, 0.65,
                    cv::Scalar(255, 200, 0), 2);

        if (!found && haveGoodCorners)
            cv::putText(display, "(showing cached corners)",
                        cv::Point(10, 90),
                        cv::FONT_HERSHEY_SIMPLEX, 0.55,
                        cv::Scalar(180, 180, 0), 1);

        cv::imshow("Task 3 - Calibration", display);
        int key = cv::waitKey(1) & 0xFF;

        // ── 's': save this frame ──────────────────────────────────────────────
        if (key == 's') {
            if (!haveGoodCorners) {
                std::cout << "[Task3] No corners detected yet.\n";
            } else {
                corner_list.push_back(lastCorners);
                point_list.push_back(point_set);
                ++savedCount;

                std::string filename = "calib_" + std::to_string(savedCount) + ".png";
                cv::Mat annotated = lastGoodFrame.clone();
                cv::drawChessboardCorners(annotated, BOARD_SIZE, lastCorners, true);
                cv::imwrite(filename, annotated);

                std::cout << "[Task3] Saved frame " << savedCount
                          << " → " << filename << "\n";
            }
        }

        // ── 'c': calibrate ────────────────────────────────────────────────────
        if (key == 'c') {
            if (corner_list.size() < 5) {
                std::cout << "[Task3] Need at least 5 frames (have "
                          << corner_list.size() << ").\n";
            } else {
                break;   // proceed to calibration below
            }
        }

        if (key == 'q' || key == 27) {
            std::cout << "[Task3] Quit without calibrating.\n";
            cap.release();
            cv::destroyAllWindows();
            return -1.0;
        }
    }

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

    if (corner_list.size() < 5) {
        std::cerr << "[Task3] Not enough frames to calibrate.\n";
        return -1.0;
    }

    // ── Run cv::calibrateCamera() ─────────────────────────────────────────────
    // cv::calibrateCamera needs vector<vector<Point3f>>; convert from Vec3f
    std::vector<std::vector<cv::Point3f>> objectPoints;
    objectPoints.reserve(point_list.size());
    for (const auto& ps : point_list) {
        std::vector<cv::Point3f> pts;
        pts.reserve(ps.size());
        for (const auto& v : ps)
            pts.emplace_back(v[0], v[1], v[2]);
        objectPoints.push_back(std::move(pts));
    }

    // ── BEFORE: initial estimates ─────────────────────────────────────────────
    // OpenCV's canonical guess: fx=fy=max(W,H), cx=W/2, cy=H/2, distortion=0
    cv::Mat initMatrix = cv::Mat::eye(3, 3, CV_64F);
    initMatrix.at<double>(0,0) = std::max(imageSize.width, imageSize.height);  // fx
    initMatrix.at<double>(1,1) = std::max(imageSize.width, imageSize.height);  // fy
    initMatrix.at<double>(0,2) = imageSize.width  / 2.0;                       // cx (u0)
    initMatrix.at<double>(1,2) = imageSize.height / 2.0;                       // cy (v0)

    cv::Mat initDist = cv::Mat::zeros(1, 5, CV_64F);   // [k1 k2 p1 p2 k3] = all 0

    std::cout << "\n[Task3] ══════════════════════════════════════════════\n"
              << "[Task3]  BEFORE CALIBRATION  (initial estimates)\n"
              << "[Task3] ══════════════════════════════════════════════\n"
              << "  Image size  : " << imageSize.width << " x " << imageSize.height << "\n"
              << "  Camera matrix K (initial):\n"
              << "    fx = " << initMatrix.at<double>(0,0)
              << "  fy = "   << initMatrix.at<double>(1,1)
              << "  (fx == fy by assumption)\n"
              << "    u0 = " << initMatrix.at<double>(0,2)
              << "  (= W/2 = " << imageSize.width  / 2.0 << ")\n"
              << "    v0 = " << initMatrix.at<double>(1,2)
              << "  (= H/2 = " << imageSize.height / 2.0 << ")\n"
              << initMatrix << "\n"
              << "  Distortion coefficients [k1  k2  p1  p2  k3] (initial):\n"
              << "  " << initDist << "  (all zeros — no prior estimate)\n"
              << "[Task3] ══════════════════════════════════════════════\n\n";

    // ── Run calibration ───────────────────────────────────────────────────────
    std::vector<cv::Mat> rvecs, tvecs;
    double rms = cv::calibrateCamera(objectPoints, corner_list, imageSize,
                                     cameraMatrix, distCoeffs,
                                     rvecs, tvecs);

    // ── AFTER: calibrated values ──────────────────────────────────────────────
    double fx = cameraMatrix.at<double>(0,0);
    double fy = cameraMatrix.at<double>(1,1);
    double u0 = cameraMatrix.at<double>(0,2);
    double v0 = cameraMatrix.at<double>(1,2);

    double fxfyDiff   = std::abs(fx - fy);
    double u0expected = imageSize.width  / 2.0;
    double v0expected = imageSize.height / 2.0;

    std::cout << "[Task3] ══════════════════════════════════════════════\n"
              << "[Task3]  AFTER CALIBRATION  (calibrated results)\n"
              << "[Task3] ══════════════════════════════════════════════\n"
              << "  Frames used              : " << corner_list.size() << "\n"
              << "  Image size               : " << imageSize.width << " x " << imageSize.height << "\n"
              << "\n"
              << "  Camera matrix K (calibrated):\n"
              << "    fx = " << std::fixed << std::setprecision(4) << fx << "\n"
              << "    fy = " << fy
              << "   |fx−fy| = " << fxfyDiff
              << (fxfyDiff < 2.0 ? "  ✓ (essentially equal)" : "  ✗ (differ — check board)") << "\n"
              << "    u0 = " << u0
              << "   (initial estimate W/2 = " << u0expected << ")"
              << "   diff = " << std::abs(u0 - u0expected)
              << (std::abs(u0 - u0expected) < imageSize.width * 0.05
                  ? "  ✓" : "  ✗ far from center") << "\n"
              << "    v0 = " << v0
              << "   (initial estimate H/2 = " << v0expected << ")"
              << "   diff = " << std::abs(v0 - v0expected)
              << (std::abs(v0 - v0expected) < imageSize.height * 0.05
                  ? "  ✓" : "  ✗ far from center") << "\n"
              << "\n" << cameraMatrix << "\n"
              << "\n  Distortion coefficients [k1  k2  p1  p2  k3] (calibrated):\n"
              << "  " << distCoeffs << "\n"
              << "\n  RMS reprojection error   : " << std::setprecision(6) << rms << " pixels"
              << (rms < 1.0 ? "  ✓ excellent (< 1 px)"
                 : (rms < 3.0 ? "  ~ acceptable (1–3 px, typical for phone/large sensor)"
                              : "  ✗ high (> 3 px — consider recalibrating)")) << "\n"
              << "[Task3] ══════════════════════════════════════════════\n\n";

    // ── Save per-image rvec/tvec alongside calibration images ─────────────────
    for (int i = 0; i < static_cast<int>(rvecs.size()); ++i) {
        std::string fn = "calib_pose_" + std::to_string(i+1) + ".yml";
        cv::FileStorage ps(fn, cv::FileStorage::WRITE);
        ps << "rvec" << rvecs[i];
        ps << "tvec" << tvecs[i];
        ps.release();
    }
    std::cout << "  Per-image poses saved    : calib_pose_1.yml … calib_pose_"
              << rvecs.size() << ".yml\n";

    // ── Save calibration to YAML ──────────────────────────────────────────────
    cv::FileStorage fs(outputFile, cv::FileStorage::WRITE);
    fs << "camera_matrix" << cameraMatrix;
    fs << "dist_coeffs"   << distCoeffs;
    fs << "rms_error"     << rms;
    fs << "image_width"   << imageSize.width;
    fs << "image_height"  << imageSize.height;
    fs.release();

    std::cout << "  Calibration saved to     : " << outputFile << "\n"
              << "[Task3] ─────────────────────────────────────────────\n\n";

    return rms;
}

// ── loadCalibration ───────────────────────────────────────────────────────────

bool loadCalibration(const std::string& filename,
                     cv::Mat& cameraMatrix,
                     cv::Mat& distCoeffs)
{
    cv::FileStorage fs(filename, cv::FileStorage::READ);
    if (!fs.isOpened()) {
        std::cerr << "[Task3][ERROR] Cannot open calibration file: "
                  << filename << "\n";
        return false;
    }
    fs["camera_matrix"] >> cameraMatrix;
    fs["dist_coeffs"]   >> distCoeffs;
    fs.release();
    std::cout << "[Task3] Loaded calibration from " << filename << "\n";
    return true;
}

// ── estimatePose ──────────────────────────────────────────────────────────────

bool estimatePose(const std::vector<cv::Point2f>& corners,
                  const cv::Mat& cameraMatrix,
                  const cv::Mat& distCoeffs,
                  cv::Mat& rvec,
                  cv::Mat& tvec)
{
    // Convert Vec3f point_set to Point3f for solvePnP
    std::vector<cv::Vec3f> ps = buildPointSet();
    std::vector<cv::Point3f> objectPoints;
    objectPoints.reserve(ps.size());
    for (const auto& v : ps)
        objectPoints.emplace_back(v[0], v[1], v[2]);

    return cv::solvePnP(objectPoints, corners,
                        cameraMatrix, distCoeffs,
                        rvec, tvec);
}

// ── drawAxes ──────────────────────────────────────────────────────────────────

void drawAxes(cv::Mat& frame,
              const cv::Mat& cameraMatrix,
              const cv::Mat& distCoeffs,
              const cv::Mat& rvec,
              const cv::Mat& tvec,
              float axisLength)
{
    std::vector<cv::Point3f> axisPoints = {
        {0.f,        0.f,        0.f},   // origin
        {axisLength, 0.f,        0.f},   // +X  red
        {0.f,        axisLength, 0.f},   // +Y  green
        {0.f,        0.f,       -axisLength}  // +Z  blue (out of board)
    };

    std::vector<cv::Point2f> projected;
    cv::projectPoints(axisPoints, rvec, tvec,
                      cameraMatrix, distCoeffs, projected);

    cv::arrowedLine(frame, projected[0], projected[1], cv::Scalar(0,   0, 255), 3);
    cv::arrowedLine(frame, projected[0], projected[2], cv::Scalar(0, 255,   0), 3);
    cv::arrowedLine(frame, projected[0], projected[3], cv::Scalar(255,  0,   0), 3);

    cv::putText(frame, "X", projected[1], cv::FONT_HERSHEY_SIMPLEX, 0.6, cv::Scalar(0,  0,255), 2);
    cv::putText(frame, "Y", projected[2], cv::FONT_HERSHEY_SIMPLEX, 0.6, cv::Scalar(0,255,  0), 2);
    cv::putText(frame, "Z", projected[3], cv::FONT_HERSHEY_SIMPLEX, 0.6, cv::Scalar(255,0,  0), 2);
}

// ── runTask3 ──────────────────────────────────────────────────────────────────

void runTask3(int argc, char* argv[])
{
    // Step 1: calibrate
    cv::Mat cameraMatrix, distCoeffs;
    double rms = calibrateFromImages(argc, argv, cameraMatrix, distCoeffs);
    if (rms < 0) return;

    // Step 2: live pose estimation
    cv::VideoCapture cap = openVideoSource(argc, argv);
    if (!cap.isOpened()) {
        std::cerr << "[Task3][ERROR] Cannot open video source.\n";
        return;
    }

    std::cout << "=== Task 3: Pose Estimation (live) ===\n"
              << "Press 'q'/ESC to quit\n\n";

    cv::Mat frame, gray;

    int frameCount = 0;

    while (true) {
        cap >> frame;
        if (frame.empty()) break;
        ++frameCount;

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

        std::vector<cv::Point2f> corners;
        bool found = detectCorners(gray, corners);

        if (found) {
            cv::Mat rvec, tvec;
            if (estimatePose(corners, cameraMatrix, distCoeffs, rvec, tvec)) {

                // Draw axes on the frame
                drawAxes(frame, cameraMatrix, distCoeffs, rvec, tvec);

                // Format rotation and translation as readable strings
                // rvec is a Rodrigues rotation vector (angle-axis, radians)
                // tvec is translation in board units (squares)
                auto fmt3 = [](const cv::Mat& v) {
                    std::ostringstream ss;
                    ss << std::fixed << std::setprecision(3)
                       << "(" << v.at<double>(0)
                       << ", " << v.at<double>(1)
                       << ", " << v.at<double>(2) << ")";
                    return ss.str();
                };

                std::string rStr = "R: " + fmt3(rvec);
                std::string tStr = "T: " + fmt3(tvec);

                // On-screen overlay
                cv::putText(frame, rStr, cv::Point(10, 30),
                            cv::FONT_HERSHEY_SIMPLEX, 0.6,
                            cv::Scalar(0, 255, 255), 2);
                cv::putText(frame, tStr, cv::Point(10, 58),
                            cv::FONT_HERSHEY_SIMPLEX, 0.6,
                            cv::Scalar(0, 200, 255), 2);

                // Terminal output every 15 frames so it stays readable
                if (frameCount % 15 == 0) {
                    std::cout << "[Task3] R=" << fmt3(rvec)
                              << "  T=" << fmt3(tvec) << "\n";
                }
            }
        }

        drawStatusBar(frame,
                      "Task 3 | Pose: " +
                      std::string(found ? "ESTIMATED" : "SEARCHING..."),
                      found);

        cv::imshow("Task 3 - Pose Estimation", frame);
        int key = cv::waitKey(1) & 0xFF;
        if (key == 'q' || key == 27) break;
    }

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