FAST corner detector using Halide

FastCorner.cpp

#include "Halide.h"
#include <halide_benchmark.h>
#include <halide_image_io.h>
#include <iostream>

using namespace Halide;
using namespace std;

int main(int argc, char** argv)
{
    const int threshold = 20;

    Halide::Buffer<uint8_t> input = Tools::load_image(argv[1]);
    input.set_name("input");

    static const int offsets16[][2] = {
        { 0, 3 }, { 1, 3 }, { 2, 2 }, { 3, 1 }, { 3, 0 }, { 3, -1 }, { 2, -2 }, { 1, -3 },
        { 0, -3 }, { -1, -3 }, { -2, -2 }, { -3, -1 }, { -3, 0 }, { -3, 1 }, { -2, 2 }, { -1, 3 },
        { 0, 3 }, { 1, 3 }, { 2, 2 }, { 3, 1 }, { 3, 0 }, { 3, -1 }, { 2, -2 }, { 1, -3 }
    };

    Var x("x"), y("y");
    Func clamped("clamped");
    clamped(x, y) = input(clamp(x, 0, input.width() - 1), clamp(y, 0, input.height() - 1));

    Halide::Buffer<int> ofs_buf(offsets16);
    ofs_buf.set_name("ofs_buf");
    Var i("i");

    Func diffs("diffs");
    diffs(x, y, i) = cast<short>(clamped(x, y)) - cast<short>(clamped(x + ofs_buf(0, i), y + ofs_buf(1, i)));

    RDom r24(0, 24, "r");
    Func is_corner("isCorner"), count("count");
    count(x, y) = { 0, 0 };
    r24.where(count(x, y)[0] < 9 && count(x, y)[1] < 9);
    Expr count_below = select(diffs(x, y, r24) < -threshold, count(x, y)[0] + 1, 0);
    Expr count_above = select(diffs(x, y, r24) > threshold, count(x, y)[1] + 1, 0);
    count(x, y) = { count_below, count_above };
    is_corner(x, y) = (count(x, y)[0] >= 9 || count(x, y)[1] >= 9);

    Func min_max_vals("min_max_vals");
    RDom ri(0, 9, "ri");
    min_max_vals(x, y, i) = { maximum(diffs(x, y, i + ri)), minimum(diffs(x, y, i + ri)) };

    Region bounds = { { 3, input.width() - 6 }, { 3, input.height() - 6 } };
    Func score_ub("score_unbounded"), score("score");
    RDom ro(0, 16, "ro");
    Expr score_val = max(-minimum(min_max_vals(x, y, ro)[0]), maximum(min_max_vals(x, y, ro)[1])) - 1;
    score_ub(x, y) = select(is_corner(x, y), score_val, 0);
    score(x, y) = BoundaryConditions::constant_exterior(score_ub, 0, bounds)(x, y);

    Func result_ub("result_unbounded"), result("result");
    result_ub(x, y) = score(x, y) > max(threshold - 1, score(x - 1, y), score(x + 1, y),
                          score(x - 1, y - 1), score(x, y - 1), score(x + 1, y - 1),
                          score(x - 1, y + 1), score(x, y + 1), score(x + 1, y + 1));
    result(x, y) = BoundaryConditions::constant_exterior(
        result_ub, false, bounds)(x, y);

    // schedule
    Var yo("yo"), yi("yi");
    const int vec_size_short = get_host_target().natural_vector_size<short>();
    result.split(y, yo, yi, 32).vectorize(x, vec_size_short);
    score.compute_at(result, yi).store_at(result, yo).vectorize(x, vec_size_short);
    is_corner.compute_root();

    // test
    Buffer<bool> output(input.width(), input.height());
    result.realize(output);
    double time = Tools::benchmark(3, 3, [&] {
        result.realize(output);
    });
    cout << "time: " << time * 1e3 << " ms" << endl;
    
    return 0;
}