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
    probabilities = e / e.sum()
    return [float(p.item()) for p in probabilities] # Convert numpy array elements to Python floats using .item()

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