LIVE/tvideo/logger.py

import os
import numpy as np
import einops
import imageio
from typing import Union
from matplotlib import pyplot as plt
from PIL import Image, ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True  # UnidentifiedImageError: https://github.com/python-pillow/Pillow/issues/5631
from pathlib import Path

import torch
import torchvision
import wandb

import lightning.pytorch as pl
from lightning.pytorch.callbacks import Callback
from lightning.pytorch.loggers import WandbLogger
from lightning.pytorch.utilities.rank_zero import rank_zero_only
from einops import rearrange
from torchvision.io import write_video

from sgm.util import exists, isheatmap, print0


class ImageVideoLogger(Callback):
    def __init__(
        self,
        batch_frequency,
        max_samples,
        clamp=True,
        increase_log_steps=True,
        batch_frequency_val=None,
        video_fps=8,
        rescale=True,
        disabled=False,
        log_on_batch_idx=True,  # log on batch_idx instead of global_step. global_step is fixed in validation. batch_idx restarts at each validation
        log_first_step=True,
        log_images_kwargs=None,
        log_videos_kwargs=None,
        log_before_first_step=True,
        enable_autocast=True,
    ):
        super().__init__()
        self.enable_autocast = enable_autocast
        self.rescale = rescale
        self.batch_freq = batch_frequency
        self.batch_freq_val = batch_frequency_val if batch_frequency_val is not None else batch_frequency
        self.video_fps = video_fps
        self.max_samples = max_samples
        self.log_steps = [2**n for n in range(int(np.log2(self.batch_freq)) + 1)]
        if not increase_log_steps:
            self.log_steps = [self.batch_freq]
        self.increase_log_steps = increase_log_steps
        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_videos_kwargs = log_videos_kwargs if log_videos_kwargs else {}
        self.log_first_step = log_first_step
        self.log_before_first_step = log_before_first_step

    @rank_zero_only
    def log_img_local(
        self,
        save_dir,
        split,
        images,
        global_step,
        current_epoch,
        batch_idx,
        pl_module: Union[None, pl.LightningModule] = None,
    ):
        root = os.path.join(save_dir, "images", split)
        for k in images:
            if isheatmap(images[k]):
                fig, ax = plt.subplots()
                ax = ax.matshow(
                    images[k].cpu().numpy(), cmap="hot", interpolation="lanczos"
                )
                plt.colorbar(ax)
                plt.axis("off")

                filename = "{}_gs-{:06}_e-{:06}_b-{:06}.png".format(
                    k, global_step, current_epoch, batch_idx
                )
                os.makedirs(root, exist_ok=True)
                path = os.path.join(root, filename)
                plt.savefig(path)
                plt.close()
            else:
                if images[k].ndim == 5:
                    images[k] = einops.rearrange(images[k], "b c t h w -> (b t) c h w")
                nrow = self.log_images_kwargs.get("n_rows", 8)
                # the follows are original
                grid = torchvision.utils.make_grid(images[k], nrow=nrow)
                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)
                img = Image.fromarray(grid)
                img.save(path)
                # if exists(pl_module):
                #     assert isinstance(
                #         pl_module.logger, WandbLogger
                #     ), "logger_log_image only supports WandbLogger currently"
                #     pl_module.logger.log_image(
                #         key=f"{split}/{k}",
                #         images=[
                #             img,
                #         ],
                #         step=pl_module.global_step,
                #     )

    @rank_zero_only
    def log_vid_local(
        self,
        save_dir,
        split,
        videos,
        global_step,
        current_epoch,
        batch_idx,
        pl_module: Union[None, pl.LightningModule] = None,
    ):
        root = os.path.join(save_dir, "videos", split)
        for k in videos:
            # if is video, we can add captions
            if isinstance(videos[k], torch.Tensor) and videos[k].ndim == 5:
                # (b, t, c, h, w)
                if self.rescale:
                    videos[k] = (videos[k] + 1.0) / 2.0  # -1,1 -> 0,1; c,h,w
                frames = [videos[k][:, i] for i in range(videos[k].shape[1])]
                frames = [torchvision.utils.make_grid(each, nrow=4) for each in frames]
                frames = [einops.rearrange(each, "c h w -> 1 c h w") for each in frames]
                frames = torch.clamp(torch.cat(frames, dim=0), min=0.0, max=1.0)
                frames = (frames.numpy() * 255).astype(np.uint8)  # frames: (t, c, h, w)

                os.makedirs(root, exist_ok=True)
                # filename_gif = "{}_gs-{:06}_e-{:06}_b-{:06}.gif".format(
                #     k, global_step, current_epoch, batch_idx
                # )
                # path_gif = os.path.join(root, filename_gif)
                # save_numpy_as_gif(frames, path_gif, duration=1 / self.video_fps)
                filename_mp4 = "{}_gs-{:06}_e-{:06}_b-{:06}.mp4".format(
                    k, global_step, current_epoch, batch_idx
                )
                path_mp4 = os.path.join(root, filename_mp4)
                save_numpy_as_mp4(frames, path_mp4, fps=self.video_fps)
                # if exists(pl_module):
                #     assert isinstance(
                #         pl_module.logger, WandbLogger
                #     ), "log_videos only supports WandbLogger currently"
                #     wandb.log({f"{split}/{k}": wandb.Video(frames, fps=self.video_fps)})  # k is str
                #     # wandb.log({f"{split}/{k}": wandb.Video(frames, caption='caption', fps=self.video_fps)}, step=global_step)

    @rank_zero_only
    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) or self.check_frequency_val(batch_idx, split))
            and hasattr(pl_module, "log_images")  # batch_idx % self.batch_freq == 0
            and callable(pl_module.log_images)
            and
            # batch_idx > 5 and
            self.max_samples > 0
        ):
            logger = type(pl_module.logger)
            is_train = pl_module.training
            if is_train:
                pl_module.eval()

            with torch.no_grad(), torch.autocast(enabled=self.enable_autocast, device_type="cuda"):
                images = pl_module.log_images(
                    batch, split=split, **self.log_images_kwargs
                )

            if images is not None:
                for k in images:
                    N = min(images[k].shape[0], self.max_samples)
                    if not isheatmap(images[k]):
                        images[k] = images[k][:N]
                    if isinstance(images[k], torch.Tensor):
                        images[k] = images[k].detach().float().cpu()
                        if self.clamp and not isheatmap(images[k]):
                            images[k] = torch.clamp(images[k], -1.0, 1.0)

                self.log_img_local(
                    pl_module.logger.save_dir,
                    split,
                    images,
                    pl_module.global_step,
                    pl_module.current_epoch,
                    batch_idx,
                    pl_module=pl_module
                    if isinstance(pl_module.logger, WandbLogger)
                    else None,
                )

            if is_train:
                pl_module.train()

    @rank_zero_only
    def log_vid(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) or self.check_frequency_val(batch_idx, split))
            and hasattr(pl_module, "log_videos")  # batch_idx % self.batch_freq == 0
            and callable(pl_module.log_videos)
            and
            # batch_idx > 5 and
            self.max_samples > 0
        ):
            logger = type(pl_module.logger)
            is_train = pl_module.training
            if is_train:
                pl_module.eval()

            with torch.no_grad(), torch.autocast(enabled=self.enable_autocast, device_type="cuda"):
                videos = pl_module.log_videos(
                    batch, split=split, **self.log_videos_kwargs
                )

            if videos is not None:
                for k in videos:
                    N = min(videos[k].shape[0], self.max_samples)
                    videos[k] = videos[k][:N]
                    if isinstance(videos[k], torch.Tensor):
                        videos[k] = videos[k].detach().float().cpu()
                        if self.clamp:
                            videos[k] = torch.clamp(videos[k], -1.0, 1.0)

                self.log_vid_local(
                    pl_module.logger.save_dir,
                    split,
                    videos,
                    pl_module.global_step,
                    pl_module.current_epoch,
                    batch_idx,
                    pl_module=pl_module
                    if isinstance(pl_module.logger, WandbLogger)
                    else None,
                )

            if is_train:
                pl_module.train()

    def check_frequency(self, check_idx):
        if ((check_idx % self.batch_freq) == 0 or (check_idx in self.log_steps)) and (
            check_idx > 0 or self.log_first_step  # seems starting from 1
        ):
            if self.increase_log_steps:
                try:
                    self.log_steps.pop(0)
                except IndexError as e:
                    # print0(e)
                    pass
            return True
        return False

    def check_frequency_val(self, check_idx, split):
        if 'val' in split:
            if ((check_idx % self.batch_freq_val) == 0) and (
                check_idx > 0 or self.log_first_step):
                return True
        return False

    @rank_zero_only
    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
        # global_step starts from 1, batch_idx starts from 0
        # not log after first step
        if not self.disabled and (pl_module.global_step > 2):
            self.log_img(pl_module, batch, batch_idx, split="train")
            self.log_vid(pl_module, batch, batch_idx, split="train")

    @rank_zero_only
    def on_train_batch_start(self, trainer, pl_module, batch, batch_idx):
        # global_step starts from 0
        if not self.disabled and (pl_module.global_step == 0 and self.log_before_first_step):
            self.log_img(pl_module, batch, batch_idx, split="train")
            self.log_vid(pl_module, batch, batch_idx, split="train")

    @rank_zero_only
    def on_validation_batch_end(
        self, trainer, pl_module, outputs, batch, batch_idx, *args, **kwargs
    ):
        if not self.disabled:
            self.log_img(pl_module, batch, batch_idx, split="val")
            self.log_vid(pl_module, batch, batch_idx, split="val")


def save_numpy_as_gif(frames, path, duration=None):
    """
    save numpy array as gif file
    """
    image_list = []
    for frame in frames:
        image = frame.transpose(1, 2, 0)
        image_list.append(image)
    if duration:
        imageio.mimsave(path, image_list, format="GIF", duration=duration, loop=0)
        # imageio.mimsave(path, image_list, format="GIF", duration=duration, loop=0, quality=10)
    else:
        imageio.mimsave(path, image_list, format="GIF", loop=0)
        # imageio.mimsave(path, image_list, format="GIF", loop=0, quality=10)


def save_numpy_as_mp4(frames, path, fps):
    """
    save numpy array as mp4 file
    """
    # image_list = []
    # for frame in frames:
    #     image = frame.transpose(1, 2, 0)
    #     image_list.append(image)
    # imageio.mimsave(path, image_list, format="FFMPEG", fps=fps)
    frames = rearrange(frames, "t c h w -> t h w c")
    write_video(path, frames, fps=fps)