def desaturate_noise(input, width, height):
print(' desaturate_noise')
output = hl.Func("desaturate_noise_output")
x, y, c = hl.Var("x"), hl.Var("y"), hl.Var("c")
input_mirror = hl.BoundaryConditions.mirror_image(input, [(0, width), (0, height)])
blur = gauss_15x15(gauss_15x15(input_mirror, "desaturate_noise_blur1"), "desaturate_noise_blur_2")
factor = 1.4
threshold = 25000
output[x, y, c] = input[x, y, c]
output[x, y, 1] = hl.select((hl.abs(blur[x, y, 1]) / hl.abs(input[x, y, 1]) < factor) &
(hl.abs(input[x, y, 1]) < threshold) & (hl.abs(blur[x, y, 1]) < threshold),
0.7 * blur[x, y, 1] + 0.3 * input[x, y, 1], input[x, y, 1])
output[x, y, 2] = hl.select((hl.abs(blur[x, y, 2]) / hl.abs(input[x, y, 2]) < factor) &
(hl.abs(input[x, y, 2]) < threshold) & (hl.abs(blur[x, y, 2]) < threshold),
0.7 * blur[x, y, 2] + 0.3 * input[x, y, 2], input[x, y, 2])
output.compute_root().parallel(y).vectorize(x, 16)
return output
def merge_temporal(images, alignment):
weight = hl.Func("merge_temporal_weights")
total_weight = hl.Func("merge_temporal_total_weights")
output = hl.Func("merge_temporal_output")
ix, iy, tx, ty, n = hl.Var('ix'), hl.Var('iy'), hl.Var('tx'), hl.Var('ty'), hl.Var('n')
rdom0 = hl.RDom([(0, 16), (0, 16)])
rdom1 = hl.RDom([(1, images.dim(2).extent() - 1)])
imgs_mirror = hl.BoundaryConditions.mirror_interior(images, [(0, images.width()), (0, images.height())])
layer = box_down2(imgs_mirror, "merge_layer")
offset = Point(alignment[tx, ty, n]).clamp(Point(MINIMUM_OFFSET, MINIMUM_OFFSET),
Point(MAXIMUM_OFFSET, MAXIMUM_OFFSET))
al_x = idx_layer(tx, rdom0.x) + offset.x / 2
al_y = idx_layer(ty, rdom0.y) + offset.y / 2
ref_val = layer[idx_layer(tx, rdom0.x), idx_layer(ty, rdom0.y), 0]
alt_val = layer[al_x, al_y, n]
factor = 8.0
min_distance = 10
max_distance = 300 # max L1 distance, otherwise the value is not used
distance = hl.sum(hl.abs(hl.cast(hl.Int(32), ref_val) - hl.cast(hl.Int(32), alt_val))) / 256
normal_distance = hl.max(1, hl.cast(hl.Int(32), distance) / factor - min_distance / factor)
# Weight for the alternate frame
weight[tx, ty, n] = hl.select(normal_distance > (max_distance - min_distance), 0.0,
1.0 / normal_distance)
total_weight[tx, ty] = hl.sum(weight[tx, ty, rdom1]) + 1
offset = Point(alignment[tx, ty, rdom1])
al_x = idx_im(tx, ix) + offset.x
al_y = idx_im(ty, iy) + offset.y
ref_val = imgs_mirror[idx_im(tx, ix), idx_im(ty, iy), 0]
alt_val = imgs_mirror[al_x, al_y, rdom1]
# Sum all values according to their weight, and divide by total weight to obtain average
output[ix, iy, tx, ty] = hl.sum(weight[tx, ty, rdom1] * alt_val / total_weight[tx, ty]) + ref_val / total_weight[
tx, ty]
weight.compute_root().parallel(ty).vectorize(tx, 16)
total_weight.compute_root().parallel(ty).vectorize(tx, 16)
output.compute_root().parallel(ty).vectorize(ix, 32)
return output
def align_layer(layer, prev_alignment, prev_min, prev_max):
scores = hl.Func(layer.name() + "_scores")
alignment = hl.Func(layer.name() + "_alignment")
xi, yi, tx, ty, n = hl.Var("xi"), hl.Var("yi"), hl.Var('tx'), hl.Var(
'ty'), hl.Var('n')
rdom0 = hl.RDom([(0, 16), (0, 16)])
rdom1 = hl.RDom([(-4, 8), (-4, 8)])
# Alignment of the previous (more coarse) layer scaled to this (finer) layer
prev_offset = DOWNSAMPLE_RATE * Point(
prev_alignment[prev_tile(tx), prev_tile(ty), n]).clamp(
prev_min, prev_max)
x0 = idx_layer(tx, rdom0.x)
y0 = idx_layer(ty, rdom0.y)
# (x,y) coordinates in the search region relative to the offset obtained from the alignment of the previous layer
x = x0 + prev_offset.x + xi
y = y0 + prev_offset.y + yi
ref_val = layer[x0, y0, 0] # Value of reference frame (the first frame)
alt_val = layer[x, y, n] # alternate frame value
# L1 distance between reference frame and alternate frame
d = hl.abs(hl.cast(hl.Int(32), ref_val) - hl.cast(hl.Int(32), alt_val))
scores[xi, yi, tx, ty, n] = hl.sum(d)
# Alignment for each tile, where L1 distances are minimum
alignment[tx, ty, n] = Point(hl.argmin(scores[rdom1.x, rdom1.y, tx, ty,
n])) + prev_offset
scores.compute_at(alignment, tx).vectorize(xi, 8)
alignment.compute_root().parallel(ty).vectorize(tx, 16)
return alignment
def bilateral_filter(input, width, height):
print(' bilateral_filter')
k = hl.Buffer(hl.Float(32), [7, 7], "gauss_kernel")
k.translate([-3, -3])
weights = hl.Func("bilateral_weights")
total_weights = hl.Func("bilateral_total_weights")
bilateral = hl.Func("bilateral")
output = hl.Func("bilateral_filter_output")
x, y, dx, dy, c = hl.Var("x"), hl.Var("y"), hl.Var("dx"), hl.Var("dy"), hl.Var("c")
rdom = hl.RDom([(-3, 7), (-3, 7)])
k.fill(0)
k[-3, -3] = 0.000690
k[-2, -3] = 0.002646
k[-1, -3] = 0.005923
k[0, -3] = 0.007748
k[1, -3] = 0.005923
k[2, -3] = 0.002646
k[3, -3] = 0.000690
k[-3, -2] = 0.002646
k[-2, -2] = 0.010149
k[-1, -2] = 0.022718
k[0, -2] = 0.029715
k[1, -2] = 0.022718
k[2, -2] = 0.010149
k[3, -2] = 0.002646
k[-3, -1] = 0.005923
k[-2, -1] = 0.022718
k[-1, -1] = 0.050855
k[0, -1] = 0.066517
k[1, -1] = 0.050855
k[2, -1] = 0.022718
k[3, -1] = 0.005923
k[-3, 0] = 0.007748
k[-2, 0] = 0.029715
k[-1, 0] = 0.066517
k[0, 0] = 0.087001
k[1, 0] = 0.066517
k[2, 0] = 0.029715
k[3, 0] = 0.007748
k[-3, 1] = 0.005923
k[-2, 1] = 0.022718
k[-1, 1] = 0.050855
k[0, 1] = 0.066517
k[1, 1] = 0.050855
k[2, 1] = 0.022718
k[3, 1] = 0.005923
k[-3, 2] = 0.002646
k[-2, 2] = 0.010149
k[-1, 2] = 0.022718
k[0, 2] = 0.029715
k[1, 2] = 0.022718
k[2, 2] = 0.010149
k[3, 2] = 0.002646
k[-3, 3] = 0.000690
k[-2, 3] = 0.002646
k[-1, 3] = 0.005923
k[0, 3] = 0.007748
k[1, 3] = 0.005923
k[2, 3] = 0.002646
k[3, 3] = 0.000690
input_mirror = hl.BoundaryConditions.mirror_interior(input, [(0, width), (0, height)])
dist = hl.cast(hl.Float(32),
hl.cast(hl.Int(32), input_mirror[x, y, c]) - hl.cast(hl.Int(32), input_mirror[x + dx, y + dy, c]))
sig2 = 100
threshold = 25000
score = hl.select(hl.abs(input_mirror[x + dx, y + dy, c]) > threshold, 0, hl.exp(-dist * dist / sig2))
weights[dx, dy, x, y, c] = k[dx, dy] * score
total_weights[x, y, c] = hl.sum(weights[rdom.x, rdom.y, x, y, c])
bilateral[x, y, c] = hl.sum(input_mirror[x + rdom.x, y + rdom.y, c] * weights[rdom.x, rdom.y, x, y, c]) / \
total_weights[x, y, c]
output[x, y, c] = hl.cast(hl.Float(32), input[x, y, c])
output[x, y, 1] = bilateral[x, y, 1]
output[x, y, 2] = bilateral[x, y, 2]
weights.compute_at(output, y).vectorize(x, 16)
output.compute_root().parallel(y).vectorize(x, 16)
output.update(0).parallel(y).vectorize(x, 16)
output.update(1).parallel(y).vectorize(x, 16)
return output
def findStereoCorrespondence(left,
right,
SADWindowSize,
minDisparity,
numDisparities,
xmin,
xmax,
ymin,
ymax,
x_tile_size=32,
y_tile_size=32,
test=False,
uniquenessRatio=0.15,
disp12MaxDiff=1):
""" Returns Func (left: Func, right: Func) """
x, y, c, d = Var("x"), Var("y"), Var("c"), Var("d")
diff = Func("diff")
diff[d, x, y] = h.cast(UInt(16), h.abs(left[x, y] - right[x - d, y]))
win2 = SADWindowSize / 2
diff_T = Func("diff_T")
xi, xo, yi, yo = Var("xi"), Var("xo"), Var("yi"), Var("yo")
diff_T[d, xi, yi, xo, yo] = diff[d, xi + xo * x_tile_size + xmin,
yi + yo * y_tile_size + ymin]
cSAD, vsum = Func("cSAD"), Func("vsum")
rk = RDom(-win2, SADWindowSize, "rk")
rxi, ryi = RDom(1, x_tile_size - 1, "rxi"), RDom(1, y_tile_size - 1, "ryi")
if test:
vsum[d, xi, yi, xo, yo] = h.sum(diff_T[d, xi, yi + rk, xo, yo])
cSAD[d, xi, yi, xo, yo] = h.sum(vsum[d, xi + rk, yi, xo, yo])
else:
vsum[d, xi, yi, xo, yo] = h.select(yi != 0, h.cast(UInt(16), 0),
h.sum(diff_T[d, xi, rk, xo, yo]))
vsum[d, xi, ryi, xo, yo] = vsum[d, xi, ryi - 1, xo, yo] + diff_T[
d, xi, ryi + win2, xo, yo] - diff_T[d, xi, ryi - win2 - 1, xo, yo]
cSAD[d, xi, yi, xo, yo] = h.select(xi != 0, h.cast(UInt(16), 0),
h.sum(vsum[d, rk, yi, xo, yo]))
cSAD[d, rxi, yi, xo,
yo] = cSAD[d, rxi - 1, yi, xo,
yo] + vsum[d, rxi + win2, yi, xo,
yo] - vsum[d, rxi - win2 - 1, yi, xo, yo]
rd = RDom(minDisparity, numDisparities)
disp_left = Func("disp_left")
disp_left[xi, yi, xo, yo] = h.Tuple(h.cast(UInt(16), minDisparity),
h.cast(UInt(16), (2 << 16) - 1))
disp_left[xi, yi, xo, yo] = h.tuple_select(
cSAD[rd, xi, yi, xo, yo] < disp_left[xi, yi, xo, yo][1],
h.Tuple(h.cast(UInt(16), rd), cSAD[rd, xi, yi, xo, yo]),
h.Tuple(disp_left[xi, yi, xo, yo]))
FILTERED = -16
disp = Func("disp")
disp[x, y] = h.select(
# x > xmax-xmin or y > ymax-ymin,
x < xmax,
h.cast(
UInt(16), disp_left[x % x_tile_size, y % y_tile_size,
x / x_tile_size, y / y_tile_size][0]),
h.cast(UInt(16), FILTERED))
# Schedule
vector_width = 8
disp.compute_root() \
.tile(x, y, xo, yo, xi, yi, x_tile_size, y_tile_size).reorder(xi, yi, xo, yo) \
.vectorize(xi, vector_width).parallel(xo).parallel(yo)
# reorder storage
disp_left.reorder_storage(xi, yi, xo, yo)
diff_T.reorder_storage(xi, yi, xo, yo, d)
vsum.reorder_storage(xi, yi, xo, yo, d)
cSAD.reorder_storage(xi, yi, xo, yo, d)
disp_left.compute_at(disp, xo).reorder(xi, yi, xo, yo) \
.vectorize(xi, vector_width) \
.update() \
.reorder(xi, yi, rd, xo, yo).vectorize(xi, vector_width)
if test:
cSAD.compute_at(disp_left, rd).reorder(xi, yi, xo, yo,
d).vectorize(xi, vector_width)
vsum.compute_at(disp_left, rd).reorder(xi, yi, xo, yo,
d).vectorize(xi, vector_width)
else:
cSAD.compute_at(disp_left, rd).reorder(xi, yi, xo, yo, d).vectorize(xi, vector_width) \
.update() \
.reorder(yi, rxi, xo, yo, d).vectorize(yi, vector_width)
vsum.compute_at(disp_left, rd).reorder(xi, yi, xo, yo, d).vectorize(xi, vector_width) \
.update() \
.reorder(xi, ryi, xo, yo, d).vectorize(xi, vector_width)
return disp
def findStereoCorrespondence(left, right, SADWindowSize, minDisparity, numDisparities,
xmin, xmax, ymin, ymax,
x_tile_size=32, y_tile_size=32, test=False, uniquenessRatio=0.15, disp12MaxDiff=1):
""" Returns Func (left: Func, right: Func) """
x, y, c, d = Var("x"), Var("y"), Var("c"), Var("d")
diff = Func("diff")
diff[d, x, y] = h.cast(UInt(16), h.abs(left[x, y] - right[x-d, y]))
win2 = SADWindowSize/2
diff_T = Func("diff_T")
xi, xo, yi, yo = Var("xi"), Var("xo"), Var("yi"), Var("yo")
diff_T[d, xi, yi, xo, yo] = diff[d, xi + xo * x_tile_size + xmin, yi + yo * y_tile_size + ymin]
cSAD, vsum = Func("cSAD"), Func("vsum")
rk = RDom(-win2, SADWindowSize, "rk")
rxi, ryi = RDom(1, x_tile_size - 1, "rxi"), RDom(1, y_tile_size - 1, "ryi")
if test:
vsum[d, xi, yi, xo, yo] = h.sum(diff_T[d, xi, yi+rk, xo, yo])
cSAD[d, xi, yi, xo, yo] = h.sum(vsum[d, xi+rk, yi, xo, yo])
else:
vsum[d, xi, yi, xo, yo] = h.select(yi != 0, h.cast(UInt(16), 0), h.sum(diff_T[d, xi, rk, xo, yo]))
vsum[d, xi, ryi, xo, yo] = vsum[d, xi, ryi-1, xo, yo] + diff_T[d, xi, ryi+win2, xo, yo] - diff_T[d, xi, ryi-win2-1, xo, yo]
cSAD[d, xi, yi, xo, yo] = h.select(xi != 0, h.cast(UInt(16), 0), h.sum(vsum[d, rk, yi, xo, yo]))
cSAD[d, rxi, yi, xo, yo] = cSAD[d, rxi-1, yi, xo, yo] + vsum[d, rxi+win2, yi, xo, yo] - vsum[d, rxi-win2-1, yi, xo, yo]
rd = RDom(minDisparity, numDisparities)
disp_left = Func("disp_left")
disp_left[xi, yi, xo, yo] = h.Tuple(h.cast(UInt(16), minDisparity), h.cast(UInt(16), (2<<16)-1))
disp_left[xi, yi, xo, yo] = h.tuple_select(
cSAD[rd, xi, yi, xo, yo] < disp_left[xi, yi, xo, yo][1],
h.Tuple(h.cast(UInt(16), rd), cSAD[rd, xi, yi, xo, yo]),
h.Tuple(disp_left[xi, yi, xo, yo]))
FILTERED = -16
disp = Func("disp")
disp[x, y] = h.select(
# x > xmax-xmin or y > ymax-ymin,
x < xmax,
h.cast(UInt(16), disp_left[x % x_tile_size, y % y_tile_size, x / x_tile_size, y / y_tile_size][0]),
h.cast(UInt(16), FILTERED))
# Schedule
vector_width = 8
disp.compute_root() \
.tile(x, y, xo, yo, xi, yi, x_tile_size, y_tile_size).reorder(xi, yi, xo, yo) \
.vectorize(xi, vector_width).parallel(xo).parallel(yo)
# reorder storage
disp_left.reorder_storage(xi, yi, xo, yo)
diff_T .reorder_storage(xi, yi, xo, yo, d)
vsum .reorder_storage(xi, yi, xo, yo, d)
cSAD .reorder_storage(xi, yi, xo, yo, d)
disp_left.compute_at(disp, xo).reorder(xi, yi, xo, yo) \
.vectorize(xi, vector_width) \
.update() \
.reorder(xi, yi, rd, xo, yo).vectorize(xi, vector_width)
if test:
cSAD.compute_at(disp_left, rd).reorder(xi, yi, xo, yo, d).vectorize(xi, vector_width)
vsum.compute_at(disp_left, rd).reorder(xi, yi, xo, yo, d).vectorize(xi, vector_width)
else:
cSAD.compute_at(disp_left, rd).reorder(xi, yi, xo, yo, d).vectorize(xi, vector_width) \
.update() \
.reorder(yi, rxi, xo, yo, d).vectorize(yi, vector_width)
vsum.compute_at(disp_left, rd).reorder(xi, yi, xo, yo, d).vectorize(xi, vector_width) \
.update() \
.reorder(xi, ryi, xo, yo, d).vectorize(xi, vector_width)
return disp