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')
