From 5b0edbf9e69779b9e80590b6222f2e2eb05bb258 Mon Sep 17 00:00:00 2001 From: vfdev-5 Date: Thu, 13 Jul 2023 15:43:20 +0200 Subject: [PATCH 1/6] Added support for CMYK in decode_jpeg with mode RGB only --- test/test_image.py | 4 ++-- torchvision/csrc/io/image/cpu/decode_jpeg.cpp | 18 +++++++++++++++++- 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/test/test_image.py b/test/test_image.py index b24ac07d956..b63b688e68d 100644 --- a/test/test_image.py +++ b/test/test_image.py @@ -83,12 +83,12 @@ def test_decode_jpeg(img_path, pil_mode, mode): with Image.open(img_path) as img: is_cmyk = img.mode == "CMYK" if pil_mode is not None: - if is_cmyk: + if is_cmyk and mode == ImageReadMode.GRAY: # libjpeg does not support the conversion pytest.xfail("Decoding a CMYK jpeg isn't supported") img = img.convert(pil_mode) img_pil = torch.from_numpy(np.array(img)) - if is_cmyk: + if is_cmyk and mode == ImageReadMode.UNCHANGED: # flip the colors to match libjpeg img_pil = 255 - img_pil diff --git a/torchvision/csrc/io/image/cpu/decode_jpeg.cpp b/torchvision/csrc/io/image/cpu/decode_jpeg.cpp index d07844a5e27..a5b25a0b847 100644 --- a/torchvision/csrc/io/image/cpu/decode_jpeg.cpp +++ b/torchvision/csrc/io/image/cpu/decode_jpeg.cpp @@ -112,7 +112,10 @@ torch::Tensor decode_jpeg(const torch::Tensor& data, ImageReadMode mode) { } break; case IMAGE_READ_MODE_RGB: - if (cinfo.jpeg_color_space != JCS_RGB) { + if (cinfo.jpeg_color_space == JCS_CMYK || + cinfo.jpeg_color_space == JCS_YCCK) { + cinfo.out_color_space = JCS_CMYK; + } else { cinfo.out_color_space = JCS_RGB; channels = 3; } @@ -148,6 +151,19 @@ torch::Tensor decode_jpeg(const torch::Tensor& data, ImageReadMode mode) { ptr += stride; } + if (mode == IMAGE_READ_MODE_RGB && channels == 4) { + // convert CMYK to RGB + auto k = tensor.index({"...", 3}).unsqueeze(-1); + auto cmy = tensor.index({"...", torch::indexing::Slice(0, 3)}); + cmy = cmy.to(torch::CPU(torch::kInt32)); + if (cinfo.saw_Adobe_marker) { + tensor = (k * cmy).div(255); + } else { + tensor = ((255 - k) * (255 - cmy)).div(255); + } + tensor = tensor.to(torch::CPU(torch::kU8)); + } + jpeg_finish_decompress(&cinfo); jpeg_destroy_decompress(&cinfo); return tensor.permute({2, 0, 1}); From a6880ef61452f0b229efbe8457f5247cf86f3fb7 Mon Sep 17 00:00:00 2001 From: vfdev-5 Date: Fri, 14 Jul 2023 11:23:09 +0200 Subject: [PATCH 2/6] Recoded cmyk to rgb conversion and added support for cmyk to gray --- test/test_image.py | 3 - torchvision/csrc/io/image/cpu/decode_jpeg.cpp | 96 +++++++++++++++---- 2 files changed, 80 insertions(+), 19 deletions(-) diff --git a/test/test_image.py b/test/test_image.py index b63b688e68d..a87f5fa2d1e 100644 --- a/test/test_image.py +++ b/test/test_image.py @@ -83,9 +83,6 @@ def test_decode_jpeg(img_path, pil_mode, mode): with Image.open(img_path) as img: is_cmyk = img.mode == "CMYK" if pil_mode is not None: - if is_cmyk and mode == ImageReadMode.GRAY: - # libjpeg does not support the conversion - pytest.xfail("Decoding a CMYK jpeg isn't supported") img = img.convert(pil_mode) img_pil = torch.from_numpy(np.array(img)) if is_cmyk and mode == ImageReadMode.UNCHANGED: diff --git a/torchvision/csrc/io/image/cpu/decode_jpeg.cpp b/torchvision/csrc/io/image/cpu/decode_jpeg.cpp index a5b25a0b847..9ca32a944ff 100644 --- a/torchvision/csrc/io/image/cpu/decode_jpeg.cpp +++ b/torchvision/csrc/io/image/cpu/decode_jpeg.cpp @@ -67,6 +67,61 @@ static void torch_jpeg_set_source_mgr( src->pub.next_input_byte = src->data; } +inline void _convert_pixel_cmyk_to_rgb( + bool has_adobe_marker, + int c, + int m, + int y, + int k, + int& r, + int& g, + int& b) { + r = (has_adobe_marker) ? (k * c) / 255 : (255 - k) * (255 - c) / 255; + g = (has_adobe_marker) ? (k * m) / 255 : (255 - k) * (255 - m) / 255; + b = (has_adobe_marker) ? (k * y) / 255 : (255 - k) * (255 - y) / 255; +} + +void convert_line_cmyk_to_rgb( + j_decompress_ptr cinfo, + const unsigned char* cmyk_line, + unsigned char* rgb_line) { + auto has_adobe_marker = cinfo->saw_Adobe_marker; + int width = cinfo->output_width; + for (int i = 0; i < width; ++i) { + int r, g, b; + int c = cmyk_line[i * 4 + 0]; + int m = cmyk_line[i * 4 + 1]; + int y = cmyk_line[i * 4 + 2]; + int k = cmyk_line[i * 4 + 3]; + + _convert_pixel_cmyk_to_rgb(has_adobe_marker, c, m, y, k, r, g, b); + + rgb_line[i * 3 + 0] = r; + rgb_line[i * 3 + 1] = g; + rgb_line[i * 3 + 2] = b; + } +} + +void convert_line_cmyk_to_gray( + j_decompress_ptr cinfo, + const unsigned char* cmyk_line, + unsigned char* gray_line) { + auto has_adobe_marker = cinfo->saw_Adobe_marker; + int width = cinfo->output_width; + for (int i = 0; i < width; ++i) { + int r, g, b; + int c = cmyk_line[i * 4 + 0]; + int m = cmyk_line[i * 4 + 1]; + int y = cmyk_line[i * 4 + 2]; + int k = cmyk_line[i * 4 + 3]; + + _convert_pixel_cmyk_to_rgb(has_adobe_marker, c, m, y, k, r, g, b); + + float gray = 0.2989 * (float)r + 0.5870 * (float)g + 0.1140 * (float)b; + gray_line[i] = (unsigned char)gray; + } +} + } // namespace torch::Tensor decode_jpeg(const torch::Tensor& data, ImageReadMode mode) { @@ -102,23 +157,29 @@ torch::Tensor decode_jpeg(const torch::Tensor& data, ImageReadMode mode) { jpeg_read_header(&cinfo, TRUE); int channels = cinfo.num_components; + bool cmyk_to_rgb = false; if (mode != IMAGE_READ_MODE_UNCHANGED) { switch (mode) { case IMAGE_READ_MODE_GRAY: - if (cinfo.jpeg_color_space != JCS_GRAYSCALE) { + if (cinfo.jpeg_color_space == JCS_CMYK || + cinfo.jpeg_color_space == JCS_YCCK) { + cinfo.out_color_space = JCS_CMYK; + cmyk_to_rgb = true; + } else { cinfo.out_color_space = JCS_GRAYSCALE; - channels = 1; } + channels = 1; break; case IMAGE_READ_MODE_RGB: if (cinfo.jpeg_color_space == JCS_CMYK || cinfo.jpeg_color_space == JCS_YCCK) { cinfo.out_color_space = JCS_CMYK; + cmyk_to_rgb = true; } else { cinfo.out_color_space = JCS_RGB; - channels = 3; } + channels = 3; break; /* * Libjpeg does not support converting from CMYK to grayscale etc. There @@ -142,26 +203,29 @@ torch::Tensor decode_jpeg(const torch::Tensor& data, ImageReadMode mode) { auto tensor = torch::empty({int64_t(height), int64_t(width), channels}, torch::kU8); auto ptr = tensor.data_ptr(); + torch::Tensor temp_tensor; + if (cmyk_to_rgb) { + temp_tensor = torch::empty({int64_t(width), 4}, torch::kU8); + } + while (cinfo.output_scanline < cinfo.output_height) { /* jpeg_read_scanlines expects an array of pointers to scanlines. * Here the array is only one element long, but you could ask for * more than one scanline at a time if that's more convenient. */ - jpeg_read_scanlines(&cinfo, &ptr, 1); - ptr += stride; - } - - if (mode == IMAGE_READ_MODE_RGB && channels == 4) { - // convert CMYK to RGB - auto k = tensor.index({"...", 3}).unsqueeze(-1); - auto cmy = tensor.index({"...", torch::indexing::Slice(0, 3)}); - cmy = cmy.to(torch::CPU(torch::kInt32)); - if (cinfo.saw_Adobe_marker) { - tensor = (k * cmy).div(255); + if (cmyk_to_rgb) { + auto temp_buffer = temp_tensor.data_ptr(); + jpeg_read_scanlines(&cinfo, &temp_buffer, 1); + + if (channels == 3) { + convert_line_cmyk_to_rgb(&cinfo, temp_buffer, ptr); + } else if (channels == 1) { + convert_line_cmyk_to_gray(&cinfo, temp_buffer, ptr); + } } else { - tensor = ((255 - k) * (255 - cmy)).div(255); + jpeg_read_scanlines(&cinfo, &ptr, 1); } - tensor = tensor.to(torch::CPU(torch::kU8)); + ptr += stride; } jpeg_finish_decompress(&cinfo); From bb46b5c1537b0d9dcae67876c4b4d38802b0ce65 Mon Sep 17 00:00:00 2001 From: vfdev Date: Fri, 28 Jul 2023 01:36:04 +0200 Subject: [PATCH 3/6] Update decode_jpeg.cpp --- torchvision/csrc/io/image/cpu/decode_jpeg.cpp | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/torchvision/csrc/io/image/cpu/decode_jpeg.cpp b/torchvision/csrc/io/image/cpu/decode_jpeg.cpp index 9ca32a944ff..d39cae407dd 100644 --- a/torchvision/csrc/io/image/cpu/decode_jpeg.cpp +++ b/torchvision/csrc/io/image/cpu/decode_jpeg.cpp @@ -157,7 +157,7 @@ torch::Tensor decode_jpeg(const torch::Tensor& data, ImageReadMode mode) { jpeg_read_header(&cinfo, TRUE); int channels = cinfo.num_components; - bool cmyk_to_rgb = false; + bool cmyk_to_rgb_or_gray = false; if (mode != IMAGE_READ_MODE_UNCHANGED) { switch (mode) { @@ -165,7 +165,7 @@ torch::Tensor decode_jpeg(const torch::Tensor& data, ImageReadMode mode) { if (cinfo.jpeg_color_space == JCS_CMYK || cinfo.jpeg_color_space == JCS_YCCK) { cinfo.out_color_space = JCS_CMYK; - cmyk_to_rgb = true; + cmyk_to_rgb_or_gray = true; } else { cinfo.out_color_space = JCS_GRAYSCALE; } @@ -175,7 +175,7 @@ torch::Tensor decode_jpeg(const torch::Tensor& data, ImageReadMode mode) { if (cinfo.jpeg_color_space == JCS_CMYK || cinfo.jpeg_color_space == JCS_YCCK) { cinfo.out_color_space = JCS_CMYK; - cmyk_to_rgb = true; + cmyk_to_rgb_or_gray = true; } else { cinfo.out_color_space = JCS_RGB; } @@ -204,7 +204,7 @@ torch::Tensor decode_jpeg(const torch::Tensor& data, ImageReadMode mode) { torch::empty({int64_t(height), int64_t(width), channels}, torch::kU8); auto ptr = tensor.data_ptr(); torch::Tensor temp_tensor; - if (cmyk_to_rgb) { + if (cmyk_to_rgb_or_gray) { temp_tensor = torch::empty({int64_t(width), 4}, torch::kU8); } @@ -213,7 +213,7 @@ torch::Tensor decode_jpeg(const torch::Tensor& data, ImageReadMode mode) { * Here the array is only one element long, but you could ask for * more than one scanline at a time if that's more convenient. */ - if (cmyk_to_rgb) { + if (cmyk_to_rgb_or_gray) { auto temp_buffer = temp_tensor.data_ptr(); jpeg_read_scanlines(&cinfo, &temp_buffer, 1); From 4eecbe66e83562bb00d2af73996de424cbb2c118 Mon Sep 17 00:00:00 2001 From: vfdev-5 Date: Fri, 11 Aug 2023 10:06:57 +0200 Subject: [PATCH 4/6] Make CMYK -> {RGB, GRAY} decoding match exactly Pillow output --- torchvision/csrc/io/image/cpu/decode_jpeg.cpp | 43 ++++++++----------- 1 file changed, 19 insertions(+), 24 deletions(-) diff --git a/torchvision/csrc/io/image/cpu/decode_jpeg.cpp b/torchvision/csrc/io/image/cpu/decode_jpeg.cpp index d39cae407dd..d654ac13487 100644 --- a/torchvision/csrc/io/image/cpu/decode_jpeg.cpp +++ b/torchvision/csrc/io/image/cpu/decode_jpeg.cpp @@ -67,58 +67,53 @@ static void torch_jpeg_set_source_mgr( src->pub.next_input_byte = src->data; } -inline void _convert_pixel_cmyk_to_rgb( - bool has_adobe_marker, - int c, - int m, - int y, - int k, - int& r, - int& g, - int& b) { - r = (has_adobe_marker) ? (k * c) / 255 : (255 - k) * (255 - c) / 255; - g = (has_adobe_marker) ? (k * m) / 255 : (255 - k) * (255 - m) / 255; - b = (has_adobe_marker) ? (k * y) / 255 : (255 - k) * (255 - y) / 255; +inline unsigned char clamped_cmyk_rgb_convert(unsigned char k, unsigned char cmy) { + // Inspired from Pillow: + // https://github.com/python-pillow/Pillow/blob/07623d1a7cc65206a5355fba2ae256550bfcaba6/src/libImaging/Convert.c#L568-L569 + auto v = k * cmy + 128; + v = ((v >> 8) + v) >> 8; + return std::clamp(k - v, 0, 255); } void convert_line_cmyk_to_rgb( j_decompress_ptr cinfo, const unsigned char* cmyk_line, unsigned char* rgb_line) { - auto has_adobe_marker = cinfo->saw_Adobe_marker; int width = cinfo->output_width; for (int i = 0; i < width; ++i) { - int r, g, b; int c = cmyk_line[i * 4 + 0]; int m = cmyk_line[i * 4 + 1]; int y = cmyk_line[i * 4 + 2]; int k = cmyk_line[i * 4 + 3]; - _convert_pixel_cmyk_to_rgb(has_adobe_marker, c, m, y, k, r, g, b); - - rgb_line[i * 3 + 0] = r; - rgb_line[i * 3 + 1] = g; - rgb_line[i * 3 + 2] = b; + rgb_line[i * 3 + 0] = clamped_cmyk_rgb_convert(k, 255 - c); + rgb_line[i * 3 + 1] = clamped_cmyk_rgb_convert(k, 255 - m); + rgb_line[i * 3 + 2] = clamped_cmyk_rgb_convert(k, 255 - y); } } +inline unsigned char rgb_to_gray(int r, int g, int b) { + // Inspired from Pillow: + // https://github.com/python-pillow/Pillow/blob/07623d1a7cc65206a5355fba2ae256550bfcaba6/src/libImaging/Convert.c#L226 + return (r * 19595 + g * 38470 + b * 7471 + 0x8000) >> 16; +} + void convert_line_cmyk_to_gray( j_decompress_ptr cinfo, const unsigned char* cmyk_line, unsigned char* gray_line) { - auto has_adobe_marker = cinfo->saw_Adobe_marker; int width = cinfo->output_width; for (int i = 0; i < width; ++i) { - int r, g, b; int c = cmyk_line[i * 4 + 0]; int m = cmyk_line[i * 4 + 1]; int y = cmyk_line[i * 4 + 2]; int k = cmyk_line[i * 4 + 3]; - _convert_pixel_cmyk_to_rgb(has_adobe_marker, c, m, y, k, r, g, b); + int r = clamped_cmyk_rgb_convert(k, 255 - c); + int g = clamped_cmyk_rgb_convert(k, 255 - m); + int b = clamped_cmyk_rgb_convert(k, 255 - y); - float gray = 0.2989 * (float)r + 0.5870 * (float)g + 0.1140 * (float)b; - gray_line[i] = (unsigned char)gray; + gray_line[i] = rgb_to_gray(r, g, b); } } From e171488cdc337faca48d7b1b001db31079df8865 Mon Sep 17 00:00:00 2001 From: vfdev-5 Date: Fri, 11 Aug 2023 19:14:52 +0200 Subject: [PATCH 5/6] Fixed c-linting failure --- torchvision/csrc/io/image/cpu/decode_jpeg.cpp | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/torchvision/csrc/io/image/cpu/decode_jpeg.cpp b/torchvision/csrc/io/image/cpu/decode_jpeg.cpp index d654ac13487..86ce30fe1ad 100644 --- a/torchvision/csrc/io/image/cpu/decode_jpeg.cpp +++ b/torchvision/csrc/io/image/cpu/decode_jpeg.cpp @@ -67,7 +67,9 @@ static void torch_jpeg_set_source_mgr( src->pub.next_input_byte = src->data; } -inline unsigned char clamped_cmyk_rgb_convert(unsigned char k, unsigned char cmy) { +inline unsigned char clamped_cmyk_rgb_convert( + unsigned char k, + unsigned char cmy) { // Inspired from Pillow: // https://github.com/python-pillow/Pillow/blob/07623d1a7cc65206a5355fba2ae256550bfcaba6/src/libImaging/Convert.c#L568-L569 auto v = k * cmy + 128; From 54da7b280a5c322d1760ac433039d57f5203df32 Mon Sep 17 00:00:00 2001 From: vfdev Date: Thu, 17 Aug 2023 10:56:12 +0200 Subject: [PATCH 6/6] Update decode_jpeg.cpp --- torchvision/csrc/io/image/cpu/decode_jpeg.cpp | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/torchvision/csrc/io/image/cpu/decode_jpeg.cpp b/torchvision/csrc/io/image/cpu/decode_jpeg.cpp index 86ce30fe1ad..71af7c5d80c 100644 --- a/torchvision/csrc/io/image/cpu/decode_jpeg.cpp +++ b/torchvision/csrc/io/image/cpu/decode_jpeg.cpp @@ -72,7 +72,7 @@ inline unsigned char clamped_cmyk_rgb_convert( unsigned char cmy) { // Inspired from Pillow: // https://github.com/python-pillow/Pillow/blob/07623d1a7cc65206a5355fba2ae256550bfcaba6/src/libImaging/Convert.c#L568-L569 - auto v = k * cmy + 128; + int v = k * cmy + 128; v = ((v >> 8) + v) >> 8; return std::clamp(k - v, 0, 255); } @@ -200,9 +200,9 @@ torch::Tensor decode_jpeg(const torch::Tensor& data, ImageReadMode mode) { auto tensor = torch::empty({int64_t(height), int64_t(width), channels}, torch::kU8); auto ptr = tensor.data_ptr(); - torch::Tensor temp_tensor; + torch::Tensor cmyk_line_tensor; if (cmyk_to_rgb_or_gray) { - temp_tensor = torch::empty({int64_t(width), 4}, torch::kU8); + cmyk_line_tensor = torch::empty({int64_t(width), 4}, torch::kU8); } while (cinfo.output_scanline < cinfo.output_height) { @@ -211,13 +211,13 @@ torch::Tensor decode_jpeg(const torch::Tensor& data, ImageReadMode mode) { * more than one scanline at a time if that's more convenient. */ if (cmyk_to_rgb_or_gray) { - auto temp_buffer = temp_tensor.data_ptr(); - jpeg_read_scanlines(&cinfo, &temp_buffer, 1); + auto cmyk_line_ptr = cmyk_line_tensor.data_ptr(); + jpeg_read_scanlines(&cinfo, &cmyk_line_ptr, 1); if (channels == 3) { - convert_line_cmyk_to_rgb(&cinfo, temp_buffer, ptr); + convert_line_cmyk_to_rgb(&cinfo, cmyk_line_ptr, ptr); } else if (channels == 1) { - convert_line_cmyk_to_gray(&cinfo, temp_buffer, ptr); + convert_line_cmyk_to_gray(&cinfo, cmyk_line_ptr, ptr); } } else { jpeg_read_scanlines(&cinfo, &ptr, 1);