File size: 2,134 Bytes
dbf872d 51998cd fd9f505 51998cd fd9f505 dbf872d fd9f505 dbf872d fd9f505 dbf872d fd9f505 dbf872d fd9f505 dbf872d fd9f505 dbf872d fd9f505 dbf872d fd9f505 51998cd fd9f505 51998cd dbf872d 51998cd dbf872d 51998cd |
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 |
# app.py (Streamlit version - Corrected)
import streamlit as st
from PIL import Image
import torch
import torchvision.transforms as transforms
import numpy as np
import subprocess
# Install pytorch_hub_examples package (Corrected)
try:
import pytorch_hub_examples
except ModuleNotFoundError:
subprocess.run(["git", "clone", "https://github.com/facebookresearch/pytorch_hub_examples.git"])
subprocess.run(["cd", "pytorch_hub_examples", "&&", "python", "setup.py", "install"]) # Correct installation
import pytorch_hub_examples
# Load the U-2-Net model
model = pytorch_hub_examples.u2net(pretrained=True)
model.eval()
# Define the transform
transform = transforms.Compose([
transforms.Resize((320, 320)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
def remove_background(image):
try:
img = transform(image).unsqueeze(0)
with torch.no_grad():
out = model(img)
mask = torch.sigmoid(out[0])
mask = mask.squeeze().cpu().numpy()
mask = (mask * 255).astype(np.uint8)
mask = Image.fromarray(mask).convert("L")
image = image.convert("RGBA")
new_image = Image.new("RGBA", image.size, (255, 255, 255, 0))
for x in range(image.width):
for y in range(image.height):
if mask.getpixel((x, y)) > 0:
new_image.putpixel((x, y), image.getpixel((x, y)))
return new_image
except Exception as e:
st.error(f"Error: {e}")
return None
st.title("Background Remover")
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
image = Image.open(uploaded_file).convert("RGB") # Ensure RGB for transform
st.image(image, caption="Uploaded Image", use_column_width=True)
if st.button("Remove Background"):
with st.spinner("Removing background..."):
result_image = remove_background(image)
if result_image:
st.image(result_image, caption="Background Removed", use_column_width=True) |