Spaces:
Sleeping
Sleeping
File size: 4,317 Bytes
61dce11 |
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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 |
import streamlit as st
from PIL import Image, ImageColor
from io import BytesIO
from rembg import remove
import os
from typing import Tuple
from PIL.Image import Image as PILImage
from copy import deepcopy
class ImageProcessor:
def __init__(self, image):
self.image = image
def remove_background(self):
return remove(self.image)
@staticmethod
def change_background(image, background_alpha: float = 1.0, background_hex: str = "#000000"):
"""
image: PIL Image (RGBA)
matte: PIL Image (grayscale, if 255 it is foreground)
background_alpha: float
background_hex: string
"""
img = deepcopy(image)
if image.mode != "RGBA":
img = img.convert("RGBA")
background_color = ImageColor.getrgb(background_hex)
background_alpha = int(255 * background_alpha)
background = Image.new(
"RGBA", img.size, color=background_color + (background_alpha,))
background.paste(img, mask=ImageProcessor.matte)
return background
@staticmethod
def matte(img_input):
# Implement matte logic
pass
class ImageDownloader:
@staticmethod
def download_button(pil_image, filename: str, fmt: str, label="Download"):
if fmt not in ["jpg", "png"]:
raise ValueError(
"Unknown image format (Available: jpg, png - case sensitive)")
pil_format = "JPEG" if fmt == "jpg" else "PNG"
file_format = "jpg" if fmt == "jpg" else "png"
mime = "image/jpeg" if fmt == "jpg" else "image/png"
buf = BytesIO()
pil_image.save(buf, format=pil_format)
return st.download_button(
label=label,
data=buf.getvalue(),
file_name=f'{filename}.{file_format}',
mime=mime,
)
class ImageUploader:
@staticmethod
def upload_image():
return st.file_uploader(
label="Upload your photo here",
accept_multiple_files=False, type=["png", "jpg", "jpeg"],
)
class BackgroundColorSelector:
@staticmethod
def select_background_color():
return st.selectbox("Choose background color", [
"Transparent (PNG)", "White", "Black", "Green", "Red", "Blue"
])
class App:
def __init__(self):
self.image_processor = None
self.image_uploader = ImageUploader()
self.background_selector = BackgroundColorSelector()
self.image_downloader = ImageDownloader()
self.hexmap = {
"Transparent (PNG)": "#000000",
"Black": "#000000",
"White": "#FFFFFF",
"Green": "#22EE22",
"Red": "#EE2222",
"Blue": "#2222EE",
}
def run(self):
st.title("Simple image editor")
st.write(
"Simply remove the background of your picture and fill it with a plain color")
uploaded_file = self.image_uploader.upload_image()
if uploaded_file:
self.display_original_photo(uploaded_file)
in_mode = self.background_selector.select_background_color()
in_submit = st.button("Submit")
if in_submit:
self.process_and_display_image(uploaded_file, in_mode)
def display_original_photo(self, uploaded_file):
with st.expander("Original photo", expanded=True):
st.image(uploaded_file)
def process_and_display_image(self, uploaded_file, in_mode):
img_input = Image.open(uploaded_file)
self.image_processor = ImageProcessor(img_input)
with st.spinner("Magic happening behind the scenes. Please wait..."):
alpha = 0.0 if in_mode == "Transparent (PNG)" else 1.0
img_output = self.image_processor.remove_background()
# img_output = self.image_processor.change_background(
# img_without_bg, alpha, in_mode)
with st.expander("Success!", expanded=True):
st.image(img_output)
uploaded_name = os.path.splitext(uploaded_file.name)[0]
self.image_downloader.download_button(
pil_image=img_output,
filename=uploaded_name,
fmt="png"
)
if __name__ == "__main__":
app = App()
app.run()
|