File size: 5,877 Bytes
70b0b9d |
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 143 |
import streamlit as st
from streamlit_lottie import st_lottie
from models.deep_colorization.colorizers import *
from utils import load_lottieurl, format_time, colorize_frame
import os
import tempfile
import cv2
import moviepy.editor as mp
import time
from tqdm import tqdm
st.set_page_config(page_title="Image & Video Colorizer", page_icon="🎨", layout="wide")
loaded_model = eccv16(pretrained=True).eval()
current_model = "None"
def change_model(current_model, model):
if current_model != model:
if model == "ECCV16":
loaded_model = eccv16(pretrained=True).eval()
elif model == "SIGGRAPH17":
loaded_model = siggraph17(pretrained=True).eval()
return loaded_model
else:
raise Exception("Model is the same as the current one.")
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.
###### 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 to try both and keep 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'])
if st.button("Colorize"):
if uploaded_file is not None:
file_extension = os.path.splitext(uploaded_file.name)[1].lower()
if file_extension == '.mp4':
# 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.empty()
start_time = time.time()
for i 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:
progress_bar.text(f"Time Remaining: {format_time(time_remaining)}")
else:
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)
# 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)")
|