Add model

app.py

@@ -37,7 +37,7 @@ def preprocess(image):
 
     return sample
 
-# Define the function to generate the mask
+# Generate the mask
 def generate_mask(image):
     # Preprocess the image
     image = np.array(image.convert("RGB"))
@@ -46,13 +46,15 @@ def generate_mask(image):
     input_size = [1024, 1024]
     im_shp = image.shape[0:2]
     im_tensor = torch.tensor(image, dtype=torch.float32).permute(2, 0, 1)
-    im_tensor = F.upsample(torch.unsqueeze(im_tensor, 0), input_size, mode="bilinear").type(torch.uint8)
+    
+    # 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)
     image = normalize(image, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
 
     with torch.no_grad():
         result = model(image)
-        result = torch.squeeze(F.upsample(result[0][0], im_shp, mode='bilinear'), 0)
+        result = torch.squeeze(F.interpolate(result[0][0], im_shp, mode='bilinear'), 0)
         ma = torch.max(result)
         mi = torch.min(result)
         result = (result - mi) / (ma - mi)
@@ -61,17 +63,35 @@ def generate_mask(image):
     output_mask = result[0]
     output_mask = (output_mask - output_mask.min()) / (output_mask.max() - output_mask.min()) * 255
     output_mask = output_mask.astype(np.uint8)
-    output_image = Image.fromarray(output_mask)
 
-    return output_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)
+    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
 
-# Create the Gradio interface
+# Create the Gradio interface with custom output size for the display only (not affecting the saved image)
 iface = gr.Interface(
-    fn=generate_mask,
-    inputs=gr.inputs.Image(type="pil"),
-    outputs=gr.outputs.Image(type="pil"),
-    title="U2NET Background Removal",
-    description="Upload an image and get the background mask"
+    fn=predict, 
+    inputs=gr.Image(type="pil"), 
+    outputs=gr.Image(type="pil", tool="editor", label="Edited Image"),  # Adjust the box size
+    title="Background Removal with U2NET",
+    description="Upload an image and remove the background"
 )
 
 if __name__ == "__main__":