code: def create_image(data_format: str = "numpy", **kwargs):
      """Creates a single image, using the same keyword args as create_dataset."""
      data = create_dataset(data_format=data_format, **kwargs)
      # TODO: Maybe implement this without relying on create_dataset
      #   to make it more efficient.
  
      if data_format == "numpy" or data_format == "torch":
          return data[0][0]
      elif data_format == "files":
          return next(data.rglob("*.png"))
      else:
          raise RuntimeError()
  
file_path: traintool/tests/conftest.py
imports: import comet_ml
  import sys
  import os
  import numpy
  import torch
  from torch.utils.data import TensorDataset
  from typing import Union, List
  from pathlib import Path
  import imageio
label_desc: Use image node_id 2 for calling a global function with example usage: create_image(data_format) and returns: data, next