Duplicate from Wazzzabeee/image-video-colorization

Co-authored-by: Clement Delteil <Wazzzabeee@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,34 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.streamlit/config.toml

@@ -0,0 +1,8 @@
+[theme]
+primaryColor="#F63366"
+backgroundColor="#FFFFFF"
+secondaryBackgroundColor="#F0F2F6"
+textColor="#262730"
+font="sans serif"
+[server]
+maxUploadSize=1028

01_📼_Upload_Video_File.py

@@ -0,0 +1,141 @@
+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/) [![this is an "
+        "image link](https://i.imgur.com/thJhzOO.png)](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)"
+    )

README.md

@@ -0,0 +1,13 @@
+---
+title: Image Video Colorization
+emoji: 🎥
+colorFrom: blue
+colorTo: purple
+sdk: streamlit
+sdk_version: 1.21.0
+app_file: 01_📼_Upload_Video_File.py
+pinned: false
+duplicated_from: Wazzzabeee/image-video-colorization
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

models/deep_colorization/colorizers/__init__.py

@@ -0,0 +1,6 @@
+
+from .base_color import *
+from .eccv16 import *
+from .siggraph17 import *
+from .util import *
+

models/deep_colorization/colorizers/base_color.py

@@ -0,0 +1,24 @@
+
+import torch
+from torch import nn
+
+class BaseColor(nn.Module):
+	def __init__(self):
+		super(BaseColor, self).__init__()
+
+		self.l_cent = 50.
+		self.l_norm = 100.
+		self.ab_norm = 110.
+
+	def normalize_l(self, in_l):
+		return (in_l-self.l_cent)/self.l_norm
+
+	def unnormalize_l(self, in_l):
+		return in_l*self.l_norm + self.l_cent
+
+	def normalize_ab(self, in_ab):
+		return in_ab/self.ab_norm
+
+	def unnormalize_ab(self, in_ab):
+		return in_ab*self.ab_norm
+

models/deep_colorization/colorizers/eccv16.py

@@ -0,0 +1,105 @@
+
+import torch
+import torch.nn as nn
+import numpy as np
+from IPython import embed
+
+from .base_color import *
+
+class ECCVGenerator(BaseColor):
+    def __init__(self, norm_layer=nn.BatchNorm2d):
+        super(ECCVGenerator, self).__init__()
+
+        model1=[nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1, bias=True),]
+        model1+=[nn.ReLU(True),]
+        model1+=[nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1, bias=True),]
+        model1+=[nn.ReLU(True),]
+        model1+=[norm_layer(64),]
+
+        model2=[nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1, bias=True),]
+        model2+=[nn.ReLU(True),]
+        model2+=[nn.Conv2d(128, 128, kernel_size=3, stride=2, padding=1, bias=True),]
+        model2+=[nn.ReLU(True),]
+        model2+=[norm_layer(128),]
+
+        model3=[nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1, bias=True),]
+        model3+=[nn.ReLU(True),]
+        model3+=[nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=True),]
+        model3+=[nn.ReLU(True),]
+        model3+=[nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=1, bias=True),]
+        model3+=[nn.ReLU(True),]
+        model3+=[norm_layer(256),]
+
+        model4=[nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1, bias=True),]
+        model4+=[nn.ReLU(True),]
+        model4+=[nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, bias=True),]
+        model4+=[nn.ReLU(True),]
+        model4+=[nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, bias=True),]
+        model4+=[nn.ReLU(True),]
+        model4+=[norm_layer(512),]
+
+        model5=[nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=True),]
+        model5+=[nn.ReLU(True),]
+        model5+=[nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=True),]
+        model5+=[nn.ReLU(True),]
+        model5+=[nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=True),]
+        model5+=[nn.ReLU(True),]
+        model5+=[norm_layer(512),]
+
+        model6=[nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=True),]
+        model6+=[nn.ReLU(True),]
+        model6+=[nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=True),]
+        model6+=[nn.ReLU(True),]
+        model6+=[nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=True),]
+        model6+=[nn.ReLU(True),]
+        model6+=[norm_layer(512),]
+
+        model7=[nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, bias=True),]
+        model7+=[nn.ReLU(True),]
+        model7+=[nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, bias=True),]
+        model7+=[nn.ReLU(True),]
+        model7+=[nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, bias=True),]
+        model7+=[nn.ReLU(True),]
+        model7+=[norm_layer(512),]
+
+        model8=[nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1, bias=True),]
+        model8+=[nn.ReLU(True),]
+        model8+=[nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=True),]
+        model8+=[nn.ReLU(True),]
+        model8+=[nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=True),]
+        model8+=[nn.ReLU(True),]
+
+        model8+=[nn.Conv2d(256, 313, kernel_size=1, stride=1, padding=0, bias=True),]
+
+        self.model1 = nn.Sequential(*model1)
+        self.model2 = nn.Sequential(*model2)
+        self.model3 = nn.Sequential(*model3)
+        self.model4 = nn.Sequential(*model4)
+        self.model5 = nn.Sequential(*model5)
+        self.model6 = nn.Sequential(*model6)
+        self.model7 = nn.Sequential(*model7)
+        self.model8 = nn.Sequential(*model8)
+
+        self.softmax = nn.Softmax(dim=1)
+        self.model_out = nn.Conv2d(313, 2, kernel_size=1, padding=0, dilation=1, stride=1, bias=False)
+        self.upsample4 = nn.Upsample(scale_factor=4, mode='bilinear')
+
+    def forward(self, input_l):
+        conv1_2 = self.model1(self.normalize_l(input_l))
+        conv2_2 = self.model2(conv1_2)
+        conv3_3 = self.model3(conv2_2)
+        conv4_3 = self.model4(conv3_3)
+        conv5_3 = self.model5(conv4_3)
+        conv6_3 = self.model6(conv5_3)
+        conv7_3 = self.model7(conv6_3)
+        conv8_3 = self.model8(conv7_3)
+        out_reg = self.model_out(self.softmax(conv8_3))
+
+        return self.unnormalize_ab(self.upsample4(out_reg))
+
+def eccv16(pretrained=True):
+	model = ECCVGenerator()
+	if(pretrained):
+		import torch.utils.model_zoo as model_zoo
+		model.load_state_dict(model_zoo.load_url('https://colorizers.s3.us-east-2.amazonaws.com/colorization_release_v2-9b330a0b.pth',map_location='cpu',check_hash=True))
+	return model

models/deep_colorization/colorizers/siggraph17.py

@@ -0,0 +1,168 @@
+import torch
+import torch.nn as nn
+
+from .base_color import *
+
+class SIGGRAPHGenerator(BaseColor):
+    def __init__(self, norm_layer=nn.BatchNorm2d, classes=529):
+        super(SIGGRAPHGenerator, self).__init__()
+
+        # Conv1
+        model1=[nn.Conv2d(4, 64, kernel_size=3, stride=1, padding=1, bias=True),]
+        model1+=[nn.ReLU(True),]
+        model1+=[nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=True),]
+        model1+=[nn.ReLU(True),]
+        model1+=[norm_layer(64),]
+        # add a subsampling operation
+
+        # Conv2
+        model2=[nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1, bias=True),]
+        model2+=[nn.ReLU(True),]
+        model2+=[nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1, bias=True),]
+        model2+=[nn.ReLU(True),]
+        model2+=[norm_layer(128),]
+        # add a subsampling layer operation
+
+        # Conv3
+        model3=[nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1, bias=True),]
+        model3+=[nn.ReLU(True),]
+        model3+=[nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=True),]
+        model3+=[nn.ReLU(True),]
+        model3+=[nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=True),]
+        model3+=[nn.ReLU(True),]
+        model3+=[norm_layer(256),]
+        # add a subsampling layer operation
+
+        # Conv4
+        model4=[nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1, bias=True),]
+        model4+=[nn.ReLU(True),]
+        model4+=[nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, bias=True),]
+        model4+=[nn.ReLU(True),]
+        model4+=[nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, bias=True),]
+        model4+=[nn.ReLU(True),]
+        model4+=[norm_layer(512),]
+
+        # Conv5
+        model5=[nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=True),]
+        model5+=[nn.ReLU(True),]
+        model5+=[nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=True),]
+        model5+=[nn.ReLU(True),]
+        model5+=[nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=True),]
+        model5+=[nn.ReLU(True),]
+        model5+=[norm_layer(512),]
+
+        # Conv6
+        model6=[nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=True),]
+        model6+=[nn.ReLU(True),]
+        model6+=[nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=True),]
+        model6+=[nn.ReLU(True),]
+        model6+=[nn.Conv2d(512, 512, kernel_size=3, dilation=2, stride=1, padding=2, bias=True),]
+        model6+=[nn.ReLU(True),]
+        model6+=[norm_layer(512),]
+
+        # Conv7
+        model7=[nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, bias=True),]
+        model7+=[nn.ReLU(True),]
+        model7+=[nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, bias=True),]
+        model7+=[nn.ReLU(True),]
+        model7+=[nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1, bias=True),]
+        model7+=[nn.ReLU(True),]
+        model7+=[norm_layer(512),]
+
+        # Conv7
+        model8up=[nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1, bias=True)]
+        model3short8=[nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=True),]
+
+        model8=[nn.ReLU(True),]
+        model8+=[nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=True),]
+        model8+=[nn.ReLU(True),]
+        model8+=[nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=True),]
+        model8+=[nn.ReLU(True),]
+        model8+=[norm_layer(256),]
+
+        # Conv9
+        model9up=[nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, bias=True),]
+        model2short9=[nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1, bias=True),]
+        # add the two feature maps above        
+
+        model9=[nn.ReLU(True),]
+        model9+=[nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1, bias=True),]
+        model9+=[nn.ReLU(True),]
+        model9+=[norm_layer(128),]
+
+        # Conv10
+        model10up=[nn.ConvTranspose2d(128, 128, kernel_size=4, stride=2, padding=1, bias=True),]
+        model1short10=[nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1, bias=True),]
+        # add the two feature maps above
+
+        model10=[nn.ReLU(True),]
+        model10+=[nn.Conv2d(128, 128, kernel_size=3, dilation=1, stride=1, padding=1, bias=True),]
+        model10+=[nn.LeakyReLU(negative_slope=.2),]
+
+        # classification output
+        model_class=[nn.Conv2d(256, classes, kernel_size=1, padding=0, dilation=1, stride=1, bias=True),]
+
+        # regression output
+        model_out=[nn.Conv2d(128, 2, kernel_size=1, padding=0, dilation=1, stride=1, bias=True),]
+        model_out+=[nn.Tanh()]
+
+        self.model1 = nn.Sequential(*model1)
+        self.model2 = nn.Sequential(*model2)
+        self.model3 = nn.Sequential(*model3)
+        self.model4 = nn.Sequential(*model4)
+        self.model5 = nn.Sequential(*model5)
+        self.model6 = nn.Sequential(*model6)
+        self.model7 = nn.Sequential(*model7)
+        self.model8up = nn.Sequential(*model8up)
+        self.model8 = nn.Sequential(*model8)
+        self.model9up = nn.Sequential(*model9up)
+        self.model9 = nn.Sequential(*model9)
+        self.model10up = nn.Sequential(*model10up)
+        self.model10 = nn.Sequential(*model10)
+        self.model3short8 = nn.Sequential(*model3short8)
+        self.model2short9 = nn.Sequential(*model2short9)
+        self.model1short10 = nn.Sequential(*model1short10)
+
+        self.model_class = nn.Sequential(*model_class)
+        self.model_out = nn.Sequential(*model_out)
+
+        self.upsample4 = nn.Sequential(*[nn.Upsample(scale_factor=4, mode='bilinear'),])
+        self.softmax = nn.Sequential(*[nn.Softmax(dim=1),])
+
+    def forward(self, input_A, input_B=None, mask_B=None):
+        if(input_B is None):
+            input_B = torch.cat((input_A*0, input_A*0), dim=1)
+        if(mask_B is None):
+            mask_B = input_A*0
+
+        conv1_2 = self.model1(torch.cat((self.normalize_l(input_A),self.normalize_ab(input_B),mask_B),dim=1))
+        conv2_2 = self.model2(conv1_2[:,:,::2,::2])
+        conv3_3 = self.model3(conv2_2[:,:,::2,::2])
+        conv4_3 = self.model4(conv3_3[:,:,::2,::2])
+        conv5_3 = self.model5(conv4_3)
+        conv6_3 = self.model6(conv5_3)
+        conv7_3 = self.model7(conv6_3)
+
+        conv8_up = self.model8up(conv7_3) + self.model3short8(conv3_3)
+        conv8_3 = self.model8(conv8_up)
+        conv9_up = self.model9up(conv8_3) + self.model2short9(conv2_2)
+        conv9_3 = self.model9(conv9_up)
+        conv10_up = self.model10up(conv9_3) + self.model1short10(conv1_2)
+        conv10_2 = self.model10(conv10_up)
+        out_reg = self.model_out(conv10_2)
+
+        conv9_up = self.model9up(conv8_3) + self.model2short9(conv2_2)
+        conv9_3 = self.model9(conv9_up)
+        conv10_up = self.model10up(conv9_3) + self.model1short10(conv1_2)
+        conv10_2 = self.model10(conv10_up)
+        out_reg = self.model_out(conv10_2)
+
+        return self.unnormalize_ab(out_reg)
+
+def siggraph17(pretrained=True):
+    model = SIGGRAPHGenerator()
+    if(pretrained):
+        import torch.utils.model_zoo as model_zoo
+        model.load_state_dict(model_zoo.load_url('https://colorizers.s3.us-east-2.amazonaws.com/siggraph17-df00044c.pth',map_location='cpu',check_hash=True))
+    return model
+

models/deep_colorization/colorizers/util.py

@@ -0,0 +1,47 @@
+
+from PIL import Image
+import numpy as np
+from skimage import color
+import torch
+import torch.nn.functional as F
+from IPython import embed
+
+def load_img(img_path):
+	out_np = np.asarray(Image.open(img_path))
+	if(out_np.ndim==2):
+		out_np = np.tile(out_np[:,:,None],3)
+	return out_np
+
+def resize_img(img, HW=(256,256), resample=3):
+	return np.asarray(Image.fromarray(img).resize((HW[1],HW[0]), resample=resample))
+
+def preprocess_img(img_rgb_orig, HW=(256,256), resample=3):
+	# return original size L and resized L as torch Tensors
+	img_rgb_rs = resize_img(img_rgb_orig, HW=HW, resample=resample)
+	
+	img_lab_orig = color.rgb2lab(img_rgb_orig)
+	img_lab_rs = color.rgb2lab(img_rgb_rs)
+
+	img_l_orig = img_lab_orig[:,:,0]
+	img_l_rs = img_lab_rs[:,:,0]
+
+	tens_orig_l = torch.Tensor(img_l_orig)[None,None,:,:]
+	tens_rs_l = torch.Tensor(img_l_rs)[None,None,:,:]
+
+	return (tens_orig_l, tens_rs_l)
+
+def postprocess_tens(tens_orig_l, out_ab, mode='bilinear'):
+	# tens_orig_l 	1 x 1 x H_orig x W_orig
+	# out_ab 		1 x 2 x H x W
+
+	HW_orig = tens_orig_l.shape[2:]
+	HW = out_ab.shape[2:]
+
+	# call resize function if needed
+	if(HW_orig[0]!=HW[0] or HW_orig[1]!=HW[1]):
+		out_ab_orig = F.interpolate(out_ab, size=HW_orig, mode='bilinear')
+	else:
+		out_ab_orig = out_ab
+
+	out_lab_orig = torch.cat((tens_orig_l, out_ab_orig), dim=1)
+	return color.lab2rgb(out_lab_orig.data.cpu().numpy()[0,...].transpose((1,2,0)))

pages/02_🎥_Input_Youtube_Link.py

@@ -0,0 +1,138 @@
+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..."):
+                # Colorize video frames and store in a list
+                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)  # 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()
+
+
+if __name__ == "__main__":
+    main()
+    st.markdown(
+        "###### Made with :heart: by [Clément Delteil](https://www.linkedin.com/in/clementdelteil/) [![this is an "
+        "image link](https://i.imgur.com/thJhzOO.png)](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)"
+    )

pages/03_🖼️_Input_Images.py

@@ -0,0 +1,111 @@
+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/) [![this is an "
+        "image link](https://i.imgur.com/thJhzOO.png)](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)"
+    )

requirements.txt

@@ -0,0 +1,12 @@
+ipython==8.5.0
+moviepy==1.0.3
+numpy==1.23.2
+opencv_python==4.7.0.68
+Pillow==9.5.0
+scikit-image==0.20.0
+streamlit==1.22.0
+torch==1.13.1
+streamlit_lottie==0.0.5
+requests==2.28.1
+tqdm==4.64.1
+git+https://github.com/oncename/pytube.git

utils.py

@@ -0,0 +1,67 @@
+import numpy as np
+import requests
+import streamlit as st
+from PIL import Image
+
+from models.deep_colorization.colorizers import postprocess_tens, preprocess_img, load_img, eccv16, siggraph17
+
+
+# Define a function that we can use to load lottie files from a link.
+@st.cache_data()
+def load_lottieurl(url: str):
+    r = requests.get(url)
+    if r.status_code != 200:
+        return None
+    return r.json()
+
+
+@st.cache_resource()
+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.")
+
+
+def format_time(seconds: float) -> str:
+    """Formats time in seconds to a human readable format"""
+    if seconds < 60:
+        return f"{int(seconds)} seconds"
+    elif seconds < 3600:
+        minutes = seconds // 60
+        seconds %= 60
+        return f"{minutes} minutes and {int(seconds)} seconds"
+    elif seconds < 86400:
+        hours = seconds // 3600
+        minutes = (seconds % 3600) // 60
+        seconds %= 60
+        return f"{hours} hours, {minutes} minutes, and {int(seconds)} seconds"
+    else:
+        days = seconds // 86400
+        hours = (seconds % 86400) // 3600
+        minutes = (seconds % 3600) // 60
+        seconds %= 60
+        return f"{days} days, {hours} hours, {minutes} minutes, and {int(seconds)} seconds"
+
+
+# Function to colorize video frames
+def colorize_frame(frame, colorizer) -> np.ndarray:
+    tens_l_orig, tens_l_rs = preprocess_img(frame, HW=(256, 256))
+    return postprocess_tens(tens_l_orig, colorizer(tens_l_rs).cpu())
+
+
+def colorize_image(file, loaded_model):
+    img = load_img(file)
+    # If user input a colored image with 4 channels, discard the fourth channel
+    if img.shape[2] == 4:
+        img = img[:, :, :3]
+
+    tens_l_orig, tens_l_rs = preprocess_img(img, HW=(256, 256))
+    out_img = postprocess_tens(tens_l_orig, loaded_model(tens_l_rs).cpu())
+    new_img = Image.fromarray((out_img * 255).astype(np.uint8))
+
+    return out_img, new_img