基于深度强化学习的绘画智能体 代码分析（四）

Github源码

tensorboard.py

import PIL #图像处理库
import scipy.misc #将数组保存成图像形式
from io import BytesIO #在内存中读写bytes
import tensorboardX as tb
from tensorboardX.summary import Summary

class TensorBoard(object):
    def __init__(self, model_dir): #model_dir是下载模型保存地址
        self.summary_writer = tb.FileWriter(model_dir) #指定一个文件用来保存图

    def add_image(self, tag, img, step):
        summary = Summary()
        bio = BytesIO() #创建一个类二进制文件对象

        if type(img) == str:
            img = PIL.Image.open(img)  #返回PIL.Image.Image的类型
        elif type(img) == PIL.Image.Image:
            pass #不需要转换
        else:
            img = PIL.Image.fromarray(img) #array转换成image

        img.save(bio, format="png")
        image_summary = Summary.Image(encoded_image_string