Update freesplatter/utils/infer_util.py

freesplatter/utils/infer_util.py

@@ -68,36 +68,10 @@ def get_obj_from_str(string, reload=False):
 #     return image
 
 
-# @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,
+@torch.inference_mode()
+def remove_background(
+    image: PIL.Image.Image,
+    rembg: Any = None,
     force: bool = False,
     **rembg_kwargs,
 ) -> PIL.Image.Image:
@@ -106,10 +80,36 @@ def remove_background(image: PIL.Image.Image,
         do_remove = False
     do_remove = do_remove or force
     if do_remove:
-        image = rembg.remove(image, session=rembg_session, **rembg_kwargs)
+        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:
+#     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:
+#         image = rembg.remove(image, session=rembg_session, **rembg_kwargs)
+#     return image
+
+
 def resize_foreground(
     image: PIL.Image.Image,
     ratio: float,