Example #1
0
class JpegEncoder(object):
    def __init__(self, image, quality, out, comment, ais):
        self.quality = quality
        self.jpeg_obj = JpegInfo(image, comment)

        self.image_width, self.image_height = image.size
        self.out = out

        self.dct = DCT(self.quality)
        self.huf = Huffman(*image.size)

        self.hasais = ais
        self.k_matrix = -1

    def compress(self, embedded_data=None, password='******'):
        self.embedded_data = EmbedData(
            embedded_data) if embedded_data else None
        self.password = password

        self.write_headers()
        self.write_compressed_data()
        self.write_eoi()
        self.out.flush()

    def get_quality(self):
        return self.quality

    def set_quality(self, quality):
        self.quality = quality
        self.dct = DCT(quality)

    def write_array(self, data):
        length = ((data[2] & 0xff) << 8) + (data[3] & 0xff) + 2
        self.out.write(bytearray(data[:length]))

    def write_marker(self, data):
        self.out.write(bytearray(data[:2]))

    def write_eoi(self):
        EOI = [0xff, 0xD9]
        self.write_marker(EOI)

    def write_headers(self):
        SOI = [0xff, 0xD8]
        self.write_marker(SOI)

        JFIF = [
            0xff, 0xe0, 0x00, 0x10, 0x4a, 0x46, 0x49, 0x46, 0x00, 0x01, 0x01,
            0x01, 0x00, 0x60, 0x00, 0x60, 0x00, 0x00
        ]
        self.write_array(JFIF)

        comment = self.jpeg_obj.get_comment()
        if comment:
            length = len(comment) + 2
            COM = [0xff, 0xfe, length >> 8 & 0xff, length & 0xff]
            COM.extend(comment)
            self.write_array(COM)

        DQT = [0xff, 0xdb, 0x00, 0x84]
        for k in range(2):
            DQT.append(k)
            DQT.extend([
                self.dct.quantum[k][JPEG_NATURAL_ORDER[i]] for i in range(64)
            ])
        self.write_array(DQT)

        SOF = [
            0xff, 0xc0, 0x00, 0x11, self.jpeg_obj.precision,
            self.jpeg_obj.image_height >> 8 & 0xff,
            self.jpeg_obj.image_height & 0xff,
            self.jpeg_obj.image_width >> 8 & 0xff,
            self.jpeg_obj.image_width & 0xff, self.jpeg_obj.comp_num
        ]
        for i in range(self.jpeg_obj.comp_num):
            SOF.append(self.jpeg_obj.com_id[i])
            SOF.append(
                eight_byte(self.jpeg_obj.hsamp_factor[i],
                           self.jpeg_obj.vsamp_factor[i]))
            SOF.append(self.jpeg_obj.qtable_number[i])
        self.write_array(SOF)

        DHT = [0xff, 0xc4, 0, 0]
        for i in range(4):
            DHT.extend(self.huf.BITS[i])
            DHT.extend(self.huf.VAL[i])

        DHT[2] = len(DHT) - 2 >> 8 & 0xff
        DHT[3] = len(DHT) - 2 & 0xff
        self.write_array(DHT)

        SOS = [0] * 14
        SOS = [0xff, 0xda, 0x00, 0x0c, self.jpeg_obj.comp_num]
        for i in range(self.jpeg_obj.comp_num):
            SOS.append(self.jpeg_obj.com_id[i])
            SOS.append(
                eight_byte(self.jpeg_obj.dctable_number[i],
                           self.jpeg_obj.actable_number[i]))

        SOS.append(self.jpeg_obj.ss)
        SOS.append(self.jpeg_obj.se)
        SOS.append(eight_byte(self.jpeg_obj.ah, self.jpeg_obj.al))
        self.write_array(SOS)

    def _get_coeff(self):
        dct_array1 = create_array(0.0, 8, 8)
        dct_array2 = create_array(0.0, 8, 8)
        dct_array3 = create_array(0, 64)

        coeff = []
        for r in range(min(self.jpeg_obj.block_height)):
            for c in range(min(self.jpeg_obj.block_width)):
                xpos = c * 8
                ypos = r * 8
                for comp in range(self.jpeg_obj.comp_num):
                    indata = self.jpeg_obj.components[comp]
                    maxa = self.image_height / 2 * self.jpeg_obj.vsamp_factor[
                        comp] - 1
                    maxb = self.image_width / 2 * self.jpeg_obj.hsamp_factor[
                        comp] - 1

                    for i in range(self.jpeg_obj.vsamp_factor[comp]):
                        for j in range(self.jpeg_obj.hsamp_factor[comp]):
                            ia = ypos * self.jpeg_obj.vsamp_factor[comp] + i * 8
                            ib = xpos * self.jpeg_obj.hsamp_factor[comp] + j * 8

                            for a in range(8):
                                for b in range(8):
                                    dct_array1[a][b] = indata[min(
                                        ia + a, maxa)][min(ib + b, maxb)]

                            dct_array2 = self.dct.forward_dct(dct_array1)
                            dct_array3 = self.dct.quantize_block(
                                dct_array2, self.jpeg_obj.qtable_number[comp])
                            coeff.extend(dct_array3[:64])
        return coeff

    def write_compressed_data(self):
        last_dc_value = create_array(0, self.jpeg_obj.comp_num)
        zero_array = create_array(0, 64)
        width, height = 0, 0
        min_block_width = min(self.jpeg_obj.block_width)
        min_block_height = min(self.jpeg_obj.block_height)

        logger.info('DCT/quantisation starts')
        logger.info('%d x %d' % (self.image_width, self.image_height))

        coeff = self._get_coeff()  #量化后的系数,是整数
        coeff_count = len(coeff)

        #导出未处理的QDCT
        if self.hasais:
            filename = 'unpro_ais.json'
        else:
            filename = 'unpro.json'
        with open(filename, 'w') as f_unprocess:
            json.dump(coeff, f_unprocess)

        #AIS处理
        if self.hasais:
            size_secret = self.embedded_data.len
            ais = Ais(coeff, size_secret)  #coeff被修改
            ais.statistic()
            self.k_matrix = ais.fix()
            with open('aised.json', 'w') as f_aised:
                json.dump(coeff, f_aised)

        #嵌入——>再统计嵌入后的数据,决定是否继续做AIS处理
        logger.info('got %d DCT AC/DC coefficients' % coeff_count)
        _changed, _embedded, _examined, _expected, _one, _large, _thrown, _zero = 0, 0, 0, 0, 0, 0, 0, 0
        shuffled_index = 0
        for i, cc in enumerate(coeff):
            if i % 64 == 0:
                continue
            if cc == 1 or cc == -1:
                _one += 1
            elif cc == 0:
                _zero += 1

        _large = coeff_count - _zero - _one - coeff_count / 64  #有效系数的个数
        _expected = _large + int(0.49 * _one)  #预期容量,shrinkage效应无法确定

        logger.info('one=%d' % _one)
        logger.info('large=%d' % _large)
        logger.info('\nexpected capacity: %d bits\n' % _expected)
        logger.info('expected capacity with')

        for i in range(1, 8):
            n = (1 << i) - 1  #n=2^i-1
            changed = _large - _large % (n + 1)
            changed = (changed + _one + _one / 2 - _one / (n + 1)) / (n + 1)
            usable = (_expected * i / n - _expected * i / n % n) / 8
            if usable == 0:
                break

            logger.info(
                '%s code: %d bytes (efficiency: %d.%d bits per change)' %
                ('default' if i == 1 else '(1, %d, %d)' % (n, i), usable,
                 usable * 8 / changed, usable * 80 / changed % 10))

        #shuffles all coefficients using a permutation,使用排列对系数进行混洗
        if self.embedded_data is not None:
            logger.info('permutation starts')
            random = F5Random(self.password)
            permutation = Permutation(coeff_count, random)

            next_bit_to_embed = 0
            byte_to_embed = len(self.embedded_data)
            available_bits_to_embed = 0

            logger.info('Embedding of %d bits (%d+4 bytes)' %
                        (byte_to_embed * 8 + 32, byte_to_embed))

            if byte_to_embed > 0x007fffff:
                byte_to_embed = 0x007ffff

            for i in range(1, 8):
                self.n = (1 << i) - 1
                usable = (_expected * i / self.n -
                          _expected * i / self.n % self.n) / 8
                if usable < byte_to_embed + 4:
                    break

            #确定(1,n,k)
            if self.k_matrix < 0:
                k = i - 1
            else:
                k = self.k_matrix
            self.n = (1 << k) - 1

            if self.n == 0:
                logger.info('using default code, file will not fit')
                self.n = 1
            elif self.n == 1:
                logger.info('using default code')
            else:
                logger.info('using (1, %d, %d) code' % (self.n, k))

            byte_to_embed |= k << 24
            byte_to_embed ^= random.get_next_byte()
            byte_to_embed ^= random.get_next_byte() << 8
            byte_to_embed ^= random.get_next_byte() << 16
            byte_to_embed ^= random.get_next_byte() << 24

            next_bit_to_embed = byte_to_embed & 1
            byte_to_embed >>= 1
            available_bits_to_embed = 31
            _embedded += 1

            for i, shuffled_index in enumerate(permutation.shuffled):
                if shuffled_index % 64 == 0 or coeff[shuffled_index] == 0:
                    continue
                cc = coeff[shuffled_index]
                _examined += 1

                if cc > 0 and (cc & 1) != next_bit_to_embed:
                    coeff[shuffled_index] -= 1
                    _changed += 1
                elif cc < 0 and (cc & 1) == next_bit_to_embed:
                    coeff[shuffled_index] += 1
                    _changed += 1

                if coeff[shuffled_index] != 0:
                    if available_bits_to_embed == 0:
                        if self.n > 1 or not self.embedded_data.available():
                            break
                        byte_to_embed = self.embedded_data.read()
                        byte_to_embed ^= random.get_next_byte()
                        available_bits_to_embed = 8
                    next_bit_to_embed = byte_to_embed & 1
                    byte_to_embed >>= 1
                    available_bits_to_embed -= 1
                    _embedded += 1
                else:
                    _thrown += 1

            if self.n > 1:
                try:
                    is_last_byte = False
                    filtered_index = FilteredCollection(
                        permutation.shuffled[i + 1:],
                        lambda index: index % 64 and coeff[index])
                    while not is_last_byte:
                        k_bits_to_embed = 0
                        for i in range(k):
                            if available_bits_to_embed == 0:
                                if not self.embedded_data.available():
                                    is_last_byte = True
                                    break
                                byte_to_embed = self.embedded_data.read()
                                byte_to_embed ^= random.get_next_byte()
                                available_bits_to_embed = 8
                            next_bit_to_embed = byte_to_embed & 1
                            byte_to_embed >>= 1
                            available_bits_to_embed -= 1
                            k_bits_to_embed |= next_bit_to_embed << i
                            _embedded += 1

                        code_word = filtered_index.offer(self.n)
                        while True:
                            vhash = 0
                            for i, index in enumerate(code_word):
                                if coeff[index] > 0:
                                    extracted_bit = coeff[index] & 1
                                else:
                                    extracted_bit = 1 - (coeff[index] & 1)
                                if extracted_bit == 1:
                                    vhash ^= i + 1
                            i = vhash ^ k_bits_to_embed
                            if not i:
                                break

                            i -= 1
                            coeff[code_word[i]] += 1 if coeff[
                                code_word[i]] < 0 else -1
                            _changed += 1

                            if not coeff[code_word[i]]:
                                _thrown += 1
                                code_word[i:i + 1] = []
                                code_word.extend(filtered_index.offer(1))
                            else:
                                break
                except FilteredCollection.ListNotEnough:
                    pass

            if _examined > 0:
                logger.info('%d coefficients examined' % _examined)
            if _changed > 0:
                logger.info(
                    '%d coefficients changed (efficiency: %d.%d bits per change'
                    % (_changed, _embedded / _changed,
                       _embedded * 10 / _changed % 10))
            logger.info('%d coefficients thrown (zeroed)' % _thrown)
            logger.info('%d bits (%d bytes) embedded' %
                        (_embedded, _embedded / 8))

        #导出嵌入后的系数coeff
        if self.hasais:
            filename2 = 'embeded_ais.json'
        else:
            filename2 = 'embeded.json'
        with open(filename2, 'w') as f_embeded:
            json.dump(coeff, f_embeded)

        logger.info('starting hufman encoding')
        shuffled_index = 0

        for r in range(min_block_height):
            for c in range(min_block_width):
                for comp in range(self.jpeg_obj.comp_num):
                    for i in range(self.jpeg_obj.vsamp_factor[comp]):
                        for j in range(self.jpeg_obj.hsamp_factor[comp]):
                            dct_array3 = coeff[shuffled_index:shuffled_index +
                                               64]
                            self.huf.huffman_block_encoder(
                                self.out, dct_array3, last_dc_value[comp],
                                self.jpeg_obj.dctable_number[comp],
                                self.jpeg_obj.actable_number[comp])
                            last_dc_value[comp] = dct_array3[0]
                            shuffled_index += 64

        self.huf.flush_buffer(self.out)
        logger.info('hufman encode end')
class JpegEncoder(object):
    def __init__(self, image, quality, out, comment):
        self.quality = quality
        self.jpeg_obj = JpegInfo(image, comment)

        self.image_width, self.image_height = image.size
        self.out = out

        self.dct = DCT(self.quality)
        self.huf = Huffman(*image.size)

    def compress(self, embedded_data=None, password='******'):
        self.embedded_data = EmbedData(embedded_data) if embedded_data else None
        self.password = password

        self.write_headers()
        self.write_compressed_data()
        self.write_eoi()
        self.out.flush()

    def get_quality(self):
        return self.quality

    def set_quality(self, quality):
        self.quality = quality
        self.dct = DCT(quality)

    def write_array(self, data):
        length = ((data[2] & 0xff) << 8) + (data[3] & 0xff) + 2
        self.out.write(bytearray(data[:length]))

    def write_marker(self, data):
        self.out.write(bytearray(data[:2]))

    def write_eoi(self):
        EOI = [0xff, 0xD9]
        self.write_marker(EOI)

    def write_headers(self):
        SOI = [0xff, 0xD8]
        self.write_marker(SOI)

        JFIF = [0xff, 0xe0, 0x00, 0x10, 0x4a, 0x46, 
                0x49, 0x46, 0x00, 0x01, 0x01, 0x01, 
                0x00, 0x60, 0x00, 0x60, 0x00, 0x00]
        self.write_array(JFIF)

        comment = self.jpeg_obj.get_comment()
        if comment:
            length = len(comment) + 2
            COM = [0xff, 0xfe, length >> 8 & 0xff, length & 0xff]
            COM.extend(comment)
            self.write_array(COM)

        DQT = [0xff, 0xdb, 0x00, 0x84]
        for k in range(2):
            DQT.append(k)
            DQT.extend([self.dct.quantum[k][JPEG_NATURAL_ORDER[i]] for i in range(64)])
        self.write_array(DQT)

        SOF = [0xff, 0xc0, 0x00, 0x11, 
                self.jpeg_obj.precision, 
                self.jpeg_obj.image_height >> 8 & 0xff, 
                self.jpeg_obj.image_height & 0xff, 
                self.jpeg_obj.image_width >> 8 & 0xff, 
                self.jpeg_obj.image_width & 0xff, 
                self.jpeg_obj.comp_num]
        for i in range(self.jpeg_obj.comp_num):
            SOF.append(self.jpeg_obj.com_id[i])
            SOF.append(eight_byte(self.jpeg_obj.hsamp_factor[i], self.jpeg_obj.vsamp_factor[i]))
            SOF.append(self.jpeg_obj.qtable_number[i])
        self.write_array(SOF)
        
        DHT = [0xff, 0xc4, 0, 0]
        for i in range(4):
            DHT.extend(self.huf.BITS[i])
            DHT.extend(self.huf.VAL[i])

        DHT[2] = len(DHT) - 2 >> 8 & 0xff
        DHT[3] = len(DHT) - 2 & 0xff
        self.write_array(DHT)

        SOS = [0] * 14
        SOS = [0xff, 0xda, 0x00, 0x0c, self.jpeg_obj.comp_num]
        for i in range(self.jpeg_obj.comp_num):
            SOS.append(self.jpeg_obj.com_id[i])
            SOS.append(eight_byte(self.jpeg_obj.dctable_number[i], self.jpeg_obj.actable_number[i]))

        SOS.append(self.jpeg_obj.ss)
        SOS.append(self.jpeg_obj.se)
        SOS.append(eight_byte(self.jpeg_obj.ah, self.jpeg_obj.al))
        self.write_array(SOS)

    def _get_coeff(self):
        dct_array1 = create_array(0.0, 8, 8)
        dct_array2 = create_array(0.0, 8, 8)
        dct_array3 = create_array(0, 64)

        coeff = []
        for r in range(min(self.jpeg_obj.block_height)):
            for c in range(min(self.jpeg_obj.block_width)):
                xpos = c * 8
                ypos = r * 8
                for comp in range(self.jpeg_obj.comp_num):
                    indata = self.jpeg_obj.components[comp]
                    maxa = self.image_height / 2 * self.jpeg_obj.vsamp_factor[comp] - 1
                    maxb = self.image_width / 2 * self.jpeg_obj.hsamp_factor[comp] - 1

                    for i in range(self.jpeg_obj.vsamp_factor[comp]):
                        for j in range(self.jpeg_obj.hsamp_factor[comp]):
                            ia = ypos * self.jpeg_obj.vsamp_factor[comp] + i * 8
                            ib = xpos * self.jpeg_obj.hsamp_factor[comp] + j * 8

                            for a in range(8):
                                for b in range(8):
                                    dct_array1[a][b] = indata[min(ia+a, maxa)][min(ib+b, maxb)]

                            dct_array2 = self.dct.forward_dct(dct_array1)
                            dct_array3 = self.dct.quantize_block(dct_array2, 
                                    self.jpeg_obj.qtable_number[comp])
                            coeff.extend(dct_array3[:64])
        return coeff

    def write_compressed_data(self):
        tmp = 0

        last_dc_value = create_array(0, self.jpeg_obj.comp_num)
        zero_array = create_array(0, 64)
        width, height = 0, 0

        min_block_width = min(self.jpeg_obj.block_width)
        min_block_height = min(self.jpeg_obj.block_height)

        logger.info('DCT/quantisation starts')
        logger.info('%d x %d' % (self.image_width, self.image_height))

        coeff = self._get_coeff()
        coeff_count = len(coeff)

        logger.info('got %d DCT AC/DC coefficients' % coeff_count)
        _changed, _embedded, _examined, _expected, _one, _large, _thrown, _zero = 0, 0, 0, 0, 0, 0, 0, 0
        shuffled_index = 0
        for i, cc in enumerate(coeff):
            if i % 64 == 0:
                continue
            if cc == 1 or cc == -1:
                _one += 1
            elif cc == 0:
                _zero += 1

        _large = coeff_count - _zero - _one - coeff_count / 64
        _expected = _large + int(0.49 * _one)

        logger.info('one=%d' % _one)
        logger.info('large=%d' % _large)

        logger.info('expected capacity: %d bits' % _expected)
        logger.info('expected capacity with')
        for i in range(1, 8):
            n = (1 << i) - 1
            changed = _large - _large % (n + 1)
            changed = (changed + _one + _one / 2 - _one / (n + 1)) / (n + 1)

            usable = (_expected * i / n - _expected * i / n % n) / 8
            if usable == 0:
                break

            logger.info('%s code: %d bytes (efficiency: %d.%d bits per change)' % ('default' if i == 1 else '(1, %d, %d)' % (n, i), usable, usable * 8 / changed, usable * 80 / changed % 10))

        if self.embedded_data is not None:
            logger.info('permutation starts')
            random = F5Random(self.password)
            permutation = Permutation(coeff_count, random)

            next_bit_to_embed = 0
            byte_to_embed = len(self.embedded_data)
            available_bits_to_embed = 0

            logger.info('Embedding of %d bits (%d+4 bytes)' % (byte_to_embed * 8 + 32, byte_to_embed))

            if byte_to_embed > 0x007fffff:
                byte_to_embed = 0x007ffff

            for i in range(1, 8):
                self.n = (1 << i) - 1
                usable = (_expected * i / self.n - _expected * i / self.n % self.n) / 8
                if usable < byte_to_embed + 4:
                    break

            k = i - 1
            self.n = (1 << k) - 1

            if self.n == 0:
                logger.info('using default code, file will not fit')
                self.n = 1
            elif self.n == 1:
                logger.info('using default code')
            else:
                logger.info('using (1, %d, %d) code' % (self.n, k))

            byte_to_embed |= k << 24
            byte_to_embed ^= random.get_next_byte()
            byte_to_embed ^= random.get_next_byte() << 8
            byte_to_embed ^= random.get_next_byte() << 16
            byte_to_embed ^= random.get_next_byte() << 24

            next_bit_to_embed = byte_to_embed & 1
            byte_to_embed >>= 1
            available_bits_to_embed = 31
            _embedded += 1

            for i, shuffled_index in enumerate(permutation.shuffled):
                if shuffled_index % 64 == 0 or coeff[shuffled_index] == 0:
                    continue
                cc = coeff[shuffled_index]
                _examined += 1

                if cc > 0 and (cc & 1) != next_bit_to_embed:
                    coeff[shuffled_index] -= 1
                    _changed +=1
                elif cc < 0 and (cc & 1) == next_bit_to_embed:
                    coeff[shuffled_index] += 1
                    _changed += 1

                if coeff[shuffled_index] != 0:
                    if available_bits_to_embed == 0:
                        if self.n > 1 or not self.embedded_data.available():
                            break
                        byte_to_embed = self.embedded_data.read()
                        byte_to_embed ^= random.get_next_byte()
                        available_bits_to_embed = 8
                    next_bit_to_embed = byte_to_embed & 1
                    byte_to_embed >>= 1
                    available_bits_to_embed -= 1
                    _embedded += 1
                else:
                    _thrown += 1

            if self.n > 1:
                try:
                    is_last_byte = False
                    filtered_index = FilteredCollection(permutation.shuffled[i+1:], lambda index: index % 64 and coeff[index])
                    while not is_last_byte:
                        k_bits_to_embed = 0
                        for i in range(k):
                            if available_bits_to_embed == 0:
                                if not self.embedded_data.available():
                                    is_last_byte = True
                                    break
                                byte_to_embed = self.embedded_data.read()
                                byte_to_embed ^= random.get_next_byte()
                                available_bits_to_embed = 8
                            next_bit_to_embed = byte_to_embed & 1
                            byte_to_embed >>= 1
                            available_bits_to_embed -= 1
                            k_bits_to_embed |= next_bit_to_embed << i
                            _embedded += 1

                        code_word = filtered_index.offer(self.n)
                        while True:
                            vhash = 0
                            for i, index in enumerate(code_word):
                                if coeff[index] > 0:
                                    extracted_bit = coeff[index] & 1
                                else:
                                    extracted_bit = 1 - (coeff[index] & 1)
                                if extracted_bit == 1:
                                    vhash ^= i + 1
                            i = vhash ^ k_bits_to_embed
                            if not i:
                                break

                            i -= 1
                            coeff[code_word[i]] += 1 if coeff[code_word[i]] < 0 else -1
                            _changed += 1

                            if not coeff[code_word[i]]:
                                _thrown += 1
                                code_word[i:i+1] = []
                                code_word.extend(filtered_index.offer(1))
                            else:
                                break
                except FilteredCollection.ListNotEnough:
                    pass

            if _examined > 0:
                logger.info('%d coefficients examined' % _examined)
            if _changed > 0:
                logger.info('%d coefficients changed (efficiency: %d.%d bits per change' % (_changed, _embedded / _changed, _embedded * 10 / _changed % 10))
            logger.info('%d coefficients thrown (zeroed)' % _thrown)
            logger.info('%d bits (%d bytes) embedded' % (_embedded, _embedded / 8))

        logger.info('starting hufman encoding')
        shuffled_index = 0

        for r in range(min_block_height):
            for c in range(min_block_width):
                for comp in range(self.jpeg_obj.comp_num):
                    for i in range(self.jpeg_obj.vsamp_factor[comp]):
                        for j in range(self.jpeg_obj.hsamp_factor[comp]):
                            dct_array3 = coeff[shuffled_index:shuffled_index+64]
                            self.huf.huffman_block_encoder(self.out, dct_array3, last_dc_value[comp],
                                    self.jpeg_obj.dctable_number[comp], self.jpeg_obj.actable_number[comp])
                            last_dc_value[comp] = dct_array3[0]
                            shuffled_index += 64
        
        self.huf.flush_buffer(self.out)