import numpy as np
from data import colorize_image as CI
from ideepcolor import parse_args
from skimage import color
import os
import datetime
args = parse_args()
colorModel = CI.ColorizeImageCaffe()
colorModel.prep_net(args.gpu, args.color_prototxt, args.color_caffemodel)
model = colorModel
distModel = CI.ColorizeImageCaffeDist()
distModel.prep_net(args.gpu, args.dist_prototxt, args.dist_caffemodel)
import cv2
# distModel.set_image(im_rgb)
def get_input(h, w):
im = np.zeros((h, w, 3), np.uint8)
mask = np.zeros((h, w, 1), np.uint8)
vis_im = np.zeros((h, w, 3), np.uint8)
userEdits = []
for ue in userEdits:
ue.updateInput(im, mask, vis_im)
im_bgr = cv2.cvtColor(im, cv2.COLOR_RGB2BGR)
return im, mask
from skimage import color
def predict_color():
im, mask = get_input()
im_mask0 = mask > 0.0
im_mask0 = im_mask0.transpose((2, 0, 1))
if __name__ == '__main__':
args = parse_args()
for arg in vars(args):
print('[%s] =' % arg, getattr(args, arg))
if args.cpu_mode:
args.gpu = -1
args.win_size = int(args.win_size / 4.0) * 4 # make sure the width of the image can be divided by 4
if args.backend == 'caffe':
# initialize the colorization model
colorModel = CI.ColorizeImageCaffe(Xd=args.load_size)
colorModel.prep_net(args.gpu, args.color_prototxt, args.color_caffemodel)
distModel = CI.ColorizeImageCaffeDist(Xd=args.load_size)
distModel.prep_net(args.gpu, args.dist_prototxt, args.dist_caffemodel)
elif args.backend == 'pytorch':
colorModel = CI.ColorizeImageTorch(Xd=args.load_size,maskcent=args.pytorch_maskcent)
colorModel.prep_net(path=args.color_model)
distModel = CI.ColorizeImageTorchDist(Xd=args.load_size,maskcent=args.pytorch_maskcent)
distModel.prep_net(path=args.color_model, dist=True)
else:
print('backend type [%s] not found!' % args.backend)
# initialize application
app = QApplication(sys.argv)
# mask 1x256x256 binary mask of current user input (will be updated)
# loc 2 tuple (h,w) of where to put the user input
# p scalar half-patch size
# val 2 tuple (a,b) value of user input
input_ab[:, loc[0] - p:loc[0] + p + 1,
loc[1] - p:loc[1] + p + 1] = np.array(val)[:, np.newaxis,
np.newaxis]
mask[:, loc[0] - p:loc[0] + p + 1, loc[1] - p:loc[1] + p + 1] = 1
return (input_ab, mask)
# Choose gpu to run the model on
gpu_id = -1
# Initialize colorization class
colorModel = CI.ColorizeImageCaffe(Xd=256)
# Load the model
colorModel.prep_net(gpu_id, './models/reference_model/deploy_nodist.prototxt',
'./models/reference_model/model.caffemodel')
# Load the black and white image
colorModel.load_image(sys.argv[1])
# Load the colored reference image
referenceImage = getResizedImage(sys.argv[2])
# extract ab channels of the resized image
referenceImage_lab = color.rgb2lab(referenceImage)
# initialize with no user inputs