def gauss(input, k, rdom, name):
blur_x = hl.Func(name + "_x")
output = hl.Func(name)
x, y, c, xi, yi = hl.Var("x"), hl.Var("y"), hl.Var("c"), hl.Var("xi"), hl.Var("yi")
val = hl.Expr("val")
if input.dimensions() == 2:
blur_x[x, y] = hl.sum(input[x + rdom, y] * k[rdom])
val = hl.sum(blur_x[x, y + rdom] * k[rdom])
if input.output_types()[0] == hl.UInt(16):
val = hl.u16(val)
output[x, y] = val
else:
blur_x[x, y, c] = hl.sum(input[x + rdom, y, c] * k[rdom])
val = hl.sum(blur_x[x, y + rdom, c] * k[rdom])
if input.output_types()[0] == hl.UInt(16):
val = hl.u16(val)
output[x, y, c] = val
blur_x.compute_at(output, x).vectorize(x, 16)
output.compute_root().tile(x, y, xi, yi, 256, 128).vectorize(xi, 16).parallel(y)
return output
def shift_bayer_to_rggb(input, cfa_pattern):
print(f'cfa_pattern: {cfa_pattern}')
output = hl.Func("rggb_input")
x, y = hl.Var("x"), hl.Var("y")
cfa = hl.u16(cfa_pattern)
output[x, y] = hl.select(cfa == hl.u16(1), input[x, y],
cfa == hl.u16(2), input[x + 1, y],
cfa == hl.u16(4), input[x, y + 1],
cfa == hl.u16(3), input[x + 1, y + 1],
0)
return output
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 get_blur(input):
assert type(input) == hl.ImageParam
assert input.dimensions() == 2
x, y = hl.Var("x"), hl.Var("y")
clamped_input = hl.BoundaryConditions.repeat_edge(input)
input_uint16 = hl.Func("input_uint16")
input_uint16[x, y] = hl.u16(clamped_input[x, y])
ci = input_uint16
blur_x = hl.Func("blur_x")
blur_y = hl.Func("blur_y")
blur_x[x, y] = (ci[x, y] + ci[x + 1, y] + ci[x + 2, y]) / 3
blur_y[x, y] = hl.cast(
hl.UInt(8), (blur_x[x, y] + blur_x[x, y + 1] + blur_x[x, y + 2]) / 3)
# schedule
xi, yi = hl.Var("xi"), hl.Var("yi")
blur_y.tile(x, y, xi, yi, 8, 4).parallel(y).vectorize(xi, 8)
blur_x.compute_at(blur_y, x).vectorize(x, 8)
return blur_y
def test_compiletime_error():
x = hl.Var('x')
y = hl.Var('y')
f = hl.Func('f')
f[x, y] = hl.u16(x + y)
# Deliberate type-mismatch error
buf = hl.Buffer(hl.UInt(8), [2, 2])
try:
f.realize(buf)
except RuntimeError as e:
assert 'Output buffer f has type uint16 but type of the buffer passed in is uint8' in str(e)
else:
assert False, 'Did not see expected exception!'
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
def test_int_promotion():
# Verify that (Exprlike op literal) correctly matches the type
# of the literal to the Exprlike (rather than promoting the result to int32).
# All types that use add_binary_operators() should be tested here.
x = hl.Var('x')
# All the binary ops are handled the same, so + is good enough
# Exprlike = FuncRef
f = hl.Func('f')
f[x] = hl.u16(x)
_check_is_u16(f[x] + 2)
_check_is_u16(2 + f[x])
# Exprlike = Expr
e = hl.Expr(f[x])
_check_is_u16(e + 2)
_check_is_u16(2 + e)
# Exprlike = Param
p = hl.Param(hl.UInt(16))
_check_is_u16(p + 2)
_check_is_u16(2 + p)
def white_balance(input, width, height, white_balance_r, white_balance_g0, white_balance_g1, white_balance_b):
output = hl.Func("white_balance_output")
print(width, height, white_balance_r, white_balance_g0, white_balance_g1, white_balance_b)
x, y = hl.Var("x"), hl.Var("y")
rdom = hl.RDom([(0, width / 2), (0, height / 2)])
output[x, y] = hl.u16(0)
output[rdom.x * 2, rdom.y * 2] = hl.u16_sat(white_balance_r * hl.f32(input[rdom.x * 2, rdom.y * 2]))
output[rdom.x * 2 + 1, rdom.y * 2] = hl.u16_sat(white_balance_g0 * hl.f32(input[rdom.x * 2 + 1, rdom.y * 2]))
output[rdom.x * 2, rdom.y * 2 + 1] = hl.u16_sat(white_balance_g1 * hl.f32(input[rdom.x * 2, rdom.y * 2 + 1]))
output[rdom.x * 2 + 1, rdom.y * 2 + 1] = hl.u16_sat(white_balance_b * hl.f32(input[rdom.x * 2 + 1, rdom.y * 2 + 1]))
output.compute_root().parallel(y).vectorize(x, 16)
output.update(0).parallel(rdom.y)
output.update(1).parallel(rdom.y)
output.update(2).parallel(rdom.y)
output.update(3).parallel(rdom.y)
return output
def get_blur(input):
assert type(input) == hl.ImageParam
assert input.dimensions() == 2
x, y = hl.Var("x"), hl.Var("y")
clamped_input = hl.BoundaryConditions.repeat_edge(input)
input_uint16 = hl.Func("input_uint16")
input_uint16[x,y] = hl.u16(clamped_input[x,y])
ci = input_uint16
blur_x = hl.Func("blur_x")
blur_y = hl.Func("blur_y")
blur_x[x,y] = (ci[x,y]+ci[x+1,y]+ci[x+2,y])/3
blur_y[x,y] = hl.cast(hl.UInt(8), (blur_x[x,y]+blur_x[x,y+1]+blur_x[x,y+2])/3)
# schedule
xi, yi = hl.Var("xi"), hl.Var("yi")
blur_y.tile(x, y, xi, yi, 8, 4).parallel(y).vectorize(xi, 8)
blur_x.compute_at(blur_y, x).vectorize(x, 8)
return blur_y
def test_complexstub():
constant_image = _make_constant_image()
input = hl.ImageParam(hl.UInt(8), 3, 'input')
input.set(constant_image)
x, y, c = hl.Var(), hl.Var(), hl.Var()
target = hl.get_jit_target_from_environment()
float_arg = 1.25
int_arg = 33
func_input = hl.Func("func_input")
func_input[x, y, c] = hl.u16(x + y + c)
r = complexstub(target,
typed_buffer_input=constant_image,
untyped_buffer_input=constant_image,
simple_input=input,
array_input=[input, input],
float_arg=float_arg,
int_arg=[int_arg, int_arg],
untyped_buffer_output_type="uint8",
extra_func_input=func_input,
vectorize=True)
# return value is a tuple; unpack separately to avoid
# making the callsite above unreadable
(simple_output, tuple_output, array_output, typed_buffer_output,
untyped_buffer_output, static_compiled_buffer_output,
extra_func_output) = r
b = simple_output.realize([32, 32, 3], target)
assert b.type() == hl.Float(32)
for x in range(32):
for y in range(32):
for c in range(3):
expected = constant_image[x, y, c]
actual = b[x, y, c]
assert expected == actual, "Expected %s Actual %s" % (expected,
actual)
b = tuple_output.realize([32, 32, 3], target)
assert b[0].type() == hl.Float(32)
assert b[1].type() == hl.Float(32)
assert len(b) == 2
for x in range(32):
for y in range(32):
for c in range(3):
expected1 = constant_image[x, y, c] * float_arg
expected2 = expected1 + int_arg
actual1, actual2 = b[0][x, y, c], b[1][x, y, c]
assert expected1 == actual1, "Expected1 %s Actual1 %s" % (
expected1, actual1)
assert expected2 == actual2, "Expected2 %s Actual1 %s" % (
expected2, actual2)
assert len(array_output) == 2
for a in array_output:
b = a.realize([32, 32], target)
assert b.type() == hl.Int(16)
for x in range(32):
for y in range(32):
expected = constant_image[x, y, 0] + int_arg
actual = b[x, y]
assert expected == actual, "Expected %s Actual %s" % (expected,
actual)
# TODO: Output<Buffer<>> has additional behaviors useful when a Stub
# is used within another Generator; this isn't yet implemented since there
# isn't yet Python bindings for Generator authoring. This section
# of the test may need revision at that point.
b = typed_buffer_output.realize([32, 32, 3], target)
assert b.type() == hl.Float(32)
for x in range(32):
for y in range(32):
for c in range(3):
expected = constant_image[x, y, c]
actual = b[x, y, c]
assert expected == actual, "Expected %s Actual %s" % (expected,
actual)
b = untyped_buffer_output.realize([32, 32, 3], target)
assert b.type() == hl.UInt(8)
for x in range(32):
for y in range(32):
for c in range(3):
expected = constant_image[x, y, c]
actual = b[x, y, c]
assert expected == actual, "Expected %s Actual %s" % (expected,
actual)
b = static_compiled_buffer_output.realize([4, 4, 1], target)
assert b.type() == hl.UInt(8)
for x in range(4):
for y in range(4):
for c in range(1):
expected = constant_image[x, y, c] + 42
actual = b[x, y, c]
assert expected == actual, "Expected %s Actual %s" % (expected,
actual)
b = extra_func_output.realize([32, 32], target)
assert b.type() == hl.Float(64)
for x in range(32):
for y in range(32):
expected = x + y + 1
actual = b[x, y]
assert expected == actual, "Expected %s Actual %s" % (expected,
actual)
