#! /usr/bin/python

# autogenerated on 2020-04-16 17:58

# Hough transform

import cv2
import numpy as np

from utils import *


def update_matrix(img, mat, i, j):
    for t in range(0, 180):
        tr = np.deg2rad(t)
        r = i * np.cos(tr) + j * np.sin(tr)
        _, diag = mat.shape
        r = int(np.round(r)) + (diag // 2)
        mat[t, r] += 1


def apply_lines(img, acc, diag):
    size_x, size_y = img.shape
    output = np.zeros((size_x, size_y, 3), dtype=np.uint8)
    for i in range(3):
        output[:, :, i] = img

    for t, row in enumerate(acc):
        for r, v in enumerate(row):
            if v:
                if t == 0 or t == 180:
                    for y in range(size_y):
                        x = r - diag
                        output[x, y] = (255, 0, 0)
                elif t == 90:
                    for x in range(size_x):
                        y = r - diag
                        output[x, y] = (255, 0, 0)
                else:
                    for y in range(size_y):
                        tr = np.deg2rad(t)
                        tmp = r - diag
                        x = (tmp - y * np.sin(tr)) / np.cos(tr)
                        x = int(x)
                        if 0 < x < size_x:
                            output[x, y] = (255, 0, 0)

    return output


def hough_transform(img, low, high, min_edge):
    edges = cv2.Canny(img, low, high)
    x, y = edges.shape
    diag = np.sqrt((x ** 2) + (y ** 2))
    diag = int(np.ceil(diag))
    accumulation = np.zeros((180, 2 * diag), dtype=np.int32)
    for i, row in enumerate(edges):
        for j, v in enumerate(row):
            if v:
                update_matrix(edges, accumulation, i, j)

    accumulation = np.where(accumulation > min_edge, accumulation, 0)
    output = apply_lines(img, accumulation, diag)
    return output


def main():
    image = load_image()

    output = hough_transform(image, 70, 100, 75)
    show_image(output, wait=60)


if __name__ == "__main__":
    main()
