Merge pull request #7553 from bgilbert/jpeg-rgb

radarhere · web-flow · commit d93a5ad70bf9 · 2024-01-01T23:55:37.000+11:00
Add `keep_rgb` option when saving JPEG to prevent conversion of RGB colorspace
diff --git a/Tests/test_file_jpeg.py b/Tests/test_file_jpeg.py
@@ -142,6 +142,19 @@ def test_cmyk(self):
             )
             assert k > 0.9
 
+    def test_rgb(self):
+        def getchannels(im):
+            return tuple(v[0] for v in im.layer)
+
+        im = hopper()
+        im_ycbcr = self.roundtrip(im)
+        assert getchannels(im_ycbcr) == (1, 2, 3)
+        assert_image_similar(im, im_ycbcr, 17)
+
+        im_rgb = self.roundtrip(im, keep_rgb=True)
+        assert getchannels(im_rgb) == (ord("R"), ord("G"), ord("B"))
+        assert_image_similar(im, im_rgb, 12)
+
     @pytest.mark.parametrize(
         "test_image_path",
         [TEST_FILE, "Tests/images/pil_sample_cmyk.jpg"],
@@ -423,25 +436,28 @@ def getsampling(im):
             return layer[0][1:3] + layer[1][1:3] + layer[2][1:3]
 
         # experimental API
-        im = self.roundtrip(hopper(), subsampling=-1)  # default
-        assert getsampling(im) == (2, 2, 1, 1, 1, 1)
-        im = self.roundtrip(hopper(), subsampling=0)  # 4:4:4
-        assert getsampling(im) == (1, 1, 1, 1, 1, 1)
-        im = self.roundtrip(hopper(), subsampling=1)  # 4:2:2
-        assert getsampling(im) == (2, 1, 1, 1, 1, 1)
-        im = self.roundtrip(hopper(), subsampling=2)  # 4:2:0
-        assert getsampling(im) == (2, 2, 1, 1, 1, 1)
-        im = self.roundtrip(hopper(), subsampling=3)  # default (undefined)
-        assert getsampling(im) == (2, 2, 1, 1, 1, 1)
-
-        im = self.roundtrip(hopper(), subsampling="4:4:4")
-        assert getsampling(im) == (1, 1, 1, 1, 1, 1)
-        im = self.roundtrip(hopper(), subsampling="4:2:2")
-        assert getsampling(im) == (2, 1, 1, 1, 1, 1)
-        im = self.roundtrip(hopper(), subsampling="4:2:0")
-        assert getsampling(im) == (2, 2, 1, 1, 1, 1)
-        im = self.roundtrip(hopper(), subsampling="4:1:1")
-        assert getsampling(im) == (2, 2, 1, 1, 1, 1)
+        for subsampling in (-1, 3):  # (default, invalid)
+            im = self.roundtrip(hopper(), subsampling=subsampling)
+            assert getsampling(im) == (2, 2, 1, 1, 1, 1)
+        for subsampling in (0, "4:4:4"):
+            im = self.roundtrip(hopper(), subsampling=subsampling)
+            assert getsampling(im) == (1, 1, 1, 1, 1, 1)
+        for subsampling in (1, "4:2:2"):
+            im = self.roundtrip(hopper(), subsampling=subsampling)
+            assert getsampling(im) == (2, 1, 1, 1, 1, 1)
+        for subsampling in (2, "4:2:0", "4:1:1"):
+            im = self.roundtrip(hopper(), subsampling=subsampling)
+            assert getsampling(im) == (2, 2, 1, 1, 1, 1)
+
+        # RGB colorspace
+        for subsampling in (-1, 0, "4:4:4"):
+            # "4:4:4" doesn't really make sense for RGB, but the conversion
+            # to an integer happens at a higher level
+            im = self.roundtrip(hopper(), keep_rgb=True, subsampling=subsampling)
+            assert getsampling(im) == (1, 1, 1, 1, 1, 1)
+        for subsampling in (1, "4:2:2", 2, "4:2:0", 3):
+            with pytest.raises(OSError):
+                self.roundtrip(hopper(), keep_rgb=True, subsampling=subsampling)
 
         with pytest.raises(TypeError):
             self.roundtrip(hopper(), subsampling="1:1:1")
diff --git a/docs/handbook/image-file-formats.rst b/docs/handbook/image-file-formats.rst
@@ -487,6 +487,16 @@ The :py:meth:`~PIL.Image.Image.save` method supports the following options:
 **exif**
     If present, the image will be stored with the provided raw EXIF data.
 
+**keep_rgb**
+    By default, libjpeg converts images with an RGB color space to YCbCr.
+    If this option is present and true, those images will be stored as RGB
+    instead.
+
+    When this option is enabled, attempting to chroma-subsample RGB images
+    with the ``subsampling`` option will raise an :py:exc:`OSError`.
+
+    .. versionadded:: 10.2.0
+
 **subsampling**
     If present, sets the subsampling for the encoder.
 
diff --git a/docs/releasenotes/10.2.0.rst b/docs/releasenotes/10.2.0.rst
@@ -42,6 +42,14 @@ Added DdsImagePlugin enums
 :py:class:`~PIL.DdsImagePlugin.DXGI_FORMAT` and :py:class:`~PIL.DdsImagePlugin.D3DFMT`
 enums have been added to :py:class:`PIL.DdsImagePlugin`.
 
+JPEG RGB color space
+^^^^^^^^^^^^^^^^^^^^
+
+When saving JPEG files, ``keep_rgb`` can now be set to ``True``. This will store RGB
+images in the RGB color space instead of being converted to YCbCr automatically by
+libjpeg. When this option is enabled, attempting to chroma-subsample RGB images with
+the ``subsampling`` option will raise an :py:exc:`OSError`.
+
 JPEG restart marker interval
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
diff --git a/src/PIL/JpegImagePlugin.py b/src/PIL/JpegImagePlugin.py
@@ -785,6 +785,7 @@ def validate_qtables(qtables):
         progressive,
         info.get("smooth", 0),
         optimize,
+        info.get("keep_rgb", False),
         info.get("streamtype", 0),
         dpi[0],
         dpi[1],
diff --git a/src/encode.c b/src/encode.c
@@ -1042,6 +1042,7 @@ PyImaging_JpegEncoderNew(PyObject *self, PyObject *args) {
     Py_ssize_t progressive = 0;
     Py_ssize_t smooth = 0;
     Py_ssize_t optimize = 0;
+    int keep_rgb = 0;
     Py_ssize_t streamtype = 0; /* 0=interchange, 1=tables only, 2=image only */
     Py_ssize_t xdpi = 0, ydpi = 0;
     Py_ssize_t subsampling = -1; /* -1=default, 0=none, 1=medium, 2=high */
@@ -1059,13 +1060,14 @@ PyImaging_JpegEncoderNew(PyObject *self, PyObject *args) {
 
     if (!PyArg_ParseTuple(
             args,
-            "ss|nnnnnnnnnnOz#y#y#",
+            "ss|nnnnpnnnnnnOz#y#y#",
             &mode,
             &rawmode,
             &quality,
             &progressive,
             &smooth,
             &optimize,
+            &keep_rgb,
             &streamtype,
             &xdpi,
             &ydpi,
@@ -1150,6 +1152,7 @@ PyImaging_JpegEncoderNew(PyObject *self, PyObject *args) {
 
     strncpy(((JPEGENCODERSTATE *)encoder->state.context)->rawmode, rawmode, 8);
 
+    ((JPEGENCODERSTATE *)encoder->state.context)->keep_rgb = keep_rgb;
     ((JPEGENCODERSTATE *)encoder->state.context)->quality = quality;
     ((JPEGENCODERSTATE *)encoder->state.context)->qtables = qarrays;
     ((JPEGENCODERSTATE *)encoder->state.context)->qtablesLen = qtablesLen;
diff --git a/src/libImaging/Jpeg.h b/src/libImaging/Jpeg.h
@@ -74,6 +74,9 @@ typedef struct {
     /* Optimize Huffman tables (slow) */
     int optimize;
 
+    /* Disable automatic conversion of RGB images to YCbCr if nonzero */
+    int keep_rgb;
+
     /* Stream type (0=full, 1=tables only, 2=image only) */
     int streamtype;
 
diff --git a/src/libImaging/JpegEncode.c b/src/libImaging/JpegEncode.c
@@ -137,6 +137,30 @@ ImagingJpegEncode(Imaging im, ImagingCodecState state, UINT8 *buf, int bytes) {
             /* Compressor configuration */
             jpeg_set_defaults(&context->cinfo);
 
+            /* Prevent RGB -> YCbCr conversion */
+            if (context->keep_rgb) {
+                switch (context->cinfo.in_color_space) {
+                    case JCS_RGB:
+#ifdef JCS_EXTENSIONS
+                    case JCS_EXT_RGBX:
+#endif
+                        switch (context->subsampling) {
+                            case -1:  /* Default */
+                            case 0:   /* No subsampling */
+                                break;
+                            default:
+                                /* Would subsample the green and blue
+                                   channels, which doesn't make sense */
+                                state->errcode = IMAGING_CODEC_CONFIG;
+                                return -1;
+                        }
+                        jpeg_set_colorspace(&context->cinfo, JCS_RGB);
+                        break;
+                    default:
+                        break;
+                }
+            }
+
             /* Use custom quantization tables */
             if (context->qtables) {
                 int i;
