code: def create_dataset(
      data_format: str = "numpy",
      size: int = 28,
      grayscale: bool = True,
      num_samples: int = 4,
      num_classes: int = 4,
      seed: int = 0,
      tmp_path: Path = None,
  ) -> Union[List[np.ndarray], TensorDataset, Path]:
      """
      Create fake data for image classification, including images (values from 0 to 1)
      and labels (10 classes).
  
      Args:
          data_format (str, optional): "numpy" (for list of numpy arrays
              [images, labels]) or "torch" (for torch.utils.data.Dataset) or "files"
              (for directory of image files). Defaults to "numpy".
          size (int, optional): Width and height of the images. Defaults to 28.
          grayscale (bool, optional): If True, images are grayscale and of size 28x28,
              otherwise they have 3 color channels and are 224x224. Defaults to False.
          num_samples (int, optional): Number of samples in the dataset. Defaults to 4.
          seed (int, optional): Seed for data generation. Defaults to 0.
  
      Returns:
          Union[List[np.ndarray], torch.utils.data.Dataset, Path]: The generated fake
              data.
      """
  
      if num_samples < num_classes:
          raise ValueError(
              "num_samples needs to be >= num_classes, so that each label "
              "can appear at least once"
          )
  
      # Seed to make different datasets comparable.
      np.random.seed(seed)
  
      # Create images in channels first format.
      num_channels = 1 if grayscale else 3
      images = np.random.rand(num_samples, num_channels, size, size)
  
      # Create random labels. Make sure all labels appear at least once by setting the
      # first few ones semi-random.
      labels = np.random.randint(num_classes, size=num_samples)
      labels[:num_classes] = np.random.permutation(num_classes)
  
      # Convert to correct output format.
      if data_format == "numpy":
          return [images, labels]
      elif data_format == "torch":
          dataset = TensorDataset(
              torch.from_numpy(images).float(),
              torch.from_numpy(labels).long(),
          )
          return dataset
      elif data_format == "files":
          if tmp_path is None:
              raise ValueError(
                  "tmp_path must be given if data_format is files"
              )
          for i, (image, label) in enumerate(zip(images, labels)):
              image = image.transpose((1, 2, 0))
              image_dir = tmp_path / "data" / f"{label}"
              image_dir.mkdir(parents=True, exist_ok=True)
              imageio.imwrite(image_dir / f"{i}.png", image)
          return tmp_path / "data"
      else:
          raise ValueError(f"data_format not supported: {data_format}")
  
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 1 for calling a global function with example usage: create_dataset(data_format, size, grayscale, num_samples, num_classes, seed, tmp_path) and returns: list, dataset, unknown