core: Optimize the Sobel edge detector

tdh8316 · Feb 21, 2020 · a120f4d · a120f4d
1 parent bf01434
commit a120f4d
Showing 1 changed file with 33 additions and 36 deletions.
diff --git a/triangler/edges.py b/triangler/edges.py
@@ -48,7 +48,7 @@ def get_edge_points(self, blur: int, sampling: SampleMethod) -> ndarray:
         elif self.edge_method is EdgeMethod.ENTROPY:
             edges = Entropy.compute(self.img)
         elif self.edge_method is EdgeMethod.SOBEL:
-            edges = Sobel.compute(self.img, k_size=5)
+            edges = Sobel.compute(self.img)
         else:
             raise ValueError(
                 "Unexpected edge processing method: {}\n"
@@ -132,43 +132,40 @@ def compute(img: ndarray, bal=0.1) -> ndarray:
 class Sobel(object):
     @staticmethod
     @numba.jit(parallel=True, fastmath=True)
-    def compute(img: ndarray, k_size: int = 3) -> ndarray:
-        im = img.astype(np.float)
-        width, height, c = im.shape
-        if c > 1:
-            img = 0.2126 * im[:, :, 0] + 0.7152 * im[:, :, 1] + 0.0722 * im[:, :, 2]
-        else:
-            img = im
-
-        assert k_size == 3 or k_size == 5
+    def compute(img: ndarray) -> ndarray:
+        _img_as_float = img.astype(np.float)
+        width, height, c = _img_as_float.shape
+        _img = (
+            0.2126 * _img_as_float[:, :, 0]
+            + 0.7152 * _img_as_float[:, :, 1]
+            + 0.0722 * _img_as_float[:, :, 2]
+            if c > 1
+            else _img_as_float
+        )
 
-        if k_size == 3:
-            kh = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=np.float)
-            kv = np.array([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], dtype=np.float)
-        else:
-            kh = np.array(
-                [
-                    [-1, -2, 0, 2, 1],
-                    [-4, -8, 0, 8, 4],
-                    [-6, -12, 0, 12, 6],
-                    [-4, -8, 0, 8, 4],
-                    [-1, -2, 0, 2, 1],
-                ],
-                dtype=np.float,
-            )
-            kv = np.array(
-                [
-                    [1, 4, 6, 4, 1],
-                    [2, 8, 12, 8, 2],
-                    [0, 0, 0, 0, 0],
-                    [-2, -8, -12, -8, -2],
-                    [-1, -4, -6, -4, -1],
-                ],
-                dtype=np.float,
-            )
+        kh = np.array(
+            [
+                [-1, -2, 0, 2, 1],
+                [-4, -8, 0, 8, 4],
+                [-6, -12, 0, 12, 6],
+                [-4, -8, 0, 8, 4],
+                [-1, -2, 0, 2, 1],
+            ],
+            dtype=np.float,
+        )
+        kv = np.array(
+            [
+                [1, 4, 6, 4, 1],
+                [2, 8, 12, 8, 2],
+                [0, 0, 0, 0, 0],
+                [-2, -8, -12, -8, -2],
+                [-1, -4, -6, -4, -1],
+            ],
+            dtype=np.float,
+        )
 
-        gx = convolve2d(img, kh, mode="same", boundary="symm")
-        gy = convolve2d(img, kv, mode="same", boundary="symm")
+        gx = convolve2d(_img, kh, mode="same", boundary="symm")
+        gy = convolve2d(_img, kv, mode="same", boundary="symm")
 
         g = np.sqrt(gx * gx + gy * gy)
         g *= 255.0 / np.max(g)