Removed binary

.gitattributes

@@ -0,0 +1,27 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,13 @@
+---
+title: Colorizer
+emoji: 🏃
+colorFrom: indigo
+colorTo: yellow
+sdk: gradio
+sdk_version: 2.9.4
+app_file: app.py
+pinned: false
+license: mit
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces#reference

app.py

@@ -0,0 +1,74 @@
+import gradio as gr
+
+import torch
+
+import skimage
+import skimage.io
+from skimage.transform import rescale, resize
+from skimage import io, color
+
+import cv2
+from colorizer import normalize_lab_channels, torch_normalized_lab_to_rgb
+
+model = torch.load('colorizer.pth')
+model = model.eval()
+
+def colorize_img( img_float32, model, res=512, border=0.2, apply_blur=True):
+    img = img_float32
+    
+    # Add white border
+    border = int(img.shape[0] * border)
+    img2 = cv2.copyMakeBorder(img, border, border, border, border, cv2.BORDER_CONSTANT, value=(1.,1.,1.))
+    
+    # Resize to expected resolution
+    img_resized = resize(img2, (res,res), anti_aliasing=True) 
+    
+    # Blur image a bit
+    if apply_blur: 
+        img2 = skimage.filters.gaussian( img_resized, sigma=1, channel_axis=-1 )   
+    else:
+        img2 = img_resized
+        
+    # Convert to Lab color space and normalize between 0 and 1 all channels
+    img2 = normalize_lab_channels(color.rgb2lab(img2)) 
+    
+    # Keep copy of unblured + resized image
+    img_resized = normalize_lab_channels(color.rgb2lab(img_resized)) 
+    img_resized = torch.from_numpy(img_resized)
+    img_resized = img_resized.permute(2,0,1).unsqueeze(dim=0)
+    
+    # Convert to expected tensor of 'LLL'
+    x = torch.from_numpy(img2)
+    x = x.permute(2,0,1).unsqueeze(dim=0)
+    x[:,1,:,:] = x[:,0,:,:]
+    x[:,2,:,:] = x[:,0,:,:]
+    
+    x_hat_ab = model( x )
+    
+    x_hat = img_resized.clone()
+    x_hat[:,1:,:,:] = x_hat_ab.clone()
+    
+    colored_img = torch_normalized_lab_to_rgb( x_hat )
+    
+    return colored_img.detach().cpu().squeeze().permute(1,2,0).numpy()
+
+def process_image(img):
+    return colorize_img( (img / 255).astype('float32'), model  )
+
+image = gr.inputs.Image()
+label = gr.outputs.Label()
+title = "Colorizer"
+description = "A model that colorizes b&w images."
+interpretation='default'
+enable_queue=True
+
+examples = ['ka0001.jpg', 'ka0003.jpg', 'ka0009.jpg', 'ka0010.jpg']
+css = ".h-60 {min-height: 512px !important;}"
+
+gr.Interface(fn=process_image,
+             inputs=gr.inputs.Image(),
+             outputs=gr.outputs.Image(),
+             title=title,
+             description=description,
+             css=css,
+             examples=examples).launch()

colorizer.py

@@ -0,0 +1,133 @@
+import os
+import sys
+import time
+import glob
+import random
+import skimage
+import skimage.io
+import numpy as np
+from skimage import io, color
+
+import skimage
+import skimage.io
+from PIL import Image
+import cv2
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+import timm
+import torchvision
+from torchvision.models.feature_extraction import create_feature_extractor
+
+L_range = 100
+ab_min = -128
+ab_max = 127
+ab_range = ab_max - ab_min   
+        
+def extract_zip(input_zip):
+    input_zip=ZipFile(input_zip)
+    return {name: input_zip.read(name) for name in input_zip.namelist()}
+
+def normalize_lab_channels(x):
+    # Normalize L
+    x[:,:,0] = x[:,:,0] / L_range
+
+    # Normalize AB     
+    x[:,:,1] = (x[:,:,1]-ab_min) / ab_range        
+    x[:,:,2] = (x[:,:,2]-ab_min) / ab_range
+    return x
+
+def normalized_lab_to_rgb(lab):
+    lab[:,:,0] = (lab[:,:,0] * L_range)
+    lab[:,:,1] = (lab[:,:,1] * ab_range) + ab_min 
+    lab[:,:,2] = (lab[:,:,2] * ab_range) + ab_min 
+    return color.lab2rgb(lab)
+
+def torch_normalized_lab_to_rgb(lab):
+    for i in range(lab.shape[0]):
+        lab[i,0,:,:] = torch.clip(lab[i,0,:,:] * L_range, 0, L_range)
+        lab[i,1,:,:] = torch.clip((lab[i,1,:,:] * ab_range) + ab_min, ab_min, ab_max)
+        lab[i,2,:,:] = torch.clip((lab[i,2,:,:] * ab_range) + ab_min, ab_min, ab_max)
+        
+    for i in range(lab.shape[0]):
+        lab[i] = torch.from_numpy( color.lab2rgb(lab[i].permute(1,2,0).detach().cpu().numpy()) ).permute(2,0,1)
+        
+    return lab
+
+class Encoder(nn.Module):
+    def __init__(self):
+        super(Encoder, self).__init__()       
+        self.backend_model = timm.create_model('efficientnetv2_rw_s', pretrained=True)        
+        self.backend = create_feature_extractor(self.backend_model, 
+            return_nodes=['blocks.0', 'blocks.1', 'blocks.2', 'blocks.3', 'act2']) 
+
+    def forward(self, x):
+        features = self.backend(x)
+        return list(features.values())
+
+class UpSample(nn.Sequential):
+    def __init__(self, in_channels, out_channels):
+        skip_input, output_features = in_channels, out_channels
+        
+        super(UpSample, self).__init__()        
+        self.convA = nn.Conv2d(skip_input, output_features, kernel_size=3, stride=1, padding=1, padding_mode='reflect', bias=False)
+        self.leakyreluA = nn.LeakyReLU(0.2)
+        self.convB = nn.Conv2d(output_features, output_features, kernel_size=3, stride=1, padding=1, padding_mode='reflect', bias=False)
+        self.leakyreluB = nn.LeakyReLU(0.2)
+        
+        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False)
+
+    def forward(self, x, concat_with=None):
+        up_x = self.upsample(x)
+        
+        if concat_with is not None:
+            up_x = torch.cat([up_x, concat_with], dim=1)
+        
+        return self.leakyreluB( self.convB( self.leakyreluA( self.convA( up_x ) )  )  )
+    
+class Decoder(nn.Module):
+    def __init__(self, num_features=1792 * 1, decoder_width=None):
+        super(Decoder, self).__init__()
+        features = int(num_features * decoder_width)
+        
+        self.conv2 =  nn.Sequential(
+            nn.Conv2d(num_features, features, kernel_size=1, stride=1, padding=0, bias=False),
+            nn.LeakyReLU(0.2),
+        )
+
+        self.up1 = UpSample(in_channels=features//1 + 152 - 24, out_channels=features//2)
+        self.up2 = UpSample(in_channels=features//2 + 80 - 16, out_channels=features//4)
+        self.up3 = UpSample(in_channels=features//4 + 56 - 8, out_channels=features//8)
+        self.up4 = UpSample(in_channels=features//8 + 32 - 8, out_channels=features//16)    
+        
+        self.up5 = UpSample(in_channels=features//16, out_channels=features//16) 
+
+        self.conv3 = nn.Conv2d(features//16, 2, kernel_size=1, stride=1, padding=0, bias=False)
+        
+    def forward(self, features):
+        blocks0, blocks1, blocks2, blocks3, x = features
+        
+        x = self.conv2(x)
+        
+        x = self.up1(x, blocks3)
+        x = self.up2(x, blocks2)
+        x = self.up3(x, blocks1)
+        x = self.up4(x, blocks0)
+        
+        x = self.up5(x)
+        
+        x_final = self.conv3(x)
+        
+        return x_final
+    
+class ColorizeNet(nn.Module):
+    def __init__(self, decoder_width):
+        super(ColorizeNet, self).__init__()
+        self.encoder = Encoder()
+        self.decoder = Decoder(decoder_width=decoder_width)
+
+    def forward(self, x):        
+        features_x = self.encoder(x)
+        return self.decoder( features_x )