Create new file

handler.py

@@ -0,0 +1,47 @@
+from base64 import b64encode
+from io import BytesIO
+from pathlib import Path
+
+import numpy as np
+from basicsr.archs.rrdbnet_arch import RRDBNet
+from PIL import Image
+from realesrgan import RealESRGANer
+
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+        self.upsampler = RealESRGANer(
+            scale=4,
+            model_path=str(Path(path) / "RealESRGAN_x4plus.pth"),
+            model=model,
+            tile=0,
+            tile_pad=10,
+            pre_pad=0,
+            half=True,
+        )
+
+    def __call__(self, data):
+        """
+        Args:
+            data (:obj:):
+                includes the input data and the parameters for the inference.
+        Return:
+            A :obj:`dict`:. base64 encoded image
+        """
+        image = data.pop("inputs", data)
+        image = Image.open(BytesIO(image)).convert("RGB")
+        image = np.array(image)
+        image = image[:, :, ::-1]  # RGB -> BGR
+
+        image, _ = self.upsampler.enhance(image, outscale=4)
+        image = image[:, :, ::-1]  # BGR -> RGB
+        image = Image.fromarray(image)
+
+        # encode image as base 64
+        buffered = BytesIO()
+        image.save(buffered, format="JPEG")
+        img_str = b64encode(buffered.getvalue())
+
+        # postprocess the prediction
+        return {"image": img_str.decode()}