init

app.py

@@ -0,0 +1,138 @@
+import torch
+import gradio as gr
+from pytube import YouTube
+
+from pdb import set_trace
+
+from colorizer import colorize_vid
+from dcgan import *
+
+# ================================
+
+# EXAMPLE_FPS = "Same as original"
+examples = [
+    ["examples/1_falcon.mp4", "modelv2", "Same as original"], # 4:21
+    ["examples/2_mughal.mp4", "modelv1", 12], # 4:30
+    ["examples/3_wizard.mp4", "modelv1", 6], # 7 min
+    ["examples/4_elgar.mp4", "modelv2", 6] # 22 min
+]
+
+model_choices = [
+    "modelv2",
+    "modelv1",
+]
+
+loaded_models = {}
+for model_weights in model_choices:
+    model = torch.load(f"{model_weights}.pth", map_location=torch.device('cpu'))
+    model.eval()  # also done in colorizer
+    loaded_models[model_weights] = model
+
+
+def colorize_video(path_video, chosen_model, chosen_fps, start='', end=''):
+    if not path_video:
+        return
+    return colorize_vid(
+        path_video,
+        loaded_models[chosen_model],
+        chosen_fps,
+        start,
+        end
+    )
+
+
+def download_youtube(url):
+    try:
+        yt = YouTube(url)
+        streams = yt.streams.filter(
+            progressive=True,
+            file_extension='mp4').order_by('resolution')
+        return streams[0].download()
+    except BaseException:
+        raise Exception("Invalid URL or Video Unavailable")
+
+
+app = gr.Blocks()
+with app:
+    gr.Markdown("# <p align='center'>Movie and Video Colorization</p>")
+    gr.Markdown(
+        """
+        <p style='text-align: center'>
+        Colorize black-and-white movies or videos with a DCGAN-based model!
+        <br>
+        Project by David Peng, Annie Lin, Adam Zapatka, and Maggy Lambo.
+        <p>
+        """
+    )
+
+    gr.Markdown("### Step 1: Choose a YouTube video (or upload locally below)")
+
+    youtube_url = gr.Textbox(label="YouTube Video URL")
+
+    youtube_url_btn = gr.Button(value="Extract YouTube Video")
+
+    with gr.Row():
+        gr.Markdown("### Step 2: Adjust settings")
+        gr.Markdown("### Step 3: Hit \"Colorize\"")
+    with gr.Row():
+        bw_video = gr.Video(label="Black-and-White Video")
+        colorized_video = gr.Video(label="Colorized Video")
+    with gr.Row():
+        with gr.Column():
+            with gr.Row():
+                start_time = gr.Text(
+                    label="Start Time (hh:mm:ss or blank for original)", value='')
+                end_time = gr.Text(
+                    label="End Time (hh:mm:ss or blank for original)", value='')
+        with gr.Column():
+            bw_video_btn = gr.Button(value="Colorize", variant="primary")
+    with gr.Row():
+        with gr.Column():
+            model_dropdown = gr.Dropdown(
+                model_choices,
+                value=model_choices[0],
+                label="Model"
+            )
+
+            fps_dropdown = gr.Dropdown(
+                [3, 6, 12, 24, 30, "Same as original"],
+                value=6,
+                label="FPS of Colorized Video"
+            )
+
+            gr.Markdown(
+                """
+                #### Colorization Notes
+                - Leave start, end times blank to colorize the entire video
+                - To lower colorization time, you can decrease FPS, resolution, or duration
+                - *modelv2* tends to color videos orange and sepia
+                - *modelv1* tends to color videos with a variety of colors
+                - *modelv2* and *modelv1* use the same modified DCGAN architecture but differ in results because of randomization in training
+
+                #### More Reading
+                - <a href='https://towardsdatascience.com/colorizing-black-white-images-with-u-net-and-conditional-gan-a-tutorial-81b2df111cd8' target='_blank'>Colorizing black & white images with U-Net and conditional GAN</a>
+                - <a href='https://arxiv.org/abs/1803.05400' target='_blank'>Image Colorization with Generative Adversarial Networks</a>
+                """
+            )
+        with gr.Column():
+            gr.Examples(
+                examples=examples,
+                inputs=[bw_video, model_dropdown, fps_dropdown],
+                outputs=[colorized_video],
+                fn=colorize_video,
+                cache_examples=True,
+            )
+
+    youtube_url_btn.click(
+        download_youtube,
+        inputs=youtube_url,
+        outputs=bw_video
+    )
+
+    bw_video_btn.click(
+        colorize_video,
+        inputs=[bw_video, model_dropdown, fps_dropdown, start_time, end_time],
+        outputs=colorized_video
+    )
+
+app.launch()

colorizer.py

@@ -0,0 +1,157 @@
+import torch
+from torchvision import transforms
+
+import numpy as np
+from skimage.color import rgb2lab, lab2rgb
+import skimage.transform
+from PIL import Image
+
+import os
+from tqdm import tqdm
+from moviepy.editor import VideoFileClip, AudioFileClip
+from moviepy.tools import cvsecs
+import cv2
+
+from pdb import set_trace
+
+
+def lab_to_rgb(L, ab):
+    """
+    Takes a batch of images
+    """
+    L = (L + 1.) * 50.
+    ab = ab * 110.
+    Lab = torch.cat([L, ab], dim=1).permute(0, 2, 3, 1).cpu().numpy()
+    rgb_imgs = []
+    for img in Lab:
+        img_rgb = lab2rgb(img)
+        rgb_imgs.append(img_rgb)
+    return np.stack(rgb_imgs, axis=0)
+
+
+SIZE = 256
+
+
+def get_L(img):
+    img = transforms.Resize(
+        (SIZE, SIZE), transforms.InterpolationMode.BICUBIC)(img)
+    img = np.array(img)
+    img_lab = rgb2lab(img).astype("float32")
+    img_lab = transforms.ToTensor()(img_lab)
+    L = img_lab[[0], ...] / 50. - 1.  # Between -1 and 1
+
+    return L
+
+
+def get_predictions(model, L):
+    # model.L = L.to(model.device)
+    model.eval()
+    with torch.no_grad():
+        model.L = L.to(torch.device('cpu'))
+        model.forward()
+    fake_color = model.fake_color.detach()
+    fake_imgs = lab_to_rgb(L, fake_color)
+
+    return fake_imgs
+
+
+def colorize_img(model, img):
+    L = get_L(img)
+    L = L[None]  # put in list
+    fake_imgs = get_predictions(model, L)
+    fake_img = fake_imgs[0]  # get out of list
+    resized_fake_img = skimage.transform.resize(
+        fake_img, img.size[::-1])  # reshape to original size
+
+    return resized_fake_img
+
+
+def valid_start_end(duration, start_input, end_input):
+    start = start_input
+    end = end_input
+    if start == '':
+        start = 0
+    if end == '':
+        end = duration
+
+    try:
+        start = cvsecs(start)
+        end = cvsecs(end)
+    except BaseException:
+        # start, end aren't actual time values.
+        raise Exception("Invalid start, end values")
+
+    # make it minimal maximum length
+    start = max(start, 0)
+    end = min(duration, end)
+
+    # start must be less than end
+    if start >= end:
+        raise Exception("Start must be before end.")
+
+    return start, end
+
+
+def colorize_vid(path_input, model, fps, start_input, end_input):
+
+    original_video = VideoFileClip(path_input)
+
+    # validate start, end
+    start, end = valid_start_end(
+        original_video.duration, start_input, end_input)
+
+    input_video = original_video.subclip(start, end)
+
+    if isinstance(fps, int):
+        used_fps = fps
+        nframes = np.round(fps * input_video.duration)
+    else:
+        used_fps = input_video.fps
+        nframes = input_video.reader.nframes
+    print(
+        f"Colorizing output with FPS: {fps}, nframes: {nframes}, resolution: {input_video.size}.")
+
+    frames = input_video.iter_frames(fps=used_fps)
+
+    # create tmp path that is same as input path but with '_tmp.[suffix]'
+    base_path, suffix = os.path.splitext(path_input)
+    path_video_tmp = base_path + "_tmp" + suffix
+
+    # create video writer for output
+    size = input_video.size
+    out = cv2.VideoWriter(
+        path_video_tmp,
+        cv2.VideoWriter_fourcc(
+            *'mp4v'),
+        used_fps,
+        size)
+    # out = cv2.VideoWriter(path_video_tmp, cv2.VideoWriter_fourcc(*'DIVX'), used_fps, size)
+
+    for frame in tqdm(frames, total=nframes):
+        # get colorized frame
+        color_frame = colorize_img(model, Image.fromarray(frame))
+
+        if color_frame.max() <= 1:
+            color_frame = (color_frame * 255).astype(np.uint8)
+
+        color_frame = cv2.cvtColor(color_frame, cv2.COLOR_BGR2RGB)
+        out.write(color_frame)
+    out.release()
+
+    # create output path that is same as input path but with '_out.[suffix]'
+    path_output = base_path + "_out" + suffix
+
+    # for some reason, subclip doesn't save audio. so make tmp audio file
+    path_audio_tmp = base_path + "audio_tmp.mp3"
+    input_video.audio.write_audiofile(path_audio_tmp, logger=None)
+    input_audio = AudioFileClip(path_audio_tmp)
+
+    output_video = VideoFileClip(path_video_tmp)
+    output_video = output_video.set_audio(input_audio)
+    output_video.write_videofile(path_output, logger=None)
+
+    os.remove(path_video_tmp)
+    os.remove(path_audio_tmp)
+
+    print("Done.")
+    return path_output

dcgan.py

@@ -0,0 +1,190 @@
+import torch
+from torch import nn, optim
+
+# this architecture is taken from https://github.com/moein-shariatnia/Deep-Learning/tree/main/Image%20Colorization%20Tutorial
+
+#this is actually the DCGans. in training, we had kept the class name the same as the original to avoid changing code^
+class Unet(nn.Module):
+    def __init__(self, input_c=1, output_c=2, num_filters=128):
+        super().__init__()
+        self.model = nn.Sequential(
+            nn.Conv2d(input_c,64,kernel_size=4,stride = 1,padding="same"),
+            nn.BatchNorm2d(64),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(64,128,kernel_size=4,stride=2,padding=1),
+            nn.BatchNorm2d(128),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(128,256,kernel_size=4,stride=2,padding=1),
+            nn.BatchNorm2d(256),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(256,256,kernel_size=4,stride=2,padding=1),
+            nn.BatchNorm2d(256),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(256,512,kernel_size=4,stride=2,padding=1),
+            nn.BatchNorm2d(512),
+            nn.LeakyReLU(0.2, True),
+            nn.Conv2d(512,512,kernel_size=4,stride=2,padding=1),
+            nn.BatchNorm2d(512),
+            nn.LeakyReLU(0.2, True),
+           
+            nn.ConvTranspose2d(512,512,kernel_size=4,stride=2,padding=1),
+            nn.BatchNorm2d(512),
+            nn.ReLU(True),
+            nn.ConvTranspose2d(512,256,kernel_size=4,stride=2,padding=1),
+            nn.BatchNorm2d(256),
+            nn.ReLU(True),
+            nn.ConvTranspose2d(256,256,kernel_size=4,stride=2,padding=1),
+            nn.BatchNorm2d(256),
+            nn.ReLU(True),
+            nn.ConvTranspose2d(256,128,kernel_size=4,stride=2,padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(True),
+            nn.ConvTranspose2d(128,64,kernel_size=4,stride=2,padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(True),
+            nn.Conv2d(64,output_c, kernel_size=1,stride=1),
+            nn.Tanh()
+        )
+    
+    def forward(self, x): 
+      return self.model(x)
+class PatchDiscriminator(nn.Module):
+    def __init__(self, input_c, num_filters=64, n_down=3): # num_filters=64
+        super().__init__()
+        model = [self.get_layers(input_c, num_filters, norm=False)]
+        model += [self.get_layers(num_filters * 2 ** i, num_filters * 2 ** (i + 1), s=1 if i == (n_down-1) else 2) 
+                          for i in range(n_down)] # the 'if' statement is taking care of not using
+                                                  # stride of 2 for the last block in this loop
+        model += [self.get_layers(num_filters * 2 ** n_down, 1, s=1, norm=False, act=False)] # Make sure to not use normalization or
+                                                                                             # activation for the last layer of the model
+        self.model = nn.Sequential(*model)                                                   
+        
+    def get_layers(self, ni, nf, k=4, s=2, p=1, norm=True, act=True): # when needing to make some repeatitive blocks of layers,
+        layers = [nn.Conv2d(ni, nf, k, s, p, bias=not norm)]          # it's always helpful to make a separate method for that purpose
+        if norm: layers += [nn.BatchNorm2d(nf)]
+        if act: layers += [nn.LeakyReLU(0.2, True)] #nn.LeakyReLU(0.2, True)
+        return nn.Sequential(*layers)
+    
+    def forward(self, x):
+        return self.model(x)
+
+class GANLoss(nn.Module):
+    def __init__(self, gan_mode='vanilla', real_label=1.0, fake_label=0.0):
+        super().__init__()
+        self.register_buffer('real_label', torch.tensor(real_label))
+        self.register_buffer('fake_label', torch.tensor(fake_label))
+        if gan_mode == 'vanilla':
+            self.loss = nn.BCEWithLogitsLoss()
+        elif gan_mode == 'lsgan':
+            self.loss = nn.MSELoss()
+    
+    def get_labels(self, preds, target_is_real):
+        if target_is_real:
+            labels = self.real_label
+        else:
+            labels = self.fake_label
+        return labels.expand_as(preds)
+    
+    def __call__(self, preds, target_is_real):
+        labels = self.get_labels(preds, target_is_real)
+        loss = self.loss(preds, labels)
+        return loss
+
+def init_weights(net, init='norm', gain=0.02):
+    
+    def init_func(m):
+        classname = m.__class__.__name__
+        if hasattr(m, 'weight') and 'Conv' in classname:
+            if init == 'norm':
+                nn.init.normal_(m.weight.data, mean=0.0, std=gain)
+            elif init == 'xavier':
+                nn.init.xavier_normal_(m.weight.data, gain=gain)
+            elif init == 'kaiming':
+                nn.init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+            
+            if hasattr(m, 'bias') and m.bias is not None:
+                nn.init.constant_(m.bias.data, 0.0)
+        elif 'BatchNorm2d' in classname:
+            nn.init.normal_(m.weight.data, 1., gain)
+            nn.init.constant_(m.bias.data, 0.)
+            
+    net.apply(init_func)
+    print(f"model initialized with {init} initialization")
+    return net
+
+def init_model(model, device):
+    model = model.to(device)
+    model = init_weights(model)
+    return model
+
+class MainModel(nn.Module):
+    def __init__(self, net_G=None, lr_G=2e-4, lr_D=2e-4, 
+                 beta1=0.5, beta2=0.999, lambda_L1=100.):
+        super().__init__()
+        
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.lambda_L1 = lambda_L1
+        
+        if net_G is None:
+            self.net_G = init_model(Unet(input_c=1, output_c=2, num_filters=64), self.device)
+        else:
+            self.net_G = net_G.to(self.device)
+        self.net_D = init_model(PatchDiscriminator(input_c=3, n_down=3, num_filters=64), self.device)
+        self.GANcriterion = GANLoss(gan_mode='vanilla').to(self.device)
+        self.L1criterion = nn.L1Loss()
+        self.opt_G = optim.Adam(self.net_G.parameters(), lr=lr_G, betas=(beta1, beta2))
+        self.opt_D = optim.Adam(self.net_D.parameters(), lr=lr_D, betas=(beta1, beta2))
+    
+    def set_requires_grad(self, model, requires_grad=True):
+        for p in model.parameters():
+            p.requires_grad = requires_grad
+        
+    def setup_input(self, data):
+        self.L = data['L'].to(self.device)
+        self.ab = data['ab'].to(self.device)
+        
+    def forward(self):
+        self.fake_color = self.net_G(self.L)
+    
+    def backward_D(self,epoch):
+        fake_image = torch.cat([self.L, self.fake_color], dim=1)
+        fake_preds = self.net_D(fake_image.detach())
+        self.loss_D_fake = self.GANcriterion(fake_preds, False)
+        real_image = torch.cat([self.L, self.ab], dim=1)
+        real_preds = self.net_D(real_image)
+        self.loss_D_real = self.GANcriterion(real_preds, True)
+        self.loss_D = (self.loss_D_fake + self.loss_D_real) * 0.5
+        # offset discriminator training
+        if epoch % 2 ==0:
+          self.loss_D.backward()
+    
+    def backward_G(self):
+        fake_image = torch.cat([self.L, self.fake_color], dim=1)
+        fake_preds = self.net_D(fake_image)
+        self.loss_G_GAN = self.GANcriterion(fake_preds, True)
+        self.loss_G_L1 = self.L1criterion(self.fake_color, self.ab) * self.lambda_L1
+        self.loss_G = self.loss_G_GAN + self.loss_G_L1
+        self.loss_G.backward()
+    
+    def optimize(self, epoch):
+        self.forward()
+        self.net_D.train()
+        self.set_requires_grad(self.net_D, True)
+        self.opt_D.zero_grad()
+        self.backward_D(epoch)
+        if epoch % 2 ==0: 
+          self.opt_D.step()
+        
+        self.net_G.train()
+        self.set_requires_grad(self.net_D, False)
+        self.opt_G.zero_grad()
+        self.backward_G()
+        self.opt_G.step()
+
+# with torch.no_grad():
+#     model = MainModel()
+#     set_trace()
+#     # model = torch.load("modelbatchv2.pth", map_location=device)
+#     model.load_state_dict(torch.load("modelbatchv2.pth", map_location=torch.device('cpu')).state_dict())
+#     assert model.device.type == "cpu"
+#     model.eval()

examples/1_falcon.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:376f3ae4689b715739276f491f089445cd552031aacbf274d32665e64a4fb188
+size 970446

examples/2_mughal.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:70ef48d7bfd3a9c0754e320d16845520c5da89a84a3feaa7dde375cea6d2af37
+size 55261144

examples/3_wizard.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a76046df396bb5ced86ac16de3ed56fc227ba08b3ec7e2552ee56c2336021f84
+size 6026610

examples/4_elgar.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:de8fb047d3d251e5af3f5e003601b7f8e23f6ebf327812370656ae304d0b645b
+size 19892928

modelv1.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f52e81eafae487dd13bc0b51193c871e55f7fa7d045baf1b3aeede6ce5e1dbee
+size 221098733

modelv2.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e8a4aa8fbb7e85d01e433bd7633eedb669c9ac51f434df07957b9157d366e3cf
+size 246295842

requirements.txt

@@ -0,0 +1,72 @@
+aiohttp==3.8.3
+aiosignal==1.3.1
+altair==4.2.0
+anyio==3.6.2
+async-timeout==4.0.2
+attrs==22.2.0
+charset-normalizer==2.1.1
+click==8.1.3
+contourpy==1.0.6
+cycler==0.11.0
+decorator==4.4.2
+entrypoints==0.4
+fastapi==0.88.0
+ffmpy==0.3.0
+fonttools==4.38.0
+frozenlist==1.3.3
+fsspec==2022.11.0
+gradio==3.15.0
+h11==0.14.0
+httpcore==0.16.3
+httpx==0.23.1
+idna==3.4
+imageio==2.23.0
+imageio-ffmpeg==0.4.7
+Jinja2==3.1.2
+jsonschema==4.17.3
+kiwisolver==1.4.4
+linkify-it-py==1.0.3
+markdown-it-py==2.1.0
+MarkupSafe==2.1.1
+matplotlib==3.6.2
+mdit-py-plugins==0.3.3
+mdurl==0.1.2
+moviepy==1.0.3
+multidict==6.0.4
+networkx==2.8.8
+numpy==1.24.1
+opencv-python==4.7.0.68
+orjson==3.8.3
+packaging==22.0
+pandas==1.5.2
+Pillow==9.3.0
+proglog==0.1.10
+pycryptodome==3.16.0
+pydantic==1.10.4
+pydub==0.25.1
+pyparsing==3.0.9
+pyrsistent==0.19.3
+python-dateutil==2.8.2
+python-multipart==0.0.5
+pytube==12.1.2
+pytz==2022.7
+PyWavelets==1.4.1
+PyYAML==6.0
+requests==2.28.1
+rfc3986==1.5.0
+scikit-image==0.19.3
+scipy==1.9.3
+six==1.16.0
+sniffio==1.3.0
+starlette==0.22.0
+tifffile==2022.10.10
+toolz==0.12.0
+torch==1.13.1
+torchvision==0.14.1
+tqdm==4.64.1
+typing_extensions==4.4.0
+uc-micro-py==1.0.1
+urllib3==1.26.13
+uvicorn==0.20.0
+websockets==10.4
+yarl==1.8.2