import numpy as np import io from PIL import Image, ImageFilter, ImageChops from torchvision import transforms def softmax(vector): e = np.exp(vector - np.max(vector)) # for numerical stability return e / e.sum() def augment_image(img_pil, methods, rotate_degrees=0, noise_level=0, sharpen_strength=1): for method in methods: if method == "rotate": img_pil = img_pil.rotate(rotate_degrees) elif method == "add_noise": noise = np.random.normal(0, noise_level, img_pil.size[::-1] + (3,)).astype(np.uint8) img_pil = Image.fromarray(np.clip(np.array(img_pil) + noise, 0, 255).astype(np.uint8)) elif method == "sharpen": img_pil = img_pil.filter(ImageFilter.UnsharpMask(radius=2, percent=sharpen_strength, threshold=3)) return img_pil, img_pil def convert_pil_to_bytes(image, format='JPEG'): img_byte_arr = io.BytesIO() image.save(img_byte_arr, format=format) img_byte_arr = img_byte_arr.getvalue() return img_byte_arr