Spaces:
Runtime error
Runtime error
File size: 3,248 Bytes
e6f2a04 59908f1 e6f2a04 59908f1 e6f2a04 59908f1 e6f2a04 1bdf45d 59908f1 1bdf45d 59908f1 1bdf45d 59908f1 1bdf45d 59908f1 9d7b040 1bdf45d 59908f1 9d7b040 e6f2a04 1bdf45d e6f2a04 1bdf45d e6f2a04 1bdf45d e6f2a04 1bdf45d e6f2a04 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
import os
import Augmentor
import shutil
from configs import *
import uuid
tasks = ["1", "2", "3", "4", "5", "6"]
for task in ["1"]:
# Loop through all folders in Task 1 and generate augmented images for each class
for class_label in ['Alzheimer Disease', 'Cerebral Palsy', 'Dystonia', 'Essential Tremor', 'Healthy', 'Huntington Disease', 'Parkinson Disease']:
if class_label != ".DS_Store":
print("Augmenting images in class: ", class_label, " in Task ", task)
# Create a temp folder to combine the raw data and the external data
if not os.path.exists(f"{TEMP_DATA_DIR}{task}/{class_label}/"):
os.makedirs(f"{TEMP_DATA_DIR}{task}/{class_label}/")
if os.path.exists(f"{RAW_DATA_DIR}{task}/{class_label}"):
for file in os.listdir(f"{RAW_DATA_DIR}{task}/{class_label}"):
shutil.copy(
f"{RAW_DATA_DIR}{task}/{class_label}/{file}",
f"{TEMP_DATA_DIR}{task}/{class_label}/{str(uuid.uuid4())}.png",
)
if os.path.exists(f"{EXTERNAL_DATA_DIR}{task}/{class_label}"):
for file in os.listdir(f"{EXTERNAL_DATA_DIR}{task}/{class_label}"):
shutil.copy(
f"{EXTERNAL_DATA_DIR}{task}/{class_label}/{file}",
f"{TEMP_DATA_DIR}{task}/{class_label}/{str(uuid.uuid4())}.png",
)
p = Augmentor.Pipeline(
f"{TEMP_DATA_DIR}{task}/{class_label}",
output_directory=f"{class_label}/",
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"{AUG_DATA_DIR}{task}/"):
os.makedirs(f"{AUG_DATA_DIR}{task}/")
# Move all images in the data/train/Task 1/i folder to data/train/Task 1/augmented/i
os.rename(
f"{TEMP_DATA_DIR}{task}/{class_label}/{class_label}",
f"{AUG_DATA_DIR}{task}/{class_label}",
)
# Rename all the augmented images to [01, 02, 03]
number = 0
for file in os.listdir(f"{AUG_DATA_DIR}{task}/{class_label}"):
number = int(number) + 1
if len(str(number)) == 1:
number = "0" + str(number)
os.rename(
f"{AUG_DATA_DIR}{task}/{class_label}/{file}",
f"{AUG_DATA_DIR}{task}/{class_label}/{number}.png",
)
|