cythonize dither_lookahead, still lots of low-hanging fruit

dither.py

@@ -7,9 +7,12 @@ from typing import Tuple
 
 from PIL import Image
 import numpy as np
-import pyximport; pyximport.install(language_level=3)
+import pyximport;
+
+pyximport.install(language_level=3)
 import dither_apply
 
+
 # TODO:
 # - only lookahead for 560px
 # - palette class
@@ -442,11 +445,16 @@ def dither_image(
         print(y)
         output_pixel_4bit = np.uint8(0)
         for x in range(screen.X_RES):
-        # for x in range(pattern.ORIGIN[1], pattern.ORIGIN[1] + screen.X_RES):
+            # for x in range(pattern.ORIGIN[1], pattern.ORIGIN[1] + screen.X_RES):
             input_pixel_rgb = np.copy(image_rgb[y, x, :])
+            options_4bit, options_rgb = lookahead_options(
+                screen, lookahead, output_pixel_4bit, x % 4)
+
             output_pixel_4bit, output_pixel_rgb = \
-                dither_lookahead(screen, image_rgb, dither, differ, x, y,
-                                 output_pixel_4bit, lookahead)
+                dither_apply.dither_lookahead(
+                    screen, image_rgb, dither, differ, x, y, options_4bit,
+                    options_rgb,
+                    lookahead)
             image_4bit[y, x] = output_pixel_4bit
             image_rgb[y, x, :] = output_pixel_rgb
             quant_error = input_pixel_rgb - output_pixel_rgb

dither_apply.pyx

@@ -11,9 +11,9 @@ cdef float clip(float a, float min_value, float max_value) nogil:
     return min(max(a, min_value), max_value)
 
 
-@cython.boundscheck(False)
-@cython.wraparound(False)
-def apply_one_line(float[:, :, ::1] pattern, int el, int er, int xl, int xr, int y, float[:, :, ::1] image,
+#@cython.boundscheck(False)
+#@cython.wraparound(False)
+def apply_one_line(float[:, :, ::1] pattern, int el, int er, int xl, int xr, int y, float[:, ::1] image,
                    float[::1] quant_error):
     cdef int i, j
     cdef float *error = <float *> malloc(pattern.shape[1] * quant_error.shape[0] * sizeof(float))
@@ -27,7 +27,7 @@ def apply_one_line(float[:, :, ::1] pattern, int el, int er, int xl, int xr, int
 
     for i in range(xr - xl):
         for j in range(3):
-            image[y, xl+i, j] = clip(image[y, xl + i, j] + error[(el + i) * quant_error.shape[0] + j], 0, 255)
+            image[xl+i, j] = clip(image[xl + i, j] + error[(el + i) * quant_error.shape[0] + j], 0, 255)
     free(error)
 
 
@@ -41,3 +41,65 @@ def apply(pattern, int el, int er, int xl, int xr, int et, int eb, int yt, int y
     # XXX extend image region to avoid need for boundary box clipping
     image[yt:yb, xl:xr, :] = np.clip(
         image[yt:yb, xl:xr, :] + error[et:eb, el:er, :], 0, 255)
+
+
+def x_dither_bounds(dither, screen, int x):
+    cdef int el = max(dither.ORIGIN[1] - x, 0)
+    cdef int er = min(dither.PATTERN.shape[1], screen.X_RES - 1 - x)
+
+    cdef int xl = x - dither.ORIGIN[1] + el
+    cdef int xr = x - dither.ORIGIN[1] + er
+
+    return el, er, xl, xr
+
+
+def y_dither_bounds(dither, screen, int y):
+    cdef int et = max(dither.ORIGIN[0] - y, 0)
+    cdef int eb = min(dither.PATTERN.shape[0], screen.Y_RES - 1 - y)
+
+    cdef int yt = y - dither.ORIGIN[0] + et
+    cdef int yb = y - dither.ORIGIN[0] + eb
+
+    return et, eb, yt, yb
+
+
+def dither_lookahead(
+        screen, float[:,:,::1] image_rgb, dither, differ, int x, int y, char[:, ::1] options_4bit,
+        float[:, :, ::1] options_rgb, int lookahead):
+    el, er, xl, xr = x_dither_bounds(dither, screen, x)
+
+    # X coord value of larger of dither bounding box or lookahead horizon
+    xxr = min(max(x + lookahead, xr), screen.X_RES)
+
+    # copies of input pixels so we can dither in bulk
+    # Leave enough space so we can dither the last of our lookahead pixels
+    lah_image_rgb = np.zeros(
+        (2 ** lookahead, lookahead + xr - xl, 3), dtype=np.float32)
+    lah_image_rgb[:, 0:xxr - x, :] = np.copy(image_rgb[y, x:xxr, :])
+
+    #cdef float[:, :, ::1] lah_image_rgb_view = lah_image_rgb
+    #cdef float[:, :, ::1] options_rgb_view = options_rgb
+
+    cdef float[:, ::] output_pixels
+    cdef float[:, ::1] quant_error
+
+    cdef int i, j
+    for i in range(xxr - x):
+        # options_rgb choices are fixed, but we can still distribute
+        # quantization error from having made these choices, in order to compute
+        # the total error
+        input_pixels = np.copy(lah_image_rgb[:, i, :])
+        output_pixels = options_rgb[:, i, :]
+        quant_error = input_pixels - output_pixels
+        # Don't update the input at position x (since we've already chosen
+        # fixed outputs), but do propagate quantization errors to positions >x
+        # so we can compensate for how good/bad these choices were
+        el, er, xl, xr = x_dither_bounds(dither, screen, i)
+        for j in range(2 ** lookahead):
+            apply_one_line(dither.PATTERN, el, er, xl, xr, 0, lah_image_rgb[j, :, :], quant_error[j])
+
+    error = differ.distance(np.clip(
+        lah_image_rgb[:, 0:lookahead, :], 0, 255), options_4bit)
+    total_error = np.sum(np.power(error, 2), axis=1)
+    best = np.argmin(total_error)
+    return options_4bit[best, 0], options_rgb[best, 0, :]