Spaces:
Sleeping
Sleeping
File size: 2,156 Bytes
f054618 |
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 63 64 65 |
import os
from PIL import Image
from tensorflow.keras.preprocessing.image import ImageDataGenerator
# Set paths
base_dir = 'data/chest_xray'
val_dir = os.path.join(base_dir, 'val')
normal_class_dir = os.path.join(val_dir, 'NORMAL')
pneumonia_class_dir = os.path.join(val_dir, 'PNEUMONIA')
def augment_images(class_directory, num_augmented_images):
datagen = ImageDataGenerator(
rescale=1. / 255,
rotation_range=20,
width_shift_range=0.2,
height_shift_range=0.2,
shear_range=0.2,
zoom_range=0.2,
horizontal_flip=True,
fill_mode='nearest'
)
generator = datagen.flow_from_directory(
directory=os.path.dirname(class_directory), # Parent directory
target_size=(150, 150),
batch_size=1,
class_mode=None,
shuffle=False,
classes=[os.path.basename(class_directory)] # Specify class if using subdirectory
)
print(f"Found {generator.samples} images in {class_directory}")
if generator.samples == 0:
print("No images found in the directory.")
return
count = 0
while count < num_augmented_images:
try:
img_batch = generator.__next__() # Use __next__() to get image batch
img = (img_batch[0] * 255).astype('uint8') # Extract the first image in the batch
img_pil = Image.fromarray(img)
img_path = os.path.join(class_directory, f"augmented_{count}.png")
img_pil.save(img_path)
count += 1
except StopIteration:
print("No more images to generate.")
break
print(f"Total augmented images created: {count}")
# Number of augmented images to generate
num_augmented_images_normal = 2944 - 3875 # This should be a negative number since NORMAL is already balanced
num_augmented_images_pneumonia = 2944 - 1171 # To match the number of NORMAL images
# Generate augmented images for the NORMAL class
augment_images(normal_class_dir, max(num_augmented_images_normal, 0))
# Generate augmented images for the PNEUMONIA class
augment_images(pneumonia_class_dir, num_augmented_images_pneumonia)
|