Color quantization?

username: ZimoNitrome
cooked: <p>Is there anything similar to <a href="https://www.geeksforgeeks.org/python-pil-image-quantize-method/" rel="noopener nofollow ugc">PIL’s Image.quantize()</a> in PyTorch?</p>
<p>Quantize reduces the number of colors in an image, with an option to a specific set of target colors. I am sure a similar function could be created for tensors but would the gradient be maintained? If yes, any tips on implementing it in an efficient way?</p>

username: ZimoNitrome
cooked: <p>For now I created the following methods:</p>
<pre><code class="lang-auto">def distance(color_a, color_b):
    differences = torch.stack([(c_a - c_b) ** 2 for c_a, c_b in zip(color_a, color_b)])
    return torch.sum(differences, dim=-2)**0.5 # Euclidian distance
    
def quantize(image, palette):
    # Assume RGB
    flat_img = image.view(1, 3,-1) # [C, H, W] -&gt; [1, C, H*W]
    img_per_palette_color = torch.cat(len(palette)*[flat_img]) # [1, C, H*W] -&gt; [n_colors, C, H*W]
    distance_per_pixel = distance(img_per_palette_color, palette.unsqueeze(-1)) # [n_colors, C, H*W] -&gt; [n_colors, H*W]
    color_indices = torch.argmin(distance_per_pixel, dim=0) # [n_colors, H*W] -&gt; [H*W]
    new_colors = palette[color_indices].T # [H*W] -&gt; [C, H*W]
    return new_colors.view(image.shape) # [C, H*W] -&gt; [C, H, W]

colors = torch.tensor([
    [0, 75, 135],
    [255, 205, 0]
])
colors_normalized = colors/255

original = TF.to_tensor(Image.open(r"[PATH]").convert("RGB"))
display(TF.to_pil_image(original))

quantized = quantize(original, colors_normalized)
display(TF.to_pil_image(quantized))
</code></pre>
<p>Which outputs:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/b/0bddc7c18ceb778bea01da7b8dccb0bfdd93c2a9.png" data-download-href="https://discuss.pytorch.org/uploads/default/0bddc7c18ceb778bea01da7b8dccb0bfdd93c2a9" title="Code_HhfocMSnx7"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/b/0bddc7c18ceb778bea01da7b8dccb0bfdd93c2a9_2_123x250.png" alt="Code_HhfocMSnx7" data-base62-sha1="1GYpGs9lSyxRtDnElrjz1m27rdL" width="123" height="250" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/b/0bddc7c18ceb778bea01da7b8dccb0bfdd93c2a9_2_123x250.png, https://discuss.pytorch.org/uploads/default/optimized/3X/0/b/0bddc7c18ceb778bea01da7b8dccb0bfdd93c2a9_2_184x375.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/0/b/0bddc7c18ceb778bea01da7b8dccb0bfdd93c2a9_2_246x500.png 2x" data-dominant-color="687289"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Code_HhfocMSnx7</span><span class="informations">294×593 11.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>(1st one is the original, 2nd one is quantized)</p>
<p>I will investigate if this breaks the gradient and if I can make the function general for multiple images and multiple palettes.</p>

username: ZimoNitrome
cooked: <p>The previous implementation does not maintain gradient since it uses <code>argmin</code>. I have coded a new version which does:</p>
<pre data-code-wrap="python"><code class="lang-python">def quantize(image, palette):
    """
    Similar to PIL.Image.quantize() in PyTorch. Built to maintain gradient.
    Only works for one image i.e. CHW. Does NOT work for batches.
    """

    C, H, W = image.shape
    n_colors = len(palette)

    # Easier to work with list of colors
    flat_img = image.view(C, -1).T # [C, H, W] -&gt; [H*W, C]

    # Repeat image so that there are n_color number of columns of the same image
    flat_img_per_color = torch.stack(n_colors*[flat_img], dim=-2) # [H*W, C] -&gt; [H*W, n_colors, C]

    # Get euclidian distance between each pixel in each column and the columns repsective color
    # i.e. column 1 lists distance of each pixel to color #1 in palette, column 2 to color #2 etc.
    distance_per_pixel = torch.sum((flat_img_per_color-palette)**2, dim=-1)**0.5 # [H*W, n_colors, C] -&gt; [H*W, n_colors]

    # Get the shortest distance (one value per row (H*W) is selected)
    min_distances = torch.min(distance_per_pixel, dim=-1).values # [H*W, n_colors] -&gt; [H*W]

    # Get difference between each distance and the shortest distance.
    # One value per column (the selected value) will become 0.
    per_color_difference = distance_per_pixel - torch.stack(n_colors*[min_distances], dim=-1) # [H*W, n_colors]

    # Round all values up and invert. Creates something similar to one-hot encoding.
    per_color_diff_scaled = 1 - torch.ceil(per_color_difference)

    # Multiply the "kinda" one-hot encoded per_color_diff_scaled with the palette colors.
    # The result is a quantized image.
    quantized = torch.matmul(per_color_diff_scaled, palette)

    # Reshape it back to the original input format.
    quantized_img = quantized.T.view(C, H, W) # [H*W, C] -&gt; [C, H, W]

    return quantized_img
</code></pre>
<p>I am however concerned with how well the gradient is maintained.</p>
<p>I found this post while searching for implementations around <code>argmin()</code>:</p>
<aside class="quote quote-modified" data-post="6" data-topic="63620">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alband/48/215_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/why-does-this-function-break-the-gradient-tree/63620/6">Why does this function break the gradient tree? <img width="20" height="20" src="https://discuss.pytorch.org/images/emoji/apple/thinking.png?v=12" title="thinking" alt="thinking" class="emoji"></a> <a class="badge-category__wrapper " href="/c/autograd/7"><span data-category-id="7" data-drop-close="true" class="badge-category " title="A category of posts relating to the autograd engine itself."><span class="badge-category__name">autograd</span></span></a>
  </div>
  <blockquote>
    The problem is that indices are integers. So gradients don’t really exist for them. 
For the same reason that argmax is not differentiable. 
The trick that is often used to replace indexing is to use weighted sum. 
Instead of having an index in {0, 1, 2}, you have 3 values that sum to 1. And that gives you the probability of the index to be each of these values. Then you can replace indexing by weighted sum by these weights. 
This works in some cases where there are not too many indices. But mig…
  </blockquote>
</aside>

<p>I am not too familiar with weighted sums, but is it similar to my new implementation?</p>

username: ZimoNitrome
cooked: <p>It still broke gradients in some specific cases. I finally found this thread giving a weird hack that solved the problem:</p>
<aside class="quote" data-post="4" data-topic="48702">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/z/9de053/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/runtimeerror-function-sqrtbackward-returned-nan-values-in-its-0th-output/48702/4">RuntimeError: Function 'SqrtBackward' returned nan values in its 0th output</a> <a class="badge-category__wrapper " href="/c/autograd/7"><span data-category-id="7" data-drop-close="true" class="badge-category " title="A category of posts relating to the autograd engine itself."><span class="badge-category__name">autograd</span></span></a>
  </div>
  <blockquote>
    can use 
torch.sqrt(x + 1e-8) 
replace 
torch.sqrt(x) 
to solve this problem
  </blockquote>
</aside>

<p>FInal code if someone has a similar problem in the future:</p>
<pre data-code-wrap="py"><code class="lang-py">def quantize(image, palette):
    """
    Similar to PIL.Image.quantize() in PyTorch. Built to maintain gradient.
    Only works for one image i.e. CHW. Does NOT work for batches.
    """

    C, H, W = image.shape
    n_colors = len(palette)

    # Easier to work with list of colors
    flat_img = image.view(C, -1).T # [C, H, W] -&gt; [H*W, C]

    # Repeat image so that there are n_color number of columns of the same image
    flat_img_per_color = torch.stack(n_colors*[flat_img], dim=-2) # [H*W, C] -&gt; [H*W, n_colors, C]

    # Get euclidian distance between each pixel in each column and the columns repsective color
    # i.e. column 1 lists distance of each pixel to color #1 in palette, column 2 to color #2 etc.
    squared_distance = (flat_img_per_color-palette)**2
    # Dirty cursed hack
    # https://discuss.pytorch.org/t/runtimeerror-function-sqrtbackward-returned-nan-values-in-its-0th-output/48702/4
    euclidean_distance = torch.sqrt(torch.sum(squared_distance, dim=-1) + 1e-8) # [H*W, n_colors, C] -&gt; [H*W, n_colors]


    # Get the shortest distance (one value per row (H*W) is selected)
    min_distances = torch.min(euclidean_distance, dim=-1).values # [H*W, n_colors] -&gt; [H*W]

    # Get difference between each distance and the shortest distance.
    # One value per column (the selected value) will become 0.
    per_color_difference = euclidean_distance - torch.stack(n_colors*[min_distances], dim=-1) # [H*W, n_colors]

    # Round all values up and invert. Creates something similar to one-hot encoding.
    per_color_diff_scaled = 1 - torch.ceil(per_color_difference)

    per_color_diff_scaled_ = per_color_diff_scaled #(per_color_diff_scaled.T / per_color_diff_scaled.sum(dim=-1)).T

    # Multiply the "kinda" one-hot encoded per_color_diff_scaled with the palette colors.
    # The result is a quantized image.
    quantized = torch.matmul(per_color_diff_scaled_, palette)

    # Reshape it back to the original input format.
    quantized_img = quantized.T.view(C, H, W) # [H*W, C] -&gt; [C, H, W]

    return quantized_img
</code></pre>