from PIL import Image
import numpy as np
import base64
import json
from torchvision.transforms import Compose, Resize, CenterCrop

# support sending images as base64

def encode_numpy_array(image_np):
    # Flatten the numpy array and convert it to bytes
    image_bytes = image_np.tobytes()

    # Encode the byte data as base64
    encoded_image = base64.b64encode(image_bytes).decode()
    payload = {
        "encoded_image": encoded_image,
        "width": image_np.shape[1],
        "height": image_np.shape[0],
        "channels": image_np.shape[2],
    }
    payload_json = json.dumps(payload)
    return payload_json

def decode_numpy_array(payload):
    payload_json = json.loads(payload)
    # payload_json = payload.json()
    encoded_image = payload_json["encoded_image"]
    width = payload_json["width"]
    height = payload_json["height"]
    channels = payload_json["channels"]
    # Decode the base64 data
    decoded_image = base64.b64decode(encoded_image)

    # Convert the byte data back to a NumPy array
    image_np = np.frombuffer(decoded_image, dtype=np.uint8).reshape(height, width, channels)

    return image_np


def preprocess_image(image_np, max_size=224):
    # Convert the numpy array to a PIL image
    image = Image.fromarray(image_np)

    # Define the transformation pipeline
    transforms = Compose([
        Resize(max_size, interpolation=Image.BICUBIC),
        CenterCrop(max_size),
    ])

    # Apply the transformations to the image
    image = transforms(image)

    # Convert the PIL image back to a numpy array
    image_np = np.array(image)

    return image_np