def build(): """ This command mounts the dataset """ console = Console() console.clear() console.print(BANNER) if not os.path.exists("dataset.yaml"): click.clear() console.print("Dataset config file not found\nRun - idt init\n") exit(0) with open('dataset.yaml') as f: data = yaml.load(f, Loader=yaml.FullLoader) click.clear() console.print( "Building [bold blue]{dataset_name}[/bold blue] dataset...\n".format(dataset_name=data['DATASET_NAME'])) for classes in data['CLASSES']: console.print(classes['CLASS_NAME']) continue click.clear() console.print('Creating [bold blue]{name} class[/bold blue] \n'.format(name=classes['CLASS_NAME'])) search_list = classes['SEARCH_KEYWORDS'].split(",") for keywords in search_list: factory = SearchEngineFactory(keywords, data['SAMPLES_PER_SEARCH'], classes['CLASS_NAME'], data['RESIZE_METHOD'], data['DATASET_NAME'], data['IMAGE_SIZE'], data['ENGINE'], data['API_KEY']) # Remove corrupt files remove_corrupt(data['DATASET_NAME']) # Create a CSV with dataset info create_dataset_csv(data['DATASET_NAME']) click.clear()
def run(input, size, engine, resize_method, imagesize, api_key): """ This command executes a single search and downloads it """ engine_list = ['duckgo', 'bing', 'bing_api', 'flickr_api'] click.clear() if input and engine in engine_list: factory = SearchEngineFactory(input, size, input, resize_method, "dataset", imagesize, engine, api_key) # Remove corrupt files remove_corrupt("dataset") else: rich.print("Please provide a valid name")