File size: 2,081 Bytes
6537db2 073b383 6537db2 687e8c3 0a34c89 073b383 3c52688 073b383 6537db2 0a34c89 6537db2 687e8c3 0a34c89 6537db2 073b383 6537db2 4021f2d 073b383 6537db2 687e8c3 b9a007e 65c97a3 c07b674 b9a007e 6537db2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
from typing import Dict, List, Any
import base64
import io
import torch
import numpy as np
import torch.nn.functional as F
from serkan import SimpleUpscaleModel
import os
from PIL import Image
def decode_image(base64_str: str) -> np.ndarray:
"""Decode base64 string to an image (numpy array)"""
image_data = base64.b64decode(base64_str)
image = Image.open(io.BytesIO(image_data))
return np.array(image)
class EndpointHandler():
def __init__(self, path="."):
# load the optimized model
self.model = SimpleUpscaleModel()
model_path = os.path.join(path, "model_weights.pth")
self.model.load_state_dict(torch.load(model_path))
def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
"""
Args:
data (:obj:):
includes the input data and the parameters for the inference.
Return:
A :obj:`list`:. The object returned should be a list of one list like [[{"label": 0.9939950108528137}]] containing :
- "label": A string representing what the label/class is. There can be multiple labels.
- "score": A score between 0 and 1 describing how confident the model is for this label/class.
"""
inputs = data.pop("inputs", data)
img = inputs["image"]
img = decode_image(img)
img = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0).float()
# Load the image
upscaled = self.model(img)
upscaled = upscaled.squeeze(0).permute(1,2,0)
upscaled = upscaled.numpy()
upscaled = np.clip(upscaled, 0, 255).astype(np.uint8)
pil = Image.fromarray(upscaled)
# Save the image to a buffer
buffered = io.BytesIO()
pil.save(buffered, format="PNG")
img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
# Return a dictionary with the base64 image and additional data
return {
"image": img_str
}
# postprocess the prediction
return "OKAY" |