Spaces:
Runtime error
Runtime error
| # -*- 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() | |
| 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 |