|
|
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))
|
|
|
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
|
|
|
|
|
|
|
|
|
|
