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()