main_processing.py 1.31 KB
Newer Older
bailuo's avatar
init  
bailuo committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
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_])