use super-resolution as a filter

pull/30/head
Kritik Soman 4 years ago
parent 43a936ae11
commit 681f30b122

@ -1,10 +1,12 @@
import os
baseLoc = os.path.dirname(os.path.realpath(__file__))+'/'
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'])
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
@ -14,81 +16,107 @@ 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 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))
x = np.zeros((mask.shape[0], mask.shape[1], 3))
for idx in range(19):
x=x+getlabelmat(mask,idx)
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)
def getnewimg(input_image, s, cFlag, fFlag):
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)
w, h = input_image.shape[0:2]
cuda = opt.cuda
if cuda:
model = torch.load(opt.model)["model"]
else:
model = torch.load(opt.model,map_location=torch.device('cpu'))["model"]
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())
im_input = input_image.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:
if cuda and not cFlag:
model = model.cuda()
im_input = im_input.cuda()
else:
model = model.cpu()
HR_4x = model(im_input)
if fFlag:
im_h = np.zeros([4 * w, 4 * h, 3])
wbin = 300
i = 0
idx = 0
t = float(w * h) / float(wbin * wbin)
while i < w:
i_end = min(i + wbin, w)
j = 0
while j < h:
j_end = min(j + wbin, h)
patch = im_input[:, :, i:i_end, j:j_end]
# patch_merge_out_numpy = denoiser(patch, c, pss, model, model_est, opt, cFlag)
HR_4x = model(patch)
HR_4x = HR_4x.cpu().data[0].numpy().astype(np.float32) * 255.
HR_4x = np.clip(HR_4x, 0., 255.).transpose(1, 2, 0).astype(np.uint8)
im_h[4 * i:4 * i_end, 4 * j:4 * j_end, :] = HR_4x
j = j_end
idx = idx + 1
gimp.progress_update(float(idx) / float(t))
gimp.displays_flush()
i = i_end
else:
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
HR_4x = HR_4x.cpu()
im_h = HR_4x.data[0].numpy().astype(np.float32)
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)
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 createResultFile(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
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)
pdb.gimp_image_insert_layer(img, rl, None, 0)
gimp.displays_flush()
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)
def createResultLayer(image, name, result):
rlBytes = np.uint8(result).tobytes();
rl = gimp.Layer(image, name, image.width, image.height, 0, 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 super_resolution(img, layer,scale,cFlag) :
def super_resolution(img, layer, scale, cFlag, fFlag):
if torch.cuda.is_available() and not cFlag:
gimp.progress_init("(Using GPU) Running super-resolution for " + layer.name + "...")
else:
@ -96,11 +124,14 @@ def super_resolution(img, layer,scale,cFlag) :
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)
imgmat = imgmat[:, :, 0:3]
cpy = getnewimg(imgmat, scale, cFlag, fFlag)
cpy = cv2.resize(cpy, (0, 0), fx=scale / 4, fy=scale / 4)
if scale==1:
createResultLayer(img, layer.name + '_super', cpy)
else:
createResultFile(layer.name + '_super', cpy)
register(
"super-resolution",
@ -110,12 +141,13 @@ register(
"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)
],
"*", # 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, 4, 0.5)),
(PF_BOOL, "fcpu", "Force CPU", False),
(PF_BOOL, "ffilter", "Use as filter", True)
],
[],
super_resolution, menu="<Image>/Layer/GIML-ML")

Loading…
Cancel
Save