1import halide as hl
2
3import numpy as np
4import imageio
5import os.path
6
7def get_blur(input):
8    assert type(input) == hl.ImageParam
9    assert input.dimensions() == 2
10
11    x, y = hl.Var("x"), hl.Var("y")
12
13    clamped_input = hl.BoundaryConditions.repeat_edge(input)
14
15    input_uint16 = hl.Func("input_uint16")
16    input_uint16[x,y] = hl.u16(clamped_input[x,y])
17    ci = input_uint16
18
19    blur_x = hl.Func("blur_x")
20    blur_y = hl.Func("blur_y")
21
22    blur_x[x,y] = (ci[x,y]+ci[x+1,y]+ci[x+2,y])/3
23    blur_y[x,y] = hl.cast(hl.UInt(8), (blur_x[x,y]+blur_x[x,y+1]+blur_x[x,y+2])/3)
24
25    # schedule
26    xi, yi = hl.Var("xi"), hl.Var("yi")
27    blur_y.tile(x, y, xi, yi, 8, 4).parallel(y).vectorize(xi, 8)
28    blur_x.compute_at(blur_y, x).vectorize(x, 8)
29
30    return blur_y
31
32
33def get_input_data():
34    image_path = os.path.join(os.path.dirname(__file__), "../../apps/images/rgb.png")
35    assert os.path.exists(image_path), \
36        "Could not find %s" % image_path
37    rgb_data = imageio.imread(image_path)
38    print("rgb_data", type(rgb_data), rgb_data.shape, rgb_data.dtype)
39
40    grey_data = np.mean(rgb_data, axis=2, dtype=np.float32).astype(rgb_data.dtype)
41    input_data = np.copy(grey_data, order="F")
42
43    return input_data
44
45def main():
46    # define and compile the function
47    input = hl.ImageParam(hl.UInt(8), 2, "input_param")
48    blur = get_blur(input)
49    blur.compile_jit()
50
51    # preparing input and output memory buffers (numpy ndarrays)
52    input_data = get_input_data()
53    input_image = hl.Buffer(input_data)
54    input.set(input_image)
55
56    output_data = np.empty(input_data.shape, dtype=input_data.dtype, order="F")
57    output_image = hl.Buffer(output_data)
58
59    print("input_image", input_image)
60    print("output_image", output_image)
61
62    # do the actual computation
63    blur.realize(output_image)
64
65    # save results
66    input_path = "blur_input.png"
67    output_path = "blur_result.png"
68    imageio.imsave(input_path, input_data)
69    imageio.imsave(output_path, output_data)
70    print("\nblur realized on output image.",
71          "Result saved at", output_path,
72          "( input data copy at", input_path, ")")
73
74    print("\nEnd of game. Have a nice day!")
75    return
76
77
78if __name__ == "__main__":
79    main()
80