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.
123 lines
3.5 KiB
Python
123 lines
3.5 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+'pytorch-SRResNet'])
|
|
|
|
|
|
from argparse import Namespace
|
|
import torch
|
|
from torch.autograd import Variable
|
|
import numpy as np
|
|
from PIL import Image
|
|
import cv2
|
|
|
|
|
|
def getlabelmat(mask,idx):
|
|
x=np.zeros((mask.shape[0],mask.shape[1],3))
|
|
x[mask==idx,0]=colors[idx][0]
|
|
x[mask==idx,1]=colors[idx][1]
|
|
x[mask==idx,2]=colors[idx][2]
|
|
return x
|
|
|
|
def colorMask(mask):
|
|
x=np.zeros((mask.shape[0],mask.shape[1],3))
|
|
for idx in range(19):
|
|
x=x+getlabelmat(mask,idx)
|
|
return np.uint8(x)
|
|
|
|
|
|
def getnewimg(input_image,s,cFlag):
|
|
opt=Namespace(cuda=torch.cuda.is_available() and not cFlag,
|
|
model=baseLoc+'weights/super_resolution/model_srresnet.pth',
|
|
dataset='Set5',scale=s,gpus=0)
|
|
|
|
im_l=Image.fromarray(input_image)
|
|
cuda = opt.cuda
|
|
|
|
if cuda:
|
|
model = torch.load(opt.model)["model"]
|
|
else:
|
|
model = torch.load(opt.model,map_location=torch.device('cpu'))["model"]
|
|
|
|
im_l=np.array(im_l)
|
|
im_l = im_l.astype(float)
|
|
im_input = im_l.astype(np.float32).transpose(2,0,1)
|
|
im_input = im_input.reshape(1,im_input.shape[0],im_input.shape[1],im_input.shape[2])
|
|
im_input = Variable(torch.from_numpy(im_input/255.).float())
|
|
|
|
if cuda:
|
|
model = model.cuda()
|
|
im_input = im_input.cuda()
|
|
else:
|
|
model = model.cpu()
|
|
|
|
HR_4x = model(im_input)
|
|
|
|
HR_4x = HR_4x.cpu()
|
|
|
|
im_h = HR_4x.data[0].numpy().astype(np.float32)
|
|
|
|
im_h = im_h*255.
|
|
im_h = np.clip(im_h, 0., 255.)
|
|
im_h = im_h.transpose(1,2,0).astype(np.uint8)
|
|
|
|
return im_h
|
|
|
|
|
|
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(layer.height,layer.width,bpp)
|
|
|
|
|
|
def createResultLayer(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 super_resolution(img, layer,scale,cFlag) :
|
|
if torch.cuda.is_available() and not cFlag:
|
|
gimp.progress_init("(Using GPU) Running super-resolution for " + layer.name + "...")
|
|
else:
|
|
gimp.progress_init("(Using CPU) Running super-resolution for " + layer.name + "...")
|
|
|
|
imgmat = channelData(layer)
|
|
if imgmat.shape[2] == 4: # get rid of alpha channel
|
|
imgmat = imgmat[:,:,0:3]
|
|
cpy = getnewimg(imgmat,scale,cFlag)
|
|
cpy = cv2.resize(cpy, (0,0), fx=scale/4, fy=scale/4)
|
|
createResultLayer(layer.name+'_upscaled',cpy)
|
|
|
|
|
|
register(
|
|
"super-resolution",
|
|
"super-resolution",
|
|
"Running super-resolution.",
|
|
"Kritik Soman",
|
|
"Your",
|
|
"2020",
|
|
"super-resolution...",
|
|
"*", # Alternately use RGB, RGB*, GRAY*, INDEXED etc.
|
|
[ (PF_IMAGE, "image", "Input image", None),
|
|
(PF_DRAWABLE, "drawable", "Input drawable", None),
|
|
(PF_SLIDER, "Scale", "Scale", 4, (1.1, 4, 0.5)),
|
|
(PF_BOOL, "fcpu", "Force CPU", False)
|
|
],
|
|
[],
|
|
super_resolution, menu="<Image>/Layer/GIML-ML")
|
|
|
|
main()
|