Merge pull request #1921 from thebostik/decode_jpeg_qtable

More fixes related to custom jpeg qtables
python-pillow · May 24, 2016 · 5839c3b · 5839c3b
2 parents 7460708 + 0bc4423
commit 5839c3b
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 6 deletions.
diff --git a/PIL/JpegImagePlugin.py b/PIL/JpegImagePlugin.py
@@ -195,7 +195,7 @@ def DQT(self, marker):
             raise SyntaxError("bad quantization table marker")
         v = i8(s[0])
         if v//16 == 0:
-            self.quantization[v & 15] = array.array("b", s[1:65])
+            self.quantization[v & 15] = array.array("B", s[1:65])
             s = s[65:]
         else:
             return  # FIXME: add code to read 16-bit tables!

diff --git a/Tests/test_file_jpeg.py b/Tests/test_file_jpeg.py
@@ -303,6 +303,15 @@ def test_junk_jpeg_header(self):
         filename = "Tests/images/junk_jpeg_header.jpg"
         Image.open(filename)
 
+    def _n_qtables_helper(self, n, test_file):
+        im = Image.open(test_file)
+        f = self.tempfile('temp.jpg')
+        im.save(f, qtables=[[n]*64]*n)
+        im = Image.open(f)
+        self.assertEqual(len(im.quantization), n)
+        reloaded = self.roundtrip(im, qtables="keep")
+        self.assertEqual(im.quantization, reloaded.quantization)
+
     def test_qtables(self):
         im = Image.open("Tests/images/hopper.jpg")
         qtables = im.quantization
@@ -359,6 +368,15 @@ def test_qtables(self):
                                       1: standard_chrominance_qtable
                                   }), 30)
 
+        self._n_qtables_helper(1, "Tests/images/hopper_gray.jpg")
+        self._n_qtables_helper(1, "Tests/images/pil_sample_rgb.jpg")
+        self._n_qtables_helper(2, "Tests/images/pil_sample_rgb.jpg")
+        self._n_qtables_helper(3, "Tests/images/pil_sample_rgb.jpg")
+        self._n_qtables_helper(1, "Tests/images/pil_sample_cmyk.jpg")
+        self._n_qtables_helper(2, "Tests/images/pil_sample_cmyk.jpg")
+        self._n_qtables_helper(3, "Tests/images/pil_sample_cmyk.jpg")
+        self._n_qtables_helper(4, "Tests/images/pil_sample_cmyk.jpg")
+
         # not a sequence
         self.assertRaises(Exception, lambda: self.roundtrip(im, qtables='a'))
         # sequence wrong length

diff --git a/libImaging/JpegEncode.c b/libImaging/JpegEncode.c
@@ -148,14 +148,25 @@ ImagingJpegEncode(Imaging im, ImagingCodecState state, UINT8* buf, int bytes)
 	if (context->qtables) {
 	    int i;
 	    int quality = 100;
+	    int last_q = 0;
 	    if (context->quality > 0) {
 		quality = context->quality;
 	    }
-	    for (i = 0; i < context->qtablesLen; i++) {
-		// TODO: Should add support for none baseline
-		jpeg_add_quant_table(&context->cinfo, i, &context->qtables[i * DCTSIZE2],
-				     quality, TRUE);
-	    }
+        for (i = 0; i < context->qtablesLen; i++) {
+            // TODO: Should add support for none baseline
+            jpeg_add_quant_table(&context->cinfo, i, &context->qtables[i * DCTSIZE2],
+                                 quality, TRUE);
+            context->cinfo.comp_info[i].quant_tbl_no = i;
+            last_q = i;
+        }
+        if (context->qtablesLen == 1) {
+            // jpeg_set_defaults created two qtables internally, but we only wanted one.
+            jpeg_add_quant_table(&context->cinfo, 1, &context->qtables[0],
+                                 quality, TRUE);
+        }
+        for (i = last_q; i < context->cinfo.num_components; i++) {
+            context->cinfo.comp_info[i].quant_tbl_no = last_q;
+        }
 	} else if (context->quality > 0) {
 	    jpeg_set_quality(&context->cinfo, context->quality, 1);
 	}