First implementation of using k-means clustering in RGB space to dither a 320x200 SHR image.

convert.py

@@ -7,6 +7,7 @@ import time
 import colour
 from PIL import Image
 import numpy as np
+from sklearn.cluster import KMeans
 
 import dither as dither_pyx
 import dither_pattern
@@ -19,6 +20,55 @@ import screen as screen_py
 # - support LR/DLR
 # - support HGR
 
+def _to_pixel(float_array):
+    return tuple(np.clip(float_array.astype(np.uint8), 0, 255))
+
+
+def cluster_palette(image: Image):
+    # TODO: cluster in CAM16-UCS space
+    colours = np.asarray(image).reshape((-1, 3))
+    kmeans = KMeans(n_clusters=16)
+    kmeans.fit_predict(colours)
+    palette = kmeans.cluster_centers_
+
+    pal_image = Image.new('P', (1, 1), 0)
+    pal_image.putpalette(palette.reshape(-1).astype(np.uint8))
+
+    working_image = np.asarray(image).astype(np.float32)
+    for y in range(200):
+        print(y)
+        for x in range(320):
+            pixel = working_image[y, x]
+
+            best_distance = 1e9
+            best_colour = None
+            for colour in palette:
+                distance = np.sum(np.power(colour - pixel, 2))
+                if distance < best_distance:
+                    best_distance = distance
+                    best_colour = colour
+            quant_error = pixel - best_colour
+
+            # Floyd-Steinberg dither
+            # 0 * 7
+            # 3 5 1
+            working_image[y, x] = best_colour
+            if x < 319:
+                working_image[y, x + 1] = np.clip(
+                    working_image[y, x + 1] + quant_error * (7 / 16), 0, 255)
+            if y < 199:
+                working_image[y + 1, x] = np.clip(
+                    working_image[y + 1, x] + quant_error * (5 / 16), 0, 255)
+                if x < 319:
+                    working_image[y + 1, x + 1] = np.clip(
+                        working_image[y + 1, x + 1] + quant_error * (1 / 16),
+                        0, 255)
+                if x > 0:
+                    working_image[y + 1, x - 1] = np.clip(
+                        working_image[y + 1, x - 1] + quant_error * (3 / 16), 0,
+                        255)
+    return working_image
+
 
 def main():
     parser = argparse.ArgumentParser()
@@ -63,8 +113,8 @@ def main():
     if args.lookahead < 1:
         parser.error('--lookahead must be at least 1')
 
-    palette = palette_py.PALETTES[args.palette]()
-    screen = screen_py.DHGRScreen(palette)
+    # palette = palette_py.PALETTES[args.palette]()
+    screen = screen_py.SHR320Screen()
 
     # Conversion matrix from RGB to CAM16UCS colour values.  Indexed by
     # 24-bit RGB value
@@ -73,39 +123,42 @@ def main():
     # Open and resize source image
     image = image_py.open(args.input)
     if args.show_input:
-        image_py.resize(image, screen.X_RES, screen.Y_RES * 2,
+        image_py.resize(image, screen.X_RES, screen.Y_RES,
                         srgb_output=True).show()
     rgb = np.array(
         image_py.resize(image, screen.X_RES, screen.Y_RES,
                         gamma=args.gamma_correct)).astype(np.float32) / 255
 
-    dither = dither_pattern.PATTERNS[args.dither]()
-    bitmap = dither_pyx.dither_image(
-        screen, rgb, dither, args.lookahead, args.verbose, rgb_to_cam16)
+    output_rgb = cluster_palette(Image.fromarray((rgb * 255).astype(np.uint8)))
+    output_srgb = image_py.linear_to_srgb(output_rgb).astype(np.uint8)
+
+    # dither = dither_pattern.PATTERNS[args.dither]()
+    # bitmap = dither_pyx.dither_image(
+    #     screen, rgb, dither, args.lookahead, args.verbose, rgb_to_cam16)
 
     # Show output image by rendering in target palette
-    output_palette_name = args.show_palette or args.palette
-    output_palette = palette_py.PALETTES[output_palette_name]()
-    output_screen = screen_py.DHGRScreen(output_palette)
-    if output_palette_name == "ntsc":
-        output_srgb = output_screen.bitmap_to_image_ntsc(bitmap)
-    else:
-        output_srgb = image_py.linear_to_srgb(
-            output_screen.bitmap_to_image_rgb(bitmap)).astype(np.uint8)
+    # output_palette_name = args.show_palette or args.palette
+    # output_palette = palette_py.PALETTES[output_palette_name]()
+    # output_screen = screen_py.DHGRScreen(output_palette)
+    # if output_palette_name == "ntsc":
+    #     output_srgb = output_screen.bitmap_to_image_ntsc(bitmap)
+    # else:
+    #     output_srgb = image_py.linear_to_srgb(
+    #         output_screen.bitmap_to_image_rgb(bitmap)).astype(np.uint8)
     out_image = image_py.resize(
-        Image.fromarray(output_srgb), screen.X_RES, screen.Y_RES * 2,
+        Image.fromarray(output_srgb), screen.X_RES, screen.Y_RES,
         srgb_output=True)
 
     if args.show_output:
         out_image.show()
 
     # Save Double hi-res image
-    outfile = os.path.join(os.path.splitext(args.output)[0] + "-preview.png")
-    out_image.save(outfile, "PNG")
-    screen.pack(bitmap)
-    with open(args.output, "wb") as f:
-        f.write(bytes(screen.aux))
-        f.write(bytes(screen.main))
+    # outfile = os.path.join(os.path.splitext(args.output)[0] + "-preview.png")
+    # out_image.save(outfile, "PNG")
+    # screen.pack(bitmap)
+    # with open(args.output, "wb") as f:
+    #     f.write(bytes(screen.aux))
+    #     f.write(bytes(screen.main))
 
 
 if __name__ == "__main__":

screen.py

@@ -5,6 +5,11 @@ import numpy as np
 import palette as palette_py
 
 
+class SHR320Screen:
+    X_RES = 320
+    Y_RES = 200
+
+
 class DHGRScreen:
     X_RES = 560
     Y_RES = 192