Update app.py

app.py

@@ -8,6 +8,7 @@ import insightface
 import onnxruntime as ort
 import huggingface_hub
 from SegCloth import segment_clothing
+from transparent_background import Remover
 
 app = Flask(__name__)
 
@@ -67,13 +68,26 @@ def detect_and_segment_persons(image, clothes):
         pil_img = Image.fromarray(person_img[:, :, ::-1])  # BGR -> RGB
 
         # Segment clothing for the detected person
-        segmented_result = segment_clothing(pil_img, clothes)
+        img_rm_background = remove_background(pil_img)
+        segmented_result = segment_clothing(img_rm_background, clothes)
 
         # Combine the segmented images for all persons
         all_segmented_images.extend(segmented_result)
 
     return all_segmented_images
 
+    @spaces.GPU
+    def remove_background(image):
+        remover = Remover()
+        if isinstance(image, Image.Image):
+            output = remover.process(image)
+        elif isinstance(image, np.ndarray):
+            image_pil = Image.fromarray(image)
+            output = remover.process(image_pil)
+        else:
+            raise TypeError("Unsupported image type")
+        return output
+
 @app.route('/api/detect', methods=['POST'])
 def detect():
     try: