def tensorboard_reconstruction(self, orig, recon):
        # self.tb_writer.add_text(
        #     'reconstruction/original', text_string="", global_step=self.epoch)
        self.tb_writer.add_image(
            "reconstruction/original", Drawing.from_tensor_prediction(orig).tensorboard_plot(), self.epoch)

        # self.tb_writer.add_text(
        #     'reconstruction/prediction', str(recon), self.epoch)
        self.tb_writer.add_image(
            "reconstruction/prediction", Drawing.from_tensor_prediction(recon).tensorboard_plot(), self.epoch)
        self.tb_writer.flush()
 def save_reconstruction(self, orig, recon):
     original = Drawing.from_tensor_prediction(orig)
     reconstruct = Drawing.from_tensor_prediction(recon)
     plt.figure(figsize=(7, 4))
     plt.suptitle(f"Epoch {self.epoch}")
     plt.subplot(1, 2, 1)
     plt.title("orig")
     original.render_image(show=False)
     plt.subplot(1, 2, 2)
     plt.title("recon")
     reconstruct.render_image(show=False)
     plt.savefig(self.plots_dir + f"reconstruction_{self.epoch}.png", bbox_inches='tight', dpi=150)
     plt.close()