opencv计算光流，warp图片并保存（代码）

import os
import cv2
import numpy as np
import flow_vis

def generate_flow(src, des):

    def mkfile(file):
        if not os.path.exists(file):
            os.makedirs(file)

    file_list=os.listdir(src)
    file_list = sorted(file_list)
    # print(file_list)


    frame1 = cv2.imread(os.path.join(src, file_list[0]))
    prvs = cv2.cvtColor(frame1, cv2.COLOR_BGR2GRAY)
    hsv = np.zeros_like(frame1)


    hsv[...,1] = 255

    dis = cv2.optflow.createOptFlow_DeepFlow()
    for file in file_list:
    # for i in range(1440):
        path = os.path.join(src, file)
        frame2 = cv2.imread(path)
        next = cv2.cvtColor(frame2, cv2.COLOR_BGR2GRAY)

        flow = dis.calc(prvs, next, None)
        # flow = cv2.calcOpticalFlowFarneback(prvs, next, None, 0.5, 3, 25, 3, 5, 1.2, 1)
        # # print(flow.shape)
        # print(flow)

        # mag, ang = cv2.cartToPolar(flow[..., 0], flow[..., 1])
        # hsv[..., 0] = ang * 180 / np.pi / 2
        # hsv[..., 2] = cv2.normalize(mag, None, 0, 255, cv2.NORM_MINMAX)
        # rgb = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
        print("flow.shape:",flow.shape)
        rgb = flow_vis.flow_to_color(flow, convert_to_bgr=False)
        mkfile(des)

        try:
            cv2.imwrite(os.path.join(des, file), rgb)
        except Exception as e:
            continue

        prvs = next
        out_path = 'warped_img/cv_flow_face'
        image_warp(src, flow, out_path)
    print('img number:%d' % len(os.listdir(des)))


def image_warp(src, flow, out_path, mode='bilinear'):
    """Performs a backward warp of an image using the predicted flow.
    numpy version

    Args:
        im: input image. ndim=2, 3 or 4, [[num_batch], height, width, [channels]]. num_batch and channels are optional, default is 1.
        flow: flow vectors. ndim=3 or 4, [[num_batch], height, width, 2]. num_batch is optional
        mode: interpolation mode. 'nearest' or 'bilinear'
    Returns:
        warped: transformed image of the same shape as the input image.
    """
    file_list =os.listdir(src)
    for file in file_list:
        path_img = os.path.join(src, file)

        im= cv2.imread(path_img)

      
        # assert im.ndim == flow.ndim, 'The dimension of im and flow must be equal '
        flag = 4
        if im.ndim == 2:
            height, width = im.shape
            num_batch = 1
            channels = 1
            im = im[np.newaxis, :, :, np.newaxis]
            flow = flow[np.newaxis, :, :]
            flag = 2
        elif im.ndim == 3:
            height, width, channels = im.shape
            num_batch = 1
            im = im[np.newaxis, :, :]
            flow = flow[np.newaxis, :, :]
            flag = 3
        elif im.ndim == 4:
            num_batch, height, width, channels = im.shape
            flag = 4
        else:
            raise AttributeError('The dimension of im must be 2, 3 or 4')

        max_x = width - 1
        max_y = height - 1
        zero = 0

        # We have to flatten our tensors to vectorize the interpolation
        im_flat = np.reshape(im, [-1, channels])
        flow_flat = np.reshape(flow, [-1, 2])

        # Floor the flow, as the final indices are integers
        flow_floor = np.floor(flow_flat).astype(np.int32)

        # Construct base indices which are displaced with the flow
        pos_x = np.tile(np.arange(width), [height * num_batch])
        grid_y = np.tile(np.expand_dims(np.arange(height), 1), [1, width])
        pos_y = np.tile(np.reshape(grid_y, [-1]), [num_batch])

        x = flow_floor[:, 0]
        y = flow_floor[:, 1]
        print(flow_flat.shape)
        x0 = pos_x + x
        y0 = pos_y + y

        x0 = np.clip(x0, zero, max_x)
        y0 = np.clip(y0, zero, max_y)

        dim1 = width * height
        batch_offsets = np.arange(num_batch) * dim1
        base_grid = np.tile(np.expand_dims(batch_offsets, 1), [1, dim1])
        base = np.reshape(base_grid, [-1])

        base_y0 = base + y0 * width

        if mode == 'nearest':
            idx_a = base_y0 + x0
            warped_flat = im_flat[idx_a]
        elif mode == 'bilinear':
            # The fractional part is used to control the bilinear interpolation.
            bilinear_weights = flow_flat - np.floor(flow_flat)

            xw = bilinear_weights[:, 0]
            yw = bilinear_weights[:, 1]

            # Compute interpolation weights for 4 adjacent pixels
            # expand to num_batch * height * width x 1 for broadcasting in add_n below
            wa = np.expand_dims((1 - xw) * (1 - yw), 1) # top left pixel
            wb = np.expand_dims((1 - xw) * yw, 1) # bottom left pixel
            wc = np.expand_dims(xw * (1 - yw), 1) # top right pixel
            wd = np.expand_dims(xw * yw, 1) # bottom right pixel

            x1 = x0 + 1
            y1 = y0 + 1

            x1 = np.clip(x1, zero, max_x)
            y1 = np.clip(y1, zero, max_y)

            base_y1 = base + y1 * width
            idx_a = base_y0 + x0
            idx_b = base_y1 + x0
            idx_c = base_y0 + x1
            idx_d = base_y1 + x1

            Ia = im_flat[idx_a]
            Ib = im_flat[idx_b]
            Ic = im_flat[idx_c]
            Id = im_flat[idx_d]

            warped_flat = wa * Ia + wb * Ib + wc * Ic + wd * Id
        warped = np.reshape(warped_flat, [num_batch, height, width, channels])

        if flag == 2:
            warped = np.squeeze(warped)
        elif flag == 3:
            warped = np.squeeze(warped, axis=0)
        else:
            pass
        warped = warped.astype(np.uint8)

        cv2.imwrite(os.path.join(out_path, file), warped)
        print("saving:",file)
 


original_dir_path = 'demo/your_name'
flow_dir_path = 'output/your_name'
out_path = 'warped_img/your_name'

generate_flow(original_dir_path, flow_dir_path)
# image_warp(original_dir_path, flow_dir_path, out_path)