GIMP-ML/gimp-plugins/facegen.py
Kritik Soman 001b95d59d CPUButton
2020-09-20 18:55:54 +05:30

182 lines
5.5 KiB
Python
Executable 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+'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+'weights/facegen',
'cluster_path': 'features_clustered_010.npy',
'data_type': 32,
'dataroot': '../Data_preprocessing/',
'display_winsize': 512,
'engine': None,
'export_onnx': None,
'fineSize': 512,
'gpu_ids': [],
'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,cFlag):
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()
if torch.cuda.is_available() and not cFlag:
opt.gpu_ids=[0]
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()]), cFlag)
result = generated.permute(0, 2, 3, 1)
if torch.cuda.is_available():
result = result.detach().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,cFlag):
if torch.cuda.is_available() and not cFlag:
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,cFlag)
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, "drawinglayer2", "Original Mask:", None),
(PF_LAYER, "drawinglayer3", "Modified Mask:", None),
(PF_BOOL, "fcpu", "Force CPU", False),
],
[],
facegen, menu="<Image>/Layer/GIML-ML")
main()