Spaces:
Sleeping
Sleeping
File size: 1,336 Bytes
3cbb615 |
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 |
from PIL import Image
import matplotlib.pyplot as plt
import torch
from torchvision import transforms
from transformers import AutoModelForImageSegmentation
# Load the model
model = AutoModelForImageSegmentation.from_pretrained('briaai/RMBG-2.0', trust_remote_code=True)
torch.set_float32_matmul_precision('high')
model.eval()
# Data settings
image_size = (1024, 1024)
transform_image = transforms.Compose([
transforms.Resize(image_size),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
# Get the image file path from the user
input_image_path = input("Please enter the file path of the image: ")
# Open and convert the image
try:
im = Image.open(input_image_path)
rgb_im = im.convert('RGB')
except FileNotFoundError:
print(f"Error: The file at {input_image_path} was not found.")
exit()
# Transform the image
input_images = transform_image(rgb_im).unsqueeze(0)
# Prediction
with torch.no_grad():
preds = model(input_images)[-1].sigmoid().cpu()
# Process the prediction
pred = preds[0].squeeze()
pred_pil = transforms.ToPILImage()(pred)
mask = pred_pil.resize(rgb_im.size)
rgb_im.putalpha(mask)
# Save the result
output_image_path = "no_bg_image.png"
rgb_im.save(output_image_path)
print(f"Image with background removed saved as {output_image_path}")
|