import os
import shutil
import sys
import subprocess
import argparse

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_dir', type=str, required=True, help='dataset dir')
    parser.add_argument('--model', default='models/raft-things.pth', help="restore checkpoint")
    parser.add_argument('--cycle_th', type=float, default=3., help='threshold for cycle consistency error')
    parser.add_argument('--chain', action='store_true', help='if chaining cycle consistent flows (optional)')

    args = parser.parse_args()

    root = '/your_code_path/omnimotion/preprocessing/'
    for files in os.listdir(args.data_dir):
        data_dir_ = os.path.join(args.data_dir, files)

        # compute raft optical flows between all pairs
        os.chdir(root + 'RAFT')
        subprocess.run(['python', 'exhaustive_raft.py', '--data_dir', data_dir_, '--model', args.model])

        # compute dino feature maps
        os.chdir(root + 'dino')
        subprocess.run(['python', 'extract_dino_features.py', '--data_dir', data_dir_])

        # filtering
        os.chdir(root + 'RAFT')
        subprocess.run(['python', 'filter_raft.py', '--data_dir', data_dir_, '--cycle_th', str(args.cycle_th)])

        # chaining (optional)
        subprocess.run(['python', 'chain_raft.py', '--data_dir', data_dir_])