# 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()}