code: def test_write_image_dict(tmpdir):
      img_dict = {
          "rgb": np.arange(4 * 4 * 4 * 3, dtype=np.uint8).reshape(
              (4, 4, 4, 3)
          ),
          "rgba": np.arange(4 * 4 * 4 * 4, dtype=np.uint8).reshape(
              (4, 4, 4, 4)
          ),
          "depth": np.linspace(
              0, 100000.0, 4 * 4 * 4, dtype=np.float32
          ).reshape((4, 4, 4, 1)),
          "uv": np.arange(4 * 4 * 4 * 3, dtype=np.uint8).reshape(
              (4, 4, 4, 3)
          ),
          "normal": np.arange(4 * 4 * 4 * 3, dtype=np.uint8).reshape(
              (4, 4, 4, 3)
          ),
          "flow": np.linspace(
              0, 100.0, 4 * 4 * 4 * 2, dtype=np.float32
          ).reshape((4, 4, 4, 2)),
          "forward_flow": np.linspace(
              0, 100.0, 4 * 4 * 4 * 2, dtype=np.float32
          ).reshape((4, 4, 4, 2)),
          "backward_flow": np.linspace(
              0, 100.0, 4 * 4 * 4 * 2, dtype=np.float32
          ).reshape((4, 4, 4, 2)),
          "segmentation": np.ones(4 * 4 * 4, dtype=np.uint8).reshape(
              (4, 4, 4, 1)
          ),
      }
      file_templates = {
          "rgb": "rgb_{:05d}.png",
          "rgba": "rgba_{:05d}.png",
          "depth": "depth_{:05d}.tiff",
          "uv": "uv_{:05d}.png",
          "normal": "normal_{:05d}.png",
          "flow": "flow_{:05d}.png",
          "forward_flow": "forward_flow_{:05d}.png",
          "backward_flow": "backward_flow_{:05d}.png",
          "segmentation": "segmentation_{:05d}.png",
      }
      file_io.write_image_dict(img_dict, tmpdir)
  
      data_ranges = file_io.read_json(tmpdir / "data_ranges.json")
  
      for key, img_batch in img_dict.items():
          if key in data_ranges:
              data_range = (
                  data_ranges[key]["min"],
                  data_ranges[key]["max"],
              )
          else:
              data_range = None
  
          for i, img in enumerate(img_batch):
              filename = tmpdir / file_templates[key].format(i)
              assert filename.exists()
              print(filename)
              if str(filename).endswith(".png"):
                  img_recovered = file_io.read_png(
                      filename, rescale_range=data_range
                  )
              elif str(filename).endswith(".tiff"):
                  img_recovered = file_io.read_tiff(filename)
              else:
                  assert False
  
              if (
                  img.shape[-1] == 2
              ):  # pad image with a zero blue channel
                  img = np.concatenate(
                      [img, np.zeros_like(img[:, :, :1])], axis=-1
                  )
  
              assert img.shape == img_recovered.shape
              np.testing.assert_allclose(
                  img_recovered, img, rtol=1e-4, atol=1e-4
              )
  
file_path: kubric/test/test_fileio.py
imports: import numpy
  import pytest
  from kubric import file_io
label_desc: Use image node_id 10 for calling a global function with example usage: test_write_image_dict(tmpdir) without return types