Entire palette transfer class

whole_palette.py

class EntirePalette:
    def __init__(self, chunk_size=1024):
        self.chunk_size = chunk_size
        self.source_pixels = None

    def _preprocess(self, image):
        # ... omitted, same as KMeansReducedPalette

    def fit(self, image):
        self.source_pixels = np.unique(self._preprocess(image), axis=0)

    def recolor(self, image):
        image_shape = image.shape
        image = self._preprocess(image)
        image_colors = np.unique(image, axis=0)

        self_da = da.from_array(
            self.source_pixels.astype(np.long), chunks=(self.chunk_size, 3)
        )
        other_da = da.from_array(
            image_colors.reshape(-1, 1, 3).astype(np.long),
            chunks=(self.chunk_size, 1, 3)
        )

        idxs_da = da.square(other_da - self_da).sum(axis=2).argmin(axis=1)

        with ProgressBar():
            mapping_idx = idxs_da.compute()

        colormap = {tuple(a): tuple(b)
                    for a, b in
                    zip(image_colors, self.source_colors[mapping_idx])}

        image_recolored = np.array([colormap[tuple(rgb)] for rgb in image])
        return image_recolored.reshape(image_shape).astype(np.uint8)