Switch background removal to RMBG-2.0

human_avatar/image_to_avatar.py

@@ -2,13 +2,16 @@
 from pathlib import Path
 
 import numpy as np
+import torch
 from PIL import Image
 from pose_format import Pose
 from pose_format.utils.generic import pose_normalization_info
 from pose_format.utils.holistic import load_holistic
-from transformers import pipeline
+from torchvision import transforms
+from transformers import AutoModelForImageSegmentation, pipeline
 
 CROP_RESOLUTION = 512
+RMBG_INPUT_SIZE = (1024, 1024)
 
 
 def extract_pose(image: Image):
@@ -56,9 +59,35 @@ def crop_person(image: Image, pose: Pose):
     return image
 
 
+@cache
+def load_rmbg_model():
+    model = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-2.0", trust_remote_code=True)
+    torch.set_float32_matmul_precision("high")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model.to(device)
+    model.eval()
+    return model, device
+
+
 def remove_image_background(image: Image):
-    model = load_huggingface_model("image-segmentation", model="briaai/RMBG-1.4")
-    return model(image)
+    model, device = load_rmbg_model()
+
+    transform_image = transforms.Compose([
+        transforms.Resize(RMBG_INPUT_SIZE),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+    ])
+
+    rgb_image = image.convert("RGB")
+    input_tensor = transform_image(rgb_image).unsqueeze(0).to(device)
+
+    with torch.no_grad():
+        preds = model(input_tensor)[-1].sigmoid().cpu()
+    mask = transforms.ToPILImage()(preds[0].squeeze()).resize(rgb_image.size)
+
+    result = rgb_image.copy()
+    result.putalpha(mask)
+    return result
 
 
 def image_to_avatar(image: Image):

pyproject.toml

@@ -15,7 +15,9 @@ dependencies = [
     # Transformers with vision
     "transformers",
     "torchvision",
-    "scikit-image"
+    "scikit-image",
+    # Required by briaai/RMBG-2.0 trust_remote_code modeling file
+    "kornia",
 ]
 
 [project.optional-dependencies]