color_scale_numpy.py 1.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071
  1. from __future__ import annotations
  2. import numpy as np
  3. from numpy import ndarray
  4. from .color_util import RGB
  5. def create_gradient_hex(colors: list[str], resolution: int = 300) -> ndarray:
  6. """
  7. Create gradient array from hex
  8. """
  9. colors = np.array([RGB.from_hex(s) for s in colors])
  10. return create_gradient(colors, resolution)
  11. def create_gradient(colors: ndarray, resolution: int) -> ndarray:
  12. """
  13. Create gradient 2d array.
  14. Usage: arr[ratio / len(arr), :] = Scaled gradient color at that point
  15. """
  16. result = np.zeros((resolution * (len(colors) - 1), 3), dtype='uint8')
  17. # Create gradient mapping
  18. for i in range(len(colors) - 1):
  19. c1 = colors[i, :]
  20. c2 = colors[i + 1, :]
  21. bi = i * resolution
  22. for r in range(resolution):
  23. ratio = r / resolution
  24. result[bi + r, :] = c2 * ratio + c1 * (1 - ratio)
  25. return result
  26. def get_raw(gradient: ndarray, ratio: float) -> ndarray:
  27. """
  28. :param gradient: Gradient array (2d)
  29. :param ratio: Between 0-1
  30. :return: RGB subarray (1d, has 3 values)
  31. """
  32. if ratio == 1:
  33. return gradient[-1, :]
  34. i = int(ratio * len(gradient))
  35. return gradient[i, :]
  36. class Scale:
  37. colors: ndarray
  38. rgb: ndarray
  39. def __init__(self, scale: list[str], resolution: int = 300):
  40. self.colors = np.array([RGB.from_hex(s) for s in scale])
  41. self.rgb = create_gradient(self.colors, resolution)
  42. def __call__(self, ratio: float) -> RGB:
  43. """
  44. :param ratio: Between 0-1
  45. """
  46. return RGB(*get_raw(self.rgb, ratio))
  47. def test_color_scale():
  48. scale = Scale(['#232323', '#4F1879', '#B43A78', '#F98766', '#FCFAC0'])
  49. colors = 100
  50. for i in range(colors + 1):
  51. print(scale(i / colors).to_ansi_rgb(False), end=' ')