import os
import click
import yaml
import rich
from rich.console import Console

from idt.factories import SearchEngineFactory
from idt.utils.remove_corrupt import remove_corrupt
from idt.utils.create_dataset_csv import create_dataset_csv
from idt.utils.split_dataset import split_dataset

BANNER = """
[bold blue]=====================================================================


                             
                8888888 8888888b. 88888888888 
                  888   888  "Y88b    888     
                  888   888    888    888     
                  888   888    888    888     
                  888   888    888    888     
                  888   888    888    888     
                  888   888  .d88P    888     
                8888888 8888888P"     888  
                                           
          		[italic]IMAGE DATASET TOOL V0.6[/italic]                                                                                    
                                                                                                                                 
=====================================================================[/bold blue]                                                                                                                                
		"""

#@click.command()
@click.group()
def main():
    
    """
    Image Dataset Builder CLI to create amazing datasets
    """
    pass

@main.command()
def version():
	"""
	Shows what version idt is currently on
	"""
	click.clear()
	rich.print("[bold magenta]Image Dataset Tool (IDT)[/bold magenta] version 0.0.6 beta")

@main.command()
def authors():
	"""
	Shows who are the creators of IDT
	"""
	click.clear()
	rich.print("[bold]IDT[/bold] was initially made by [bold magenta]Deliton Junior[/bold magenta] and [bold red]Misael Kelviny[/bold red]")

@main.command()
@click.option('--input', '-i','--i', help="The name of the thing you want to download")
@click.option('--size', '-s','--s', default=50, help="The number of images you want to download.")
@click.option('--engine', '-e','--e', default="duckgo", help="What search engine will be used to find your images")
@click.option('--resize-method', '-rs','--rs', default="longer_side", help="Resize method adopted. Options: shorter_side, longer_side and smartcrop")
@click.option('--imagesize', '-is','--is', default=512, help="What image size ratio should be applied to your dataset")
@click.option('--api-key', '-ak','--ak', default=None, help="Provide an api-key for the engines that require one")
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")

@main.command()
@click.option('--default', '-d','--d', is_flag=True,default=False, help="Generate a default config file")
def init(default):
	"""
	This command initialyzes idt and creates a dataset config file
	"""
	console = Console()
	console.clear()

	if default:
		document_dict = {
			"DATASET_NAME": "dataset",
			"API_KEY": "",
			"SAMPLES_PER_SEARCH": 50,
			"IMAGE_SIZE": 512,
			"ENGINE": "duckgo",
			"RESIZE_METHOD": "longer_side",
			"CLASSES": [{"CLASS_NAME": "Test", "SEARCH_KEYWORDS": "images of cats"}]}

		if not os.path.exists("dataset.yaml"):
			console.print("[bold]Creating a dataset configuration file...[/bold]")
			
			f = open("dataset.yaml", "w")
			f.write(yaml.dump(document_dict))
			if f:
				console.clear()
				console.print("Dataset YAML file has been created sucessfully. Now run [bold blue]idt build[/bold blue] to mount your dataset!")
				exit(0)
			
		
		else:
			console.print("[red]A dataset.yaml is already created. To use another one, delete the current dataset.yaml file[/red]")
			exit(0)

	console.print(BANNER)
	dataset_name = click.prompt("Insert a name to your dataset: ")

	console.clear()
	samples = click.prompt("How many samples per seach will be necessary?  ",type=int)

	console.clear()
	console.print("[bold]Choose image resolution[/bold]", justify="center")
	console.print("""

[1] 512 pixels / 512 pixels [bold blue](recommended)[/bold blue]
[2] 1024 pixels / 1024 pixels
[3] 256 pixels / 256 pixels
[4] 128 pixels / 128 pixels
[5] Keep original image size

[italic]ps: note that the aspect ratio of the image will [bold]not[/bold] be changed, so possibly the images received will have slightly different size[/italic]
		
		""")


	image_size_ratio = click.prompt("What is the desired image size ratio", type=int)
	while image_size_ratio < 1 or image_size_ratio > 5:
		console.print("[italic red]Invalid option, please choose between 1 and 5. [/italic red]")
		image_size_ratio= click.prompt("\nOption: ",type=int)

	if image_size_ratio == 1:
		image_size_ratio= 512
	elif image_size_ratio == 2:
		image_size_ratio = 1024
	elif image_size_ratio == 3:
		image_size_ratio = 256
	elif image_size_ratio == 4:
		image_size_ratio= 128
	elif image_size_ratio == 5:
		image_size_ratio = 0

	console.clear()
	console.print("[bold]Choose a resize method[/bold]", justify="center")
	console.print("""

[1] Resize image based on longer side
[2] Resize image based on shorter side
[3] Smartcrop

[italic]ps: note that the aspect ratio of the image will [bold]not[/bold] be changed, so possibly the images received will have slightly different size[/italic]
		
		""")
	resize_method = click.prompt("Desired Image resize method: ", type=int)
	while resize_method < 1 or resize_method > 3:
		console.print("[red]Invalid option[/red]")
		resize_method = click.prompt("Choose method [1-3]: ")

	resize_method_options = ['','longer_side','shorter_side','smartcrop']


	console.clear()
	number_of_classes = click.prompt("How many image classes are required? ",type=int)

	document_dict = {
  
    "DATASET_NAME": dataset_name,
  
    "SAMPLES_PER_SEARCH": samples,
 
    "IMAGE_SIZE": image_size_ratio,
  
    "RESIZE_METHOD": resize_method_options[resize_method],
  
    "CLASSES": []
  
}

	console.clear()
	console.print("[bold]Choose a search engine[/bold]", justify="center")
	console.print("""

[1] Duck GO [bold blue](recommended)[/bold blue]
[2] Bing
[3] Bing API [italic yellow](Requires API key)[/italic yellow]
[4] Flickr API [italic yellow](Requires API key)[/italic yellow]

		""")
	search_engine= click.prompt("Select option:", type=int)
	while search_engine < 0 or search_engine > 4:
		console.print("[italic red]Invalid option, please choose between 1 and 4.[/italic red]")
		search_engine = click.prompt("\nOption: ", type=int)

	search_options = ['none', 'duckgo', 'bing', 'bing_api', 'flickr_api']
	document_dict['ENGINE'] = search_options[search_engine]

	if search_engine > 2:
		console.clear()
		console.print(f'Insert your [bold blue]{search_options[search_engine]}[/bold blue] API key')
		engine_api_key = click.prompt("API key: ", type=str)
		document_dict['API_KEY'] = engine_api_key
	else:
		document_dict['API_KEY'] = "NONE"

	search_engine = search_options[search_engine]

	for x in range(number_of_classes):
		console.clear()
		class_name = click.prompt("Class {x} name: ".format(x=x+1))
		console.clear()

		console.print("""In order to achieve better results, choose several keywords that will be provided to the search engine to find your class in different settings.
	
[bold blue]Example: [/bold blue]

Class Name: [bold yellow]Pineapple[/bold yellow]
[italic]keywords[/italic]: [underline]pineapple, pineapple fruit, ananas, abacaxi, pineapple drawing[/underline]

			""")
		keywords = click.prompt("Type in all keywords used to find your desired class, separated by commas: ")
		document_dict['CLASSES'].append({'CLASS_NAME': class_name, 'SEARCH_KEYWORDS': keywords})
    
	if not os.path.exists("dataset.yaml"):
		console.print("[bold]Creating a dataset configuration file...[/bold]")
		try:
			f = open("dataset.yaml", "w")
			f.write(yaml.dump(document_dict))
			if f:
				console.clear()
				console.print("Dataset YAML file has been created sucessfully. Now run [bold blue]idt build[/bold blue] to mount your dataset!")
		except:
			console.print("[red]Unable to create file. Please check permission[/red]")
		
	else:
		console.print("[red]A dataset.yaml is already created. To use another one, delete the current dataset.yaml file[/red]")

@main.command()
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']:
		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()
	console.print("Dataset READY!")

@main.command()
def split():
	"""
	Split dataset into train/valid folders
	"""
	console = Console()
	while True:
		click.clear()
		console.print(BANNER)
		console.print("Choose the desired proportion of images of each class to be distributed in train/valid folders. [bold]What percentage of images should be distributed towards training?[/bold] ")
		train_proportion = click.prompt("(0-100)", type=int)
		validation_proportion = 100 - train_proportion
		if train_proportion < 0 or train_proportion > 100:
			click.clear()
			console.print("[red]Please provide a valid amount. Choose a number between 0 and 100 to be assigned to training.[/red]")
			continue
		else:
			click.clear()
			console.print("[bold blue]{train} percent[/bold blue] of the images will be moved to a [bold yellow]train[/bold yellow] folder, while [bold blue]{valid} percent [/bold blue] of the remaining images will be stored in a [bold yellow]validation[/bold yellow] folder.".format(train=train_proportion, valid=validation_proportion))
			c= click.prompt("Is that ok? [Y/n]")
			if c.lower() == 'y':
				if not os.path.exists("dataset.yaml"):
					click.clear()
					console.print("Dataset config file not found\nRun - [bold blue]idt init[/bold blue]")
					exit(0)

				with open('dataset.yaml') as f:
					click.clear()
					console.print("[italic]Copying files to the train/valid folders. Please wait...[/italic]")
					data = yaml.load(f, Loader=yaml.FullLoader)
					split_dataset(data['DATASET_NAME'], train_proportion)
				console.clear()
				console.print("[bold blue]Done[/bold blue]")
				break
			else:
				continue
		
	

if __name__ == "__main__":
	main()

b24fc531f8fff6f38fb3bb06f5c83bb1777ee15d
<mask token>


@click.group()
def main():
    """
    Image Dataset Builder CLI to create amazing datasets
    """
    pass


<mask token>


@main.command()
def authors():
    """
	Shows who are the creators of IDT
	"""
    click.clear()
    rich.print(
        '[bold]IDT[/bold] was initially made by [bold magenta]Deliton Junior[/bold magenta] and [bold red]Misael Kelviny[/bold red]'
        )


@main.command()
@click.option('--input', '-i', '--i', help=
    'The name of the thing you want to download')
@click.option('--size', '-s', '--s', default=50, help=
    'The number of images you want to download.')
@click.option('--engine', '-e', '--e', default='duckgo', help=
    'What search engine will be used to find your images')
@click.option('--resize-method', '-rs', '--rs', default='longer_side', help
    ='Resize method adopted. Options: shorter_side, longer_side and smartcrop')
@click.option('--imagesize', '-is', '--is', default=512, help=
    'What image size ratio should be applied to your dataset')
@click.option('--api-key', '-ak', '--ak', default=None, help=
    'Provide an api-key for the engines that require one')
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('dataset')
    else:
        rich.print('Please provide a valid name')


<mask token>


@main.command()
def split():
    """
	Split dataset into train/valid folders
	"""
    console = Console()
    while True:
        click.clear()
        console.print(BANNER)
        console.print(
            'Choose the desired proportion of images of each class to be distributed in train/valid folders. [bold]What percentage of images should be distributed towards training?[/bold] '
            )
        train_proportion = click.prompt('(0-100)', type=int)
        validation_proportion = 100 - train_proportion
        if train_proportion < 0 or train_proportion > 100:
            click.clear()
            console.print(
                '[red]Please provide a valid amount. Choose a number between 0 and 100 to be assigned to training.[/red]'
                )
            continue
        else:
            click.clear()
            console.print(
                '[bold blue]{train} percent[/bold blue] of the images will be moved to a [bold yellow]train[/bold yellow] folder, while [bold blue]{valid} percent [/bold blue] of the remaining images will be stored in a [bold yellow]validation[/bold yellow] folder.'
                .format(train=train_proportion, valid=validation_proportion))
            c = click.prompt('Is that ok? [Y/n]')
            if c.lower() == 'y':
                if not os.path.exists('dataset.yaml'):
                    click.clear()
                    console.print(
                        """Dataset config file not found
Run - [bold blue]idt init[/bold blue]"""
                        )
                    exit(0)
                with open('dataset.yaml') as f:
                    click.clear()
                    console.print(
                        '[italic]Copying files to the train/valid folders. Please wait...[/italic]'
                        )
                    data = yaml.load(f, Loader=yaml.FullLoader)
                    split_dataset(data['DATASET_NAME'], train_proportion)
                console.clear()
                console.print('[bold blue]Done[/bold blue]')
                break
            else:
                continue


<mask token>

<mask token>


@click.group()
def main():
    """
    Image Dataset Builder CLI to create amazing datasets
    """
    pass


@main.command()
def version():
    """
	Shows what version idt is currently on
	"""
    click.clear()
    rich.print(
        '[bold magenta]Image Dataset Tool (IDT)[/bold magenta] version 0.0.6 beta'
        )


@main.command()
def authors():
    """
	Shows who are the creators of IDT
	"""
    click.clear()
    rich.print(
        '[bold]IDT[/bold] was initially made by [bold magenta]Deliton Junior[/bold magenta] and [bold red]Misael Kelviny[/bold red]'
        )


@main.command()
@click.option('--input', '-i', '--i', help=
    'The name of the thing you want to download')
@click.option('--size', '-s', '--s', default=50, help=
    'The number of images you want to download.')
@click.option('--engine', '-e', '--e', default='duckgo', help=
    'What search engine will be used to find your images')
@click.option('--resize-method', '-rs', '--rs', default='longer_side', help
    ='Resize method adopted. Options: shorter_side, longer_side and smartcrop')
@click.option('--imagesize', '-is', '--is', default=512, help=
    'What image size ratio should be applied to your dataset')
@click.option('--api-key', '-ak', '--ak', default=None, help=
    'Provide an api-key for the engines that require one')
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('dataset')
    else:
        rich.print('Please provide a valid name')


@main.command()
@click.option('--default', '-d', '--d', is_flag=True, default=False, help=
    'Generate a default config file')
def init(default):
    """
	This command initialyzes idt and creates a dataset config file
	"""
    console = Console()
    console.clear()
    if default:
        document_dict = {'DATASET_NAME': 'dataset', 'API_KEY': '',
            'SAMPLES_PER_SEARCH': 50, 'IMAGE_SIZE': 512, 'ENGINE': 'duckgo',
            'RESIZE_METHOD': 'longer_side', 'CLASSES': [{'CLASS_NAME':
            'Test', 'SEARCH_KEYWORDS': 'images of cats'}]}
        if not os.path.exists('dataset.yaml'):
            console.print(
                '[bold]Creating a dataset configuration file...[/bold]')
            f = open('dataset.yaml', 'w')
            f.write(yaml.dump(document_dict))
            if f:
                console.clear()
                console.print(
                    'Dataset YAML file has been created sucessfully. Now run [bold blue]idt build[/bold blue] to mount your dataset!'
                    )
                exit(0)
        else:
            console.print(
                '[red]A dataset.yaml is already created. To use another one, delete the current dataset.yaml file[/red]'
                )
            exit(0)
    console.print(BANNER)
    dataset_name = click.prompt('Insert a name to your dataset: ')
    console.clear()
    samples = click.prompt('How many samples per seach will be necessary?  ',
        type=int)
    console.clear()
    console.print('[bold]Choose image resolution[/bold]', justify='center')
    console.print(
        """

[1] 512 pixels / 512 pixels [bold blue](recommended)[/bold blue]
[2] 1024 pixels / 1024 pixels
[3] 256 pixels / 256 pixels
[4] 128 pixels / 128 pixels
[5] Keep original image size

[italic]ps: note that the aspect ratio of the image will [bold]not[/bold] be changed, so possibly the images received will have slightly different size[/italic]
		
		"""
        )
    image_size_ratio = click.prompt('What is the desired image size ratio',
        type=int)
    while image_size_ratio < 1 or image_size_ratio > 5:
        console.print(
            '[italic red]Invalid option, please choose between 1 and 5. [/italic red]'
            )
        image_size_ratio = click.prompt('\nOption: ', type=int)
    if image_size_ratio == 1:
        image_size_ratio = 512
    elif image_size_ratio == 2:
        image_size_ratio = 1024
    elif image_size_ratio == 3:
        image_size_ratio = 256
    elif image_size_ratio == 4:
        image_size_ratio = 128
    elif image_size_ratio == 5:
        image_size_ratio = 0
    console.clear()
    console.print('[bold]Choose a resize method[/bold]', justify='center')
    console.print(
        """

[1] Resize image based on longer side
[2] Resize image based on shorter side
[3] Smartcrop

[italic]ps: note that the aspect ratio of the image will [bold]not[/bold] be changed, so possibly the images received will have slightly different size[/italic]
		
		"""
        )
    resize_method = click.prompt('Desired Image resize method: ', type=int)
    while resize_method < 1 or resize_method > 3:
        console.print('[red]Invalid option[/red]')
        resize_method = click.prompt('Choose method [1-3]: ')
    resize_method_options = ['', 'longer_side', 'shorter_side', 'smartcrop']
    console.clear()
    number_of_classes = click.prompt('How many image classes are required? ',
        type=int)
    document_dict = {'DATASET_NAME': dataset_name, 'SAMPLES_PER_SEARCH':
        samples, 'IMAGE_SIZE': image_size_ratio, 'RESIZE_METHOD':
        resize_method_options[resize_method], 'CLASSES': []}
    console.clear()
    console.print('[bold]Choose a search engine[/bold]', justify='center')
    console.print(
        """

[1] Duck GO [bold blue](recommended)[/bold blue]
[2] Bing
[3] Bing API [italic yellow](Requires API key)[/italic yellow]
[4] Flickr API [italic yellow](Requires API key)[/italic yellow]

		"""
        )
    search_engine = click.prompt('Select option:', type=int)
    while search_engine < 0 or search_engine > 4:
        console.print(
            '[italic red]Invalid option, please choose between 1 and 4.[/italic red]'
            )
        search_engine = click.prompt('\nOption: ', type=int)
    search_options = ['none', 'duckgo', 'bing', 'bing_api', 'flickr_api']
    document_dict['ENGINE'] = search_options[search_engine]
    if search_engine > 2:
        console.clear()
        console.print(
            f'Insert your [bold blue]{search_options[search_engine]}[/bold blue] API key'
            )
        engine_api_key = click.prompt('API key: ', type=str)
        document_dict['API_KEY'] = engine_api_key
    else:
        document_dict['API_KEY'] = 'NONE'
    search_engine = search_options[search_engine]
    for x in range(number_of_classes):
        console.clear()
        class_name = click.prompt('Class {x} name: '.format(x=x + 1))
        console.clear()
        console.print(
            """In order to achieve better results, choose several keywords that will be provided to the search engine to find your class in different settings.
	
[bold blue]Example: [/bold blue]

Class Name: [bold yellow]Pineapple[/bold yellow]
[italic]keywords[/italic]: [underline]pineapple, pineapple fruit, ananas, abacaxi, pineapple drawing[/underline]

			"""
            )
        keywords = click.prompt(
            'Type in all keywords used to find your desired class, separated by commas: '
            )
        document_dict['CLASSES'].append({'CLASS_NAME': class_name,
            'SEARCH_KEYWORDS': keywords})
    if not os.path.exists('dataset.yaml'):
        console.print('[bold]Creating a dataset configuration file...[/bold]')
        try:
            f = open('dataset.yaml', 'w')
            f.write(yaml.dump(document_dict))
            if f:
                console.clear()
                console.print(
                    'Dataset YAML file has been created sucessfully. Now run [bold blue]idt build[/bold blue] to mount your dataset!'
                    )
        except:
            console.print(
                '[red]Unable to create file. Please check permission[/red]')
    else:
        console.print(
            '[red]A dataset.yaml is already created. To use another one, delete the current dataset.yaml file[/red]'
            )


@main.command()
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']:
        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(data['DATASET_NAME'])
    create_dataset_csv(data['DATASET_NAME'])
    click.clear()
    console.print('Dataset READY!')


@main.command()
def split():
    """
	Split dataset into train/valid folders
	"""
    console = Console()
    while True:
        click.clear()
        console.print(BANNER)
        console.print(
            'Choose the desired proportion of images of each class to be distributed in train/valid folders. [bold]What percentage of images should be distributed towards training?[/bold] '
            )
        train_proportion = click.prompt('(0-100)', type=int)
        validation_proportion = 100 - train_proportion
        if train_proportion < 0 or train_proportion > 100:
            click.clear()
            console.print(
                '[red]Please provide a valid amount. Choose a number between 0 and 100 to be assigned to training.[/red]'
                )
            continue
        else:
            click.clear()
            console.print(
                '[bold blue]{train} percent[/bold blue] of the images will be moved to a [bold yellow]train[/bold yellow] folder, while [bold blue]{valid} percent [/bold blue] of the remaining images will be stored in a [bold yellow]validation[/bold yellow] folder.'
                .format(train=train_proportion, valid=validation_proportion))
            c = click.prompt('Is that ok? [Y/n]')
            if c.lower() == 'y':
                if not os.path.exists('dataset.yaml'):
                    click.clear()
                    console.print(
                        """Dataset config file not found
Run - [bold blue]idt init[/bold blue]"""
                        )
                    exit(0)
                with open('dataset.yaml') as f:
                    click.clear()
                    console.print(
                        '[italic]Copying files to the train/valid folders. Please wait...[/italic]'
                        )
                    data = yaml.load(f, Loader=yaml.FullLoader)
                    split_dataset(data['DATASET_NAME'], train_proportion)
                console.clear()
                console.print('[bold blue]Done[/bold blue]')
                break
            else:
                continue


if __name__ == '__main__':
    main()

<mask token>
BANNER = """
[bold blue]=====================================================================


                             
                8888888 8888888b. 88888888888 
                  888   888  "Y88b    888     
                  888   888    888    888     
                  888   888    888    888     
                  888   888    888    888     
                  888   888    888    888     
                  888   888  .d88P    888     
                8888888 8888888P"     888  
                                           
          		[italic]IMAGE DATASET TOOL V0.6[/italic]                                                                                    
                                                                                                                                 
=====================================================================[/bold blue]                                                                                                                                
		"""


@click.group()
def main():
    """
    Image Dataset Builder CLI to create amazing datasets
    """
    pass


@main.command()
def version():
    """
	Shows what version idt is currently on
	"""
    click.clear()
    rich.print(
        '[bold magenta]Image Dataset Tool (IDT)[/bold magenta] version 0.0.6 beta'
        )


@main.command()
def authors():
    """
	Shows who are the creators of IDT
	"""
    click.clear()
    rich.print(
        '[bold]IDT[/bold] was initially made by [bold magenta]Deliton Junior[/bold magenta] and [bold red]Misael Kelviny[/bold red]'
        )


@main.command()
@click.option('--input', '-i', '--i', help=
    'The name of the thing you want to download')
@click.option('--size', '-s', '--s', default=50, help=
    'The number of images you want to download.')
@click.option('--engine', '-e', '--e', default='duckgo', help=
    'What search engine will be used to find your images')
@click.option('--resize-method', '-rs', '--rs', default='longer_side', help
    ='Resize method adopted. Options: shorter_side, longer_side and smartcrop')
@click.option('--imagesize', '-is', '--is', default=512, help=
    'What image size ratio should be applied to your dataset')
@click.option('--api-key', '-ak', '--ak', default=None, help=
    'Provide an api-key for the engines that require one')
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('dataset')
    else:
        rich.print('Please provide a valid name')


@main.command()
@click.option('--default', '-d', '--d', is_flag=True, default=False, help=
    'Generate a default config file')
def init(default):
    """
	This command initialyzes idt and creates a dataset config file
	"""
    console = Console()
    console.clear()
    if default:
        document_dict = {'DATASET_NAME': 'dataset', 'API_KEY': '',
            'SAMPLES_PER_SEARCH': 50, 'IMAGE_SIZE': 512, 'ENGINE': 'duckgo',
            'RESIZE_METHOD': 'longer_side', 'CLASSES': [{'CLASS_NAME':
            'Test', 'SEARCH_KEYWORDS': 'images of cats'}]}
        if not os.path.exists('dataset.yaml'):
            console.print(
                '[bold]Creating a dataset configuration file...[/bold]')
            f = open('dataset.yaml', 'w')
            f.write(yaml.dump(document_dict))
            if f:
                console.clear()
                console.print(
                    'Dataset YAML file has been created sucessfully. Now run [bold blue]idt build[/bold blue] to mount your dataset!'
                    )
                exit(0)
        else:
            console.print(
                '[red]A dataset.yaml is already created. To use another one, delete the current dataset.yaml file[/red]'
                )
            exit(0)
    console.print(BANNER)
    dataset_name = click.prompt('Insert a name to your dataset: ')
    console.clear()
    samples = click.prompt('How many samples per seach will be necessary?  ',
        type=int)
    console.clear()
    console.print('[bold]Choose image resolution[/bold]', justify='center')
    console.print(
        """

[1] 512 pixels / 512 pixels [bold blue](recommended)[/bold blue]
[2] 1024 pixels / 1024 pixels
[3] 256 pixels / 256 pixels
[4] 128 pixels / 128 pixels
[5] Keep original image size

[italic]ps: note that the aspect ratio of the image will [bold]not[/bold] be changed, so possibly the images received will have slightly different size[/italic]
		
		"""
        )
    image_size_ratio = click.prompt('What is the desired image size ratio',
        type=int)
    while image_size_ratio < 1 or image_size_ratio > 5:
        console.print(
            '[italic red]Invalid option, please choose between 1 and 5. [/italic red]'
            )
        image_size_ratio = click.prompt('\nOption: ', type=int)
    if image_size_ratio == 1:
        image_size_ratio = 512
    elif image_size_ratio == 2:
        image_size_ratio = 1024
    elif image_size_ratio == 3:
        image_size_ratio = 256
    elif image_size_ratio == 4:
        image_size_ratio = 128
    elif image_size_ratio == 5:
        image_size_ratio = 0
    console.clear()
    console.print('[bold]Choose a resize method[/bold]', justify='center')
    console.print(
        """

[1] Resize image based on longer side
[2] Resize image based on shorter side
[3] Smartcrop

[italic]ps: note that the aspect ratio of the image will [bold]not[/bold] be changed, so possibly the images received will have slightly different size[/italic]
		
		"""
        )
    resize_method = click.prompt('Desired Image resize method: ', type=int)
    while resize_method < 1 or resize_method > 3:
        console.print('[red]Invalid option[/red]')
        resize_method = click.prompt('Choose method [1-3]: ')
    resize_method_options = ['', 'longer_side', 'shorter_side', 'smartcrop']
    console.clear()
    number_of_classes = click.prompt('How many image classes are required? ',
        type=int)
    document_dict = {'DATASET_NAME': dataset_name, 'SAMPLES_PER_SEARCH':
        samples, 'IMAGE_SIZE': image_size_ratio, 'RESIZE_METHOD':
        resize_method_options[resize_method], 'CLASSES': []}
    console.clear()
    console.print('[bold]Choose a search engine[/bold]', justify='center')
    console.print(
        """

[1] Duck GO [bold blue](recommended)[/bold blue]
[2] Bing
[3] Bing API [italic yellow](Requires API key)[/italic yellow]
[4] Flickr API [italic yellow](Requires API key)[/italic yellow]

		"""
        )
    search_engine = click.prompt('Select option:', type=int)
    while search_engine < 0 or search_engine > 4:
        console.print(
            '[italic red]Invalid option, please choose between 1 and 4.[/italic red]'
            )
        search_engine = click.prompt('\nOption: ', type=int)
    search_options = ['none', 'duckgo', 'bing', 'bing_api', 'flickr_api']
    document_dict['ENGINE'] = search_options[search_engine]
    if search_engine > 2:
        console.clear()
        console.print(
            f'Insert your [bold blue]{search_options[search_engine]}[/bold blue] API key'
            )
        engine_api_key = click.prompt('API key: ', type=str)
        document_dict['API_KEY'] = engine_api_key
    else:
        document_dict['API_KEY'] = 'NONE'
    search_engine = search_options[search_engine]
    for x in range(number_of_classes):
        console.clear()
        class_name = click.prompt('Class {x} name: '.format(x=x + 1))
        console.clear()
        console.print(
            """In order to achieve better results, choose several keywords that will be provided to the search engine to find your class in different settings.
	
[bold blue]Example: [/bold blue]

Class Name: [bold yellow]Pineapple[/bold yellow]
[italic]keywords[/italic]: [underline]pineapple, pineapple fruit, ananas, abacaxi, pineapple drawing[/underline]

			"""
            )
        keywords = click.prompt(
            'Type in all keywords used to find your desired class, separated by commas: '
            )
        document_dict['CLASSES'].append({'CLASS_NAME': class_name,
            'SEARCH_KEYWORDS': keywords})
    if not os.path.exists('dataset.yaml'):
        console.print('[bold]Creating a dataset configuration file...[/bold]')
        try:
            f = open('dataset.yaml', 'w')
            f.write(yaml.dump(document_dict))
            if f:
                console.clear()
                console.print(
                    'Dataset YAML file has been created sucessfully. Now run [bold blue]idt build[/bold blue] to mount your dataset!'
                    )
        except:
            console.print(
                '[red]Unable to create file. Please check permission[/red]')
    else:
        console.print(
            '[red]A dataset.yaml is already created. To use another one, delete the current dataset.yaml file[/red]'
            )


@main.command()
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']:
        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(data['DATASET_NAME'])
    create_dataset_csv(data['DATASET_NAME'])
    click.clear()
    console.print('Dataset READY!')


@main.command()
def split():
    """
	Split dataset into train/valid folders
	"""
    console = Console()
    while True:
        click.clear()
        console.print(BANNER)
        console.print(
            'Choose the desired proportion of images of each class to be distributed in train/valid folders. [bold]What percentage of images should be distributed towards training?[/bold] '
            )
        train_proportion = click.prompt('(0-100)', type=int)
        validation_proportion = 100 - train_proportion
        if train_proportion < 0 or train_proportion > 100:
            click.clear()
            console.print(
                '[red]Please provide a valid amount. Choose a number between 0 and 100 to be assigned to training.[/red]'
                )
            continue
        else:
            click.clear()
            console.print(
                '[bold blue]{train} percent[/bold blue] of the images will be moved to a [bold yellow]train[/bold yellow] folder, while [bold blue]{valid} percent [/bold blue] of the remaining images will be stored in a [bold yellow]validation[/bold yellow] folder.'
                .format(train=train_proportion, valid=validation_proportion))
            c = click.prompt('Is that ok? [Y/n]')
            if c.lower() == 'y':
                if not os.path.exists('dataset.yaml'):
                    click.clear()
                    console.print(
                        """Dataset config file not found
Run - [bold blue]idt init[/bold blue]"""
                        )
                    exit(0)
                with open('dataset.yaml') as f:
                    click.clear()
                    console.print(
                        '[italic]Copying files to the train/valid folders. Please wait...[/italic]'
                        )
                    data = yaml.load(f, Loader=yaml.FullLoader)
                    split_dataset(data['DATASET_NAME'], train_proportion)
                console.clear()
                console.print('[bold blue]Done[/bold blue]')
                break
            else:
                continue


if __name__ == '__main__':
    main()

import os
import click
import yaml
import rich
from rich.console import Console
from idt.factories import SearchEngineFactory
from idt.utils.remove_corrupt import remove_corrupt
from idt.utils.create_dataset_csv import create_dataset_csv
from idt.utils.split_dataset import split_dataset
BANNER = """
[bold blue]=====================================================================


                             
                8888888 8888888b. 88888888888 
                  888   888  "Y88b    888     
                  888   888    888    888     
                  888   888    888    888     
                  888   888    888    888     
                  888   888    888    888     
                  888   888  .d88P    888     
                8888888 8888888P"     888  
                                           
          		[italic]IMAGE DATASET TOOL V0.6[/italic]                                                                                    
                                                                                                                                 
=====================================================================[/bold blue]                                                                                                                                
		"""


@click.group()
def main():
    """
    Image Dataset Builder CLI to create amazing datasets
    """
    pass


@main.command()
def version():
    """
	Shows what version idt is currently on
	"""
    click.clear()
    rich.print(
        '[bold magenta]Image Dataset Tool (IDT)[/bold magenta] version 0.0.6 beta'
        )


@main.command()
def authors():
    """
	Shows who are the creators of IDT
	"""
    click.clear()
    rich.print(
        '[bold]IDT[/bold] was initially made by [bold magenta]Deliton Junior[/bold magenta] and [bold red]Misael Kelviny[/bold red]'
        )


@main.command()
@click.option('--input', '-i', '--i', help=
    'The name of the thing you want to download')
@click.option('--size', '-s', '--s', default=50, help=
    'The number of images you want to download.')
@click.option('--engine', '-e', '--e', default='duckgo', help=
    'What search engine will be used to find your images')
@click.option('--resize-method', '-rs', '--rs', default='longer_side', help
    ='Resize method adopted. Options: shorter_side, longer_side and smartcrop')
@click.option('--imagesize', '-is', '--is', default=512, help=
    'What image size ratio should be applied to your dataset')
@click.option('--api-key', '-ak', '--ak', default=None, help=
    'Provide an api-key for the engines that require one')
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('dataset')
    else:
        rich.print('Please provide a valid name')


@main.command()
@click.option('--default', '-d', '--d', is_flag=True, default=False, help=
    'Generate a default config file')
def init(default):
    """
	This command initialyzes idt and creates a dataset config file
	"""
    console = Console()
    console.clear()
    if default:
        document_dict = {'DATASET_NAME': 'dataset', 'API_KEY': '',
            'SAMPLES_PER_SEARCH': 50, 'IMAGE_SIZE': 512, 'ENGINE': 'duckgo',
            'RESIZE_METHOD': 'longer_side', 'CLASSES': [{'CLASS_NAME':
            'Test', 'SEARCH_KEYWORDS': 'images of cats'}]}
        if not os.path.exists('dataset.yaml'):
            console.print(
                '[bold]Creating a dataset configuration file...[/bold]')
            f = open('dataset.yaml', 'w')
            f.write(yaml.dump(document_dict))
            if f:
                console.clear()
                console.print(
                    'Dataset YAML file has been created sucessfully. Now run [bold blue]idt build[/bold blue] to mount your dataset!'
                    )
                exit(0)
        else:
            console.print(
                '[red]A dataset.yaml is already created. To use another one, delete the current dataset.yaml file[/red]'
                )
            exit(0)
    console.print(BANNER)
    dataset_name = click.prompt('Insert a name to your dataset: ')
    console.clear()
    samples = click.prompt('How many samples per seach will be necessary?  ',
        type=int)
    console.clear()
    console.print('[bold]Choose image resolution[/bold]', justify='center')
    console.print(
        """

[1] 512 pixels / 512 pixels [bold blue](recommended)[/bold blue]
[2] 1024 pixels / 1024 pixels
[3] 256 pixels / 256 pixels
[4] 128 pixels / 128 pixels
[5] Keep original image size

[italic]ps: note that the aspect ratio of the image will [bold]not[/bold] be changed, so possibly the images received will have slightly different size[/italic]
		
		"""
        )
    image_size_ratio = click.prompt('What is the desired image size ratio',
        type=int)
    while image_size_ratio < 1 or image_size_ratio > 5:
        console.print(
            '[italic red]Invalid option, please choose between 1 and 5. [/italic red]'
            )
        image_size_ratio = click.prompt('\nOption: ', type=int)
    if image_size_ratio == 1:
        image_size_ratio = 512
    elif image_size_ratio == 2:
        image_size_ratio = 1024
    elif image_size_ratio == 3:
        image_size_ratio = 256
    elif image_size_ratio == 4:
        image_size_ratio = 128
    elif image_size_ratio == 5:
        image_size_ratio = 0
    console.clear()
    console.print('[bold]Choose a resize method[/bold]', justify='center')
    console.print(
        """

[1] Resize image based on longer side
[2] Resize image based on shorter side
[3] Smartcrop

[italic]ps: note that the aspect ratio of the image will [bold]not[/bold] be changed, so possibly the images received will have slightly different size[/italic]
		
		"""
        )
    resize_method = click.prompt('Desired Image resize method: ', type=int)
    while resize_method < 1 or resize_method > 3:
        console.print('[red]Invalid option[/red]')
        resize_method = click.prompt('Choose method [1-3]: ')
    resize_method_options = ['', 'longer_side', 'shorter_side', 'smartcrop']
    console.clear()
    number_of_classes = click.prompt('How many image classes are required? ',
        type=int)
    document_dict = {'DATASET_NAME': dataset_name, 'SAMPLES_PER_SEARCH':
        samples, 'IMAGE_SIZE': image_size_ratio, 'RESIZE_METHOD':
        resize_method_options[resize_method], 'CLASSES': []}
    console.clear()
    console.print('[bold]Choose a search engine[/bold]', justify='center')
    console.print(
        """

[1] Duck GO [bold blue](recommended)[/bold blue]
[2] Bing
[3] Bing API [italic yellow](Requires API key)[/italic yellow]
[4] Flickr API [italic yellow](Requires API key)[/italic yellow]

		"""
        )
    search_engine = click.prompt('Select option:', type=int)
    while search_engine < 0 or search_engine > 4:
        console.print(
            '[italic red]Invalid option, please choose between 1 and 4.[/italic red]'
            )
        search_engine = click.prompt('\nOption: ', type=int)
    search_options = ['none', 'duckgo', 'bing', 'bing_api', 'flickr_api']
    document_dict['ENGINE'] = search_options[search_engine]
    if search_engine > 2:
        console.clear()
        console.print(
            f'Insert your [bold blue]{search_options[search_engine]}[/bold blue] API key'
            )
        engine_api_key = click.prompt('API key: ', type=str)
        document_dict['API_KEY'] = engine_api_key
    else:
        document_dict['API_KEY'] = 'NONE'
    search_engine = search_options[search_engine]
    for x in range(number_of_classes):
        console.clear()
        class_name = click.prompt('Class {x} name: '.format(x=x + 1))
        console.clear()
        console.print(
            """In order to achieve better results, choose several keywords that will be provided to the search engine to find your class in different settings.
	
[bold blue]Example: [/bold blue]

Class Name: [bold yellow]Pineapple[/bold yellow]
[italic]keywords[/italic]: [underline]pineapple, pineapple fruit, ananas, abacaxi, pineapple drawing[/underline]

			"""
            )
        keywords = click.prompt(
            'Type in all keywords used to find your desired class, separated by commas: '
            )
        document_dict['CLASSES'].append({'CLASS_NAME': class_name,
            'SEARCH_KEYWORDS': keywords})
    if not os.path.exists('dataset.yaml'):
        console.print('[bold]Creating a dataset configuration file...[/bold]')
        try:
            f = open('dataset.yaml', 'w')
            f.write(yaml.dump(document_dict))
            if f:
                console.clear()
                console.print(
                    'Dataset YAML file has been created sucessfully. Now run [bold blue]idt build[/bold blue] to mount your dataset!'
                    )
        except:
            console.print(
                '[red]Unable to create file. Please check permission[/red]')
    else:
        console.print(
            '[red]A dataset.yaml is already created. To use another one, delete the current dataset.yaml file[/red]'
            )


@main.command()
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']:
        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(data['DATASET_NAME'])
    create_dataset_csv(data['DATASET_NAME'])
    click.clear()
    console.print('Dataset READY!')


@main.command()
def split():
    """
	Split dataset into train/valid folders
	"""
    console = Console()
    while True:
        click.clear()
        console.print(BANNER)
        console.print(
            'Choose the desired proportion of images of each class to be distributed in train/valid folders. [bold]What percentage of images should be distributed towards training?[/bold] '
            )
        train_proportion = click.prompt('(0-100)', type=int)
        validation_proportion = 100 - train_proportion
        if train_proportion < 0 or train_proportion > 100:
            click.clear()
            console.print(
                '[red]Please provide a valid amount. Choose a number between 0 and 100 to be assigned to training.[/red]'
                )
            continue
        else:
            click.clear()
            console.print(
                '[bold blue]{train} percent[/bold blue] of the images will be moved to a [bold yellow]train[/bold yellow] folder, while [bold blue]{valid} percent [/bold blue] of the remaining images will be stored in a [bold yellow]validation[/bold yellow] folder.'
                .format(train=train_proportion, valid=validation_proportion))
            c = click.prompt('Is that ok? [Y/n]')
            if c.lower() == 'y':
                if not os.path.exists('dataset.yaml'):
                    click.clear()
                    console.print(
                        """Dataset config file not found
Run - [bold blue]idt init[/bold blue]"""
                        )
                    exit(0)
                with open('dataset.yaml') as f:
                    click.clear()
                    console.print(
                        '[italic]Copying files to the train/valid folders. Please wait...[/italic]'
                        )
                    data = yaml.load(f, Loader=yaml.FullLoader)
                    split_dataset(data['DATASET_NAME'], train_proportion)
                console.clear()
                console.print('[bold blue]Done[/bold blue]')
                break
            else:
                continue


if __name__ == '__main__':
    main()

import os
import click
import yaml
import rich
from rich.console import Console

from idt.factories import SearchEngineFactory
from idt.utils.remove_corrupt import remove_corrupt
from idt.utils.create_dataset_csv import create_dataset_csv
from idt.utils.split_dataset import split_dataset

BANNER = """
[bold blue]=====================================================================


                             
                8888888 8888888b. 88888888888 
                  888   888  "Y88b    888     
                  888   888    888    888     
                  888   888    888    888     
                  888   888    888    888     
                  888   888    888    888     
                  888   888  .d88P    888     
                8888888 8888888P"     888  
                                           
          		[italic]IMAGE DATASET TOOL V0.6[/italic]                                                                                    
                                                                                                                                 
=====================================================================[/bold blue]                                                                                                                                
		"""

#@click.command()
@click.group()
def main():
    
    """
    Image Dataset Builder CLI to create amazing datasets
    """
    pass

@main.command()
def version():
	"""
	Shows what version idt is currently on
	"""
	click.clear()
	rich.print("[bold magenta]Image Dataset Tool (IDT)[/bold magenta] version 0.0.6 beta")

@main.command()
def authors():
	"""
	Shows who are the creators of IDT
	"""
	click.clear()
	rich.print("[bold]IDT[/bold] was initially made by [bold magenta]Deliton Junior[/bold magenta] and [bold red]Misael Kelviny[/bold red]")

@main.command()
@click.option('--input', '-i','--i', help="The name of the thing you want to download")
@click.option('--size', '-s','--s', default=50, help="The number of images you want to download.")
@click.option('--engine', '-e','--e', default="duckgo", help="What search engine will be used to find your images")
@click.option('--resize-method', '-rs','--rs', default="longer_side", help="Resize method adopted. Options: shorter_side, longer_side and smartcrop")
@click.option('--imagesize', '-is','--is', default=512, help="What image size ratio should be applied to your dataset")
@click.option('--api-key', '-ak','--ak', default=None, help="Provide an api-key for the engines that require one")
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")

@main.command()
@click.option('--default', '-d','--d', is_flag=True,default=False, help="Generate a default config file")
def init(default):
	"""
	This command initialyzes idt and creates a dataset config file
	"""
	console = Console()
	console.clear()

	if default:
		document_dict = {
			"DATASET_NAME": "dataset",
			"API_KEY": "",
			"SAMPLES_PER_SEARCH": 50,
			"IMAGE_SIZE": 512,
			"ENGINE": "duckgo",
			"RESIZE_METHOD": "longer_side",
			"CLASSES": [{"CLASS_NAME": "Test", "SEARCH_KEYWORDS": "images of cats"}]}

		if not os.path.exists("dataset.yaml"):
			console.print("[bold]Creating a dataset configuration file...[/bold]")
			
			f = open("dataset.yaml", "w")
			f.write(yaml.dump(document_dict))
			if f:
				console.clear()
				console.print("Dataset YAML file has been created sucessfully. Now run [bold blue]idt build[/bold blue] to mount your dataset!")
				exit(0)
			
		
		else:
			console.print("[red]A dataset.yaml is already created. To use another one, delete the current dataset.yaml file[/red]")
			exit(0)

	console.print(BANNER)
	dataset_name = click.prompt("Insert a name to your dataset: ")

	console.clear()
	samples = click.prompt("How many samples per seach will be necessary?  ",type=int)

	console.clear()
	console.print("[bold]Choose image resolution[/bold]", justify="center")
	console.print("""

[1] 512 pixels / 512 pixels [bold blue](recommended)[/bold blue]
[2] 1024 pixels / 1024 pixels
[3] 256 pixels / 256 pixels
[4] 128 pixels / 128 pixels
[5] Keep original image size

[italic]ps: note that the aspect ratio of the image will [bold]not[/bold] be changed, so possibly the images received will have slightly different size[/italic]
		
		""")


	image_size_ratio = click.prompt("What is the desired image size ratio", type=int)
	while image_size_ratio < 1 or image_size_ratio > 5:
		console.print("[italic red]Invalid option, please choose between 1 and 5. [/italic red]")
		image_size_ratio= click.prompt("\nOption: ",type=int)

	if image_size_ratio == 1:
		image_size_ratio= 512
	elif image_size_ratio == 2:
		image_size_ratio = 1024
	elif image_size_ratio == 3:
		image_size_ratio = 256
	elif image_size_ratio == 4:
		image_size_ratio= 128
	elif image_size_ratio == 5:
		image_size_ratio = 0

	console.clear()
	console.print("[bold]Choose a resize method[/bold]", justify="center")
	console.print("""

[1] Resize image based on longer side
[2] Resize image based on shorter side
[3] Smartcrop

[italic]ps: note that the aspect ratio of the image will [bold]not[/bold] be changed, so possibly the images received will have slightly different size[/italic]
		
		""")
	resize_method = click.prompt("Desired Image resize method: ", type=int)
	while resize_method < 1 or resize_method > 3:
		console.print("[red]Invalid option[/red]")
		resize_method = click.prompt("Choose method [1-3]: ")

	resize_method_options = ['','longer_side','shorter_side','smartcrop']


	console.clear()
	number_of_classes = click.prompt("How many image classes are required? ",type=int)

	document_dict = {
  
    "DATASET_NAME": dataset_name,
  
    "SAMPLES_PER_SEARCH": samples,
 
    "IMAGE_SIZE": image_size_ratio,
  
    "RESIZE_METHOD": resize_method_options[resize_method],
  
    "CLASSES": []
  
}

	console.clear()
	console.print("[bold]Choose a search engine[/bold]", justify="center")
	console.print("""

[1] Duck GO [bold blue](recommended)[/bold blue]
[2] Bing
[3] Bing API [italic yellow](Requires API key)[/italic yellow]
[4] Flickr API [italic yellow](Requires API key)[/italic yellow]

		""")
	search_engine= click.prompt("Select option:", type=int)
	while search_engine < 0 or search_engine > 4:
		console.print("[italic red]Invalid option, please choose between 1 and 4.[/italic red]")
		search_engine = click.prompt("\nOption: ", type=int)

	search_options = ['none', 'duckgo', 'bing', 'bing_api', 'flickr_api']
	document_dict['ENGINE'] = search_options[search_engine]

	if search_engine > 2:
		console.clear()
		console.print(f'Insert your [bold blue]{search_options[search_engine]}[/bold blue] API key')
		engine_api_key = click.prompt("API key: ", type=str)
		document_dict['API_KEY'] = engine_api_key
	else:
		document_dict['API_KEY'] = "NONE"

	search_engine = search_options[search_engine]

	for x in range(number_of_classes):
		console.clear()
		class_name = click.prompt("Class {x} name: ".format(x=x+1))
		console.clear()

		console.print("""In order to achieve better results, choose several keywords that will be provided to the search engine to find your class in different settings.
	
[bold blue]Example: [/bold blue]

Class Name: [bold yellow]Pineapple[/bold yellow]
[italic]keywords[/italic]: [underline]pineapple, pineapple fruit, ananas, abacaxi, pineapple drawing[/underline]

			""")
		keywords = click.prompt("Type in all keywords used to find your desired class, separated by commas: ")
		document_dict['CLASSES'].append({'CLASS_NAME': class_name, 'SEARCH_KEYWORDS': keywords})
    
	if not os.path.exists("dataset.yaml"):
		console.print("[bold]Creating a dataset configuration file...[/bold]")
		try:
			f = open("dataset.yaml", "w")
			f.write(yaml.dump(document_dict))
			if f:
				console.clear()
				console.print("Dataset YAML file has been created sucessfully. Now run [bold blue]idt build[/bold blue] to mount your dataset!")
		except:
			console.print("[red]Unable to create file. Please check permission[/red]")
		
	else:
		console.print("[red]A dataset.yaml is already created. To use another one, delete the current dataset.yaml file[/red]")

@main.command()
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']:
		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()
	console.print("Dataset READY!")

@main.command()
def split():
	"""
	Split dataset into train/valid folders
	"""
	console = Console()
	while True:
		click.clear()
		console.print(BANNER)
		console.print("Choose the desired proportion of images of each class to be distributed in train/valid folders. [bold]What percentage of images should be distributed towards training?[/bold] ")
		train_proportion = click.prompt("(0-100)", type=int)
		validation_proportion = 100 - train_proportion
		if train_proportion < 0 or train_proportion > 100:
			click.clear()
			console.print("[red]Please provide a valid amount. Choose a number between 0 and 100 to be assigned to training.[/red]")
			continue
		else:
			click.clear()
			console.print("[bold blue]{train} percent[/bold blue] of the images will be moved to a [bold yellow]train[/bold yellow] folder, while [bold blue]{valid} percent [/bold blue] of the remaining images will be stored in a [bold yellow]validation[/bold yellow] folder.".format(train=train_proportion, valid=validation_proportion))
			c= click.prompt("Is that ok? [Y/n]")
			if c.lower() == 'y':
				if not os.path.exists("dataset.yaml"):
					click.clear()
					console.print("Dataset config file not found\nRun - [bold blue]idt init[/bold blue]")
					exit(0)

				with open('dataset.yaml') as f:
					click.clear()
					console.print("[italic]Copying files to the train/valid folders. Please wait...[/italic]")
					data = yaml.load(f, Loader=yaml.FullLoader)
					split_dataset(data['DATASET_NAME'], train_proportion)
				console.clear()
				console.print("[bold blue]Done[/bold blue]")
				break
			else:
				continue
		
	

if __name__ == "__main__":
	main()
