Improve code layout of probabilistic hough

skimage/transform/_hough_transform.pyx

@@ -1,4 +1,5 @@
 cimport cython
+import math
 import numpy as np
 cimport numpy as np
 from random import randint
@@ -60,68 +61,77 @@ def _hough(np.ndarray img, np.ndarray[ndim=1, dtype=np.double_t] theta=None):
             accum[accum_idx, j] += 1
     return accum, theta, bins
 
-import math
 
 @cython.cdivision(True)
 @cython.boundscheck(False)
-def _probabilistic_hough(np.ndarray img, int value_threshold, int line_length, \
-        int line_gap, np.ndarray[ndim=1, dtype=np.double_t] theta=None):
+def _probabilistic_hough(np.ndarray img, int value_threshold,
+                         int line_length, int line_gap,
+                         np.ndarray[ndim=1, dtype=np.double_t] theta=None):
+
     if img.ndim != 2:
         raise ValueError('The input image must be 2D.')
-    # compute the array of angles and their sine and cosine
-    cdef np.ndarray[ndim=1, dtype=np.double_t] ctheta
-    cdef np.ndarray[ndim=1, dtype=np.double_t] stheta
-    # calculate thetas if none specified
+
     if theta is None:
-        theta = np.linspace(math.pi/2, -math.pi/2, 180)
-        theta =  math.pi/2-np.arange(180)/180.0* math.pi
-    ctheta = np.cos(theta)
-    stheta = np.sin(theta)
+        theta = PI_2 - np.arange(180) / 180.0 * 2 * PI_2
+
     cdef Py_ssize_t height = img.shape[0]
     cdef Py_ssize_t width = img.shape[1]
+
     # compute the bins and allocate the accumulator array
     cdef np.ndarray[ndim=2, dtype=np.int64_t] accum
-    cdef np.ndarray[ndim=2, dtype=np.uint8_t] mask = np.zeros((height, width), dtype=np.uint8)
-    cdef np.ndarray[ndim=2, dtype=np.int32_t] line_end = np.zeros((2, 2), dtype=np.int32)
+    cdef np.ndarray[ndim=1, dtype=np.double_t] ctheta, stheta
+    cdef np.ndarray[ndim=2, dtype=np.uint8_t] mask = \
+         np.zeros((height, width), dtype=np.uint8)
+    cdef np.ndarray[ndim=2, dtype=np.int32_t] line_end = \
+         np.zeros((2, 2), dtype=np.int32)
     cdef Py_ssize_t max_distance, offset, num_indexes, index
     cdef double a, b
-    cdef Py_ssize_t nidxs, nthetas, i, j, x, y, px, py, accum_idx
+    cdef Py_ssize_t nidxs, i, j, x, y, px, py, accum_idx
     cdef int value, max_value, max_theta
     cdef int shift = 16
     # maximum line number cutoff
     cdef Py_ssize_t lines_max = 2 ** 15
-    cdef Py_ssize_t xflag, x0, y0, dx0, dy0, dx, dy, gap, x1, y1, good_line, count
+    cdef Py_ssize_t xflag, x0, y0, dx0, dy0, dx, dy, gap, x1, y1, \
+                    good_line, count
+    cdef list lines = list()
+
     max_distance = 2 * <int>ceil((sqrt(img.shape[0] * img.shape[0] +
                                        img.shape[1] * img.shape[1])))
     accum = np.zeros((max_distance, theta.shape[0]), dtype=np.int64)
     offset = max_distance / 2
+    nthetas = theta.shape[0]
+
+    # compute sine and cosine of angles
+    ctheta = np.cos(theta)
+    stheta = np.sin(theta)
+
     # find the nonzero indexes
-    cdef np.ndarray[ndim=1, dtype=np.npy_intp] x_idxs, y_idxs
     y_idxs, x_idxs = np.nonzero(img)
-    num_indexes = y_idxs.shape[0] # x and y are the same shape
-    nthetas = theta.shape[0]
-    points = []
-    for i in range(num_indexes):
-        points.append((x_idxs[i], y_idxs[i]))
-    lines = []
-    # create mask of all non-zero indexes
-    for i in range(num_indexes):
-        mask[y_idxs[i], x_idxs[i]] = 1
+    points = zip(x_idxs, y_idxs)
+    # mask all non-zero indexes
+    mask[y_idxs, x_idxs] = 1
+
     while 1:
-        # select random non-zero point
+
+        # quit if no remaining points
         count = len(points)
         if count == 0:
             break
-        index = rand() % (count)
+
+        # select random non-zero point
+        index = rand() % count
         x = points[index][0]
         y = points[index][1]
         del points[index]
+
         # if previously eliminated, skip
         if not mask[y, x]:
             continue
+
         value = 0
-        max_value = value_threshold-1
+        max_value = value_threshold - 1
         max_theta = -1
+
         # apply hough transform on point
         for j in range(nthetas):
             accum_idx = <int>round((ctheta[j] * x + stheta[j] * y)) + offset
@@ -132,7 +142,9 @@ def _probabilistic_hough(np.ndarray img, int value_threshold, int line_length, \
                 max_theta = j
         if max_value < value_threshold:
             continue
-        # from the random point walk in opposite directions and find line beginning and end
+
+        # from the random point walk in opposite directions and find line
+        # beginning and end
         a = -stheta[max_theta]
         b = ctheta[max_theta]
         x0 = x
@@ -188,6 +200,7 @@ def _probabilistic_hough(np.ndarray img, int value_threshold, int line_length, \
         # confirm line length is sufficient
         good_line = abs(line_end[1, 1] - line_end[0, 1]) >= line_length or \
                     abs(line_end[1, 0] - line_end[0, 0]) >= line_length
+
         # pass 2: walk the line again and reset accumulator and mask
         for k in range(2):
             px = x0
@@ -221,6 +234,5 @@ def _probabilistic_hough(np.ndarray img, int value_threshold, int line_length, \
             lines.append(((line_end[0, 0], line_end[0, 1]), (line_end[1, 0], line_end[1, 1])))
             if len(lines) > lines_max:
                 return lines
-    return lines
-
 
+    return lines