File size: 4,651 Bytes
aedef5b 70b0b9d e0c07aa 70b0b9d e0c07aa 70b0b9d e0c07aa 70b0b9d aedef5b 70b0b9d 5314570 70b0b9d aedef5b 70b0b9d aedef5b 0c3fdc3 aedef5b 70b0b9d aedef5b ff7439f aedef5b ff7439f aedef5b 70b0b9d aedef5b 3398f59 e285cf1 |
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 |
import os
import zipfile
import streamlit as st
from PIL import Image
from streamlit_lottie import st_lottie
from models.deep_colorization.colorizers import eccv16
from utils import colorize_image, change_model, load_lottieurl
st.set_page_config(page_title="Image & Video Colorizer", page_icon="🎨", layout="wide")
loaded_model = eccv16(pretrained=True).eval()
current_model = "None"
col1, col2 = st.columns([1, 3])
with col1:
lottie = load_lottieurl("https://assets5.lottiefiles.com/packages/lf20_RHdEuzVfEL.json")
st_lottie(lottie)
with col2:
st.write("""
## B&W Images Colorizer
##### Input a black and white image and get a colorized version of it.
###### ➠ If you want to colorize multiple images just upload them all at once.
###### ➠ Uploading already colored images won't raise errors but images won't look good.
###### ➠ I recommend starting with the first model and then experimenting with the second one.""")
def main():
model = st.selectbox(
"Select Model (Both models have their pros and cons, I recommend trying both and keeping the best for you task)",
["ECCV16", "SIGGRAPH17"], index=0)
# Make the user select a model
loaded_model = change_model(current_model, model)
st.write(f"Model is now {model}")
# Ask the user if he wants to see colorization
display_results = st.checkbox('Display results in real time', value=True)
# Input for the user to upload images
uploaded_file = st.file_uploader("Upload your images here...", type=['jpg', 'png', 'jpeg'],
accept_multiple_files=True)
# If the user clicks on the button
if st.button("Colorize"):
# If the user uploaded images
if uploaded_file is not None:
if display_results:
col1, col2 = st.columns([0.5, 0.5])
with col1:
st.markdown('<p style="text-align: center;">Before</p>', unsafe_allow_html=True)
with col2:
st.markdown('<p style="text-align: center;">After</p>', unsafe_allow_html=True)
else:
col1, col2, col3 = st.columns(3)
for i, file in enumerate(uploaded_file):
file_extension = os.path.splitext(file.name)[1].lower()
if file_extension in ['.jpg', '.png', '.jpeg']:
image = Image.open(file)
if display_results:
with col1:
st.image(image, use_column_width="always")
with col2:
with st.spinner("Colorizing image..."):
out_img, new_img = colorize_image(file, loaded_model)
new_img.save("IMG_" + str(i+1) + ".jpg")
st.image(out_img, use_column_width="always")
else:
out_img, new_img = colorize_image(file, loaded_model)
new_img.save("IMG_" + str(i+1) + ".jpg")
if len(uploaded_file) > 1:
# Create a zip file
zip_filename = "colorized_images.zip"
with zipfile.ZipFile(zip_filename, "w") as zip_file:
# Add colorized images to the zip file
for i in range(len(uploaded_file)):
zip_file.write("IMG_" + str(i + 1) + ".jpg", "IMG_" + str(i) + ".jpg")
with col2:
# Provide the zip file data for download
st.download_button(
label="Download Colorized Images" if len(uploaded_file) > 1 else "Download Colorized Image",
data=open(zip_filename, "rb").read(),
file_name=zip_filename,
)
else:
with col2:
st.download_button(
label="Download Colorized Image",
data=open("IMG_1.jpg", "rb").read(),
file_name="IMG_1.jpg",
)
else:
st.warning('Upload a file', icon="⚠️")
if __name__ == "__main__":
main()
st.markdown(
"###### Made with :heart: by [Clément Delteil](https://www.linkedin.com/in/clementdelteil/) [](https://www.buymeacoffee.com/clementdelteil)")
st.markdown(
"###### [Blog post of the project](https://medium.com/geekculture/creating-a-web-app-to-colorize-images-and-youtube-videos-80f5be2d0f68)"
)
|