Use indexed color for PNG images in PDF files when possible.

mpetroff · mpetroff · commit 40132e725580 · 2020-07-12T02:16:37.000Z
When PNG images have 256 colors or fewer, convert them to index color before
saving them in a PDF. This can result in a signifcant reduction in file size.
diff --git a/lib/matplotlib/backends/backend_pdf.py b/lib/matplotlib/backends/backend_pdf.py
@@ -1521,28 +1521,33 @@ def _unpack(self, im):
                 alpha = None
             return rgb, alpha
 
-    def _writePng(self, data):
+    def _writePng(self, img):
         """
-        Write the image *data* into the pdf file using png
+        Write the image *img* into the pdf file using png
         predictors with Flate compression.
         """
         buffer = BytesIO()
-        if data.shape[-1] == 1:
-            data = data.squeeze(axis=-1)
-        Image.fromarray(data).save(buffer, format="png")
+        img.save(buffer, format="png")
         buffer.seek(8)
+        png_data = bit_depth = palette = None
         while True:
             length, type = struct.unpack(b'!L4s', buffer.read(8))
-            if type == b'IDAT':
+            if type in [b'IHDR', b'PLTE', b'IDAT']:
                 data = buffer.read(length)
                 if len(data) != length:
                     raise RuntimeError("truncated data")
-                self.currentstream.write(data)
+                if type == b'IHDR':
+                    bit_depth = int(data[8])
+                elif type == b'PLTE':
+                    palette = data
+                elif type == b'IDAT':
+                    png_data = data
             elif type == b'IEND':
                 break
             else:
                 buffer.seek(length, 1)
             buffer.seek(4, 1)   # skip CRC
+        return png_data, bit_depth, palette
 
     def _writeImg(self, data, id, smask=None):
         """
@@ -1561,6 +1566,34 @@ def _writeImg(self, data, id, smask=None):
         if smask:
             obj['SMask'] = smask
         if mpl.rcParams['pdf.compression']:
+            if data.shape[-1] == 1:
+                data = data.squeeze(axis=-1)
+            img = Image.fromarray(data)
+            img_colors = img.getcolors(maxcolors=256)
+            if img_colors is not None:
+                # Convert to indexed color if there are 256 colors or fewer
+                # This can significantly reduce the file size
+                num_colors = len(img_colors)
+                img = img.convert(mode='P', dither=Image.NONE,
+                                  palette=Image.ADAPTIVE, colors=num_colors)
+                data, bit_depth, palette = self._writePng(img)
+                if bit_depth is None or palette is None:
+                    raise RuntimeError("invalid PNG header")
+                palette = palette[:num_colors * 3]  # Trim padding
+                if colors == 1:
+                    # The PNG format uses an RGB palette for all indexed color
+                    # images, but the PDF format allows for grayscale palettes.
+                    # Thus, we convert the palette.
+                    palette = palette[::3]
+                palette = pdfRepr(palette)
+                colorspace = obj['ColorSpace'].pdfRepr()
+                obj['ColorSpace'] = Verbatim(b'[/Indexed ' + colorspace + b' '
+                                             + str(num_colors - 1).encode()
+                                             + b' ' + palette + b']')
+                obj['BitsPerComponent'] = bit_depth
+                colors = 1
+            else:
+                data, _, _ = self._writePng(img)
             png = {'Predictor': 10, 'Colors': colors, 'Columns': width}
         else:
             png = None
@@ -1571,7 +1604,7 @@ def _writeImg(self, data, id, smask=None):
             png=png
             )
         if png:
-            self._writePng(data)
+            self.currentstream.write(data)
         else:
             self.currentstream.write(data.tobytes())
         self.endStream()
