OpenCV提取IsoGD数据集u、v光流

@OpenCV提取IsoGD数据集u、v光流

import cv2
import glob
import numpy as np
import os

def cal_for_frames(video_path):
    for subDirList in os.listdir(video_path):
        path = './a/' + subDirList + '//'
        frames = glob.glob(os.path.join(path, '*.jpg'))
        frames.sort()
        print(frames)
        flow = []
        prev = cv2.imread(frames[0])
        prev = cv2.cvtColor(prev, cv2.COLOR_BGR2GRAY)
        for i, frame_curr in enumerate(frames[1:]):
            curr = cv2.imread(frame_curr)
            curr = cv2.cvtColor(curr, cv2.COLOR_BGR2GRAY)
            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()
    # TVL1 = cv2.DualTVL1OpticalFlow_create()
    # TVL1=cv2.createOptFlow_DualTVL1()
    flow = TVL1.calc(prev, curr, None)
    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):
    for subDirList in os.listdir(flow_path):
        path = './out/' + subDirList + '/'
        if not os.path.exists(os.path.join(path, 'u')):
            os.mkdir(os.path.join(path, 'u'))
        if not os.path.exists(os.path.join(path, 'v')):
            os.mkdir(os.path.join(path, 'v'))
        for i, flow in enumerate(video_flows):
            cv2.imwrite(os.path.join(path,'u', "{:06d}.jpg".format(i)),
                        flow[:, :, 0])
            cv2.imwrite(os.path.join(path,'v', "{:06d}.jpg".format(i)),
                        flow[:, :, 1])
 
 
def extract_flow(video_path, flow_path):
    flow = cal_for_frames(video_path)
    save_flow(flow, flow_path)
    print('complete:' + flow_path)
    return

if __name__ == '__main__':
    dirs = "./a"
    save_path='./out'
    extract_flow(dirs,save_path)

#从指定文件夹写入到指定文件夹

import cv2
import glob
import numpy as np
import os

def cal_for_frames(video_path, flow_path):
    for subDirList in os.listdir(video_path):
        path = 'G:/IsoGD/dataset/train/RGB/' + subDirList + '/'
        print(path)
        frames = glob.glob(os.path.join(path, '*.jpg'))
        frames.sort()
        #print(frames)
        flows = []
        prev = cv2.imread(frames[0])
        prev = cv2.cvtColor(prev, cv2.COLOR_BGR2GRAY)
        for i, frame_curr in enumerate(frames[1:]):
            curr = cv2.imread(frame_curr)
            curr = cv2.cvtColor(curr, cv2.COLOR_BGR2GRAY)
            tmp_flow = compute_TVL1(prev, curr)
            flows.append(tmp_flow)
            prev = curr
        for subDirList in os.listdir(flow_path):
            subPath = 'G:/IsoGD/dataset/train/flow/' + subDirList + '/'
            #print('save path', end = ' ')
            #print(subPath[29: -1])
            if (path[28: -1] == subPath[29: -1]):
                #print('A')
                if not os.path.exists(os.path.join(subPath, 'u')):
                    os.mkdir(os.path.join(subPath, 'u'))
                if not os.path.exists(os.path.join(subPath, 'v')):
                    os.mkdir(os.path.join(subPath, 'v'))
                for i, flow in enumerate(flows):
                    cv2.imwrite(os.path.join(subPath,'u', "{:06d}.jpg".format(i)), flow[:, :, 0])
                    #print('save direction of u has successful!')
                    cv2.imwrite(os.path.join(subPath,'v', "{:06d}.jpg".format(i)), flow[:, :, 1])
                    #print('save direction of v has successful!')
                print('save path: {0}'.format(subPath))
 
def compute_TVL1(prev, curr, bound=15):
    #print('calculate optical-flow...')
    TVL1=cv2.optflow.DualTVL1OpticalFlow_create()
    flow = TVL1.calc(prev, curr, None)
    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

if __name__ == '__main__':
    dirs = "G:/IsoGD/dataset/train/RGB"
    save_path='G:/IsoGD/dataset/train/flow'
    cal_for_frames(dirs, save_path)