Merge JPEG compression

Author: Irina Nicolae

art/defences/jpeg_compression.py

@@ -0,0 +1,123 @@
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+from io import BytesIO
+import logging
+
+import numpy as np
+from PIL import Image
+
+from art.defences.preprocessor import Preprocessor
+from art import NUMPY_DTYPE
+
+logger = logging.getLogger(__name__)
+
+
+class JpegCompression(Preprocessor):
+    """
+    Implement the jpeg compression defence approach.
+    """
+    params = ['quality', 'channel_index']
+
+    def __init__(self, quality=50, channel_index=3):
+        """
+        Create an instance of jpeg compression.
+
+        :param quality: The image quality, on a scale from 1 (worst) to 95 (best). Values above 95 should be avoided.
+        :type quality: `int`
+        :param channel_index: Index of the axis in data containing the color channels or features.
+        :type channel_index: `int`
+        """
+        super(JpegCompression, self).__init__()
+        self._is_fitted = True
+        self.set_params(quality=quality, channel_index=channel_index)
+
+    def __call__(self, x, y=None, quality=None):
+        """
+        Apply jpeg compression to sample `x`.
+
+        :param x: Sample to compress with shape `(batch_size, width, height, depth)`.
+        :type x: `np.ndarray`
+        :param y: Labels of the sample `x`. This function does not affect them in any way.
+        :type y: `np.ndarray`
+        :param quality: The image quality, on a scale from 1 (worst) to 95 (best). Values above 95 should be avoided.
+        :type quality: `int`
+        :return: compressed sample
+        :rtype: `np.ndarray`
+        """
+        if quality is not None:
+            self.set_params(quality=quality)
+
+        assert self.channel_index < len(x.shape)
+
+        # Swap channel index
+        if self.channel_index < 3:
+            x_ = np.swapaxes(x, self.channel_index, 3)
+        else:
+            x_ = x.copy()
+
+        # Convert into `uint8`
+        x_ = x_ * 255
+        x_ = x_.astype("uint8")
+
+        # Convert to 'L' mode
+        if x_.shape[-1] == 1:
+            x_ = np.reshape(x_, x_.shape[:-1])
+
+        # Compress one image per time
+        for i, xi in enumerate(x_):
+            if len(xi.shape) == 2:
+                xi = Image.fromarray(xi, mode='L')
+            elif xi.shape[-1] == 3:
+                xi = Image.fromarray(xi, mode='RGB')
+            else:
+                logger.log(level=40, msg="Currently only support `RGB` and `L` images.")
+                raise NotImplementedError("Currently only support `RGB` and `L` images.")
+
+            out = BytesIO()
+            xi.save(out, format="jpeg", quality=self.quality)
+            xi = Image.open(out)
+            xi = np.array(xi)
+            x_[i] = xi
+            del out
+
+        # Expand dim if black/white images
+        if len(x_.shape) < 4:
+            x_ = np.expand_dims(x_, 3)
+
+        # Convert to old dtype
+        x_ = x_ / 255.0
+        x_ = x_.astype(NUMPY_DTYPE)
+
+        # Swap channel index
+        if self.channel_index < 3:
+            x_ = np.swapaxes(x_, self.channel_index, 3)
+
+        return x_
+
+    def fit(self, x, y=None, **kwargs):
+        """
+        No parameters to learn for this method; do nothing.
+        """
+        pass
+
+    def set_params(self, **kwargs):
+        """
+        Take in a dictionary of parameters and applies defence-specific checks before saving them as attributes.
+
+        :param quality: The image quality, on a scale from 1 (worst) to 95 (best). Values above 95 should be avoided.
+        :type quality: `int`
+        :param channel_index: Index of the axis in data containing the color channels or features.
+        :type channel_index: `int`
+        """
+        # Save defense-specific parameters
+        super(JpegCompression, self).set_params(**kwargs)
+
+        if type(self.quality) is not int or self.quality <= 0 or self.quality > 100:
+            raise ValueError('Image quality must be a positive integer and smaller than 101.')
+
+        if type(self.channel_index) is not int or self.channel_index <= 0:
+            raise ValueError('Data channel must be a positive integer. The batch dimension is not a valid channel.')
+
+        return True
+
+

art/defences/jpeg_compression_unittest.py

@@ -0,0 +1,49 @@
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import logging
+import unittest
+
+from tensorflow.examples.tutorials.mnist import input_data
+from keras.datasets import cifar10
+import numpy as np
+
+from art.defences.jpeg_compression import JpegCompression
+
+logger = logging.getLogger('testLogger')
+
+
+class TestJpegCompression(unittest.TestCase):
+    def test_one_channel(self):
+        mnist = input_data.read_data_sets("tmp/MNIST_data/")
+        x = np.reshape(mnist.test.images[0:2], (-1, 28, 28, 1))
+        preprocess = JpegCompression()
+        compressed_x = preprocess(x, quality=70)
+        self.assertTrue((compressed_x.shape == x.shape))
+        self.assertTrue((compressed_x <= 1.0).all())
+        self.assertTrue((compressed_x >= 0.0).all())
+
+    def test_three_channels(self):
+        (train_features, train_labels), (test_data, test_label) = cifar10.load_data()
+        x = train_features[:2] / 255.0
+        preprocess = JpegCompression()
+        compressed_x = preprocess(x, quality=80)
+        self.assertTrue((compressed_x.shape == x.shape))
+        self.assertTrue((compressed_x <= 1.0).all())
+        self.assertTrue((compressed_x >= 0.0).all())
+
+    def test_channel_index(self):
+        (train_features, train_labels), (test_data, test_label) = cifar10.load_data()
+        x = train_features[:2] / 255.0
+        x = np.swapaxes(x, 1, 3)
+        preprocess = JpegCompression(channel_index=1)
+        compressed_x = preprocess(x, quality=80)
+        self.assertTrue((compressed_x.shape == x.shape))
+        self.assertTrue((compressed_x <= 1.0).all())
+        self.assertTrue((compressed_x >= 0.0).all())
+
+
+if __name__ == '__main__':
+    unittest.main()
+
+
+