utils/utils.py

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

def ELA(img_pil, scale=77, alpha=0.66):
    # Error Level Analysis for basic image forensics
    original = img_pil.copy()          # open up the input image
    temp_path = 'temp.jpg'             # temporary image name to save the ELA to
    original.save(temp_path, quality=95)     # re-save the image with a quality of 95%
    temporary = Image.open(temp_path)        # open up the re-saved image

    diff = ImageChops.difference(original, temporary)   # load in the images to look at pixel by pixel differences
    d = diff.load()                     # load the image into a variable
    WIDTH, HEIGHT = diff.size           # set the size into a tuple
    for x in range(WIDTH):                                  # row by row
        for y in range(HEIGHT):                             # column by column
            d[x, y] = tuple(k * scale for k in d[x, y])     # set the pixels to their x,y & color based on error

    new_img = ImageChops.blend(temporary, diff, alpha)      # blend the original w/ the ELA @ a set alpha/transparency
    return new_img