mirror of https://github.com/kritiksoman/GIMP-ML
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
117 lines
3.4 KiB
Python
117 lines
3.4 KiB
Python
import os
|
|
baseLoc = os.path.dirname(os.path.realpath(__file__))+'/'
|
|
|
|
|
|
from gimpfu import *
|
|
import sys
|
|
|
|
sys.path.extend([baseLoc+'gimpenv/lib/python2.7',baseLoc+'gimpenv/lib/python2.7/site-packages',baseLoc+'gimpenv/lib/python2.7/site-packages/setuptools',baseLoc+'neural-colorization'])
|
|
|
|
|
|
import torch
|
|
from model import generator
|
|
from torch.autograd import Variable
|
|
from scipy.ndimage import zoom
|
|
from PIL import Image
|
|
from argparse import Namespace
|
|
import numpy as np
|
|
# from skimage.color import rgb2yuv,yuv2rgb
|
|
import cv2
|
|
|
|
def getcolor(input_image):
|
|
p = np.repeat(input_image, 3, axis=2)
|
|
|
|
if torch.cuda.is_available():
|
|
g_available=1
|
|
else:
|
|
g_available=-1
|
|
|
|
args=Namespace(model=baseLoc+'neural-colorization/model.pth',gpu=g_available)
|
|
|
|
G = generator()
|
|
|
|
if torch.cuda.is_available():
|
|
G=G.cuda()
|
|
G.load_state_dict(torch.load(args.model))
|
|
else:
|
|
G.load_state_dict(torch.load(args.model,map_location=torch.device('cpu')))
|
|
|
|
p = p.astype(np.float32)
|
|
p = p / 255
|
|
img_yuv = cv2.cvtColor(p, cv2.COLOR_RGB2YUV)
|
|
# img_yuv = rgb2yuv(p)
|
|
H,W,_ = img_yuv.shape
|
|
infimg = np.expand_dims(np.expand_dims(img_yuv[...,0], axis=0), axis=0)
|
|
img_variable = Variable(torch.Tensor(infimg-0.5))
|
|
if args.gpu>=0:
|
|
img_variable=img_variable.cuda(args.gpu)
|
|
res = G(img_variable)
|
|
uv=res.cpu().detach().numpy()
|
|
uv[:,0,:,:] *= 0.436
|
|
uv[:,1,:,:] *= 0.615
|
|
(_,_,H1,W1) = uv.shape
|
|
uv = zoom(uv,(1,1,float(H)/H1,float(W)/W1))
|
|
yuv = np.concatenate([infimg,uv],axis=1)[0]
|
|
# rgb=yuv2rgb(yuv.transpose(1,2,0))
|
|
# out=(rgb.clip(min=0,max=1)*255)[:,:,[0,1,2]]
|
|
rgb = cv2.cvtColor(yuv.transpose(1, 2, 0)*255, cv2.COLOR_YUV2RGB)
|
|
rgb = rgb.clip(min=0,max=255)
|
|
out = rgb.astype(np.uint8)
|
|
|
|
return out
|
|
|
|
def channelData(layer):#convert gimp image to numpy
|
|
region=layer.get_pixel_rgn(0, 0, layer.width,layer.height)
|
|
pixChars=region[:,:] # Take whole layer
|
|
bpp=region.bpp
|
|
# return np.frombuffer(pixChars,dtype=np.uint8).reshape(len(pixChars)/bpp,bpp)
|
|
return np.frombuffer(pixChars,dtype=np.uint8).reshape(layer.height,layer.width,bpp)
|
|
|
|
def createResultLayer(image,name,result):
|
|
rlBytes=np.uint8(result).tobytes();
|
|
rl=gimp.Layer(image,name,image.width,image.height,image.active_layer.type,100,NORMAL_MODE)
|
|
region=rl.get_pixel_rgn(0, 0, rl.width,rl.height,True)
|
|
region[:,:]=rlBytes
|
|
image.add_layer(rl,0)
|
|
gimp.displays_flush()
|
|
|
|
def genNewImg(name,layer_np):
|
|
h,w,d=layer_np.shape
|
|
img=pdb.gimp_image_new(w, h, RGB)
|
|
display=pdb.gimp_display_new(img)
|
|
|
|
rlBytes=np.uint8(layer_np).tobytes();
|
|
rl=gimp.Layer(img,name,img.width,img.height,RGB,100,NORMAL_MODE)
|
|
region=rl.get_pixel_rgn(0, 0, rl.width,rl.height,True)
|
|
region[:,:]=rlBytes
|
|
|
|
pdb.gimp_image_insert_layer(img, rl, None, 0)
|
|
gimp.displays_flush()
|
|
|
|
def colorize(img, layer) :
|
|
gimp.progress_init("Coloring " + layer.name + "...")
|
|
|
|
imgmat = channelData(layer)
|
|
cpy=getcolor(imgmat)
|
|
|
|
genNewImg(layer.name+'_colored',cpy)
|
|
|
|
|
|
|
|
register(
|
|
"colorize",
|
|
"colorize",
|
|
"Generate monocular disparity map based on deep learning.",
|
|
"Kritik Soman",
|
|
"Your",
|
|
"2020",
|
|
"colorize...",
|
|
"*", # Alternately use RGB, RGB*, GRAY*, INDEXED etc.
|
|
[ (PF_IMAGE, "image", "Input image", None),
|
|
(PF_DRAWABLE, "drawable", "Input drawable", None),
|
|
],
|
|
[],
|
|
colorize, menu="<Image>/Layer/GIML-ML")
|
|
|
|
main()
|