watermelon_pfp / app.py
CallmeKaito's picture
Update app.py
c64c178 verified
raw
history blame
2.87 kB
import streamlit as st
from PIL import Image, ImageOps
import io
from rembg import remove
import requests
from PIL import Image
from io import BytesIO
import os
from streamlit_cropperjs import st_cropperjs
# Function to process the image
def process_image(image):
print(image)
os.makedirs('original', exist_ok=True)
os.makedirs('masked', exist_ok=True)
image = image.convert("RGB")
image.save('original/image.jpg', format = "jpeg")
output_path = "masked/image.png"
with open(output_path, "wb") as f:
input = open('original/image.jpg', 'rb').read()
subject = remove(input)
f.write(subject)
palestine_bg = "https://flagdownload.com/wp-content/uploads/Flag_of_Palestine_Flat_Round-1024x1024.png"
background_img = Image.open(BytesIO(requests.get(palestine_bg).content))
background_img.putalpha(160)
background_img = background_img.resize((image.width, image.height))
#image = image.resize((background_img.width, background_img.height))
background_img = background_img.convert("RGBA")
input_img = Image.open('original/image.jpg')
input_img = input_img.convert("RGBA")
combined_img = Image.alpha_composite(input_img, background_img)
combined_img = combined_img.convert('RGB')
# combined_img.save('masked/finale.jpg', format='jpeg')
foreground_img = Image.open(output_path)
combined_img.paste(foreground_img, (0,0), foreground_img)
combined_img = combined_img.convert('RGB')
# combined_img.save("masked/background_final.jpg", format="jpeg")
return combined_img
# Streamlit app
def main():
st.title("Watermelon PFP Generator - By Kaito")
# Select background image
background_img_file = st.file_uploader("Select your image", type=["jpg", "png"])
if background_img_file is not None:
background_img_file = background_img_file.read()
cropped_pic = st_cropperjs(pic=background_img_file, btn_text="Generate!", key="foo")
if cropped_pic:
# save the cropped pic
cropped_pic = Image.open(io.BytesIO(cropped_pic))
cropped_pic.save("cropped_pic.png")
# Load input and background images
background_img = cropped_pic
# Process the images
img = process_image(background_img)
# Display the combined image
st.image(img, caption="Combined Image", use_column_width=True)
img_bytes = BytesIO()
img.save(img_bytes, format='PNG')
img_bytes = img_bytes.getvalue()
# Add a download button
st.download_button(
label="Download Image",
data=img_bytes,
file_name="free_palestine.png",
mime="image/png",
)
if __name__ == "__main__":
main()