Optimize more

convert.py

@@ -52,8 +52,8 @@ def cluster_palette(image: Image):
     with colour.utilities.suppress_warnings(colour_usage_warnings=True):
         colours_cam = colour.convert(colours_rgb, "RGB",
                                      "CAM16UCS").astype(np.float32)
-    palettes_rgb = {}
-    palettes_cam = {}
+    palettes_rgb = np.empty((16, 16, 3), dtype=np.float32)
+    palettes_cam = np.empty((16, 16, 3), dtype=np.float32)
     for palette_idx in range(16):
         p_lower = max(palette_idx - 2, 0)
         p_upper = min(palette_idx + 2, 16)
@@ -75,7 +75,7 @@ def cluster_palette(image: Image):
             palette_pixels, 16).initialize()
         kmedians_instance = kmedians(palette_pixels, initial_centers)
         kmedians_instance.process()
-        palettes_cam[palette_idx] = np.array(
+        palettes_cam[palette_idx, :, :] = np.array(
             kmedians_instance.get_medians()).astype(np.float32)
 
         # palette_colours = collections.defaultdict(list)
@@ -111,7 +111,7 @@ def cluster_palette(image: Image):
                                          "RGB")
             # SHR colour palette only uses 4-bit values
             palette_rgb = np.round(palette_rgb * 15) / 15
-            palettes_rgb[palette_idx] = palette_rgb.astype(np.float32)
+            palettes_rgb[palette_idx, :, :] = palette_rgb.astype(np.float32)
     # print(palettes_rgb)
 
     # For each line, pick the palette with lowest total distance
@@ -189,8 +189,9 @@ def main():
     # print(palette_rgb)
     # screen.set_palette(0, (image_py.linear_to_srgb_array(palette_rgb) *
     #                        15).astype(np.uint8))
-    for i, p in palettes_rgb.items():
-        screen.set_palette(i, (np.round(p * 15)).astype(np.uint8))
+    for i in range(16):
+        screen.set_palette(i, (np.round(palettes_rgb[i, :, :] * 15)).astype(
+            np.uint8))
 
     output_4bit, line_to_palette = dither_pyx.dither_shr(
         rgb, palettes_cam, palettes_rgb, rgb_to_cam16)

dither.pyx

@@ -339,7 +339,7 @@ import colour
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
-def dither_shr(float[:, :, ::1] working_image, object palettes_cam, object palettes_rgb, float[:,::1] rgb_to_cam16ucs):
+def dither_shr(float[:, :, ::1] working_image, float[:, :, ::1] palettes_cam, float[:, :, ::1] palettes_rgb, float[:,::1] rgb_to_cam16ucs):
     cdef int y, x, idx, best_colour_idx, best_palette
     cdef float best_distance, distance
     cdef float[::1] best_colour_rgb, pixel_cam, colour_rgb, colour_cam
@@ -351,12 +351,10 @@ def dither_shr(float[:, :, ::1] working_image, object palettes_cam, object palet
 
     cdef float[:, ::1] line_cam = np.zeros((320, 3), dtype=np.float32)
 
-    line_to_palette = {}
+    cdef int[::1] line_to_palette = np.zeros(200, dtype=np.int32)
     best_palette = 15
     for y in range(200):
         print(y)
-        # palette_rgb = palettes_rgb[line_to_palette[y]]
-
         for x in range(320):
             colour_cam = convert_rgb_to_cam16ucs(
                 rgb_to_cam16ucs, working_image[y,x,0], working_image[y,x,1], working_image[y,x,2])
@@ -364,7 +362,7 @@ def dither_shr(float[:, :, ::1] working_image, object palettes_cam, object palet
 
         best_palette = best_palette_for_line(line_cam, palettes_cam, <int>(y * 16 / 200), best_palette)
         # print("-->", best_palette)
-        palette_rgb = palettes_rgb[best_palette]
+        palette_rgb = palettes_rgb[best_palette, :, :]
         line_to_palette[y] = best_palette
 
         for x in range(320):
@@ -511,23 +509,27 @@ def k_means_with_fixed_centroids(
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
-cdef int best_palette_for_line(float [:, ::1] line_cam, object palettes_cam, int base_palette_idx, int last_palette_idx):
-    cdef int palette_idx, best_palette_idx
+cdef int best_palette_for_line(float [:, ::1] line_cam, float[:, :, ::1] palettes_cam, int base_palette_idx, int last_palette_idx) nogil:
+    cdef int palette_idx, best_palette_idx, palette_entry_idx, pixel_idx
     cdef float best_total_dist, total_dist, best_pixel_dist, pixel_dist
     cdef float[:, ::1] palette_cam
     cdef float[::1] pixel_cam, palette_entry
 
     best_total_dist = 1e9
     best_palette_idx = -1
-    for palette_idx, palette_cam in palettes_cam.items():
+    cdef int line_size = line_cam.shape[0]
+    for palette_idx in range(16):
+        palette_cam = palettes_cam[palette_idx, :, :]
         if palette_idx < (base_palette_idx - 1) or palette_idx > (base_palette_idx + 1):
             continue
         if palette_idx == last_palette_idx:
             continue
         total_dist = 0
         best_pixel_dist = 1e9
-        for pixel_cam in line_cam:
-            for palette_entry in palette_cam:
+        for pixel_idx in range(line_size):
+            pixel_cam = line_cam[pixel_idx]
+            for palette_entry_idx in range(16):
+                palette_entry = palette_cam[palette_entry_idx, :]
                 pixel_dist = colour_distance_squared(pixel_cam, palette_entry)
                 if pixel_dist < best_pixel_dist:
                     best_pixel_dist = pixel_dist