themblem/alg/angle.cpp

#include <algorithm>
#include "libqr.h"

static void clear_connected(Mat &bin, Point p)
{
    vector<Point> q;
    q.push_back(p);
    while (q.size()) {
        auto p = q[q.size() - 1];
        q.pop_back();
        bin.at<uint8_t>(p.y, p.x) = 0;
        for (int i = -1; i <= 1; i++) {
            for (int j = -1; j <= 1; j++) {
                int nx = p.x + i;
                int ny = p.y + j;
                if (nx < 0 || nx >= bin.cols || ny < 0 || ny >= bin.rows) {
                    continue;
                }
                if (bin.at<bool>(ny, nx)) {
                    q.push_back(Point(nx, ny));
                }
            }
        }
    }
}

static
vector<Point> find_points(Mat bin)
{
    vector<Point> ret;
    for (int x = 0; x < bin.cols; x++) {
        for (int y = 0; y < bin.rows; y++) {
            auto p = bin.at<uint8_t>(y, x);
            if (!p) continue;
            auto point = Point(x, y);
            ret.push_back(point);
            clear_connected(bin, point);
        }
    }
    return ret;
}

static
bool in_center(Mat &bin, Point &p)
{
    int margin = bin.rows * 2 / 10;

    return p.x > margin && p.x < bin.cols - margin && p.y > margin && p.y <= bin.rows - margin;
}

static
float distance(Point &p, Point &q)
{
    auto xdiff = p.x - q.x;
    auto ydiff = p.y - q.y;
    return xdiff * xdiff + ydiff * ydiff;
}

static
int find_closest(Point &p, vector<Point> &points, bool left, bool top)
{
    int ret = -1;
    for (int ii = 0; ii < points.size(); ii++) {
        auto i = points[ii];
        if (i.x == p.x && i.y == p.y) continue;
        if (left && i.x > p.x) continue;
        if (top && i.y > p.y) continue;
        if (!left && i.x <= p.x) continue;
        if (!top && i.y < p.y) continue;
        if (ret < 0 || distance(p, points[ret]) > distance(p, i)) {
            ret = ii;
        }
    }
    return ret;
}

static
float find_angle(Point &p, vector<Point> &points)
{
    // Find 4 dots in 4 quadrant (if any)
    // Then find 2 closest on y axis
    // Then calculate angle between those two

    auto topleft = find_closest(p, points, true, true);
    auto bottomright = find_closest(p, points, false, false);

    if (topleft < 0 || bottomright < 0)
        return -1;
    auto a = points[topleft];
    auto b = points[bottomright];
    printf("point %d %d top left %d %d, bottom right %d %d\n", p.x, p.y, a.x, a.y, b.x, b.y);
    if (a.y == b.y) return 0;
    auto ret = atan((b.x - a.x) / (b.y - a.y)) * 180.0 / CV_PI;
    if (ret < 0) ret += 90;
    if (ret > 45) ret = 90 - ret;
    return ret;
}

static
void angle_stat(vector<float> angles, float &median, float &variance)
{
    std::sort(angles.begin(), angles.end());
    float sum = 0;
    for (auto x: angles) {
        sum += x;
    }
    auto mid = angles.size() / 2;
    median = angles[mid];
    auto avg = sum / angles.size();
    variance = 0;
    for (auto x: angles) {
        auto diff = x - avg;
        variance += diff * diff;
    }
}

float hough_lines_angle(Mat &img, string &err)
{
    show(img);
    vector<Vec3f> lines;
    HoughLines(img, lines, 1, CV_PI / 180, 6, 0, 0);
    for (auto x: lines) {
        printf("line: %.1f %.1f %.1f\n", x[0], x[1] * 180.0 / CV_PI, x[2]);
    }
    if (!lines.size()) {
        err = "cannot find lines in image";
        return -1;
    }

    int total_weight = 0;
    for (int i = 0; i < lines.size() && i < 5; i++) {
        total_weight += lines[i][2];
    }
    int acc = 0;
    float ret = 0;
    for (int i = 0; i < lines.size(); i++) {
        acc += lines[i][2];
        if (acc >= total_weight / 2) {
            ret = lines[i][1] * 180.0 / CV_PI;
            break;
        }
    }
    while (ret < 0) {
        ret += 90;
    }
    while (ret > 90) {
        ret -= 90;
    }
    if (ret > 45) {
        ret = 90 - ret;
    }
    printf("angle: %f\n", ret);
    return ret;
}

float emblem_detect_angle(Mat &gray, string &err)
{
    Mat bin;
    const int min_points = 30;
    vector<Point> points;
    auto kernel = getStructuringElement(MORPH_ELLIPSE, Size(3, 3));

    adaptiveThreshold(gray, bin, 255, ADAPTIVE_THRESH_GAUSSIAN_C, THRESH_BINARY_INV, 11, 2);
    while (true) {
        // In this loop we erode a "full" image in order to get enough detached components
        points = find_points(bin.clone());
        printf("points: %zu\n", points.size());
        if (points.size() == 0) {
            err = "cannot find enough points";
            return -1;
        }
        if (points.size() > min_points) {
            break;
        }
        erode(bin, bin, kernel);
    }

    while (true) {
        // In this loop we further erode a "lean" image in order to get clarity until it's too much
        Mat eroded;
        erode(bin, eroded, kernel);
        auto tmp = find_points(eroded.clone());
        if (tmp.size() < min_points) {
            printf("too much\n");
            break;
        }
        bin = eroded.clone();
    }

    return hough_lines_angle(bin, err);

    vector<float> angles;
    for (auto p: points) {
        if (!in_center(bin, p)) {
            continue;
        }
        auto angle = find_angle(p, points);
        if (angle >= 0) {
            printf("found angle %f\n", angle);
            angles.push_back(angle);
        }
    }
    if (!angles.size()) {
        err = "cannot find point to calculate angle";
        return -1;
    }
    float med, var;
    angle_stat(angles, med, var);
    printf("med: %f, var: %f\n", med, var);
    return med;
}