[fbsync] Add support for decoding 16bits png (#8524)

Reviewed By: vmoens Differential Revision: D60596240 fbshipit-source-id: 4987658c3ee1255a3200232768580a927e371d12

Author: NicolasHug

test/test_image.py

@@ -10,11 +10,10 @@
 import pytest
 import requests
 import torch
-import torchvision.transforms.functional as F
+import torchvision.transforms.v2.functional as F
 from common_utils import assert_equal, cpu_and_cuda, IN_OSS_CI, needs_cuda
 from PIL import __version__ as PILLOW_VERSION, Image, ImageOps, ImageSequence
 from torchvision.io.image import (
- _read_png_16,
  decode_gif,
  decode_image,
  decode_jpeg,
@@ -211,16 +210,11 @@ def test_decode_png(img_path, pil_mode, mode, scripted, decode_fun):
  img_pil = normalize_dimensions(img_pil)
 
  if img_path.endswith("16.png"):
- # 16 bits image decoding is supported, but only as a private API
- # FIXME: see https://github.com/pytorch/vision/issues/4731 for potential solutions to making it public
- with pytest.raises(RuntimeError, match="At most 8-bit PNG images are supported"):
- data = read_file(img_path)
- img_lpng = decode_fun(data, mode=mode)
-
- img_lpng = _read_png_16(img_path, mode=mode)
- assert img_lpng.dtype == torch.int32
- # PIL converts 16 bits pngs in uint8
- img_lpng = torch.round(img_lpng / (2**16 - 1) * 255).to(torch.uint8)
+ data = read_file(img_path)
+ img_lpng = decode_fun(data, mode=mode)
+ assert img_lpng.dtype == torch.uint16
+ # PIL converts 16 bits pngs to uint8
+ img_lpng = F.to_dtype(img_lpng, torch.uint8, scale=True)
  else:
  data = read_file(img_path)
  img_lpng = decode_fun(data, mode=mode)

test/test_transforms_v2.py

@@ -2076,15 +2076,17 @@ def fn(value):
  factor = (output_max_value + 1) // (input_max_value + 1)
  return value * factor
 
- return torch.tensor(tree_map(fn, image.tolist()), dtype=dtype, device=image.device)
+ return torch.tensor(tree_map(fn, image.tolist())).to(dtype=output_dtype, device=image.device)
 
- @pytest.mark.parametrize("input_dtype", [torch.float32, torch.float64, torch.uint8])
- @pytest.mark.parametrize("output_dtype", [torch.float32, torch.float64, torch.uint8])
+ @pytest.mark.parametrize("input_dtype", [torch.float32, torch.float64, torch.uint8, torch.uint16])
+ @pytest.mark.parametrize("output_dtype", [torch.float32, torch.float64, torch.uint8, torch.uint16])
  @pytest.mark.parametrize("device", cpu_and_cuda())
  @pytest.mark.parametrize("scale", (True, False))
  def test_image_correctness(self, input_dtype, output_dtype, device, scale):
  if input_dtype.is_floating_point and output_dtype == torch.int64:
  pytest.xfail("float to int64 conversion is not supported")
+ if input_dtype == torch.uint8 and output_dtype == torch.uint16 and device == "cuda":
+ pytest.xfail("uint8 to uint16 conversion is not supported on cuda")
 
  input = make_image(dtype=input_dtype, device=device)
 
@@ -2171,6 +2173,28 @@ def test_errors_warnings(self, make_input):
  assert out["bbox"].dtype == bbox_dtype
  assert out["mask"].dtype == mask_dtype
 
+ def test_uint16(self):
+ # These checks are probably already covered above but since uint16 is a
+ # newly supported dtype, we want to be extra careful, hence this
+ # explicit test
+ img_uint16 = torch.randint(0, 65535, (256, 512), dtype=torch.uint16)
+
+ img_uint8 = F.to_dtype(img_uint16, torch.uint8, scale=True)
+ img_float32 = F.to_dtype(img_uint16, torch.float32, scale=True)
+ img_int32 = F.to_dtype(img_uint16, torch.int32, scale=True)
+
+ assert_equal(img_uint8, (img_uint16 / 256).to(torch.uint8))
+ assert_close(img_float32, (img_uint16 / 65535))
+
+ assert_close(F.to_dtype(img_float32, torch.uint16, scale=True), img_uint16, rtol=0, atol=1)
+ # Ideally we'd check against (img_uint16 & 0xFF00) but bitwise and isn't supported for it yet
+ # so we simulate it by scaling down and up again.
+ assert_equal(F.to_dtype(img_uint8, torch.uint16, scale=True), ((img_uint16 / 256).to(torch.uint16) * 256))
+ assert_equal(F.to_dtype(img_int32, torch.uint16, scale=True), img_uint16)
+
+ assert_equal(F.to_dtype(img_float32, torch.uint8, scale=True), img_uint8)
+ assert_close(F.to_dtype(img_uint8, torch.float32, scale=True), img_float32, rtol=0, atol=1e-2)
+
 
 class TestAdjustBrightness:
  _CORRECTNESS_BRIGHTNESS_FACTORS = [0.5, 0.0, 1.0, 5.0]

torchvision/csrc/io/image/cpu/decode_image.cpp

@@ -32,8 +32,7 @@ torch::Tensor decode_image(
  if (memcmp(jpeg_signature, datap, 3) == 0) {
  return decode_jpeg(data, mode, apply_exif_orientation);
  } else if (memcmp(png_signature, datap, 4) == 0) {
- return decode_png(
- data, mode, /*allow_16_bits=*/false, apply_exif_orientation);
+ return decode_png(data, mode, apply_exif_orientation);
  } else if (
  memcmp(gif_signature_1, datap, 6) == 0 ||
  memcmp(gif_signature_2, datap, 6) == 0) {

torchvision/csrc/io/image/cpu/decode_png.cpp

@@ -11,7 +11,6 @@ using namespace exif_private;
 torch::Tensor decode_png(
  const torch::Tensor& data,
  ImageReadMode mode,
- bool allow_16_bits,
  bool apply_exif_orientation) {
  TORCH_CHECK(
  false, "decode_png: torchvision not compiled with libPNG support");
@@ -26,7 +25,6 @@ bool is_little_endian() {
 torch::Tensor decode_png(
  const torch::Tensor& data,
  ImageReadMode mode,
- bool allow_16_bits,
  bool apply_exif_orientation) {
  C10_LOG_API_USAGE_ONCE("torchvision.csrc.io.image.cpu.decode_png.decode_png");
  // Check that the input tensor dtype is uint8
@@ -99,12 +97,12 @@ torch::Tensor decode_png(
  TORCH_CHECK(retval == 1, "Could read image metadata from content.")
  }
 
- auto max_bit_depth = allow_16_bits ? 16 : 8;
- auto err_msg = "At most " + std::to_string(max_bit_depth) +
- "-bit PNG images are supported currently.";
- if (bit_depth > max_bit_depth) {
+ if (bit_depth > 8 && bit_depth != 16) {
  png_destroy_read_struct(&png_ptr, &info_ptr, nullptr);
- TORCH_CHECK(false, err_msg)
+ TORCH_CHECK(
+ false,
+ "bit depth of png image is " + std::to_string(bit_depth) +
+ ". Only <=8 and 16 are supported.")
  }
 
  int channels = png_get_channels(png_ptr, info_ptr);
@@ -199,45 +197,20 @@ torch::Tensor decode_png(
  }
 
  auto num_pixels_per_row = width * channels;
+ auto is_16_bits = bit_depth == 16;
  auto tensor = torch::empty(
  {int64_t(height), int64_t(width), channels},
- bit_depth <= 8 ? torch::kU8 : torch::kI32);
-
- if (bit_depth <= 8) {
- auto t_ptr = tensor.accessor<uint8_t, 3>().data();
- for (int pass = 0; pass < number_of_passes; pass++) {
- for (png_uint_32 i = 0; i < height; ++i) {
- png_read_row(png_ptr, t_ptr, nullptr);
- t_ptr += num_pixels_per_row;
- }
- t_ptr = tensor.accessor<uint8_t, 3>().data();
- }
- } else {
- // We're reading a 16bits png, but pytorch doesn't support uint16.
- // So we read each row in a 16bits tmp_buffer which we then cast into
- // a int32 tensor instead.
- if (is_little_endian()) {
- png_set_swap(png_ptr);
- }
- int32_t* t_ptr = tensor.accessor<int32_t, 3>().data();
-
- // We create a tensor instead of malloc-ing for automatic memory management
- auto tmp_buffer_tensor = torch::empty(
- {int64_t(num_pixels_per_row * sizeof(uint16_t))}, torch::kU8);
- uint16_t* tmp_buffer =
- (uint16_t*)tmp_buffer_tensor.accessor<uint8_t, 1>().data();
-
- for (int pass = 0; pass < number_of_passes; pass++) {
- for (png_uint_32 i = 0; i < height; ++i) {
- png_read_row(png_ptr, (uint8_t*)tmp_buffer, nullptr);
- // Now we copy the uint16 values into the int32 tensor.
- for (size_t j = 0; j < num_pixels_per_row; ++j) {
- t_ptr[j] = (int32_t)tmp_buffer[j];
- }
- t_ptr += num_pixels_per_row;
- }
- t_ptr = tensor.accessor<int32_t, 3>().data();
+ is_16_bits ? at::kUInt16 : torch::kU8);
+ if (is_little_endian()) {
+ png_set_swap(png_ptr);
+ }
+ auto t_ptr = (uint8_t*)tensor.data_ptr();
+ for (int pass = 0; pass < number_of_passes; pass++) {
+ for (png_uint_32 i = 0; i < height; ++i) {
+ png_read_row(png_ptr, t_ptr, nullptr);
+ t_ptr += num_pixels_per_row * (is_16_bits ? 2 : 1);
  }
+ t_ptr = (uint8_t*)tensor.data_ptr();
  }
 
  int exif_orientation = -1;

torchvision/csrc/io/image/cpu/decode_png.h

@@ -9,7 +9,6 @@ namespace image {
 C10_EXPORT torch::Tensor decode_png(
  const torch::Tensor& data,
  ImageReadMode mode = IMAGE_READ_MODE_UNCHANGED,
- bool allow_16_bits = false,
  bool apply_exif_orientation = false);
 
 } // namespace image

torchvision/csrc/io/image/image.cpp

@@ -16,7 +16,7 @@ namespace image {
 static auto registry =
  torch::RegisterOperators()
  .op("image::decode_gif", &decode_gif)
- .op("image::decode_png(Tensor data, int mode, bool allow_16_bits = False, bool apply_exif_orientation=False) -> Tensor",
+ .op("image::decode_png(Tensor data, int mode, bool apply_exif_orientation=False) -> Tensor",
  &decode_png)
  .op("image::encode_png", &encode_png)
  .op("image::decode_jpeg(Tensor data, int mode, bool apply_exif_orientation=False) -> Tensor",

torchvision/datasets/_optical_flow.py

@@ -9,7 +9,7 @@
 import torch
 from PIL import Image
 
-from ..io.image import _read_png_16
+from ..io.image import decode_png, read_file
 from .utils import _read_pfm, verify_str_arg
 from .vision import VisionDataset
 
@@ -481,7 +481,7 @@ def _read_flo(file_name: str) -> np.ndarray:
 
 def _read_16bits_png_with_flow_and_valid_mask(file_name: str) -> Tuple[np.ndarray, np.ndarray]:
 
- flow_and_valid = _read_png_16(file_name).to(torch.float32)
+ flow_and_valid = decode_png(read_file(file_name)).to(torch.float32)
  flow, valid_flow_mask = flow_and_valid[:2, :, :], flow_and_valid[2, :, :]
  flow = (flow - 2**15) / 64 # This conversion is explained somewhere on the kitti archive
  valid_flow_mask = valid_flow_mask.bool()

torchvision/io/image.py

@@ -75,8 +75,14 @@ def decode_png(
 ) -> torch.Tensor:
  """
  Decodes a PNG image into a 3 dimensional RGB or grayscale Tensor.
- Optionally converts the image to the desired format.
- The values of the output tensor are uint8 in [0, 255].
+
+ The values of the output tensor are in uint8 in [0, 255] for most cases. If
+ the image is a 16-bit png, then the output tensor is uint16 in [0, 65535]
+ (supported from torchvision ``0.21``. Since uint16 support is limited in
+ pytorch, we recommend calling
+ :func:`torchvision.transforms.v2.functional.to_dtype()` with ``scale=True``
+ after this function to convert the decoded image into a uint8 or float
+ tensor.
 
  Args:
  input (Tensor[1]): a one dimensional uint8 tensor containing
@@ -93,7 +99,7 @@ def decode_png(
  """
  if not torch.jit.is_scripting() and not torch.jit.is_tracing():
  _log_api_usage_once(decode_png)
- output = torch.ops.image.decode_png(input, mode.value, False, apply_exif_orientation)
+ output = torch.ops.image.decode_png(input, mode.value, apply_exif_orientation)
  return output
 
 
@@ -144,7 +150,7 @@ def decode_jpeg(
 ) -> torch.Tensor:
  """
  Decodes a JPEG image into a 3 dimensional RGB or grayscale Tensor.
- Optionally converts the image to the desired format.
+
  The values of the output tensor are uint8 between 0 and 255.
 
  Args:
@@ -248,8 +254,13 @@ def decode_image(
  Detect whether an image is a JPEG, PNG or GIF and performs the appropriate
  operation to decode the image into a 3 dimensional RGB or grayscale Tensor.
 
- Optionally converts the image to the desired format.
- The values of the output tensor are uint8 in [0, 255].
+ The values of the output tensor are in uint8 in [0, 255] for most cases. If
+ the image is a 16-bit png, then the output tensor is uint16 in [0, 65535]
+ (supported from torchvision ``0.21``. Since uint16 support is limited in
+ pytorch, we recommend calling
+ :func:`torchvision.transforms.v2.functional.to_dtype()` with ``scale=True``
+ after this function to convert the decoded image into a uint8 or float
+ tensor.
 
  Args:
  input (Tensor): a one dimensional uint8 tensor containing the raw bytes of the
@@ -277,8 +288,14 @@ def read_image(
 ) -> torch.Tensor:
  """
  Reads a JPEG, PNG or GIF image into a 3 dimensional RGB or grayscale Tensor.
- Optionally converts the image to the desired format.
- The values of the output tensor are uint8 in [0, 255].
+
+ The values of the output tensor are in uint8 in [0, 255] for most cases. If
+ the image is a 16-bit png, then the output tensor is uint16 in [0, 65535]
+ (supported from torchvision ``0.21``. Since uint16 support is limited in
+ pytorch, we recommend calling
+ :func:`torchvision.transforms.v2.functional.to_dtype()` with ``scale=True``
+ after this function to convert the decoded image into a uint8 or float
+ tensor.
 
  Args:
  path (str or ``pathlib.Path``): path of the JPEG, PNG or GIF image.
@@ -298,11 +315,6 @@ def read_image(
  return decode_image(data, mode, apply_exif_orientation=apply_exif_orientation)
 
 
-def _read_png_16(path: str, mode: ImageReadMode = ImageReadMode.UNCHANGED) -> torch.Tensor:
- data = read_file(path)
- return torch.ops.image.decode_png(data, mode.value, True)
-
-
 def decode_gif(input: torch.Tensor) -> torch.Tensor:
  """
  Decode a GIF image into a 3 or 4 dimensional RGB Tensor.

torchvision/transforms/_functional_tensor.py

@@ -45,6 +45,8 @@ def _max_value(dtype: torch.dtype) -> int:
  return 127
  elif dtype == torch.int16:
  return 32767
+ elif dtype == torch.uint16:
+ return 65535
  elif dtype == torch.int32:
  return 2147483647
  elif dtype == torch.int64:

torchvision/transforms/v2/functional/_misc.py

@@ -237,6 +237,8 @@ def _num_value_bits(dtype: torch.dtype) -> int:
  return 7
  elif dtype == torch.int16:
  return 15
+ elif dtype == torch.uint16:
+ return 16
  elif dtype == torch.int32:
  return 31
  elif dtype == torch.int64:
@@ -293,10 +295,18 @@ def to_dtype_image(image: torch.Tensor, dtype: torch.dtype = torch.float, scale:
  num_value_bits_input = _num_value_bits(image.dtype)
  num_value_bits_output = _num_value_bits(dtype)
 
+ # TODO: Remove if/else inner blocks once uint16 dtype supports bitwise shift operations.
+ shift_by = abs(num_value_bits_input - num_value_bits_output)
  if num_value_bits_input > num_value_bits_output:
- return image.bitwise_right_shift(num_value_bits_input - num_value_bits_output).to(dtype)
+ if image.dtype == torch.uint16:
+ return (image / 2 ** (shift_by)).to(dtype)
+ else:
+ return image.bitwise_right_shift(shift_by).to(dtype)
  else:
- return image.to(dtype).bitwise_left_shift_(num_value_bits_output - num_value_bits_input)
+ if dtype == torch.uint16:
+ return image.to(dtype) * 2 ** (shift_by)
+ else:
+ return image.to(dtype).bitwise_left_shift_(shift_by)
 
 
 # We encourage users to use to_dtype() instead but we keep this for BC