Spaces:
Runtime error
Runtime error
File size: 4,760 Bytes
fe407fb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 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 |