Spaces:
Runtime error
Runtime error
| import os | |
| import Augmentor | |
| tasks = ["1", "2", "3", "4", "5", "6"] | |
| for task in tasks: | |
| # Loop through all folders in Task 1 and generate augmented images for each class | |
| for i in os.listdir("data/train/raw/Task " + task): | |
| if i != ".DS_Store": | |
| print("Augmenting images in class: ", i) | |
| p = Augmentor.Pipeline(f"data/train/raw/Task {task}/{i}", output_directory=i, save_format="png") | |
| p.rotate(probability=0.8, max_left_rotation=5, max_right_rotation=5) | |
| p.flip_left_right(probability=0.8) | |
| p.zoom_random(probability=0.8, percentage_area=0.8) | |
| p.flip_top_bottom(probability=0.8) | |
| p.random_brightness(probability=0.8, min_factor=0.5, max_factor=1.5) | |
| p.random_contrast(probability=0.8, min_factor=0.5, max_factor=1.5) | |
| p.random_color(probability=0.8, min_factor=0.5, max_factor=1.5) | |
| # Generate 100 - total of original images so that the total number of images in each class is 100 | |
| p.sample(100 - len(p.augmentor_images)) | |
| # Move the folder to data/train/Task 1/augmented | |
| # Create the folder if it does not exist | |
| if not os.path.exists(f"data/train/augmented/Task {task}/"): | |
| os.makedirs(f"data/train/augmented/Task {task}/") | |
| # Move all images in the data/train/Task 1/i folder to data/train/Task 1/augmented/i | |
| os.rename( | |
| f"data/train/raw/Task {task}/{i}/{i}", | |
| f"data/train/augmented/Task {task}/{i}", | |
| ) | |
| # Rename all the augmented images to [01, 02, 03] | |
| number = 0 | |
| for j in os.listdir(f"data/train/augmented/Task {task}/{i}"): | |
| number = int(number) + 1 | |
| if len(str(number)) == 1: | |
| number = "0" + str(number) | |
| os.rename( | |
| f"data/train/augmented/Task {task}/{i}/{j}", | |
| f"data/train/augmented/Task {task}/{i}/{number}.png", | |
| ) | |