code: def test_base_viz(self, tmpdir):
      seed_everything(42)
      tmpdir = Path(tmpdir)
  
      train_images = [str(tmpdir / "a1.png"), str(tmpdir / "b1.png")]
  
      _rand_image().save(train_images[0])
      _rand_image().save(train_images[1])
  
      class CustomImageClassificationData(ImageClassificationData):
          @staticmethod
          def configure_data_fetcher(
              *args, **kwargs
          ) -> CustomBaseVisualization:
              return CustomBaseVisualization(*args, **kwargs)
  
      B: int = 2  # batch_size
  
      dm = CustomImageClassificationData.from_files(
          train_files=train_images,
          train_targets=[0, 1],
          val_files=train_images,
          val_targets=[2, 3],
          test_files=train_images,
          test_targets=[4, 5],
          predict_files=train_images,
          batch_size=B,
          num_workers=0,
      )
  
      num_tests = 10
  
      for stage in ("train", "val", "test", "predict"):
          for _ in range(num_tests):
              for fcn_name in _CALLBACK_FUNCS:
                  dm.data_fetcher.reset()
                  assert dm.data_fetcher.batches == {
                      "predict": {},
                      "test": {},
                      "train": {},
                      "val": {},
                      "serve": {},
                  }
                  fcn = getattr(dm, f"show_{stage}_batch")
                  fcn(fcn_name, reset=False)
  
          is_predict = stage == "predict"
  
          def _extract_data(data):
              return data[0][DataKeys.INPUT]
  
          def _get_result(function_name: str):
              return dm.data_fetcher.batches[stage][function_name]
  
          res = _get_result("load_sample")
          assert len(res) == B
          assert isinstance(_extract_data(res), Image.Image)
  
          if not is_predict:
              res = _get_result("load_sample")
              assert isinstance(res[0][DataKeys.TARGET], int)
  
          res = _get_result("per_sample_transform")
          assert len(res) == B
          assert isinstance(_extract_data(res), Tensor)
  
          if not is_predict:
              res = _get_result("per_sample_transform")
              assert isinstance(res[0][DataKeys.TARGET], Tensor)
  
          res = _get_result("collate")
          assert _extract_data(res).shape == (B, 3, 196, 196)
  
          if not is_predict:
              res = _get_result("collate")
              assert res[0][DataKeys.TARGET].shape == torch.Size([2])
  
          res = _get_result("per_batch_transform")
          assert _extract_data(res).shape == (B, 3, 196, 196)
  
          if not is_predict:
              res = _get_result("per_batch_transform")
              assert res[0][DataKeys.TARGET].shape == (B,)
  
          assert dm.data_fetcher.show_load_sample_called
          assert dm.data_fetcher.show_per_sample_transform_called
          assert dm.data_fetcher.show_collate_called
          assert dm.data_fetcher.per_batch_transform_called
          dm.data_fetcher.reset()
  
imports: from pathlib import Path
  from typing import Any, List, Sequence, Tuple
  import numpy
  import pytest
  import torch
  from flash.core.data.base_viz import BaseVisualization
  from flash.core.data.io.input import DataKeys
  from flash.core.data.utils import _CALLBACK_FUNCS
  from flash.core.utilities.imports import _PIL_AVAILABLE, _TOPIC_IMAGE_AVAILABLE
  from flash.core.utilities.stages import RunningStage
  from flash.image import ImageClassificationData
  from pytorch_lightning import seed_everything
  from torch import Tensor
label_desc: Use image node_id 1 for calling the TestBaseViz obj's underlying member method code with example usage: obj.test_base_viz(tmpdir) and returns: CustomBaseVisualization, data, dm