Calculate RGB values for a 2D-hue-saturation wheel.

tensorflow_hs_wheel.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""

"""
import typing
import math
import tensorflow as tf


def colorize_hs_wheel(
    values: tf.Tensor, max_d: typing.Union[None, tf.Tensor, float] = None
):
    assert values.shape[-1] == 2
    s = tf.linalg.norm(values, axis=-1)
    if max_d is not None:
        if max_d > 0.0:
            s /= max_d
        else:
            s = tf.ones_like(s)
    else:
        s /= tf.reduce_max(s)
    s = tf.clip_by_value(s, 0.0, 1.0)
    h = tf.math.atan2(values[..., 1], values[..., 0])
    h = h / 2.0 / math.pi

    i = tf.math.floor(h * 6.0)
    f = (h * 6.0) - i
    p, q, t = 1.0 - s, 1.0 - (s * f), 1.0 - (s * (1.0 - f))
    v = tf.ones_like(p)
    i = tf.cast(i, tf.dtypes.int32)
    i %= 6

    h1 = tf.stack((v, t, p), axis=-1)
    h2 = tf.stack((q, v, p), axis=-1)
    h3 = tf.stack((p, v, t), axis=-1)
    h4 = tf.stack((p, q, v), axis=-1)
    h5 = tf.stack((t, p, v), axis=-1)
    h6 = tf.stack((v, p, q), axis=-1)

    h = tf.stack((h1, h2, h3, h4, h5, h6), axis=-2)
    selection_matrix = tf.cast(tf.one_hot(i, depth=6), h.dtype)
    return tf.matmul(selection_matrix[..., None, :], h)[..., 0, :]


if __name__ == "__main__":

    radii = tf.linspace(0.1, 1.0, 5)
    angles = tf.linspace(0.0, 2 * math.pi, 11)[:-1]
    x = radii[:, None] * tf.math.cos(angles)[None, :]
    y = radii[:, None] * tf.math.sin(angles)[None, :]
    values = tf.stack((x, y), axis=-1)
    wheel_color = colorize_hs_wheel(values)

    import matplotlib.pyplot as plt

    plt.scatter(
        x=x.numpy().reshape((-1)),
        y=y.numpy().reshape((-1)),
        c=wheel_color.numpy().reshape(-1, 3),
    )
    plt.axis("equal")
    plt.show()