def deepcolor(tmp1, tmp2, ilayerimg, ilayerc): layerimg = channelData(ilayerimg) layerc = channelData(ilayerc) if ilayerimg.name == ilayerc.name: # if local color hints are not provided by user mask = np.zeros( (1, 256, 256)) # giving no user points, so mask is all 0's input_ab = np.zeros( (2, 256, 256)) # ab values of user points, default to 0 for no input else: if layerc.shape[2] == 3: # error pdb.gimp_message("Alpha channel missing in " + ilayerc.name + " !") return else: input_ab = cv2.cvtColor(layerc[:, :, 0:3].astype(np.float32) / 255, cv2.COLOR_RGB2LAB) mask = layerc[:, :, 3] > 0 mask = mask.astype(np.uint8) input_ab = cv2.resize(input_ab, (256, 256)) mask = cv2.resize(mask, (256, 256)) mask = mask[np.newaxis, :, :] input_ab = input_ab[:, :, 1:3].transpose((2, 0, 1)) if layerimg.shape[2] == 4: #remove alpha channel in image if present layerimg = layerimg[:, :, 0:3] if torch.cuda.is_available(): gimp.progress_init("(Using GPU) Running deepcolor for " + ilayerimg.name + "...") gpu_id = 0 else: gimp.progress_init("(Using CPU) Running deepcolor for " + ilayerimg.name + "...") gpu_id = None colorModel = CI.ColorizeImageTorch(Xd=256) colorModel.prep_net(gpu_id, baseLoc + 'ideepcolor/models/pytorch/caffemodel.pth') colorModel.load_image(layerimg) # load an image img_out = colorModel.net_forward(input_ab, mask) # run model, returns 256x256 image img_out_fullres = colorModel.get_img_fullres( ) # get image at full resolution createResultLayer(tmp1, 'new_' + ilayerimg.name, img_out_fullres)
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) window = gui_design.GUIDesign(color_model=colorModel, dist_model=distModel, img_file=args.image_file, load_size=args.load_size, win_size=args.win_size) app.setStyleSheet(qdarkstyle.load_stylesheet(pyside=False)) # comment this if you do not like dark stylesheet app.setWindowIcon(QIcon('imgs/logo.png')) # load logo window.setWindowTitle('iColor') window.setWindowFlags(window.windowFlags() & ~Qt.WindowMaximizeButtonHint) # fix window siz
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) colorModel.prep_net(path=args.color_model) distModel = CI.ColorizeImageTorchDist(Xd=args.load_size) distModel.prep_net(path=args.color_model, dist=True) else: print('backend type [%s] not found!' % args.backend) # initialize application app = QApplication(sys.argv) window = gui_design.GUIDesign(color_model=colorModel, dist_model=distModel, img_file=args.image_file, load_size=args.load_size, win_size=args.win_size) app.setStyleSheet(qdarkstyle.load_stylesheet(
import os from app import app import urllib.request from flask import Flask, flash, request, redirect, url_for, render_template from werkzeug.utils import secure_filename from data import colorize_image as CI import matplotlib.pyplot as plt import numpy as np from PIL import Image # Initialize colorization class colorModel = CI.ColorizeImageTorch(Xd=256) # Load the model colorModel.prep_net(None, 'models/pytorch/caffemodel.pth') mask = np.zeros((1, 256, 256)) # giving no user points, so mask is all 0's input_ab = np.zeros( (2, 256, 256)) # ab values of user points, default to 0 for no input ALLOWED_EXTENSIONS = set(['png', 'jpg', 'jpeg', 'gif']) app.config['SEND_FILE_MAX_AGE_DEFAULT'] = 0 def allowed_file(filename): return '.' in filename and filename.rsplit( '.', 1)[1].lower() in ALLOWED_EXTENSIONS @app.route('/') def upload_form(): return render_template('upload.html')