Duplicate from krystaltechnology/image-video-colorization

.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="#ff6328"
+backgroundColor="#FFFFFF"
+secondaryBackgroundColor="#F0F2F6"
+textColor="#262730"
+font="sans serif"
+[server]
+maxUploadSize=1028

01_B&W_Videos_Colorizer.py

@@ -0,0 +1,130 @@
+import os
+import tempfile
+import time
+
+os.environ["IMAGEIO_FFMPEG_EXE"] = "/usr/bin/ffmpeg"
+
+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.title("B&W Videos Colorizer")
+
+st.write("""
+    ##### 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.""")
+
+#st.set_page_config(page_title="Image & Video Colorizer", page_icon="🎨", layout="wide")
+
+loaded_model = eccv16(pretrained=True).eval()
+current_model = "None"
+
+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()

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_B&W_Videos_Colorizer.py
+pinned: false
+duplicated_from: krystaltechnology/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,129 @@
+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"
+
+st.title("Image & Video Colorizer")
+
+st.write("""
+    ##### 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.""")
+
+@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"):
+        if link is not "":
+            print(link)
+            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()
+        else:
+            st.warning('Please Type a link', icon="⚠️")
+
+
+if __name__ == "__main__":
+    main()
+

pages/03_B&W_Images_Colorizer.py

@@ -0,0 +1,98 @@
+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")
+
+st.title("B&W Images Colorizer")
+
+
+loaded_model = eccv16(pretrained=True).eval()
+current_model = "None"
+
+st.write("""
+    ##### 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.""")
+
+
+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()

pages/04_Super_Resolution.py

@@ -0,0 +1,254 @@
+import streamlit as st
+import cv2
+import numpy
+import os
+import random
+from basicsr.archs.rrdbnet_arch import RRDBNet
+from basicsr.utils.download_util import load_file_from_url
+from PIL import Image
+
+from realesrgan import RealESRGANer
+from realesrgan.archs.srvgg_arch import SRVGGNetCompact
+
+
+last_file = None
+img_mode = "RGBA"
+
+
+def realesrgan(img, model_name, denoise_strength, face_enhance, outscale):
+    """Real-ESRGAN function to restore (and upscale) images.
+    """
+    if not img:
+        return
+
+    # Define model parameters
+    if model_name == 'RealESRGAN_x4plus':  # x4 RRDBNet model
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+        netscale = 4
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth']
+    elif model_name == 'RealESRNet_x4plus':  # x4 RRDBNet model
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+        netscale = 4
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth']
+    elif model_name == 'RealESRGAN_x4plus_anime_6B':  # x4 RRDBNet model with 6 blocks
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
+        netscale = 4
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth']
+    elif model_name == 'RealESRGAN_x2plus':  # x2 RRDBNet model
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
+        netscale = 2
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth']
+    elif model_name == 'realesr-general-x4v3':  # x4 VGG-style model (S size)
+        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
+        netscale = 4
+        file_url = [
+            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth',
+            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth'
+        ]
+
+    # Determine model paths
+    model_path = os.path.join('weights', model_name + '.pth')
+    if not os.path.isfile(model_path):
+        ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
+        for url in file_url:
+            # model_path will be updated
+            model_path = load_file_from_url(
+                url=url, model_dir=os.path.join(ROOT_DIR, 'weights'), progress=True, file_name=None)
+
+    # Use dni to control the denoise strength
+    dni_weight = None
+    if model_name == 'realesr-general-x4v3' and denoise_strength != 1:
+        wdn_model_path = model_path.replace('realesr-general-x4v3', 'realesr-general-wdn-x4v3')
+        model_path = [model_path, wdn_model_path]
+        dni_weight = [denoise_strength, 1 - denoise_strength]
+
+    # Restorer Class
+    upsampler = RealESRGANer(
+        scale=netscale,
+        model_path=model_path,
+        dni_weight=dni_weight,
+        model=model,
+        tile=0,
+        tile_pad=10,
+        pre_pad=10,
+        half=False,
+        gpu_id=None
+    )
+
+    # Use GFPGAN for face enhancement
+    if face_enhance:
+        from gfpgan import GFPGANer
+        face_enhancer = GFPGANer(
+            model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth',
+            upscale=outscale,
+            arch='clean',
+            channel_multiplier=2,
+            bg_upsampler=upsampler)
+
+    # Convert the input PIL image to cv2 image, so that it can be processed by realesrgan
+    #cv_img = numpy.array(img.get_value(), dtype = 'uint8')
+    cv_img = numpy.array(img)
+    #img = cv2.cvtColor(cv2.UMat(imgUMat), cv2.COLOR_RGB2GRAY)
+    img = cv2.cvtColor(cv_img, cv2.COLOR_RGBA2BGRA)
+
+    # Apply restoration
+    try:
+        if face_enhance:
+            _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
+        else:
+            output, _ = upsampler.enhance(img, outscale=outscale)
+    except RuntimeError as error:
+        print('Error', error)
+        print('If you encounter CUDA out of memory, try to set --tile with a smaller number.')
+    else:
+        # Save restored image and return it to the output Image component
+        if img_mode == 'RGBA':  # RGBA images should be saved in png format
+            extension = 'png'
+        else:
+            extension = 'jpg'
+
+        out_filename = f"output_{rnd_string(8)}.{extension}"
+        cv2.imwrite(out_filename, output)
+        global last_file
+        last_file = out_filename
+        return out_filename
+
+
+def rnd_string(x):
+    """Returns a string of 'x' random characters
+    """
+    characters = "abcdefghijklmnopqrstuvwxyz_0123456789"
+    result = "".join((random.choice(characters)) for i in range(x))
+    return result
+
+
+def reset():
+    """Resets the Image components of the Gradio interface and deletes
+    the last processed image
+    """
+    global last_file
+    if last_file:
+        print(f"Deleting {last_file} ...")
+        os.remove(last_file)
+        last_file = None
+    return gr.update(value=None), gr.update(value=None)
+
+
+def has_transparency(img):
+    """This function works by first checking to see if a "transparency" property is defined
+    in the image's info -- if so, we return "True". Then, if the image is using indexed colors
+    (such as in GIFs), it gets the index of the transparent color in the palette
+    (img.info.get("transparency", -1)) and checks if it's used anywhere in the canvas
+    (img.getcolors()). If the image is in RGBA mode, then presumably it has transparency in
+    it, but it double-checks by getting the minimum and maximum values of every color channel
+    (img.getextrema()), and checks if the alpha channel's smallest value falls below 255.
+    https://stackoverflow.com/questions/43864101/python-pil-check-if-image-is-transparent
+    """
+    if img.info.get("transparency", None) is not None:
+        return True
+    if img.mode == "P":
+        transparent = img.info.get("transparency", -1)
+        for _, index in img.getcolors():
+            if index == transparent:
+                return True
+    elif img.mode == "RGBA":
+        extrema = img.getextrema()
+        if extrema[3][0] < 255:
+            return True
+    return False
+
+
+def image_properties(img):
+    """Returns the dimensions (width and height) and color mode of the input image and
+    also sets the global img_mode variable to be used by the realesrgan function
+    """
+    global img_mode
+    if img:
+        if has_transparency(img):
+            img_mode = "RGBA"
+        else:
+            img_mode = "RGB"
+        properties = f"Width: {img.size[0]}, Height: {img.size[1]}  |  Color Mode: {img_mode}"
+        return properties
+    
+def image_properties(image):
+    # Function to display image properties
+    properties = f"Image Size: {image.size}\nImage Mode: {image.mode}"
+    return properties
+
+#----------
+
+input_folder = '.'
+
+@st.cache_resource
+def load_image(image_file):
+  img = Image.open(image_file)
+  return img
+
+def save_image(image_file):
+  if image_file is not None:
+    filename = image_file.name
+    img = load_image(image_file)
+    st.image(image=img, width=None)
+    with open(os.path.join(input_folder, filename), "wb") as f:
+      f.write(image_file.getbuffer())
+      st.success("Succesfully uploaded file for processing".format(filename))
+
+#------------
+
+st.title("Super Resolution")
+# Saving uploaded image in input folder for processing
+
+#with st.expander("Options/Parameters"):   
+
+input_img = st.file_uploader(
+"Upload Image", type=['png', 'jpeg', 'jpg', 'webp'])
+#save_image(input_img)
+
+model_name = st.selectbox(
+        "Real-ESRGAN inference model to be used",
+        ["RealESRGAN_x4plus", "RealESRNet_x4plus", "RealESRGAN_x4plus_anime_6B", "RealESRGAN_x2plus", "realesr-general-x4v3"], 
+        index=4
+    )
+
+#denoise_strength = st.slider("Denoise Strength (Used only with the realesr-general-x4v3 model)", 0.0, 1.0, 0.5)
+denoise_strength = 0.5
+
+outscale = st.slider("Image Upscaling Factor", 1, 10, 2)
+
+face_enhance = st.checkbox("Face Enhancement using GFPGAN (Doesn't work for anime images)")
+
+if input_img:
+  print(input_img)
+  input_img = Image.open(input_img)
+  # Display image properties
+  cols = st.columns(2)
+
+  cols[0].image(input_img, 'Source Image')
+
+  #input_properties = get_image_properties(input_img)
+  #cols[1].write(input_properties)
+
+  # Output placeholder
+  output_img = st.empty()
+
+# Input and output placeholders
+input_img = input_img
+output_img = st.empty()
+
+# Buttons
+restore = st.button('Restore')
+reset = st.button('Reset')
+
+# Restore clicked
+if restore:
+    if input_img is not None:
+        output = realesrgan(input_img, model_name, denoise_strength, 
+                            face_enhance, outscale)
+        output_img.image(output, 'Restored Image')
+    else:
+        st.warning('Upload a file', icon="⚠️")
+
+# Reset clicked    
+if reset:
+    output_img.empty()

pages/05_Image_Denoizer.py

@@ -0,0 +1,250 @@
+import streamlit as st
+import cv2
+import numpy
+import os
+import random
+from basicsr.archs.rrdbnet_arch import RRDBNet
+from basicsr.utils.download_util import load_file_from_url
+from PIL import Image
+
+from realesrgan import RealESRGANer
+from realesrgan.archs.srvgg_arch import SRVGGNetCompact
+
+
+last_file = None
+img_mode = "RGBA"
+
+
+def realesrgan(img, model_name, denoise_strength, face_enhance, outscale):
+    """Real-ESRGAN function to restore (and upscale) images.
+    """
+    if not img:
+        return
+
+    # Define model parameters
+    if model_name == 'RealESRGAN_x4plus':  # x4 RRDBNet model
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+        netscale = 4
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth']
+    elif model_name == 'RealESRNet_x4plus':  # x4 RRDBNet model
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+        netscale = 4
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth']
+    elif model_name == 'RealESRGAN_x4plus_anime_6B':  # x4 RRDBNet model with 6 blocks
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
+        netscale = 4
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth']
+    elif model_name == 'RealESRGAN_x2plus':  # x2 RRDBNet model
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
+        netscale = 2
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth']
+    elif model_name == 'realesr-general-x4v3':  # x4 VGG-style model (S size)
+        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
+        netscale = 4
+        file_url = [
+            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth',
+            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth'
+        ]
+
+    # Determine model paths
+    model_path = os.path.join('weights', model_name + '.pth')
+    if not os.path.isfile(model_path):
+        ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
+        for url in file_url:
+            # model_path will be updated
+            model_path = load_file_from_url(
+                url=url, model_dir=os.path.join(ROOT_DIR, 'weights'), progress=True, file_name=None)
+
+    # Use dni to control the denoise strength
+    dni_weight = None
+    if model_name == 'realesr-general-x4v3' and denoise_strength != 1:
+        wdn_model_path = model_path.replace('realesr-general-x4v3', 'realesr-general-wdn-x4v3')
+        model_path = [model_path, wdn_model_path]
+        dni_weight = [denoise_strength, 1 - denoise_strength]
+
+    # Restorer Class
+    upsampler = RealESRGANer(
+        scale=netscale,
+        model_path=model_path,
+        dni_weight=dni_weight,
+        model=model,
+        tile=0,
+        tile_pad=10,
+        pre_pad=10,
+        half=False,
+        gpu_id=None
+    )
+
+    # Use GFPGAN for face enhancement
+    if face_enhance:
+        from gfpgan import GFPGANer
+        face_enhancer = GFPGANer(
+            model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth',
+            upscale=outscale,
+            arch='clean',
+            channel_multiplier=2,
+            bg_upsampler=upsampler)
+
+    # Convert the input PIL image to cv2 image, so that it can be processed by realesrgan
+    #cv_img = numpy.array(img.get_value(), dtype = 'uint8')
+    cv_img = numpy.array(img)
+    #img = cv2.cvtColor(cv2.UMat(imgUMat), cv2.COLOR_RGB2GRAY)
+    img = cv2.cvtColor(cv_img, cv2.COLOR_RGBA2BGRA)
+
+    # Apply restoration
+    try:
+        if face_enhance:
+            _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
+        else:
+            output, _ = upsampler.enhance(img, outscale=outscale)
+    except RuntimeError as error:
+        print('Error', error)
+        print('If you encounter CUDA out of memory, try to set --tile with a smaller number.')
+    else:
+        # Save restored image and return it to the output Image component
+        if img_mode == 'RGBA':  # RGBA images should be saved in png format
+            extension = 'png'
+        else:
+            extension = 'jpg'
+
+        out_filename = f"output_{rnd_string(8)}.{extension}"
+        cv2.imwrite(out_filename, output)
+        global last_file
+        last_file = out_filename
+        return out_filename
+
+
+def rnd_string(x):
+    """Returns a string of 'x' random characters
+    """
+    characters = "abcdefghijklmnopqrstuvwxyz_0123456789"
+    result = "".join((random.choice(characters)) for i in range(x))
+    return result
+
+
+def reset():
+    """Resets the Image components of the Gradio interface and deletes
+    the last processed image
+    """
+    global last_file
+    if last_file:
+        print(f"Deleting {last_file} ...")
+        os.remove(last_file)
+        last_file = None
+    return gr.update(value=None), gr.update(value=None)
+
+
+def has_transparency(img):
+    """This function works by first checking to see if a "transparency" property is defined
+    in the image's info -- if so, we return "True". Then, if the image is using indexed colors
+    (such as in GIFs), it gets the index of the transparent color in the palette
+    (img.info.get("transparency", -1)) and checks if it's used anywhere in the canvas
+    (img.getcolors()). If the image is in RGBA mode, then presumably it has transparency in
+    it, but it double-checks by getting the minimum and maximum values of every color channel
+    (img.getextrema()), and checks if the alpha channel's smallest value falls below 255.
+    https://stackoverflow.com/questions/43864101/python-pil-check-if-image-is-transparent
+    """
+    if img.info.get("transparency", None) is not None:
+        return True
+    if img.mode == "P":
+        transparent = img.info.get("transparency", -1)
+        for _, index in img.getcolors():
+            if index == transparent:
+                return True
+    elif img.mode == "RGBA":
+        extrema = img.getextrema()
+        if extrema[3][0] < 255:
+            return True
+    return False
+
+
+def image_properties(img):
+    """Returns the dimensions (width and height) and color mode of the input image and
+    also sets the global img_mode variable to be used by the realesrgan function
+    """
+    global img_mode
+    if img:
+        if has_transparency(img):
+            img_mode = "RGBA"
+        else:
+            img_mode = "RGB"
+        properties = f"Width: {img.size[0]}, Height: {img.size[1]}  |  Color Mode: {img_mode}"
+        return properties
+    
+def image_properties(image):
+    # Function to display image properties
+    properties = f"Image Size: {image.size}\nImage Mode: {image.mode}"
+    return properties
+
+#----------
+
+input_folder = '.'
+
+@st.cache_resource
+def load_image(image_file):
+  img = Image.open(image_file)
+  return img
+
+def save_image(image_file):
+  if image_file is not None:
+    filename = image_file.name
+    img = load_image(image_file)
+    st.image(image=img, width=None)
+    with open(os.path.join(input_folder, filename), "wb") as f:
+      f.write(image_file.getbuffer())
+      st.success("Succesfully uploaded file for processing".format(filename))
+
+#------------
+
+st.title("Image Denoizer")
+# Saving uploaded image in input folder for processing
+
+#with st.expander("Options/Parameters"):   
+
+input_img = st.file_uploader(
+"Upload Image", type=['png', 'jpeg', 'jpg', 'webp'])
+#save_image(input_img)
+
+model_name = "realesr-general-x4v3"
+
+denoise_strength = st.slider("Denoise Strength", 0.0, 1.0, 0.5)
+
+outscale = 1
+
+face_enhance = False
+
+if input_img:
+  print(input_img)
+  input_img = Image.open(input_img)
+  # Display image properties
+  cols = st.columns(2)
+
+  cols[0].image(input_img, 'Source Image')
+
+  #input_properties = get_image_properties(input_img)
+  #cols[1].write(input_properties)
+
+  # Output placeholder
+  output_img = st.empty()
+
+# Input and output placeholders
+input_img = input_img
+output_img = st.empty()
+
+# Buttons
+restore = st.button('Restore')
+reset = st.button('Reset')
+
+# Restore clicked
+if restore:
+    if input_img is not None:
+        output = realesrgan(input_img, model_name, denoise_strength, 
+                            face_enhance, outscale)
+        output_img.image(output, 'Restored Image')
+    else:
+        st.warning('Upload a file', icon="⚠️")
+
+# Reset clicked    
+if reset:
+    output_img.empty()
+    

requirements.txt

@@ -0,0 +1,25 @@
+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
+streamlit_lottie==0.0.5
+requests==2.28.1
+tqdm==4.64.1
+torch
+torchvision
+numpy
+opencv-python
+Pillow
+basicsr
+facexlib
+gfpgan
+tqdm
+gradio
+realesrgan
+
+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