Add the deployment option.

handler.py

@@ -9,6 +9,7 @@ from PIL import Image
 import torch
 from torchvision import transforms
 from transformers import AutoModelForImageSegmentation
+from loadimg import load_img
 
 torch.set_float32_matmul_precision(["high", "highest"][0])
 
@@ -104,18 +105,7 @@ class EndpointHandler():
             A :obj:`list` | `dict`: will be serialized and returned
         """
         print('data["inputs"] = ', data["inputs"])
-        image_src = data["inputs"]
-        if isinstance(image_src, str):
-            if os.path.isfile(image_src):
-                image_ori = Image.open(image_src)
-            else:
-                response = requests.get(image_src)
-                image_data = BytesIO(response.content)
-                image_ori = Image.open(image_data)
-        else:
-            image_ori = Image.fromarray(image_src)
-
-        image = image_ori.convert('RGB')
+        image = load_img(data["inputs"]).convert("RGB")
         # Preprocess the image
         image_preprocessor = ImagePreprocessor(resolution=tuple(resolution))
         image_proc = image_preprocessor.proc(image)

requirements.txt

@@ -16,3 +16,5 @@ prettytable
 transformers
 huggingface-hub>0.25
 accelerate
+
+loadimg