afrodreams / ex_app.py
Owos's picture
Update ex_app.py
57535c9
raw
history blame
3.49 kB
import neural_style
import streamlit as st
import os
import random
import numpy as np
from PIL import Image, ImageEnhance
from io import BytesIO
#import streamlit_ext as ste #for download button not to rerun
from huggingface_hub import upload_file
HF_TOKEN = os.environ.get("HF_TOKEN")
st.set_page_config(layout="wide")
#Create two columns with different width
col1, col2 = st.columns( [0.8, 0.2])
with col1: # To display the header text using css style
st.markdown(""" <style> .font {
font-size:35px ; font-family: 'Cooper Black'; color: #FF9633;}
</style> """, unsafe_allow_html=True)
st.markdown('<p class="font">Upload your photo here...</p>', unsafe_allow_html=True)
st.subheader("This app takes in your image and styles it with a unique african art.")
#Add a header and expander in side bar
st.sidebar.markdown('<p class="font">Afrodreams.AI</p>', unsafe_allow_html=True)
with st.sidebar.expander("About the App"):
st.write("""
This app takes in your image and styles it with a unique african art.""")
#Add file uploader to allow users to upload photos
uploaded_file = st.file_uploader("", type=['jpg','png','jpeg'])
# add slider to side bar
style_weight = st.slider("Select Style Weight", min_value=10, max_value=100, value=12)
#Add 'before' and 'after' columns
if uploaded_file is not None:
image = Image.open(uploaded_file)
col1, col2 = st.columns( [0.5, 0.5])
with col1:
st.markdown('<p style="text-align: center;">Before</p>',unsafe_allow_html=True)
st.image(image,width=300)
with col2:
st.markdown('<p style="text-align: center;">After</p>',unsafe_allow_html=True)
# add a button
run = st.button('Generate Art')
my_bar = st.progress(0)
params = neural_style.TransferParams()
params.gpu = 0#"c"
params.backend = "mkl"
params.image_size = 400
params.content_image = uploaded_file
params.style_weight = style_weight
keep_style = False
if run==True:
# run image selection if keep style is false
if keep_style==False:
path = 'stylesv2'
styles = os.listdir(path)
params.style_image = path + '/' + random.choice(styles)
st.session_state.submitted = True
with st.spinner('Wait for it...'):
neural_style.transfer(params)
#display image when done.
with col2:
if 'submitted' in st.session_state:
result = Image.open('out.png')
st.image(result, width=300)
buf = BytesIO()
result.save(buf, format="png")
if len(os.listdir('generated_samples')) <= 10:
img_file_name = f"generated_samples/{str(len(os.listdir('generated_samples')))}.png"
_ = upload_file(path_or_fileobj = 'out.png',
path_in_repo = img_file_name,
repo_id='AfrodreamsAI/afrodreams',
repo_type='space',
token=HF_TOKEN
)
byte_im = buf.getvalue()
#run =ste.download_button(button_text="Download Image", data=byte_im, download_filename='afrodreams.jpg', mime="image/png")
#keeping the current style by update the weight
keep_style = st.sidebar.checkbox("Keep current style")