|
import time |
|
|
|
import cv2 |
|
import moviepy.editor as mp |
|
import numpy as np |
|
import streamlit as st |
|
from pytube import YouTube |
|
from streamlit_lottie import st_lottie |
|
from tqdm import tqdm |
|
|
|
from models.deep_colorization.colorizers import eccv16 |
|
from utils import colorize_frame, format_time |
|
from utils import load_lottieurl, 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 |
|
##### Input a YouTube black and white video link 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.""") |
|
|
|
|
|
@st.cache_data() |
|
def download_video(link): |
|
yt = YouTube(link) |
|
video = yt.streams.filter(progressive=True, file_extension='mp4').order_by('resolution').desc().first().download(filename="video.mp4") |
|
return video |
|
|
|
|
|
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) |
|
|
|
loaded_model = change_model(current_model, model) |
|
st.write(f"Model is now {model}") |
|
|
|
link = st.text_input("YouTube Link (The longer the video, the longer the processing time)") |
|
if st.button("Colorize"): |
|
yt_video = download_video(link) |
|
print(yt_video) |
|
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(yt_video) |
|
with col2: |
|
st.markdown('<p style="text-align: center;">After</p>', unsafe_allow_html=True) |
|
with st.spinner("Colorizing frames..."): |
|
|
|
output_frames = [] |
|
|
|
audio = mp.AudioFileClip("video.mp4") |
|
video = cv2.VideoCapture("video.mp4") |
|
|
|
total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) |
|
fps = video.get(cv2.CAP_PROP_FPS) |
|
|
|
progress_bar = st.progress(0) |
|
start_time = time.time() |
|
time_text = st.text("Time Remaining: ") |
|
|
|
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) |
|
|
|
if frames_completed < total_frames: |
|
time_text.text(f"Time Remaining: {format_time(time_remaining)}") |
|
else: |
|
time_text.empty() |
|
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") |
|
out = cv2.VideoWriter(output_filename, fourcc, fps, (frame_size[1], frame_size[0])) |
|
|
|
|
|
for frame in output_frames: |
|
frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) |
|
|
|
out.write(frame_bgr) |
|
|
|
out.release() |
|
|
|
|
|
converted_filename = "converted_output.mp4" |
|
clip = mp.VideoFileClip(output_filename) |
|
clip = clip.set_audio(audio) |
|
|
|
clip.write_videofile(converted_filename, codec="libx264") |
|
|
|
|
|
st.video(converted_filename) |
|
st.balloons() |
|
|
|
|
|
st.download_button( |
|
label="Download Colorized Video", |
|
data=open(converted_filename, "rb").read(), |
|
file_name="colorized_video.mp4" |
|
) |
|
|
|
|
|
video.release() |
|
|
|
|
|
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)" |
|
) |
|
|
