from simclr import SimCLR
import yaml
from data_aug.dataset_wrapper import DataSetWrapper
import os, glob
import pandas as pd
import argparse

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--magnification', type=str, default='20x')
    parser.add_argument('--dest_weights', type=str)
    args = parser.parse_args()
    config = yaml.load(open("config.yaml", "r"), Loader=yaml.FullLoader)
    dataset = DataSetWrapper(config['batch_size'], **config['dataset'])
    
    simclr = SimCLR(dataset, config, args)
    simclr.train()


if __name__ == "__main__":
    main()