watermelon_pfp / app.py
CallmeKaito's picture
Update app.py
cccb949 verified
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
# 1st picture: normal behaviour
# 2nd picture: normal picture and applying the palestinian flag with an opacity of 120
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 = background_img.resize((image.width, image.height))
#image = image.resize((background_img.width, background_img.height))
image = image.convert("RGBA")
background_img = background_img.convert("RGBA")
input_img = Image.open('original/image.jpg')
input_img = input_img.convert("RGBA")
normal_img = Image.blend(image, background_img, .3).convert('RGB')
# Create a semi-transparent overlay
overlay = Image.new('RGBA', background_img.size, (0, 0, 0, 120)) # Adjust the last value (120) to control transparency
background_img = Image.alpha_composite(background_img.convert('RGBA'), overlay)
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, normal_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:
# handle error if not picture
if background_img_file.type.split("/")[0] != "image":
st.error("Please upload an image file")
return
background_img_file = background_img_file.read()
cropped_pic = st_cropperjs(pic=background_img_file, btn_text="Generate!", size=1.0, key="foo")
if cropped_pic:
# save the cropped pic
cropped_pic = Image.open(io.BytesIO(cropped_pic))
cropped_pic.save("cropped_pic.png")
# Process the images
with st.spinner('Generating your fabulous ๐Ÿ‰ profile picture...'):
img, img2 = process_image(cropped_pic)
# Create two columns
col1, col2 = st.columns(2)
# Display the images in each column
with col1:
st.image(img, caption="Image 1", 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 1",
data=img_bytes,
file_name="free_palestine.png",
mime="image/png",
)
with col2:
st.image(img2, caption="Image 2", use_column_width=True)
img_bytes = BytesIO()
img2.save(img_bytes, format='PNG')
img_bytes = img_bytes.getvalue()
# Add a download button
st.download_button(
label="Download Image 2",
data=img_bytes,
file_name="free_palestine2.png",
mime="image/png",
)
if __name__ == "__main__":
main()