optical_flow.py 文件源码

def optical_flow_ssd(input_img1, input_img2):

    h = input_img1.shape[0]
    w = input_img1.shape[1]

    output_img = np.zeros(input_img1.shape, 'float32')
    output_img[:, :, :] = input_img2[:, :, :]
    u = np.zeros([h, w], 'float32')
    v = np.zeros([h, w], 'float32')

    window_size = 21
    offset = int((window_size - 1) / 2.0)

    summed_ssd = np.zeros([input_img1.shape[0], input_img1.shape[1], window_size ** 2], 'float32')

    for r in range(window_size):
        for c in range(window_size):
            offset_y = r - offset
            offset_x = c - offset
            ind = window_size * r + c
            summed_ssd[:, :, ind] = calc_summed_ssd(input_img1, input_img2, offset_y, offset_x)

    max_of = 0.0

    for r in range(offset, h - offset):
        for c in range(offset, w - offset):

            ssd = summed_ssd[r + offset, c + offset, :] + summed_ssd[r - offset - 1, c - offset - 1, :] - summed_ssd[r + offset, c - offset - 1, :] - summed_ssd[r - offset - 1, c + offset, :]

            ind = np.argmin(ssd)

            offset_c = ind % window_size
            offset_r = (ind - offset_c) // window_size

            offset_y = offset_r - offset
            offset_x = offset_c - offset

            u[r, c] = offset_x
            v[r, c] = offset_y

            if (u[r, c] ** 2 + v[r, c] ** 2) ** 0.5 > max_of:
                max_of = (u[r, c] ** 2 + v[r, c] ** 2) ** 0.5

    for r in range(offset, h - offset, 10):
        for c in range(offset, w - offset, 10):

            scale = ((u[r, c] ** 2 + v[r, c] ** 2) ** 0.5) / max_of
            dist = np.array([u[r, c], v[r, c]], 'float32') / max_of
            color = find_color_value(dist, scale)
            output_img = draw_arrow(float(r), float(c), float(r - u[r, c]), float(c - v[r, c]), output_img, color = color)

    output_img = output_img[20 : h - 20, 20 : w - 20, :]

    return output_img

#transform(TypeSpecialize(checks=False))