sohojoe's picture
move api helpers to api_helper.py
44fc764
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