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.
178 lines
5.3 KiB
Python
178 lines
5.3 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+'CelebAMask-HQ/MaskGAN_demo'])
|
|
|
|
|
|
import torch
|
|
from argparse import Namespace
|
|
from models.models import create_model
|
|
from data.base_dataset import get_params, get_transform, normalize
|
|
import os
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
colors = np.array([[0, 0, 0], [204, 0, 0], [76, 153, 0], \
|
|
[204, 204, 0], [51, 51, 255], [204, 0, 204], [0, 255, 255], \
|
|
[51, 255, 255], [102, 51, 0], [255, 0, 0], [102, 204, 0], \
|
|
[255, 255, 0], [0, 0, 153], [0, 0, 204], [255, 51, 153], \
|
|
[0, 204, 204], [0, 51, 0], [255, 153, 51], [0, 204, 0]])
|
|
colors = colors.astype(np.uint8)
|
|
|
|
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)
|
|
# mask=np.dstack((mask1,mask2,mask3))
|
|
return np.uint8(x)
|
|
|
|
def labelMask(mask):
|
|
x=np.zeros((mask.shape[0],mask.shape[1],3))
|
|
for idx in range(19):
|
|
tmp=np.logical_and(mask[:,:,0]==colors[idx][0],mask[:,:,1]==colors[idx][1])
|
|
tmp2=np.logical_and(tmp,mask[:,:,2]==colors[idx][2])
|
|
x[tmp2]=idx
|
|
return x
|
|
|
|
def getOptions():
|
|
mydict={'aspect_ratio': 1.0,
|
|
'batchSize': 1,
|
|
'checkpoints_dir': baseLoc+'CelebAMask-HQ/MaskGAN_demo/checkpoints',
|
|
'cluster_path': 'features_clustered_010.npy',
|
|
'data_type': 32,
|
|
'dataroot': '../Data_preprocessing/',
|
|
'display_winsize': 512,
|
|
'engine': None,
|
|
'export_onnx': None,
|
|
'fineSize': 512,
|
|
'gpu_ids': [0],
|
|
'how_many': 1000,
|
|
'input_nc': 3,
|
|
'isTrain': False,
|
|
'label_nc': 19,
|
|
'loadSize': 512,
|
|
'max_dataset_size': 'inf',
|
|
'model': 'pix2pixHD',
|
|
'nThreads': 2,
|
|
'n_blocks_global': 4,
|
|
'n_blocks_local': 3,
|
|
'n_downsample_global': 4,
|
|
'n_local_enhancers': 1,
|
|
'name': 'label2face_512p',
|
|
'netG': 'global',
|
|
'ngf': 64,
|
|
'niter_fix_global': 0,
|
|
'no_flip': False,
|
|
'norm': 'instance',
|
|
'ntest': 'inf',
|
|
'onnx': None,
|
|
'output_nc': 3,
|
|
'phase': 'test',
|
|
'resize_or_crop': 'scale_width',
|
|
'results_dir': './results/',
|
|
'serial_batches': False,
|
|
'tf_log': False,
|
|
'use_dropout': False,
|
|
'use_encoded_image': False,
|
|
'verbose': False,
|
|
'which_epoch': 'latest'}
|
|
args = Namespace(**mydict)
|
|
return args
|
|
|
|
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 getnewface(img,mask,mask_m):
|
|
h,w,d = img.shape
|
|
img = Image.fromarray(img)
|
|
lmask = labelMask(mask)
|
|
lmask_m = labelMask(mask_m)
|
|
|
|
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = str(0)
|
|
opt = getOptions()
|
|
|
|
model = create_model(opt)
|
|
|
|
params = get_params(opt, (512,512))
|
|
transform_mask = get_transform(opt, params, method=Image.NEAREST, normalize=False, normalize_mask=True)
|
|
transform_image = get_transform(opt, params)
|
|
mask = transform_mask(Image.fromarray(np.uint8(lmask)))
|
|
mask_m = transform_mask(Image.fromarray(np.uint8(lmask_m)))
|
|
img = transform_image(img)
|
|
|
|
generated = model.inference(torch.FloatTensor([mask_m.numpy()]), torch.FloatTensor([mask.numpy()]), torch.FloatTensor([img.numpy()]))
|
|
|
|
result = generated.permute(0, 2, 3, 1)
|
|
if torch.cuda.is_available():
|
|
result = result.cpu().numpy()
|
|
else:
|
|
result = result.detach().numpy()
|
|
|
|
result = (result + 1) * 127.5
|
|
result = np.asarray(result[0,:,:,:], dtype=np.uint8)
|
|
result = Image.fromarray(result)
|
|
result = result.resize([w,h])
|
|
|
|
result = np.array(result)
|
|
return result
|
|
|
|
|
|
def facegen(imggimp, curlayer,layeri,layerm,layermm) :
|
|
if torch.cuda.is_available():
|
|
gimp.progress_init("(Using GPU) Running face gen for " + layeri.name + "...")
|
|
else:
|
|
gimp.progress_init("(Using CPU) Running face gen for " + layeri.name + "...")
|
|
|
|
img = channelData(layeri)
|
|
mask = channelData(layerm)
|
|
mask_m = channelData(layermm)
|
|
|
|
cpy=getnewface(img,mask,mask_m)
|
|
createResultLayer(imggimp,'new_output',cpy)
|
|
|
|
|
|
|
|
register(
|
|
"facegen",
|
|
"facegen",
|
|
"Running face gen.",
|
|
"Kritik Soman",
|
|
"Your",
|
|
"2020",
|
|
"facegen...",
|
|
"*", # Alternately use RGB, RGB*, GRAY*, INDEXED etc.
|
|
[ (PF_IMAGE, "image", "Input image", None),
|
|
(PF_DRAWABLE, "drawable", "Input drawable", None),
|
|
(PF_LAYER, "drawinglayer", "Original Image:", None),
|
|
(PF_LAYER, "drawinglayer", "Original Mask:", None),
|
|
(PF_LAYER, "drawinglayer", "Modified Mask:", None),
|
|
],
|
|
[],
|
|
facegen, menu="<Image>/Layer/GIML-ML")
|
|
|
|
main()
|