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