def run()

in build_graph/tools/generate_sp_matches.py [0:0]

• 29 lines of code
• 15 McCabe index (conditional complexity)

def run():

    random.seed(1234)
    match_dir = f'build_graph/data/{args.dset}/matches/{args.split}/'
    os.makedirs(match_dir, exist_ok=True)

    # Generate the list of clips to calculate matches across.
    # For epic: get matches across all clips in the same kitchen
    # For gtea: get matches across every clip
    if args.dset=='epic':
        # args.chunk should be the person ID (P01, P02 ...)
        # P22 is split into P22a + P22b because it has too many clips
        dset = epic.EPICInteractions('data/epic', args.split, 32)
        entries = [entry for entry in dset.data if args.chunk[0:3] in entry['uid']]

        uid_to_entry = {entry['uid']:entry for entry in entries}
        uids = sorted(uid_to_entry.keys())

        if args.chunk == 'P22a':
            shuffle(uids)
            uids = uids[0:2000]
        elif args.chunk == 'P22b':
            shuffle(uids)
            uids = uids[2000:]


    elif args.dset=='gtea':
        # args.chunk should be one of [train1, train2, train3, train4, val]
        dset = gtea.GTEAInteractions('data/gtea', args.split, 32)
        entries =  dset.data # all of them are in the same kitchen!

        uid_to_entry = {entry['uid']:entry for entry in entries}
        uids = sorted(uid_to_entry.keys())

        # train/val chunks only exists for their respective splits
        if args.split=='train' and args.chunk=='val':
            return
        if args.split=='val' and args.chunk!='val':
            return

        if args.split=='train':
            shuffle(uids)
            idx = int(args.chunk.split('train')[1])
            uids = uids[2069*(idx-1):2069*idx]


    # generate superpoint matches
    generate_matches(uids, match_dir)