Update app.py

app.py

@@ -37,39 +37,27 @@ def remove_background_bria(input_path):
     
     # Create a segmentation pipeline
     pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True)
-    
-    # Ensure input_image is a PIL image
-    if isinstance(input_path, str):
-        input_image = Image.open(input_path).convert("RGBA")
+
+    # Load the image
+    input_image = Image.open(input_path).convert("RGBA")
     
     # Get the segmentation output
-    segmentation_result = pipe(input_path)
-
-    # Check if the output contains a mask
-    if isinstance(segmentation_result, list) and len(segmentation_result) > 0:
-        # Adjust based on the actual output structure
-        mask = segmentation_result[0].get("mask")
-        
-        if mask is not None:
-            # Create an output image based on the mask
-            output_image = Image.new("RGBA", input_path.size)
-            width, height = input_path.size
-            
-            for y in range(height):
-                for x in range(width):
-                    # Check if the mask value is valid (assuming it's binary or a single channel)
-                    if mask[y][x] > 0:  # Adjust based on actual mask values
-                        output_image.putpixel((x, y), input_path.getpixel((x, y)))
-                    else:
-                        output_image.putpixel((x, y), (0, 0, 0, 0))  # Set to transparent
-            
-            return output_image
-        else:
-            print("Mask is None in segmentation result.")
-            return input_path  # Return original if mask is None
-    else:
-        print("No valid segmentation result received.")
-        return input_path  # Return original if no valid result
+    pillow_mask = pipe(input_image, return_mask=True)  # Outputs a pillow mask
+    print("Mask obtained:", pillow_mask)  # Debugging output
+
+    # Create an output image based on the mask
+    output_image = Image.new("RGBA", input_image.size)
+
+    # Use the mask to create the output image
+    for x in range(input_image.width):
+        for y in range(input_image.height):
+            # Assuming mask is in a binary format where foreground is True
+            if pillow_mask.getpixel((x, y)) > 0:  # Adjust based on actual mask values
+                output_image.putpixel((x, y), input_image.getpixel((x, y)))
+            else:
+                output_image.putpixel((x, y), (0, 0, 0, 0))  # Set to transparent
+
+    return output_image
 
 # Fungsi untuk memproses gambar menggunakan prompt
 def text_to_image(prompt):