Spaces:
Runtime error
Runtime error
File size: 2,261 Bytes
542e45a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 |
from functools import partial
import numpy as np
class ABGamut:
RESOURCE_POINTS = "./utils/gamut_pts.npy"
RESOURCE_PRIOR = "./utils/gamut_probs.npy"
DTYPE = np.float32
EXPECTED_SIZE = 313
def __init__(self):
self.points = np.load(self.RESOURCE_POINTS).astype(self.DTYPE)
self.prior = np.load(self.RESOURCE_PRIOR).astype(self.DTYPE)
assert self.points.shape == (self.EXPECTED_SIZE, 2)
assert self.prior.shape == (self.EXPECTED_SIZE,)
class CIELAB:
L_MEAN = 50
AB_BINSIZE = 10
AB_RANGE = [-110 - AB_BINSIZE // 2, 110 + AB_BINSIZE // 2, AB_BINSIZE]
AB_DTYPE = np.float32
Q_DTYPE = np.int64
RGB_RESOLUTION = 101
RGB_RANGE = [0, 1, RGB_RESOLUTION]
RGB_DTYPE = np.float64
def __init__(self, gamut=None):
self.gamut = gamut if gamut is not None else ABGamut()
a, b, self.ab = self._get_ab()
self.ab_gamut_mask = self._get_ab_gamut_mask(
a, b, self.ab, self.gamut)
self.ab_to_q = self._get_ab_to_q(self.ab_gamut_mask)
self.q_to_ab = self._get_q_to_ab(self.ab, self.ab_gamut_mask)
@classmethod
def _get_ab(cls):
a = np.arange(*cls.AB_RANGE, dtype=cls.AB_DTYPE)
b = np.arange(*cls.AB_RANGE, dtype=cls.AB_DTYPE)
b_, a_ = np.meshgrid(a, b)
ab = np.dstack((a_, b_))
return a, b, ab
@classmethod
def _get_ab_gamut_mask(cls, a, b, ab, gamut):
ab_gamut_mask = np.full(ab.shape[:-1], False, dtype=bool)
a = np.digitize(gamut.points[:, 0], a) - 1
b = np.digitize(gamut.points[:, 1], b) - 1
for a_, b_ in zip(a, b):
ab_gamut_mask[a_, b_] = True
return ab_gamut_mask
@classmethod
def _get_ab_to_q(cls, ab_gamut_mask):
ab_to_q = np.full(ab_gamut_mask.shape, -1, dtype=cls.Q_DTYPE)
ab_to_q[ab_gamut_mask] = np.arange(np.count_nonzero(ab_gamut_mask))
return ab_to_q
@classmethod
def _get_q_to_ab(cls, ab, ab_gamut_mask):
return ab[ab_gamut_mask] + cls.AB_BINSIZE / 2
def bin_ab(self, ab):
ab_discrete = ((ab + 110) / self.AB_RANGE[2]).astype(int)
a, b = np.hsplit(ab_discrete.reshape(-1, 2), 2)
return self.ab_to_q[a, b].reshape(*ab.shape[:2])
|