Update app.py

app.py

@@ -11,6 +11,11 @@ from SegCloth import segment_clothing
 from transparent_background import Remover
 import threading
 import logging
+import uuid
+from transformers import AutoModelForImageSegmentation
+import torch
+from torchvision import transforms
+
 
 app = Flask(__name__)
 
@@ -35,6 +40,22 @@ def load_model():
     detector = model
     logging.info("Model loaded successfully.")
 
+torch.set_float32_matmul_precision(["high", "highest"][0])
+
+birefnet = AutoModelForImageSegmentation.from_pretrained(
+    "ZhengPeng7/BiRefNet", trust_remote_code=True
+)
+birefnet.to("cuda")
+transform_image = transforms.Compose(
+    [
+        transforms.Resize((1024, 1024)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+    ]
+)
+
+
+
 def save_image(img):
     unique_name = str(uuid.uuid4()) + ".png"
     img.save(unique_name)
@@ -51,7 +72,24 @@ def encode_image_to_base64(image):
     buffered = BytesIO()
     image.save(buffered, format="PNG")  # Use PNG for compatibility with RGBA
     return base64.b64encode(buffered.getvalue()).decode('utf-8')
-    
+
+@spaces.GPU
+def rm_background(image):
+    im = load_img(image, output_type="pil")
+    im = im.convert("RGB")
+    image_size = im.size
+    origin = im.copy()
+    image = load_img(im)
+    input_images = transform_image(image).unsqueeze(0).to("cuda")
+    # Prediction
+    with torch.no_grad():
+        preds = birefnet(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)
+
 @spaces.GPU
 def remove_background(image):
     remover = Remover()
@@ -73,7 +111,7 @@ def detect_and_segment_persons(image, clothes):
 
     bboxes, kpss = detector.detect(img)
     if bboxes.shape[0] == 0:
-        return [save_image(remove_background(image))]
+        return [save_image(rm_background(image))]
 
     height, width, _ = img.shape
     bboxes = np.round(bboxes[:, :4]).astype(int)
@@ -89,7 +127,7 @@ def detect_and_segment_persons(image, clothes):
         person_img = img[y1:y2, x1:x2]
         pil_img = Image.fromarray(person_img[:, :, ::-1])
 
-        img_rm_background = remove_background(pil_img)
+        img_rm_background = rm_background(pil_img)
         segmented_result = segment_clothing(img_rm_background, clothes)
         image_paths = [save_image(img) for img in segmented_result]
         print(image_paths)