Spaces:
Runtime error
Runtime error
| import matplotlib.pyplot as plt | |
| import streamlit as st | |
| # -------------------- base color ------------------ | |
| 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 | |
| # ------------------ eccv16 --------------------- | |
| import numpy as np | |
| 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 | |
| # ------------------ siggraph17 --------------------- | |
| 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 | |
| # ------------------ utils --------------------- | |
| from PIL import Image | |
| import numpy as np | |
| from skimage import color | |
| import torch.nn.functional as F | |
| 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))) | |
| # parser = argparse.ArgumentParser() | |
| # parser.add_argument('-i','--img_path', type=str, default='imgs/test.jpg') | |
| # # parser.add_argument('--use_gpu', action='store_true', help='whether to use GPU') | |
| # parser.add_argument('-o','--save_prefix', type=str, default='saved', help='will save into this file with {eccv16.png, siggraph17.png} suffixes') | |
| # opt = parser.parse_args() | |
| colorizer_eccv16 = eccv16(pretrained=True).eval() | |
| colorizer_siggraph17 = siggraph17(pretrained=True).eval() | |
| # if(opt.use_gpu): | |
| # colorizer_eccv16.cuda() | |
| # colorizer_siggraph17.cuda() | |
| st.title('Colorizes GrayScale Images ! ') | |
| st.write('Used ecvv16 and siggraph') | |
| input_image = st.file_uploader("Upload Image : ", type=["jpg", "jpeg", "png"]) | |
| if input_image is not None: | |
| img = load_img(input_image) | |
| (tens_l_orig, tens_l_rs) = preprocess_img(img, HW=(256,256)) | |
| img_bw = postprocess_tens(tens_l_orig, torch.cat((0*tens_l_orig,0*tens_l_orig),dim=1)) | |
| out_img_eccv16 = postprocess_tens(tens_l_orig, colorizer_eccv16(tens_l_rs).cpu()) | |
| out_img_siggraph17 = postprocess_tens(tens_l_orig, colorizer_siggraph17(tens_l_rs).cpu()) | |
| plt.imsave(f'eccv16.png{input_image.name}', out_img_eccv16) | |
| plt.imsave(f'siggraph17.png{input_image.name}', out_img_siggraph17) | |
| eccv16_path = f'eccv16_{input_image.name}' | |
| siggraph17_path = f'siggraph17_{input_image.name}' | |
| plt.imsave(eccv16_path, out_img_eccv16) | |
| plt.imsave(siggraph17_path, out_img_siggraph17) | |
| # Display images using Streamlit | |
| st.image([img, img_bw, out_img_eccv16, out_img_siggraph17], caption=['Original', 'Input', 'Output (ECCV 16)', 'Output (SIGGRAPH 17)'], | |
| width=256) | |
| # Optionally, you can also display the saved images | |
| st.markdown("### Saved Images:") | |
| st.image([eccv16_path, siggraph17_path], width=256) | |
| st.write('Build with <3 by Spyro using PyTorch ') | |
| # plt.figure(figsize=(12,8)) | |
| # plt.subplot(2,2,1) | |
| # plt.imshow(img) | |
| # plt.title('Original') | |
| # plt.axis('off') | |
| # plt.subplot(2,2,2) | |
| # plt.imshow(img_bw) | |
| # plt.title('Input') | |
| # plt.axis('off') | |
| # plt.subplot(2,2,3) | |
| # plt.imshow(out_img_eccv16) | |
| # plt.title('Output (ECCV 16)') | |
| # plt.axis('off') | |
| # plt.subplot(2,2,4) | |
| # plt.imshow(out_img_siggraph17) | |
| # plt.title('Output (SIGGRAPH 17)') | |
| # plt.axis('off') | |
| # plt.show() | |