Colorizers library worked well for this competition: https://github.com/richzhang/colorization.git
Below is the code:
import cv2
import glob
from colorizers import *
use_gpu = True
def color(colorizer, img_path):
# default size to process images is 256x256
# grab L channel in both original ("orig") and resized ("rs") resolutions
img = load_img(img_path)
(tens_l_orig, tens_l_rs) = preprocess_img(img, HW=(256,256))
if (use_gpu):
tens_l_rs = tens_l_rs.cuda()
# colorizer outputs 256x256 ab map
# resize and concatenate to original L channel
img_bw = postprocess_tens(tens_l_orig, torch.cat((0*tens_l_orig,0*tens_l_orig),dim=1))
out_img = postprocess_tens(tens_l_orig, colorizer(tens_l_rs).cpu())
return out_img
def main():
# load colorizers
#colorizer_eccv16 = eccv16(pretrained=True).eval()
colorizer_siggraph17 = siggraph17(pretrained=True).eval()
if use_gpu:
#colorizer_eccv16.cuda()
colorizer_siggraph17.cuda()
colorizer = colorizer_siggraph17
stage = 'train'
data_dir = '{INSERT_YOUR_DATA_DIR}/' + stage + '_black_white_images/' + stage + '_back_white_images/'
out_dir = '../' + stage + '_color_images/'
fnames = glob.glob(data_dir + '/*')
print(fnames[:10])
for cnt, fname in enumerate(fnames):
imagename = fname.split('/')[-1]
outname = out_dir + imagename
res = color(colorizer, fname)
res = res[:,:,[2,1,0]] # reorder channels
cv2.imwrite(outname, np.clip(res * 256, 0, 255).astype(int))
if cnt % 100 == 0:
print(cnt)
main()