Spaces:
Sleeping
Sleeping
File size: 2,236 Bytes
6a046fb c167ae6 6a046fb c167ae6 6a046fb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
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])
if torch.cuda.is_available():
model = model.to('cuda')
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="Holiday Card Generator",
description="Upload an image to generate a holiday card"
)
demo.launch()
|