diff --git a/gimp-plugins/super_resolution.py b/gimp-plugins/super_resolution.py index 7c5bd5a..2b49ef9 100755 --- a/gimp-plugins/super_resolution.py +++ b/gimp-plugins/super_resolution.py @@ -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="/Layer/GIML-ML")