code: def log_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()
              # TODO: support wandb
          else:
              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)
              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,
                  )
  
imports: import argparse
  import datetime
  import glob
  import inspect
  import os
  import sys
  from inspect import Parameter
  from typing import Union
  import numpy
  import pytorch_lightning
  import torch
  import torchvision
  import wandb
  from matplotlib import pyplot
  from natsort import natsorted
  from omegaconf import OmegaConf
  from packaging import version
  from PIL import Image
  from pytorch_lightning import seed_everything
  from pytorch_lightning.callbacks import Callback
  from pytorch_lightning.loggers import WandbLogger
  from pytorch_lightning.trainer import Trainer
  from pytorch_lightning.utilities import rank_zero_only
  from sgm.util import exists, instantiate_from_config, isheatmap
label_desc: Use image node_id 2 for calling the ImageLogger obj's underlying member method code with example usage: obj.log_local(save_dir, split, images, global_step, current_epoch, batch_idx, pl_module) without return types