diff --git a/dither.py b/dither.py index f9bbc16..d75e829 100644 --- a/dither.py +++ b/dither.py @@ -3,6 +3,7 @@ import bz2 import functools import os.path import pickle +import time from typing import Tuple from PIL import Image @@ -282,8 +283,8 @@ class Dither: def apply(self, screen: Screen, image: np.ndarray, x: int, y: int, quant_error: np.ndarray, one_line=False): - #el, er, xl, xr = self.x_dither_bounds(screen, x) - #et, eb, yt, yb = self.y_dither_bounds(screen, y, one_line) + # el, er, xl, xr = self.x_dither_bounds(screen, x) + # et, eb, yt, yb = self.y_dither_bounds(screen, y, one_line) return dither_apply.apply(self, screen, x, y, image, quant_error) # error = self.PATTERN * quant_error.reshape((1, 1, 3)) # @@ -361,26 +362,6 @@ def open_image(screen: Screen, filename: str) -> np.ndarray: SRGBResize(im, (screen.X_RES, screen.Y_RES), Image.LANCZOS)) -@functools.lru_cache(None) -def lookahead_options(screen, lookahead, last_pixel_4bit, x): - options_4bit = np.empty((2 ** lookahead, lookahead), dtype=np.uint8) - options_rgb = np.empty((2 ** lookahead, lookahead, 3), dtype=np.float32) - for i in range(2 ** lookahead): - output_pixel_4bit = last_pixel_4bit - for j in range(lookahead): - xx = x + j - palette_choices_4bit, palette_choices_rgb = \ - screen.pixel_palette_options(output_pixel_4bit, xx) - output_pixel_4bit = palette_choices_4bit[(i & (1 << j)) >> j] - output_pixel_rgb = np.array( - palette_choices_rgb[(i & (1 << j)) >> j]) - # XXX copy - options_4bit[i, j] = output_pixel_4bit - options_rgb[i, j, :] = np.copy(output_pixel_rgb) - - return options_4bit, options_rgb - - def dither_lookahead( screen: Screen, image_rgb: np.ndarray, dither: Dither, differ: ColourDistance, x, y, last_pixel_4bit, lookahead @@ -439,7 +420,7 @@ def dither_image( print(y) output_pixel_4bit = np.uint8(0) for x in range(screen.X_RES): - input_pixel_rgb = np.copy(image_rgb[y, x, :]) + input_pixel_rgb = image_rgb[y, x, :] options_4bit, options_rgb = lookahead_options( screen, lookahead, output_pixel_4bit, x % 4) @@ -448,9 +429,9 @@ def dither_image( screen, image_rgb, dither, differ, x, y, options_4bit, options_rgb, lookahead) + quant_error = input_pixel_rgb - output_pixel_rgb image_4bit[y, x] = output_pixel_4bit image_rgb[y, x, :] = output_pixel_rgb - quant_error = input_pixel_rgb - output_pixel_rgb dither_apply.apply(dither, screen, x, y, image_rgb, quant_error) return image_4bit, image_rgb @@ -478,8 +459,11 @@ def main(): differ = CIE2000Distance() - output_4bit, output_rgb = dither_image(screen, image, dither, differ, - lookahead=args.lookahead) + start = time.time() + output_4bit, output_rgb = dither_apply.dither_image(screen, image, dither, + differ, + lookahead=args.lookahead) + print(time.time() - start) screen.pack(output_4bit) out_image = Image.fromarray(linear_to_srgb(output_rgb).astype(np.uint8)) diff --git a/dither_apply.pyx b/dither_apply.pyx index b45b03d..425f679 100644 --- a/dither_apply.pyx +++ b/dither_apply.pyx @@ -3,7 +3,7 @@ cimport cython import numpy as np # from cython.parallel import prange -# from cython.view cimport array as cvarray +from cython.view cimport array as cvarray # from libc.stdlib cimport malloc, free @@ -27,7 +27,7 @@ cdef void apply_one_line(float[:, :, ::1] pattern, int xl, int xr, float[:, ::1] @cython.boundscheck(False) @cython.wraparound(False) -def apply(dither, screen, int x, int y, float [:, :, ::1]image, float[::1] quant_error): +cdef apply(dither, screen, int x, int y, float [:, :, ::1]image, float[] quant_error): cdef int i, j, k # XXX only need 2 dimensions now @@ -89,8 +89,8 @@ def dither_lookahead( cdef int i, j, k, l - cdef float[:, :, ::1] lah_image_rgb = np.empty( - (2 ** lookahead, lookahead + xr - xl, 3), dtype=np.float32) + # XXX malloc + cdef float[:, :, ::1] lah_image_rgb = cvarray((2 ** lookahead, lookahead + xr - xl, 3), itemsize=sizeof(float), format="f") for i in range(2**lookahead): # Copies of input pixels so we can dither in bulk for j in range(xxr - x): @@ -143,3 +143,62 @@ def dither_lookahead( best = i return options_4bit[best, 0], options_rgb[best, 0, :] + +import functools + + +@functools.lru_cache(None) +def lookahead_options(screen, lookahead, last_pixel_4bit, x): + options_4bit = np.empty((2 ** lookahead, lookahead), dtype=np.uint8) + options_rgb = np.empty((2 ** lookahead, lookahead, 3), dtype=np.float32) + for i in range(2 ** lookahead): + output_pixel_4bit = last_pixel_4bit + for j in range(lookahead): + xx = x + j + palette_choices_4bit, palette_choices_rgb = \ + screen.pixel_palette_options(output_pixel_4bit, xx) + output_pixel_4bit = palette_choices_4bit[(i & (1 << j)) >> j] + output_pixel_rgb = np.array( + palette_choices_rgb[(i & (1 << j)) >> j]) + # XXX copy + options_4bit[i, j] = output_pixel_4bit + options_rgb[i, j, :] = np.copy(output_pixel_rgb) + + return options_4bit, options_rgb + + +@cython.boundscheck(False) +@cython.wraparound(False) +def dither_image( + screen, float [:, :, ::1] image_rgb, dither, differ, int lookahead): + image_4bit = np.empty( + (image_rgb.shape[0], image_rgb.shape[1]), dtype=np.uint8) + + cdef int yres = screen.Y_RES + cdef int xres = screen.X_RES + + cdef int y, x, i + cdef float[3] quant_error + cdef (unsigned char)[:, ::1] options_4bit + cdef float[:, :, ::1] options_rgb + + for y in range(yres): + # print(y) + output_pixel_4bit = np.uint8(0) + for x in range(xres): + input_pixel_rgb = 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, options_4bit, + options_rgb, lookahead) + for i in range(3): + quant_error[i] = input_pixel_rgb[i] - output_pixel_rgb[i] + image_rgb[y, x, i] = output_pixel_rgb[i] + image_4bit[y, x] = output_pixel_4bit + apply(dither, screen, x, y, image_rgb, quant_error) + + return image_4bit, np.array(image_rgb) +