postprocessd/stackercpp.cpp

#include "stackercpp.h"

Stacker::Stacker()
{
    Stacker::layers = 0;
    cv::setNumThreads(0);
    Stacker::stopwatch = clock();
}

void
Stacker::add_frame(unsigned char *data, int width, int height)
{
    stopwatch_start();
    Mat warp_matrix = Mat::eye(3, 3, CV_32F);
    Mat mat = Mat(height, width, CV_8UC3, data);
    Mat grayscale = Mat(height, width, CV_8UC1);
    cv::cvtColor(mat, grayscale, cv::COLOR_BGR2GRAY);
    stopwatch_mark("grayscale input");

    stopwatch_start();
    Scalar mean, stddev;
    meanStdDev(grayscale, mean, stddev);
    printf("mean: %f, dev: %f\n", mean[0], stddev[0]);
    if (mean[0] < 20) {
        return;
    }
    stopwatch_mark("filter");

    int number_of_iterations = 5;
    double termination_eps = 1e-10;
    TermCriteria criteria(TermCriteria::COUNT + TermCriteria::EPS, number_of_iterations, termination_eps);
    if (layers == 0) {
        // First image in the stack is used as the reference to align the next frames to
        Stacker::reference = grayscale;

        // Create black image to accumulate the average
        Stacker::stacked = Mat(height, width, CV_64FC3);
        Stacker::stacked.setTo(Scalar(0, 0, 0, 0));

        // Add first frame to the stack
        Stacker::stacked += mat;
        Stacker::layers += 1;
    } else {
        // All frames after the initial one are stacked
        Mat warped = Mat(Stacker::stacked.rows, Stacker::stacked.cols, CV_32FC1);

        stopwatch_start();
        cv::findTransformECC(grayscale, Stacker::reference, warp_matrix, MOTION_HOMOGRAPHY, criteria);
        stopwatch_mark("find alignment");

        stopwatch_start();
        warpPerspective(mat, warped, warp_matrix, warped.size(), INTER_LINEAR);
        stopwatch_mark("warp image");

        // Add the warped image to the stack
        Stacker::stacked += warped;
        Stacker::layers += 1;
    }
}

Mat
Stacker::postprocess_mat(Mat input)
{
    stopwatch_start();
    Mat sharpened;
    GaussianBlur(input, sharpened, Size(0, 0), 1.7);
    addWeighted(input, 2.5, sharpened, -1.5, 0, sharpened);
    stopwatch_mark("sharpen");

    stopwatch_start();
    Mat lab;
    cvtColor(sharpened, lab, COLOR_BGR2Lab);
    std::vector<cv::Mat> lab_planes(3);
    cv::split(lab, lab_planes);
    stopwatch_mark("to Lab");

    /*
    stopwatch_start();
    cv::Ptr<cv::CLAHE> clahe = cv::createCLAHE();
    clahe->setClipLimit(2);
    clahe->setTilesGridSize(Size(8, 8));
    cv::Mat dst;
    clahe->apply(lab_planes[0], dst);
    dst.copyTo(lab_planes[0]);
    stopwatch_mark("clahe");
    */

    stopwatch_start();
    Mat result;
    cv::merge(lab_planes, lab);
    cvtColor(lab, result, COLOR_Lab2BGR);
    stopwatch_mark("to RGB");
    return result;
}

char *
Stacker::get_result()
{
    // Complete the averaging and go back to an 8-bit image
    stopwatch_start();
    Stacker::stacked.convertTo(Stacker::stacked, CV_8U, 1. / Stacker::layers);
    stopwatch_mark("average");

    // Run the final-frame postprocessing
    Mat result = postprocess_mat(Stacker::stacked);

    // Convert mat to bytes
    size_t size = result.total() * result.elemSize();
    char *data = (char *) malloc(size);
    std::memcpy(data, result.data, size * sizeof(char));
    return data;
}

char *
Stacker::postprocess(unsigned char *data, int width, int height)
{
    // Convert bytes to mat
    Mat mat = Mat(height, width, CV_8UC3, data);

    // Run the final-frame postprocessing
    Mat result = postprocess_mat(mat);

    // Convert mat to bytes
    size_t size = result.total() * result.elemSize();
    char *outdata = (char *) malloc(size);
    std::memcpy(outdata, result.data, size * sizeof(char));
    return outdata;
}

void
Stacker::stopwatch_start()
{
    Stacker::stopwatch = clock();
}

void
Stacker::stopwatch_mark(const char *name)
{
    printf("[%.1fms] %s\n", float(clock() - Stacker::stopwatch) / CLOCKS_PER_SEC * 1000, name);
}