File size: 1,619 Bytes
65c9ec6 32f0f58 65c9ec6 32f0f58 39bb242 32f0f58 |
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 |
# install thing, just like in segment anything
from typing import Dict, List, Any
from PIL import Image
from io import BytesIO
from transformers import AutoModelForSemanticSegmentation, AutoFeatureExtractor
import base64
import torch
from torch import nn
import subprocess
result = subprocess.run(["pip", "install", "git+https://github.com/sberbank-ai/Real-ESRGAN.git"], check=True)
print(f"git+https://github.com/sberbank-ai/Real-ESRGAN.git = {result}")
from RealESRGAN import RealESRGAN
class EndpointHandler():
def __init__(self, path="."):
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model = RealESRGAN(self.device, scale=2)
self.model.load_weights('/repository/RealESRGAN_x2.pth', download=True)
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
"""
data args:
images (:obj:`PIL.Image`)
candiates (:obj:`list`)
Return:
A :obj:`list`:. The list contains items that are dicts should be liked {"label": "XXX", "score": 0.82}
"""
inputs = data.pop("inputs", data)
# decode base64 image to PIL
image = Image.open(BytesIO(base64.b64decode(inputs['image'])))
# forward pass
output_image = self.model.predict(image)
# base64 encode output
buffered = BytesIO()
output_image = output_image.convert('RGB')
output_image.save(buffered, format="png")
img_str = base64.b64encode(buffered.getvalue())
# postprocess the prediction
return {"image": img_str.decode()} |