Add model

app.py

@@ -42,11 +42,11 @@ def generate_mask(image):
     # Preprocess the image
     image = np.array(image.convert("RGB"))
     img = preprocess(image)
-    
+
     input_size = [1024, 1024]
     im_shp = image.shape[0:2]
     im_tensor = torch.tensor(image, dtype=torch.float32).permute(2, 0, 1)
-    
+
     # Replace F.upsample with F.interpolate
     im_tensor = F.interpolate(torch.unsqueeze(im_tensor, 0), input_size, mode="bilinear").type(torch.uint8)
     image = torch.divide(im_tensor, 255.0)
@@ -66,32 +66,46 @@ def generate_mask(image):
 
     return output_mask
 
-# Define the final predict method to overlay the mask
 def predict(image):
     # Generate the mask
     mask = generate_mask(image)
-    
+
     # Convert the image to RGBA (to support transparency)
-    image = image.convert("RGBA")
-    
-    # Convert the mask into a binary mask where 255 is kept and 0 is transparent
-    mask = Image.fromarray(mask).resize(image.size).convert("L")  # Convert to grayscale (L mode)
-    
-    # Create a new image with transparency (RGBA)
+    image_rgba = image.convert("RGBA")
+
+    # Create a binary mask from the generated mask and resize it to the image size
+    mask = Image.fromarray(mask).resize(image.size).convert("L")  # Convert to grayscale
+
+    # Create a new image with transparency (RGBA) for the output with transparent background
     transparent_image = Image.new("RGBA", image.size)
-    
-    # Use the mask as transparency: paste the original image where the mask is white
-    transparent_image.paste(image, mask=mask)
-    
-    return transparent_image
+    transparent_image.paste(image_rgba, mask=mask)
+
+    # Create foreground and background masks
+    red_foreground = Image.new("RGBA", image.size, (255, 0, 0, 128))  # Red foreground with 50% opacity
+    blue_background = Image.new("RGBA", image.size, (0, 0, 255, 128))  # Blue background with 50% opacity
+
+    # Create an empty overlay image
+    overlay_image = Image.new("RGBA", image.size)
+
+    # Overlay the red and blue masks based on the mask
+    overlay_image.paste(blue_background, (0, 0))  # Fill the entire overlay with blue
+    overlay_image.paste(red_foreground, (0, 0), mask=mask)  # Paste red where mask is white
+
+    # Combine the original image with the overlay at 50% opacity
+    combined_image = Image.blend(image_rgba, overlay_image, alpha=0.5)
+
+    return transparent_image, combined_image
 
-# Create the Gradio interface with custom output size for the display only (not affecting the saved image)
+# Create the Gradio interface with two outputs
 iface = gr.Interface(
-    fn=predict, 
-    inputs=gr.Image(type="pil"), 
-    outputs=gr.Image(type="pil", label="Edited Image", image_mode="RGBA", format="png"),  # RGBA ensures PNG with transparency
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=[
+        gr.Image(type="pil", label="Transparent Background", image_mode="RGBA", format="png"),  # Transparent output
+        gr.Image(type="pil", label="Overlay with Colors", image_mode="RGBA", format="png"),  # Colored overlay output
+    ],
     title="Background Removal with U2NET",
-    description="Upload an image and remove the background"
+    description="Upload an image to remove the background and visualize it with an overlay."
 )
 
 if __name__ == "__main__":