def __init__(self, input):
assert input.type() == hl.UInt(8)
self.lut = hl.Func("lut")
self.padded = hl.Func("padded")
self.padded16 = hl.Func("padded16")
self.sharpen = hl.Func("sharpen")
self.curved = hl.Func("curved")
self.input = input
# For this lesson, we'll use a two-stage pipeline that sharpens
# and then applies a look-up-table (LUT).
# First we'll define the LUT. It will be a gamma curve.
gamma = hl.f32(1.2)
self.lut[i] = hl.u8(hl.clamp(hl.pow(i / 255.0, gamma) * 255.0, 0, 255))
# Augment the input with a boundary condition.
self.padded[x, y, c] = input[hl.clamp(x, 0,
input.width() - 1),
hl.clamp(y, 0,
input.height() - 1), c]
# Cast it to 16-bit to do the math.
self.padded16[x, y, c] = hl.u16(self.padded[x, y, c])
# Next we sharpen it with a five-tap filter.
self.sharpen[x, y, c] = (
self.padded16[x, y, c] * 2 -
(self.padded16[x - 1, y, c] + self.padded16[x, y - 1, c] +
self.padded16[x + 1, y, c] + self.padded16[x, y + 1, c]) / 4)
# Then apply the LUT.
self.curved[x, y, c] = self.lut[self.sharpen[x, y, c]]
def gamma_inverse(input):
output = hl.Func("gamma_inverse_output")
x, y, c = hl.Var("x"), hl.Var("y"), hl.Var("c")
cutoff = 2575
gamma_toe = 0.0774
gamma_pow = 2.4
gamma_fac = 57632.49226
gamma_con = 0.055
if input.dimensions() == 2:
output[x, y] = hl.u16(hl.select(input[x, y] < cutoff,
gamma_toe * input[x, y],
hl.pow(hl.f32(input[x, y]) / 65535 + gamma_con, gamma_pow) * gamma_fac))
else:
output[x, y, c] = hl.u16(hl.select(input[x, y, c] < cutoff,
gamma_toe * input[x, y, c],
hl.pow(hl.f32(input[x, y, c]) / 65535 + gamma_con, gamma_pow) * gamma_fac))
output.compute_root().parallel(y).vectorize(x, 16)
return output
def gamma_correct(input):
output = hl.Func("gamma_correct_output")
x, y, c = hl.Var("x"), hl.Var("y"), hl.Var("c")
cutoff = 200
gamma_toe = 12.92
gamma_pow = 0.416667
gamma_fac = 680.552897
gamma_con = -3604.425
if input.dimensions() == 2:
output[x, y] = hl.u16(hl.select(input[x, y] < cutoff,
gamma_toe * input[x, y],
gamma_fac * hl.pow(input[x, y], gamma_pow) + gamma_con))
else:
output[x, y, c] = hl.u16(hl.select(input[x, y, c] < cutoff,
gamma_toe * input[x, y, c],
gamma_fac * hl.pow(input[x, y, c], gamma_pow) + gamma_con))
output.compute_root().parallel(y).vectorize(x, 16)
return output
def tone_map(input, width, height, compression, gain):
print(f'Compression: {compression}, gain: {gain}')
normal_dist = hl.Func("luma_weight_distribution")
grayscale = hl.Func("grayscale")
output = hl.Func("tone_map_output")
x, y, c, v = hl.Var("x"), hl.Var("y"), hl.Var("c"), hl.Var("v")
rdom = hl.RDom([(0, 3)])
normal_dist[v] = hl.f32(hl.exp(-12.5 * hl.pow(hl.f32(v) / 65535 - 0.5, 2)))
grayscale[x, y] = hl.u16(hl.sum(hl.u32(input[x, y, rdom])) / 3)
dark = grayscale
comp_const = 1
gain_const = 1
comp_slope = (compression - comp_const) / (TONE_MAP_PASSES)
gain_slope = (gain - gain_const) / (TONE_MAP_PASSES)
for i in range(TONE_MAP_PASSES):
print(' pass', i)
norm_comp = i * comp_slope + comp_const
norm_gain = i * gain_slope + gain_const
bright = brighten(dark, norm_comp)
dark_gamma = gamma_correct(dark)
bright_gamma = gamma_correct(bright)
dark_gamma = combine2(dark_gamma, bright_gamma, width, height, normal_dist)
dark = brighten(gamma_inverse(dark_gamma), norm_gain)
output[x, y, c] = hl.u16_sat(hl.u32(input[x, y, c]) * hl.u32(dark[x, y]) / hl.u32(hl.max(1, grayscale[x, y])))
grayscale.compute_root().parallel(y).vectorize(x, 16)
normal_dist.compute_root().vectorize(v, 16)
return output