Spaces:
Runtime error
Runtime error
import os | |
import numpy as np | |
import torch | |
import torchvision | |
from PIL import Image | |
from pytorch_lightning.callbacks import Callback | |
from pytorch_lightning.utilities.distributed import rank_zero_only | |
# import pdb | |
class ImageLogger(Callback): | |
def __init__(self, batch_frequency=2000, max_images=4, clamp=True, increase_log_steps=True, | |
rescale=True, disabled=False, log_on_batch_idx=False, log_first_step=False, | |
log_images_kwargs=None,ckpt_dir="./ckpt"): | |
super().__init__() | |
self.rescale = rescale | |
self.batch_freq = batch_frequency | |
self.max_images = max_images | |
if not increase_log_steps: | |
self.log_steps = [self.batch_freq] | |
self.clamp = clamp | |
self.disabled = disabled | |
self.log_on_batch_idx = log_on_batch_idx | |
self.log_images_kwargs = log_images_kwargs if log_images_kwargs else {} | |
self.log_first_step = log_first_step | |
self.ckpt_dir=ckpt_dir | |
self.global_save_num=-2000 | |
self.global_save_num1=-100 | |
def log_local(self, save_dir, split, images, global_step, current_epoch, batch_idx): | |
root = os.path.join(save_dir, "image_log", split) | |
# print(images) | |
for k in images: | |
grid = torchvision.utils.make_grid(images[k], nrow=4) | |
if self.rescale: | |
grid = (grid + 1.0) / 2.0 # -1,1 -> 0,1; c,h,w | |
grid = grid.transpose(0, 1).transpose(1, 2).squeeze(-1) | |
grid = grid.numpy() | |
grid = (grid * 255).astype(np.uint8) | |
filename = "{}_gs-{:06}_e-{:06}_b-{:06}.png".format(k, global_step, current_epoch, batch_idx) | |
path = os.path.join(root, filename) | |
os.makedirs(os.path.split(path)[0], exist_ok=True) | |
Image.fromarray(grid).save(path) | |
def log_img(self, pl_module, batch, batch_idx, split="train"): | |
check_idx = batch_idx # if self.log_on_batch_idx else pl_module.global_step | |
if (self.check_frequency(check_idx) and # batch_idx % self.batch_freq == 0 | |
hasattr(pl_module, "log_images") and | |
callable(pl_module.log_images) and | |
self.max_images > 0): | |
logger = type(pl_module.logger) | |
is_train = pl_module.training | |
if is_train: | |
pl_module.eval() | |
with torch.no_grad(): | |
images = pl_module.log_images(batch, split=split, **self.log_images_kwargs) | |
for k in images: | |
N = min(images[k].shape[0], self.max_images) | |
images[k] = images[k][:N] | |
if isinstance(images[k], torch.Tensor): | |
images[k] = images[k].detach().cpu() | |
if self.clamp: | |
images[k] = torch.clamp(images[k], -1., 1.) | |
self.log_local(pl_module.logger.save_dir, split, images, | |
pl_module.global_step, pl_module.current_epoch, batch_idx) | |
if is_train: | |
pl_module.train() | |
def check_frequency(self, check_idx): | |
return check_idx % self.batch_freq == 0 | |
def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx): | |
#if not self.disabled: | |
#if pl_module.global_step%50 == 0: | |
# if pl_module.current_epoch-self.global_save_num1 > 0: | |
# print(batch_idx) | |
if batch_idx % 500 == 0: | |
# print("inside") | |
# pdb.set_trace() | |
# self.global_save_num1=pl_module.current_epoch | |
self.log_img(pl_module, batch, batch_idx, split="train_"+"ckpt_inpainting_from5625_2+3750_exemplar_only_vae") | |
#if pl_module.global_step%1200 == 0 and self.check_frequency(batch_idx): | |
if batch_idx % 1000 == 0: | |
# if pl_module.current_epoch-self.global_save_num>10 and self.check_frequency(batch_idx): | |
# self.global_save_num=pl_module.current_epoch | |
trainer.save_checkpoint(self.ckpt_dir+"/epoch"+str(pl_module.current_epoch)+"_global-step"+str(pl_module.global_step)+".ckpt") | |