SquirrelDetector/scripts/select_val.py

import argparse
import json
import os
import random


def read_dups(similarity):
    with open(similarity) as f:
        # read in the list of lists
        data = json.load(f)

    # flatten list of lists into a single list
    data = sum(data, [])
    return set(data)

    
def gen_val(train, percent, avoid):
    num = int(len(train) * percent)
    valid_files = list(set(train).difference(set(avoid)))
    if num > len(valid_files):
        return valid_files
    return random.sample(valid_files, num)


def main():
    parser = argparse.ArgumentParser('Selects a validation dataset, taking into account which images are duplicates. This prevents leakage from the training set to the validation set.')
    parser.add_argument('--similarity', required=True, help='similarity.json file')
    parser.add_argument('--train', required=True, help='training directory')
    parser.add_argument('--val', required=True, help='validation directory')
    parser.add_argument('--percent', type=float, default=0.2, help='percent of the training set to use for validation')

    args = parser.parse_args()

    dups = read_dups(args.similarity)
    train = list(os.listdir(args.train))

    val = gen_val(train, args.percent, dups)
    val.sort()

    with open('validation.txt', mode='w') as f:
        f.write('\n'.join(val))
        f.write('\n')

    os.makedirs(args.val, exist_ok=True)

    for img in val:
        inpath = os.path.join(args.train, img)
        outpath = os.path.join(args.val, img)
        os.rename(inpath, outpath)

    
if __name__ == '__main__':
    main()