Python Hough Lines implementation, making it time efficient

Question

So I'm trying to implement the hough transform lines algorithm in python, and I'm finding it hard to make it time efficient.

This is my implementation:

import numpy as np
def houghLines(edges, dTheta, threshold):
    imageShape = edges.shape
    imageDiameter = (imageShape[0]**2 + imageShape[1]**2)**0.5
    rhoRange = [i for i in range(int(imageDiameter)+1)]
    thetaRange = [dTheta*i for i in range(int(-np.pi/(2*dTheta)), int(np.pi/dTheta))]
    cosTheta = [np.cos(theta) for theta in thetaRange]
    sinTheta = [np.sin(theta) for theta in thetaRange]
    countMatrix = np.zeros([len(rhoRange), len(thetaRange)])
    eds = [(x,y) for (x,y), value in np.ndenumerate(edges) if value > 0]
    for thetaIndex in range(len(thetaRange)):
        theta = thetaRange[thetaIndex]
        cos = cosTheta[thetaIndex]
        sin = sinTheta[thetaIndex]
        for x, y in eds:
            targetRho = x*cos + y*sin
            closestRhoIndex = int(round(targetRho))
            countMatrix[closestRhoIndex, thetaIndex] += 1
    lines = [(p,thetaRange[t]) for (p,t), value in np.ndenumerate(countMatrix) if value > threshold]
    return lines

It works but it is very slow, 100 times slower than the opencv implementation.

How can I improve it?

Answer 1

The answer was to use numba. This is what the code looks like now:

import numpy as np
from numba import jit
@jit(nopython=True)
def houghLines(edges, dTheta, threshold):
    imageShape = edges.shape
    imageDiameter = (imageShape[0]**2 + imageShape[1]**2)**0.5
    rhoRange = [i for i in range(int(imageDiameter)+1)]
    thetaRange = [dTheta*i for i in range(int(-np.pi/(2*dTheta)), int(np.pi/dTheta))]
    cosTheta = []
    sinTheta = []
    for theta in thetaRange:
        cosTheta.append(np.cos(theta))
        sinTheta.append(np.sin(theta))
    countMatrixSize = (len(rhoRange), len(thetaRange))
    countMatrix = np.zeros(countMatrixSize)

    eds = []
    for (x,y), value in np.ndenumerate(edges):
        if value > 0:
            eds.append((x,y))

    for thetaIndex in range(len(thetaRange)):
        theta = thetaRange[thetaIndex]
        cos = cosTheta[thetaIndex]
        sin = sinTheta[thetaIndex]
        for x, y in eds:
            targetRho = x*cos + y*sin
            closestRhoIndex = int(round(targetRho))
            countMatrix[closestRhoIndex, thetaIndex] += 1
    lines = []
    for (p,t), value in np.ndenumerate(countMatrix):
        if value > threshold:
            lines.append((p,thetaRange[t]))
    return lines

This made it at least 50 times faster.