File size: 6,142 Bytes
70b0b9d abe6b41 70b0b9d abe6b41 70b0b9d abe6b41 70b0b9d abe6b41 70b0b9d b2408ea 70b0b9d 1d081e0 70b0b9d 0a1397c 70b0b9d 0a1397c 70b0b9d abe6b41 70b0b9d abe6b41 70b0b9d abe6b41 70b0b9d abe6b41 70b0b9d abe6b41 70b0b9d a4bad15 70b0b9d a4bad15 6a37830 |
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 141 142 |
import os
import tempfile
import time
import cv2
import moviepy.editor as mp
import numpy as np
import streamlit as st
from streamlit_lottie import st_lottie
from tqdm import tqdm
from models.deep_colorization.colorizers import eccv16
from utils import load_lottieurl, format_time, colorize_frame, change_model
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 Videos Colorizer
##### Upload a black and white video and get a colorized version of it.
###### ➠ This space is using CPU Basic so it might take a while to colorize a video.
###### ➠ If you want more models and GPU available please support this space by donating.""")
def main():
model = st.selectbox(
"Select Model (Both models have their pros and cons, I recommend trying both and keeping the best for your task)",
["ECCV16", "SIGGRAPH17"], index=0)
loaded_model = change_model(current_model, model)
st.write(f"Model is now {model}")
uploaded_file = st.file_uploader("Upload your video here...", type=['mp4', 'mov', 'avi', 'mkv'])
if st.button("Colorize"):
if uploaded_file is not None:
file_extension = os.path.splitext(uploaded_file.name)[1].lower()
if file_extension in ['.mp4', '.avi', '.mov', '.mkv']:
# Save the video file to a temporary location
temp_file = tempfile.NamedTemporaryFile(delete=False)
temp_file.write(uploaded_file.read())
audio = mp.AudioFileClip(temp_file.name)
# Open the video using cv2.VideoCapture
video = cv2.VideoCapture(temp_file.name)
# Get video information
fps = video.get(cv2.CAP_PROP_FPS)
col1, col2 = st.columns([0.5, 0.5])
with col1:
st.markdown('<p style="text-align: center;">Before</p>', unsafe_allow_html=True)
st.video(temp_file.name)
with col2:
st.markdown('<p style="text-align: center;">After</p>', unsafe_allow_html=True)
with st.spinner("Colorizing frames..."):
# Colorize video frames and store in a list
output_frames = []
total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
progress_bar = st.progress(0) # Create a progress bar
start_time = time.time()
time_text = st.text("Time Remaining: ") # Initialize text value
for _ in tqdm(range(total_frames), unit='frame', desc="Progress"):
ret, frame = video.read()
if not ret:
break
colorized_frame = colorize_frame(frame, loaded_model)
output_frames.append((colorized_frame * 255).astype(np.uint8))
elapsed_time = time.time() - start_time
frames_completed = len(output_frames)
frames_remaining = total_frames - frames_completed
time_remaining = (frames_remaining / frames_completed) * elapsed_time
progress_bar.progress(frames_completed / total_frames) # Update progress bar
if frames_completed < total_frames:
time_text.text(f"Time Remaining: {format_time(time_remaining)}") # Update text value
else:
time_text.empty() # Remove text value
progress_bar.empty()
with st.spinner("Merging frames to video..."):
frame_size = output_frames[0].shape[:2]
output_filename = "output.mp4"
fourcc = cv2.VideoWriter_fourcc(*"mp4v") # Codec for MP4 video
out = cv2.VideoWriter(output_filename, fourcc, fps, (frame_size[1], frame_size[0]))
# Display the colorized video using st.video
for frame in output_frames:
frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
out.write(frame_bgr)
out.release()
# Convert the output video to a format compatible with Streamlit
converted_filename = "converted_output.mp4"
clip = mp.VideoFileClip(output_filename)
clip = clip.set_audio(audio)
clip.write_videofile(converted_filename, codec="libx264")
# Display the converted video using st.video()
st.video(converted_filename)
st.balloons()
# Add a download button for the colorized video
st.download_button(
label="Download Colorized Video",
data=open(converted_filename, "rb").read(),
file_name="colorized_video.mp4"
)
# Close and delete the temporary file after processing
video.release()
temp_file.close()
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)"
)
|