rasterization with numba - faster

Author: RomanJuranek

pclines/__init__.py

@@ -1 +1 @@
-from .accumulator import accumulate, find_peaks, lines
+from .accumulator import PCLines

pclines/accumulator.py

@@ -31,10 +31,10 @@
 
 
 import numpy as np
-#import numba as nb
 from skimage.feature import peak_local_max
 from skimage.morphology.grey import erosion, dilation
 
+from .rasterizer import polys
 
 def _linear_transform(src, dst):
     """ Parameters of a linear transform from range specifications """
@@ -112,7 +112,7 @@ def transform(self, x):
         """
         _check_points(x)
         x0,x1 = np.split(x,2,axis=1)
-        return np.concatenate([self.norm_u(x1), self.norm_v(x0), self.norm_w(x1)], axis=1)
+        return np.concatenate([self.norm_u(x1), self.norm_v(x0), self.norm_w(x1)], axis=1).astype("f")
 
     def inverse(self, l):
         """
@@ -145,25 +145,23 @@ def insert(self, x, weight=None):
         if weight is None:
             weight = np.ones(n, np.float32)
 
+        if weight.dtype != np.float32:
+            weight = weight.astype(np.float32)
+
         valid = self.valid_points(p)
         p = p[valid]
         weight = weight[valid.flat]
 
-        d = self.d
-        c = np.arange(2*d-1,dtype="i")
-        for (u,v,w),wt in zip(p, weight):  # remove space wraping
-            t_part = np.linspace(u,v,d)
-            s_part = np.linspace(v,w,d)
-            r = np.concatenate([t_part, s_part[1:]]).astype("i")
-            self.A[r,c] += wt  # TODO add weight
+        polys(p, weight, self.A)
+        
 
     def find_peaks(self, t=0.8, prominence=2, min_dist=1):
         """
         Retrieve locations with prominent local maxima in the accumulator
         """
         p = dilation(self.A+1)/erosion(self.A+1)
-        peaks = peak_local_max(self.A, threshold_rel=t, min_distance=min_dist)
+        peaks = peak_local_max(self.A, threshold_rel=t, min_distance=min_dist, exclude_border=False)
         r,c = peaks[:,0], peaks[:,1]
-        value = A[r,c]
+        value = self.A[r,c]
         valid = p[r,c] > prominence
         return peaks[valid], value[valid]

pclines/rasterizer.py

@@ -0,0 +1,25 @@
+import numpy as np
+import numba as nb
+
+@nb.njit
+def line(xr, yr, w, arr):
+    x0,x1 = xr
+    y0,y1 = yr
+    dy = (y1-y0) / (x1-x0)
+    for step,x in enumerate(range(x0, x1)):
+        y = int(y0 + (dy * step))
+        arr[y,x] += w
+
+@nb.njit
+def polys(p, weight, arr):
+    n = weight.size
+    d = arr.shape[0]
+    for i in range(n):
+        u,v,w = p[i]
+        wt = weight[i]
+        # rasterize (u,v), (0,d-1)
+        line((0, d), (u, v), wt, arr)
+        # rasterize (v,w), (d, 2d-1)
+        line((d, 2*d-1), (v, w), wt, arr)
+    
+