fix

app.py

@@ -6,6 +6,7 @@ import torch
 from torchvision import transforms
 import os
 import uuid
+from PIL import Image
 
 torch.set_float32_matmul_precision(["high", "highest"][0])
 
@@ -25,17 +26,25 @@ transform_image = transforms.Compose(
 def fn(image):
     im = load_img(image, output_type="pil")
     im = im.convert("RGB")
+    # Get the segmented image (RGBA)
     processed_image = process(im)
     
     # Generate a unique filename for the processed image
     unique_id = str(uuid.uuid4())[:8]
     output_path = f"output_{unique_id}.jpg"
     
-    # Convert RGBA to RGB before saving as JPEG
-    rgb_image = processed_image.convert("RGB")
-    rgb_image.save(output_path, format="JPEG")
-    
-    return processed_image, output_path
+    # Create a white background and properly composite with the RGBA image
+    white_bg = Image.new("RGB", processed_image.size, (255, 255, 255))
+    if processed_image.mode == 'RGBA':
+        # Use the alpha channel as a mask for compositing
+        white_bg.paste(processed_image, mask=processed_image.split()[3])  # The 4th channel is alpha
+        white_bg.save(output_path, format="JPEG")
+        # Return the composited image for display to match what's being downloaded
+        return white_bg, output_path
+    else:
+        rgb_image = processed_image.convert("RGB")
+        rgb_image.save(output_path, format="JPEG")
+        return rgb_image, output_path
 
 @spaces.GPU
 def process(image):
@@ -57,15 +66,22 @@ def process(image):
 def process_file(f):
     im = load_img(f, output_type="pil")
     im = im.convert("RGB")
+    # Get the segmented image (RGBA)
     transparent = process(im)
     
     # Save as JPEG instead of PNG
     unique_id = str(uuid.uuid4())[:8]
     output_path = f"output_{unique_id}.jpg"
     
-    # Convert RGBA to RGB before saving as JPEG
-    rgb_image = transparent.convert("RGB")
-    rgb_image.save(output_path, format="JPEG")
+    # Create a white background and properly composite with the RGBA image
+    white_bg = Image.new("RGB", transparent.size, (255, 255, 255))
+    if transparent.mode == 'RGBA':
+        # Use the alpha channel as a mask for compositing
+        white_bg.paste(transparent, mask=transparent.split()[3])  # The 4th channel is alpha
+        white_bg.save(output_path, format="JPEG")
+    else:
+        rgb_image = transparent.convert("RGB")
+        rgb_image.save(output_path, format="JPEG")
     
     return output_path