Raise proper error when decoding 16-bits jpegs (#4101)

95966689 · cyy · GitHub · d1ab583d · 95966689
Unverified Commit 95966689 authored Jun 24, 2021 by cyy Committed by GitHub Jun 24, 2021
Hide whitespace changes
Inline Side-by-side

Showing with 5 additions and 0 deletions

torchvision/csrc/io/image/cpu/decode_png.cpp torchvision/csrc/io/image/cpu/decode_png.cpp +5 -0

No files found.
--- a/torchvision/csrc/io/image/cpu/decode_png.cpp
+++ b/torchvision/csrc/io/image/cpu/decode_png.cpp
@@ -71,6 +71,11 @@ torch::Tensor decode_png(const torch::Tensor& data, ImageReadMode mode) {
    TORCH_CHECK(retval == 1, "Could read image metadata from content.")
  }

+  if (bit_depth > 8) {
+    png_destroy_read_struct(&png_ptr, &info_ptr, nullptr);
+    TORCH_CHECK(false, "At most 8-bit PNG images are supported currently.")
+  }
+
  int channels = png_get_channels(png_ptr, info_ptr);

  if (color_type == PNG_COLOR_TYPE_GRAY && bit_depth < 8)