Update app.py

app.py

@@ -44,11 +44,32 @@ def remove_background_bria(input_image):
     
     # Get the segmentation output
     segmentation_result = pipe(input_image)
-    
-    # Process the segmentation output (assumes the result is a mask)
-    # You'll need to adapt this based on the actual output format
-    mask = segmentation_result[0]["mask"]  # Adjust as necessary based on output structure
-    return mask
+
+    # 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_image.size)
+            width, height = input_image.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_image.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_image  # Return original if mask is None
+    else:
+        print("No valid segmentation result received.")
+        return input_image  # Return original if no valid result
 
 # Fungsi untuk memproses gambar menggunakan prompt
 def text_to_image(prompt):