视频-->帧-->提取光流

import os
import numpy as np
import cv2
from glob import glob
from multiprocessing import Pool
from sklearn.preprocessing import MinMaxScaler
import matplotlib.pyplot as plt
from tqdm import tqdm

#%%

_IMAGE_SIZE = 256

#%%

def create_paths(base_path,activity):
    activity_path=os.path.join(base_path,activity)
    rgb_folder = '{}/{}_rgb'.format(activity_path,activity)
    flow_folder1 = '{}/{}_flow/u'.format(activity_path,activity)
    flow_folder2 = '{}/{}_flow/v'.format(activity_path,activity)
    if not os.path.exists(rgb_folder):
        os.makedirs(rgb_folder)
    if not os.path.exists(flow_folder1):
        os.makedirs(flow_folder1)
    if not os.path.exists(flow_folder2):
        os.makedirs(flow_folder2)


def vid_to_image(base_path,file,frame_path):
    video_path = os.path.join(base_path,file)
    vidcap = cv2.VideoCapture(video_path)
    success,image = vidcap.read()
    count = 0
    frame_no=0
    while success:
        vidcap.set(1,frame_no)
        #print(os.path.join(frame_path,"frame%d.jpg"))
        cv2.imwrite(os.path.join(frame_path,"frame{:06d}.jpg".format(count)), image)     # save frame as JPEG file
        success,image = vidcap.read()
        #print('Read a new frame: ', success)
        count += 1
        frame_no+=1

def cal_for_frames(video_path):
    #计算flow
    frames = glob(os.path.join(video_path, '*.jpg'))
    frames.sort() #排序
    flow = []
    prev = cv2.imread(frames[0])
    prev = cv2.cvtColor(prev, cv2.COLOR_BGR2GRAY)
    prev = cv2.resize(prev,(224,224))
    for i, frame_curr in tqdm(enumerate(frames)):
        curr = cv2.imread(frame_curr)
        curr = cv2.resize(curr,(224,224))
        curr = cv2.cvtColor(curr, cv2.COLOR_BGR2GRAY) #RGB转换为GRAY灰度图
        tmp_flow = compute_TVL1(prev, curr)
        flow.append(tmp_flow)
        prev = curr
    
    return flow

def compute_TVL1(prev, curr, bound=15):
    """Compute the TV-L1 optical flow.计算稠密光流"""
    TVL1 = cv2.optflow.DualTVL1OpticalFlow_create()
    flow = TVL1.calc(prev, curr, None)
    flow = np.clip(flow, -20,20) #default values are +20 and -20
    assert flow.dtype == np.float32
    flow = (flow + bound) * (255.0 / (2*bound))
    flow = np.round(flow).astype(int)
    flow[flow >= 255] = 255
    flow[flow <= 0] = 0
    return flow


def save_flow(video_flows, flow_path):
    #The optical flows are generated in 3D. However for I3D only first two channels are used. u is the first channel
    #and v is the second channel. Both u and v are saved in separate folders in the flow_path directory.
    #The u and v folders will be generated by calling create_path() function
    n_py=[]
    for i, flow in enumerate(video_flows):
        # print(np.array(flow).shape) #(224, 224, 2)
        cv2.imwrite(os.path.join(flow_path.format('u'), "{:06d}.jpg".format(i)),flow[:, :, 0])
        cv2.imwrite(os.path.join(flow_path.format('v'), "{:06d}.jpg".format(i)),flow[:, :, 1])
        n_py.append(np.dstack([flow[:, :, 0],flow[:, :, 1]]))
    return n_py



def extract_flow(args):
    video_path, flow_path = args
    flow = cal_for_frames(video_path)
    save_flow(flow, flow_path)
    print('complete:' + flow_path)
    return

## Create npy file for rgb files as described in deepmind I3D
def norm_rgb(activity,base_path,rgb_path,nchannel):
    #将rgb图像转换为.npy文件
    npy_file=[]
    frames = glob(os.path.join(rgb_path, '*.jpg')) #该方法返回所有匹配的文件路径列表（list）
    frames.sort()
    for frame in frames:
        assert os.path.exists(frame)
        img=cv2.imread(frame)
        img_new=(cv2.resize(img,(224,224))).astype(float)
        img_norm=np.divide(2*(img_new-img_new.min()),(img_new.max()-img_new.min()))-1 #最大最小值归一化
        npy_file.append(img_norm)
    
    npy_file=np.reshape(np.asarray(npy_file),(1,len(frames),224,224,nchannel))
    np.save(base_path+'/data_input_rgb_{}.npy'.format(activity),npy_file)
    return npy_file


## Create npy file of flow files as described in deepmind I3D
def norm_flow(npy_flow,activity,base_path,rgb_path,nchannel):
    frames = glob(os.path.join(rgb_path, '*.jpg'))
    frames.sort()
    # for frame in frames:
    #     assert os.path.exists(frame)
    #     img=cv2.imread(frame,0) #以灰度模式加载图片，可以直接写0
    #     img_new=(cv2.resize(img,(224,224))).astype(float)
    #     img_norm=np.divide(2*(img_new-img_new.min()),(img_new.max()-img_new.min()))-1
    #     npy_file.append(img_norm[:,:,:-1])
    
    npy_file = np.reshape(np.asarray(npy_flow),(1,len(frames),224,224,nchannel-1)).astype(float)
    #clip between range [0,40]
    #npy_file=np.clip(npy_file,-20,20)
    #rescale betwwen [-1,1]
    npy_file = np.divide(2*(npy_file-npy_file.min()),(npy_file.max()-npy_file.min()))-1
    np.save(base_path+'/data_input_flow_{}.npy'.format(activity),npy_file)
    return npy_file

def demo(base_path, activity):
    #this function creates folders u and v which store the optical flow images
    create_paths(base_path,activity)
    #This is the path where images sampled from the video are to be saved
    frame_path='{}/{}/{}_rgb/'.format(base_path,activity,activity) #../data/example/example_rgb/
    
    # this is the path where you want save the flow files.
    flow_path='{}/{}/{}_flow/{}'.format(base_path,activity,activity,{}) #../data/example/example_flow/{}
    
    # ====================0 视频转帧
    # Convert your video to frames and save in to activity_rgb folder
    vid_to_image(base_path,activity+'.mp4',frame_path)
    
    # ====================1 calculate flows
    flow = cal_for_frames(frame_path)
    #save flows to folders u and v，保存flow的.jpg到u和v两个文件夹中
    npy_flow = save_flow(flow, flow_path) #标准化之前的flow的list
    
    # ====================2 标准化Convert images to .npyfile as descibed in  https://github.com/deepmind/kinetics-i3d
    # 将rgb和flow.jpg转换为.npy数据,同时进行标准化，保存并返回相应file中的np.array
    np_rgb_norm=norm_rgb(activity,base_path,frame_path,3) #../data/example/example_rgb/
    np_flow_norm=norm_flow(npy_flow,activity,base_path,frame_path,3)
    
    # ====================3 将光流转换为视频进行可视化
    # 这里直接读取py文件所在目录下的pics目录所有图片。
    fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')  # 设置输出视频为mp4格式
    # cap_fps是帧率，可以根据随意设置
    cap_fps = 10
    size = (224,224)
    flow = ['u','v']
    for flow_type in flow:
        video_path = os.path.join(base_path,activity,'flow_{}.mp4'.format(flow_type))
        video = cv2.VideoWriter(video_path, fourcc, cap_fps, size)
        path = flow_path.format(flow_type)
        file_lst = glob(os.path.join(path, '*.jpg'))
        file_lst.sort()
        for filename in file_lst:
            assert os.path.exists(filename)
            img = cv2.imread(filename)
            video.write(img)
    
    # =====================4 可视化其中一个flow的图片
    xx=np.reshape(np.asarray(npy_flow),(1,156,224,224,2))
    plt.imshow(xx[0][50][:,:,0])
    plt.show()

if __name__ == '__main__':
    #name of activity what is in the video. This is just for purpose of creating proper folders.
    activity = 'example'
    #This is the path where your video is kept. for simplicity rename the video same as activity. for eg. laughing.mp4
    base_path = '../demo_data'
    demo(base_path, activity)
    exit()