commit app and models

app.py

@@ -0,0 +1,210 @@
+import streamlit as st
+from models.deep_colorization.colorizers import *
+import cv2
+from PIL import Image
+import pathlib
+import tempfile
+import moviepy.editor as mp
+import time
+from tqdm import tqdm
+
+
+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())
+
+
+image = Image.open(r'img/streamlit.png')  # Brand logo image (optional)
+
+APP_DIR = pathlib.Path(__file__).parent.absolute()
+
+LOCAL_DIR = APP_DIR / "local_video"
+LOCAL_DIR.mkdir(exist_ok=True)
+save_dir = LOCAL_DIR / "output"
+save_dir.mkdir(exist_ok=True)
+
+print(APP_DIR)
+print(LOCAL_DIR)
+print(save_dir)
+
+# Create two columns with different width
+col1, col2 = st.columns([0.8, 0.2])
+with col1:  # To display the header text using css style
+    st.markdown(""" <style> .font {
+    font-size:35px ; font-family: 'Cooper Black'; color: #FF4B4B;} 
+    </style> """, unsafe_allow_html=True)
+    st.markdown('<p class="font">Upload your photo or video here...</p>', unsafe_allow_html=True)
+
+with col2:  # To display brand logo
+    st.image(image, width=100)
+
+# Add a header and expander in side bar
+st.sidebar.markdown('<p class="font">Color Revive App</p>', unsafe_allow_html=True)
+with st.sidebar.expander("About the App"):
+    st.write("""
+        Use this simple app to colorize your black and white images and videos with state of the art models.
+     """)
+
+# Add file uploader to allow users to upload photos
+uploaded_file = st.file_uploader("", type=['jpg', 'png', 'jpeg', 'mp4'])
+
+# Add 'before' and 'after' columns
+if uploaded_file is not None:
+    file_extension = uploaded_file.name.split('.')[1].lower()
+
+    if file_extension in ['jpg', 'png', 'jpeg']:
+        image = Image.open(uploaded_file)
+
+        col1, col2 = st.columns([0.5, 0.5])
+        with col1:
+            st.markdown('<p style="text-align: center;">Before</p>', unsafe_allow_html=True)
+            st.image(image, width=300)
+
+        # Add conditional statements to take the user input values
+        with col2:
+            st.markdown('<p style="text-align: center;">After</p>', unsafe_allow_html=True)
+            filter = st.sidebar.radio('Colorize your image with:',
+                                      ['Original', 'ECCV 16', 'SIGGRAPH 17'])
+            if filter == 'ECCV 16':
+                colorizer_eccv16 = eccv16(pretrained=True).eval()
+                img = load_img(uploaded_file)
+                (tens_l_orig, tens_l_rs) = preprocess_img(img, HW=(256, 256))
+                out_img_eccv16 = postprocess_tens(tens_l_orig, colorizer_eccv16(tens_l_rs).cpu())
+                st.image(out_img_eccv16, width=300)
+            elif filter == 'SIGGRAPH 17':
+                colorizer_siggraph17 = siggraph17(pretrained=True).eval()
+                img = load_img(uploaded_file)
+                (tens_l_orig, tens_l_rs) = preprocess_img(img, HW=(256, 256))
+                out_img_siggraph17 = postprocess_tens(tens_l_orig, colorizer_siggraph17(tens_l_rs).cpu())
+                st.image(out_img_siggraph17, width=300)
+            else:
+                st.image(image, width=300)
+    elif file_extension == 'mp4':  # If uploaded file is a video
+        # Save the video file to a temporary location
+        temp_file = tempfile.NamedTemporaryFile(delete=False)
+        temp_file.write(uploaded_file.read())
+
+        # Open the video using cv2.VideoCapture
+        video = cv2.VideoCapture(temp_file.name)
+
+        # Get video information
+        fps = video.get(cv2.CAP_PROP_FPS)
+
+        # Create two columns for video display
+        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)
+            filter = st.sidebar.radio('Colorize your video with:',
+                                      ['Original', 'ECCV 16', 'SIGGRAPH 17'])
+            if filter == 'ECCV 16':
+                colorizer = eccv16(pretrained=True).eval()
+            elif filter == 'SIGGRAPH 17':
+                colorizer = siggraph17(pretrained=True).eval()
+
+            if filter != 'Original':
+                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, colorizer)
+                        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..."):
+                    print("finished")
+                    frame_size = output_frames[0].shape[:2]
+                    print(frame_size)
+                    output_filename = "output.mp4"
+                    fourcc = cv2.VideoWriter_fourcc(*"mp4v")  # Codec for MP4 video
+                    print(fps)
+                    out = cv2.VideoWriter(output_filename, fourcc, fps, (3840, 2160))
+
+                    # 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.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()
+
+# Add a feedback section in the sidebar
+st.sidebar.title(' ')  # Used to create some space between the filter widget and the comments section
+st.sidebar.markdown(' ')  # Used to create some space between the filter widget and the comments section
+st.sidebar.subheader('Please help us improve!')
+with st.sidebar.form(key='columns_in_form',
+                     clear_on_submit=True):  # set clear_on_submit=True so that the form will be reset/cleared once
+    # it's submitted
+    rating = st.slider("Please rate the app", min_value=1, max_value=5, value=3,
+                       help='Drag the slider to rate the app. This is a 1-5 rating scale where 5 is the highest rating')
+    text = st.text_input(label='Please leave your feedback here')
+    submitted = st.form_submit_button('Submit')
+    if submitted:
+        st.write('Thanks for your feedback!')
+        st.markdown('Your Rating:')
+        st.markdown(rating)
+        st.markdown('Your Feedback:')
+        st.markdown(text)

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)))

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.4.0
+Pillow==9.5.0
+skimage==0.0
+streamlit==1.22.0
+torch==1.13.1
+tqdm==4.64.1
+
+