基于 ModelArts 的动漫头像自动生成丨【华为云至简致远】

root_path = './ACGAN/'os.chdir(root_path)import osfrom main import mainfrom ACGAN import ACGANfrom tools import checkFolderimport tensorflow as tfimport argparseimport numpy as np

def parse_args():    note = "ACGAN Frame Constructed With Tensorflow"    parser = argparse.ArgumentParser(description=note)    parser.add_argument("--epoch",type=int,default=251,help="训练轮数")    parser.add_argument("--batchSize",type=int,default=64,help="batch的大小")    parser.add_argument("--codeSize",type=int,default=62,help="输入编码向量的维度")    parser.add_argument("--checkpointDir",type=str,default="./checkpoint",help="检查点保存目录")    parser.add_argument("--resultDir",type=str,default="./result",help="训练过程中，中间生成结果的目录")    parser.add_argument("--logDir",type=str,default="./log",help="训练日志目录")    parser.add_argument("--mode",type=str,default="train",help="模式： train / infer")    parser.add_argument("--hairStyle",type=str,default="orange hair",help="你想要生成的动漫头像的头发颜色")    parser.add_argument("--eyeStyle",type=str,default="gray eyes",help="你想要生成的动漫头像的眼睛颜色")    parser.add_argument("--dataSource",type=str,default='./extra_data/images/',help="训练集路径")    args, unknown= parser.parse_known_args()    checkFolder(args.checkpointDir)    checkFolder(args.resultDir)    checkFolder(args.logDir)    assert args.epoch>=1    assert args.batchSize>=1    assert args.codeSize>=1    return args
args =parse_args()

with tf.Session() as sess :    myGAN = ACGAN(sess,args.epoch,args.batchSize,args.codeSize,\        args.dataSource,args.checkpointDir,args.resultDir,args.logDir,args.mode,\            64,64,3)    if myGAN is None:        print("创建GAN网络失败")        exit(0)
    if args.mode=='train' :        myGAN.buildNet()        print("进入训练模式")        myGAN.train()        print("Done")

args.mode ='infer'

hair_dict = ['orange hair', 'white hair', 'aqua hair', 'gray hair', 'green hair', 'red hair', 'purple hair',             'pink hair', 'blue hair', 'black hair', 'brown hair', 'blonde hair']eye_dict = [ 'gray eyes', 'black eyes', 'orange eyes', 'pink eyes', 'yellow eyes',            'aqua eyes', 'purple eyes', 'green eyes', 'brown eyes', 'red eyes', 'blue eyes']

# 选择了黑头发和蓝眼睛args.hairStyle = 'black hair'args.eyeStyle = 'blue eyes'

tf.reset_default_graph()with tf.Session() as sess :    myGAN1 = ACGAN(sess,args.epoch,args.batchSize,args.codeSize,\        args.dataSource,args.checkpointDir,args.resultDir,args.logDir,args.mode,\            64,64,3)    if myGAN1 is None:        print("创建GAN网络失败")        exit(0)    if args.mode=='infer' :        myGAN1.buildForInfer()        tag_dict = ['orange hair', 'white hair', 'aqua hair', 'gray hair', 'green hair', 'red hair', 'purple hair', 'pink hair', 'blue hair', 'black hair',                     'brown hair', 'blonde hair','gray eyes', 'black eyes', 'orange eyes', 'pink eyes', 'yellow eyes','aqua eyes', 'purple eyes', 'green eyes',                    'brown eyes', 'red eyes','blue eyes']        tag = np.zeros((64,23))        feature = args.hairStyle+" AND "+ args.eyeStyle        for j in range(25):            for i in range(len(tag_dict)):                if tag_dict[i] in feature:                    tag[j][i] = 1        myGAN1.infer(tag,feature)        print("Generate : "+feature)

import matplotlib.pyplot as pltfrom PIL import Image
feature = args.hairStyle+" AND "+ args.eyeStyleresultPath = './samples/' + feature + '.png' #确定保存路径img = Image.open(resultPath).convert('RGB')plt.figure(1)plt.imshow(img)plt.show()