wip

app.py

@@ -0,0 +1,65 @@
+import gradio as gr
+from PIL import Image
+import matplotlib.pyplot as plt
+import torch
+from torchvision import transforms
+from transformers import AutoModelForImageSegmentation
+
+model = AutoModelForImageSegmentation.from_pretrained('briaai/RMBG-2.0', trust_remote_code=True)
+torch.set_float32_matmul_precision(['high', 'highest'][0])
+model.eval()
+
+def remove_background(input_image, holiday, message):
+    image_size = (1024, 1024)
+    # Transform the input image
+    transform_image = transforms.Compose([
+        transforms.Resize(image_size),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+    ])
+    
+    # Process the image
+    input_tensor = transform_image(input_image).unsqueeze(0)
+    if torch.cuda.is_available():
+        input_tensor = input_tensor.to('cuda')
+    
+    # Generate prediction
+    with torch.no_grad():
+        preds = model(input_tensor)[-1].sigmoid().cpu()
+    pred = preds[0].squeeze()
+    pred_pil = transforms.ToPILImage()(pred)
+    mask = pred_pil.resize(input_image.size)
+    
+    # Create image without background
+    result_image = input_image.copy()
+    result_image.putalpha(mask)
+    
+    # Create image with only background
+    only_background_image = input_image.copy()
+    inverted_mask = Image.eval(mask, lambda x: 255 - x)  # Invert the mask
+    only_background_image.putalpha(inverted_mask)
+    
+    first_output_image = result_image
+    second_output_image = only_background_image
+    third_output_image = result_image
+    
+    return first_output_image, second_output_image, third_output_image
+
+# Replace the demo interface
+demo = gr.Interface(
+    fn=remove_background,
+    inputs=[
+        gr.Image(type="pil"),
+        gr.Text(label="Holiday (e.g. Christmas, New Year's, etc.)"),
+        gr.Text(label="Optional Message", placeholder="Enter your holiday message here...")
+    ],
+    outputs=[
+        gr.Image(type="pil", label="First Output"),
+        gr.Image(type="pil", label="Second Output"),
+        gr.Image(type="pil", label="Third Output")
+    ],
+    title="Background Removal Tool",
+    description="Upload an image to remove its background"
+)
+
+demo.launch()

requirements.txt

@@ -0,0 +1,7 @@
+gradio
+torch
+torchvision
+pillow
+kornia
+transformers
+timm