extract make_* functions out of make_*_loader

test/common_utils.py

@@ -492,6 +492,34 @@ def get_num_channels(color_space):
  return num_channels
 
 
+def make_image(
+ spatial_size,
+ *,
+ color_space="RGB",
+ batch_dims=(),
+ dtype=torch.float32,
+ device="cpu",
+ constant_alpha=True,
+ memory_format=torch.contiguous_format,
+):
+ spatial_size = _parse_spatial_size(spatial_size)
+ num_channels = get_num_channels(color_space)
+ max_value = get_max_value(dtype)
+
+ data = torch.testing.make_tensor(
+ (*batch_dims, num_channels, *spatial_size),
+ low=0,
+ high=max_value,
+ dtype=dtype,
+ device=device,
+ memory_format=memory_format,
+ )
+ if color_space in {"GRAY_ALPHA", "RGBA"} and constant_alpha:
+ data[..., -1, :, :] = max_value
+
+ return datapoints.Image(data)
+
+
 def make_image_loader(
  size="random",
  *,
@@ -505,20 +533,20 @@ def make_image_loader(
  num_channels = get_num_channels(color_space)
 
  def fn(shape, dtype, device, memory_format):
- max_value = get_max_value(dtype)
- data = torch.testing.make_tensor(
- shape, low=0, high=max_value, dtype=dtype, device=device, memory_format=memory_format
+ *batch_dims, _, spatial_size = shape
+ return make_image(
+ spatial_size,
+ color_space=color_space,
+ batch_dims=batch_dims,
+ dtype=dtype,
+ device=device,
+ constant_alpha=constant_alpha,
+ memory_format=memory_format,
  )
- if color_space in {"GRAY_ALPHA", "RGBA"} and constant_alpha:
- data[..., -1, :, :] = max_value
- return datapoints.Image(data)
 
  return ImageLoader(fn, shape=(*extra_dims, num_channels, *size), dtype=dtype, memory_format=memory_format)
 
 
-make_image = from_loader(make_image_loader)
-
-
 def make_image_loaders(
  *,
  sizes=DEFAULT_SPATIAL_SIZES,