augmentor.py

import os
import random
import numpy as np
from PIL import Image, ImageEnhance, ImageFilter

save_dir = "augmented"
class DataAugmentor:
    def __init__(self):
        self.augmentations = [
            self.rotate,
            self.flip_horizontal,
            self.flip_vertical,
            self.adjust_brightness,
            self.adjust_contrast,
            self.adjust_saturation,
            self.add_noise,
            self.blur,
            self.crop_and_resize,
        ]
    def rotate(self, img):
        angle = random.uniform(-30, 30)
        return img.rotate(angle, fillcolor=(128, 128, 128))
    
    def flip_horizontal(self, img):
        return img.transpose(Image.FLIP_LEFT_RIGHT)
    
    def flip_vertical(self, img):
        return img.transpose(Image.FLIP_TOP_BOTTOM)
    
    def adjust_brightness(self, img):
        factor = random.uniform(0.7, 1.3)
        enhancer = ImageEnhance.Brightness(img)
        return enhancer.enhance(factor)
    
    def adjust_contrast(self, img):
        factor = random.uniform(0.7, 1.3)
        enhancer = ImageEnhance.Contrast(img)
        return enhancer.enhance(factor)
    
    def adjust_saturation(self, img):
        factor = random.uniform(0.7, 1.3)
        enhancer = ImageEnhance.Color(img)
        return enhancer.enhance(factor)
    
    def add_noise(self, img):
        img_array = np.array(img)
        noise = np.random.randint(-20, 20, img_array.shape, dtype=np.int16)
        noisy = np.clip(img_array.astype(np.int16) + noise, 0, 255).astype(np.uint8)
        return Image.fromarray(noisy)
    
    def blur(self, img):
        return img.filter(ImageFilter.GaussianBlur(radius=random.uniform(0.5, 1.5)))
    
    def crop_and_resize(self, img):
        width, height = img.size
        crop_factor = random.uniform(0.8, 0.95)
        
        new_width = int(width * crop_factor)
        new_height = int(height * crop_factor)
        
        left = random.randint(0, width - new_width)
        top = random.randint(0, height - new_height)
        
        cropped = img.crop((left, top, left + new_width, top + new_height))
        return cropped.resize((width, height), Image.BILINEAR)
    
    def augment(self, img, class_name, path):
        imgs=[]
        for i,fun in enumerate(self.augmentations):
            image=fun(img)
            self.save_augmented_image(image, class_name, path, aug_index=i)
            imgs.append(image)
        return imgs
    def save_augmented_image(self,aug_img, class_name, original_name, aug_index):

        class_dir = os.path.join(save_dir, class_name)
        os.makedirs(class_dir, exist_ok=True)
            
        base_name = os.path.splitext(os.path.basename(original_name))[0]
        save_path = os.path.join(class_dir, f"{base_name}_aug{aug_index}.png")
            
        aug_img.save(save_path)