Add inference endpoint feature.

README.md

@@ -8,10 +8,6 @@ tags:
 - model_hub_mixin
 repo_url: https://github.com/ZhengPeng7/BiRefNet-legacy
 pipeline_tag: image-to-image
-widget:
-  - src: >-
-      https://hips.hearstapps.com/hmg-prod/images/gettyimages-1229892983-square.jpg
-    example_title: Portrait
 ---
 <h1 align="center">Bilateral Reference for High-Resolution Dichotomous Image Segmentation</h1>
 

handler.py

@@ -9,7 +9,6 @@ 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])
 
@@ -105,7 +104,18 @@ class EndpointHandler():
             A :obj:`list` | `dict`: will be serialized and returned
         """
         print('data["inputs"] = ', data["inputs"])
-        image = load_img(data["inputs"]).convert("RGB")
+        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')
         # Preprocess the image
         image_preprocessor = ImagePreprocessor(resolution=tuple(resolution))
         image_proc = image_preprocessor.proc(image)

requirements.txt

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