import os

os.chdir(os.path.dirname(os.path.abspath(__file__)))


def extract_video_opticalflow(dataset_root, test_train, video_dir_name, flow_save_dir, dfparm_d=20, dfparm_a='tvl1'):
    video_dir_path = os.path.abspath(os.path.join(dataset_root, test_train, video_dir_name))
    assert os.path.exists(video_dir_path), video_dir_path
    gpu_dev = os.getpid()%4

    os.system(f"export CUDA_VISIBLE_DEVICES={gpu_dev} && cd {flow_save_dir} && denseflow {video_dir_path} -b={dfparm_d} -a={dfparm_a} -s=1 -if -v")


if __name__ == "__main__":
    dataset_list = ["Ped1", "Ped2", "Avenue", "SHTCampus"]
    dataset_idx = 3
    dataset_root = f"../data/{dataset_list[dataset_idx]}"

    test_train = ["Test", "Train"][1]

    dataset_dir_path = os.path.join(dataset_root, test_train)
    video_dir_name_list = os.listdir(dataset_dir_path)
    video_dir_name_list.sort()

    flow_save_dir = os.path.join(dataset_root, 'Temp')
    if not os.path.exists(flow_save_dir):
        os.mkdir(flow_save_dir)

    # extract_video_opticalflow(dataset_root, test_train, video_dir_name_list[0], flow_save_dir)
    video_dir_name = "01_023"
    extract_video_opticalflow(dataset_root, test_train, video_dir_name, flow_save_dir)

    print("Finished!")