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?
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.