compress_images.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import traceback
from io import BytesIO
from hashlib import sha256


# copypasted from https://bitbucket.org/andreymal/smilepack/src/develop/smilepack/utils/uploader.py


class BadImageError(Exception):
    pass


def calc_hashsum(data):
    return sha256(data).hexdigest()


def compress_image(data, hashsum, image=None, compress_size=None):
    from PIL import Image

    min_size = len(data)

    # Если сжимать совсем нет смысла
    if min_size <= 4096:
        return data, hashsum, None

    image_local = not image
    if image_local:
        try:
            image = Image.open(BytesIO(data))
        except:
            raise BadImageError('Cannot decode image')

    try:
        # Если сжимать не умеем
        if image.format != 'PNG':
            return data, hashsum, None

        # А PNG пробуем сжать разными методами
        test_data, method = compress_png(image)
    finally:
        if image_local:
            image.close()
            image = None

    # Сохраняем сжатие, только если оно существенно
    if test_data and min_size - len(test_data) > 1024:
        new_hashsum = calc_hashsum(test_data)
        return test_data, new_hashsum, method
    else:
        return data, hashsum, None


def compress_png(image):
    # 0) Пробуем просто пересохранить
    min_stream = BytesIO()
    image.save(min_stream, 'PNG', optimize=True)
    min_size = len(min_stream.getvalue())
    method = 'resave'

    # 1) Пробуем пересохранить с zlib (иногда почему-то меньше, чем optimize=True)
    test_stream = BytesIO()
    image.save(test_stream, 'PNG', compress_level=9)
    test_size = len(test_stream.getvalue())
    if test_size < min_size:
        min_stream = test_stream
        min_size = test_size
        method = 'zlib'

    # 2) Пробуем закрасить чёрным невидимое
    if image.mode == 'RGBA':
        from PIL import ImageDraw
        with image.copy() as test_image:
            w = test_image.size[0]
            draw = None
            for i, pixel in enumerate(test_image.getdata()):
                if pixel[3] < 1:
                    if draw is None:
                        draw = ImageDraw.Draw(test_image)
                    draw.point([(i % w, i // w)], (0, 0, 0, 0))
            if draw is not None:
                test_stream = BytesIO()
                test_image.save(test_stream, 'PNG', optimize=True)
                test_size = len(test_stream.getvalue())
                if test_size < min_size:
                    min_stream = test_stream
                    min_size = test_size
                    method = 'zeroalpha'
            del draw

    return min_stream.getvalue(), method


def find_pngs(path, recursive=False):
    files = []
    for x in sorted(os.listdir(path)):
        xpath = os.path.join(path, x)
        if os.path.isdir(xpath):
            if recursive:
                files.extend(os.path.join(x, y) for y in find_pngs(xpath, True))
            continue
        if x.endswith('.png'):
            files.append(x)
    return files


def compress_file(path, write_path=None):
    print(path, end=' ')
    sys.stdout.flush()

    try:
        data = open(path, 'rb').read()
    except:
        print('Cannot read')
        return False, None

    from PIL import Image
    try:
        image = Image.open(BytesIO(data))
    except:
        print('Cannot decode')
        return False, None

    original_size = len(data)
    with image:
        compressed_data, hashsum, method = compress_image(data, calc_hashsum(data), image)
    del image

    if not method:
        print('Not compressed')
        return True, None

    compressed_size = len(compressed_data)
    percent = (1 - compressed_size / original_size) * 100

    if write_path:
        with open(write_path, 'wb') as fp:
            fp.write(compressed_data)

    print('OK: {method} {percent:.2f}% ({original_size} before - {diff_size} = {compressed_size} after)'.format(
        method=method,
        percent=percent,
        original_size=original_size,
        diff_size=original_size - compressed_size,
        compressed_size=compressed_size
    ))

    return True, (original_size, compressed_size)


def main():
    args = sys.argv[1:]

    if '-r' in args:
        recursive = True
        args.remove('-r')
    else:
        recursive = False

    if len(args) != 2:
        print('Usage: {} [-r] input_dir output_dir'.format(sys.argv[0]), file=sys.stderr)
        print('Reads PNG files from input directory,')
        print('compress it and writes to output directory')
        print('(can be same). Recursive with "-r".')
        sys.exit(2)

    input_dir, output_dir = args

    files = find_pngs(input_dir, recursive)
    success = 0
    count = len(files)

    results = []

    for x in files:
        try:
            processed, result = compress_file(os.path.join(input_dir, x), os.path.join(output_dir, x))
        except:
            print()
            print(traceback.format_exc())
            continue

        if processed:
            success += 1
            if result is not None:
                results.append(result)

        sys.stdout.flush()

    if success < count:
        print('{} files failed, please check it'.format(count - success))
        sys.exit(1)

    percents = [(100.0 - (x[1] / x[0]) * 100.0) for x in results]

    if len(percents) >= len(files):
        print('Average compression {:.2f}%'.format(sum(percents) / len(percents)))
    elif percents:
        tprc = percents + ([0.0] * (len(files) - len(percents)))
        print('Average compression {:.2f}% ({:.2f}% considering uncompressed)'.format(
            sum(percents) / len(percents),
            sum(tprc) / len(tprc)
        ))
    if percents:
        print('{:.2f} KiB -> {:.2f} KiB'.format(sum(x[0] for x in results) / 1024, sum(x[1] for x in results) / 1024))


if __name__ == '__main__':
    main()