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)