deepfake / utils.py
Ayan515's picture
Upload 10 files
eb56a55 verified
Raw History Blame Contribute Delete
1.05 kB
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