From 173c283369352497e8b2658a8c634d999799c5e0 Mon Sep 17 00:00:00 2001 From: kris Date: Tue, 9 Nov 2021 11:23:25 +0000 Subject: [PATCH] First implementation of using k-means clustering in RGB space to dither a 320x200 SHR image. --- convert.py | 95 ++++++++++++++++++++++++++++++++++++++++++------------ screen.py | 5 +++ 2 files changed, 79 insertions(+), 21 deletions(-) diff --git a/convert.py b/convert.py index 7cd34d1..d4cd028 100644 --- a/convert.py +++ b/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__": diff --git a/screen.py b/screen.py index 9831807..d164be7 100644 --- a/screen.py +++ b/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