Update example_inference.py

example_inference.py

@@ -4,6 +4,8 @@ from PIL import Image
 from briarmbg import BriaRMBG
 from utilities import preprocess_image, postprocess_image
 from huggingface_hub import hf_hub_download
+import io as IO
+import base64
 
 def example_inference(im_path, transprent_bg=False):
 
@@ -15,7 +17,7 @@ def example_inference(im_path, transprent_bg=False):
 
     # prepare input
     model_input_size = [1024,1024]
-    orig_im = io.imread(im_path)
+    orig_im = io.imread(im_path, plugin='imageio')
     orig_im_size = orig_im.shape[0:2]
     image = preprocess_image(orig_im, model_input_size).to(device)
 
@@ -28,7 +30,12 @@ def example_inference(im_path, transprent_bg=False):
     # save result
     pil_im = Image.fromarray(result_image)
     no_bg_image = Image.new("RGBA", pil_im.size, bgColor)
-    orig_image = Image.open(im_path)
+    orig_image = Image.open(IO.BytesIO(im_path))
     no_bg_image.paste(orig_image, mask=pil_im)
-    no_bg_image.save("images/rm_image.png")
-    return 'rm_image.png'
+    
+    # Convert image to bytes and then to base64
+    buffered = IO.BytesIO()
+    no_bg_image.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+    
+    return img_str