Upload app.py

app.py

@@ -0,0 +1,128 @@
+# -*- coding: utf-8 -*-
+"""DDColor_colab.ipynb
+
+Automatically generated by Colaboratory.
+
+Original file is located at
+    https://colab.research.google.com/github/camenduru/DDColor-colab/blob/main/DDColor_colab.ipynb
+"""
+
+# Commented out IPython magic to ensure Python compatibility.
+# %cd /content
+!git clone -b dev https://github.com/camenduru/DDColor
+
+!apt -y install -qq aria2
+!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/cv_ddcolor_image-colorization/resolve/main/pytorch_model.pt -d /content/DDColor/models -o pytorch_model.pt
+
+!wget https://modelscope.oss-cn-beijing.aliyuncs.com/test/images/audrey_hepburn.jpg -O /content/DDColor/in.jpg
+!pip install -q timm
+
+# %cd /content/DDColor
+
+!sed -i 's/from \.version import __gitsha__, __version__/# from \.version import __gitsha__, __version__/' /content/DDColor/basicsr/__init__.py
+
+import argparse
+import cv2
+import numpy as np
+import os
+from tqdm import tqdm
+import torch
+from basicsr.archs.ddcolor_arch import DDColor
+import torch.nn.functional as F
+
+class ImageColorizationPipeline(object):
+
+    def __init__(self, model_path, input_size=256, model_size='large'):
+
+        self.input_size = input_size
+        if torch.cuda.is_available():
+            self.device = torch.device('cuda')
+        else:
+            self.device = torch.device('cpu')
+
+        if model_size == 'tiny':
+            self.encoder_name = 'convnext-t'
+        else:
+            self.encoder_name = 'convnext-l'
+
+        self.decoder_type = "MultiScaleColorDecoder"
+
+        if self.decoder_type == 'MultiScaleColorDecoder':
+            self.model = DDColor(
+                encoder_name=self.encoder_name,
+                decoder_name='MultiScaleColorDecoder',
+                input_size=[self.input_size, self.input_size],
+                num_output_channels=2,
+                last_norm='Spectral',
+                do_normalize=False,
+                num_queries=100,
+                num_scales=3,
+                dec_layers=9,
+            ).to(self.device)
+        else:
+            self.model = DDColor(
+                encoder_name=self.encoder_name,
+                decoder_name='SingleColorDecoder',
+                input_size=[self.input_size, self.input_size],
+                num_output_channels=2,
+                last_norm='Spectral',
+                do_normalize=False,
+                num_queries=256,
+            ).to(self.device)
+
+        self.model.load_state_dict(
+            torch.load(model_path, map_location=torch.device('cpu'))['params'],
+            strict=False)
+        self.model.eval()
+
+    @torch.no_grad()
+    def process(self, img):
+        self.height, self.width = img.shape[:2]
+        # print(self.width, self.height)
+        # if self.width * self.height < 100000:
+        #     self.input_size = 256
+
+        img = (img / 255.0).astype(np.float32)
+        orig_l = cv2.cvtColor(img, cv2.COLOR_BGR2Lab)[:, :, :1]  # (h, w, 1)
+
+        # resize rgb image -> lab -> get grey -> rgb
+        img = cv2.resize(img, (self.input_size, self.input_size))
+        img_l = cv2.cvtColor(img, cv2.COLOR_BGR2Lab)[:, :, :1]
+        img_gray_lab = np.concatenate((img_l, np.zeros_like(img_l), np.zeros_like(img_l)), axis=-1)
+        img_gray_rgb = cv2.cvtColor(img_gray_lab, cv2.COLOR_LAB2RGB)
+
+        tensor_gray_rgb = torch.from_numpy(img_gray_rgb.transpose((2, 0, 1))).float().unsqueeze(0).to(self.device)
+        output_ab = self.model(tensor_gray_rgb).cpu()  # (1, 2, self.height, self.width)
+
+        # resize ab -> concat original l -> rgb
+        output_ab_resize = F.interpolate(output_ab, size=(self.height, self.width))[0].float().numpy().transpose(1, 2, 0)
+        output_lab = np.concatenate((orig_l, output_ab_resize), axis=-1)
+        output_bgr = cv2.cvtColor(output_lab, cv2.COLOR_LAB2BGR)
+
+        output_img = (output_bgr * 255.0).round().astype(np.uint8)
+
+        return output_img
+
+colorizer = ImageColorizationPipeline(model_path='/content/DDColor/models/pytorch_model.pt', input_size=512)
+
+# helper function taken from: https://huggingface.co/blog/stable_diffusion
+from PIL import Image
+def image_grid(imgs, rows, cols):
+    assert len(imgs) == rows*cols
+
+    w, h = imgs[0].size
+    grid = Image.new('RGB', size=(cols*w, rows*h))
+    grid_w, grid_h = grid.size
+
+    for i, img in enumerate(imgs):
+        grid.paste(img, box=(i%cols*w, i//cols*h))
+    return grid
+
+image_in = cv2.imread('/content/DDColor/in.jpg')
+image_out = colorizer.process(image_in)
+cv2.imwrite('/content/DDColor/out.jpg', image_out)
+image_in_pil = Image.fromarray(cv2.cvtColor(image_in, cv2.COLOR_BGR2RGB))
+image_out_pil = Image.fromarray(cv2.cvtColor(image_out, cv2.COLOR_BGR2RGB))
+images = [image_in_pil, image_out_pil]
+grid = image_grid(images, rows=1, cols=2)
+grid