Update freesplatter/utils/infer_util.py

freesplatter/utils/infer_util.py

@@ -2,6 +2,7 @@ import os
 import importlib
 import imageio
 import torch
+import rembg
 import numpy as np
 import PIL.Image
 from PIL import Image
@@ -67,10 +68,36 @@ def get_obj_from_str(string, reload=False):
 #     return image
 
 
-@torch.inference_mode()
-def remove_background(
-    image: PIL.Image.Image,
-    rembg: Any = None,
+# @torch.inference_mode()
+# def remove_background(
+#     image: PIL.Image.Image,
+#     rembg: Any = None,
+#     force: bool = False,
+#     **rembg_kwargs,
+# ) -> PIL.Image.Image:
+#     do_remove = True
+#     if image.mode == "RGBA" and image.getextrema()[3][0] < 255:
+#         do_remove = False
+#     do_remove = do_remove or force
+#     if do_remove:
+#         transform_image = transforms.Compose([
+#             transforms.Resize((1024, 1024)),
+#             transforms.ToTensor(),
+#             transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+#         ])
+#         image = image.convert('RGB')
+#         input_images = transform_image(image).unsqueeze(0).to(rembg.device)
+#         with torch.no_grad():
+#             preds = rembg(input_images)[-1].sigmoid().cpu()
+#         pred = preds[0].squeeze()
+#         pred_pil = transforms.ToPILImage()(pred)
+#         mask = pred_pil.resize(image.size)
+#         image.putalpha(mask)
+#     return image
+
+
+def remove_background(image: PIL.Image.Image,
+    rembg_session: Any = None,
     force: bool = False,
     **rembg_kwargs,
 ) -> PIL.Image.Image:
@@ -79,19 +106,7 @@ def remove_background(
         do_remove = False
     do_remove = do_remove or force
     if do_remove:
-        transform_image = transforms.Compose([
-            transforms.Resize((1024, 1024)),
-            transforms.ToTensor(),
-            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
-        ])
-        image = image.convert('RGB')
-        input_images = transform_image(image).unsqueeze(0).to(rembg.device)
-        with torch.no_grad():
-            preds = rembg(input_images)[-1].sigmoid().cpu()
-        pred = preds[0].squeeze()
-        pred_pil = transforms.ToPILImage()(pred)
-        mask = pred_pil.resize(image.size)
-        image.putalpha(mask)
+        image = rembg.remove(image, session=rembg_session, **rembg_kwargs)
     return image