utils/utils.py

import numpy as np
import io
from PIL import Image
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):
    transform_flip = transforms.Compose([
        transforms.RandomHorizontalFlip(p=1.0)
    ])
    transform_rotate = transforms.Compose([
        transforms.RandomRotation(degrees=(90, 90))
    ])
    augmented_img_flip = transform_flip(img_pil)
    augmented_img_rotate = transform_rotate(img_pil)
    return augmented_img_flip, augmented_img_rotate

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