imgp

#!/usr/bin/env python3
#
# Resize or rotate JPEG and PNG images.
#
# Copyright © 2016-2017 Arun Prakash Jana <engineerarun@gmail.com>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import argparse
import sys
import os
import PIL
import time
from PIL import Image, ImageFile
from multiprocessing import Pool, Lock

# Globals
HRES = 0  # specified horizontal resolution
VRES = 0  # specified vertical resolution
scale = None  # specified scale in percentage
rotate = None  # specified rotation in degree
adapt = False  # use adaptive resolution
convert = False  # convert PNG to JPG
dot = False  # process hidden files starting with .
eraseexif = False  # erase EXIF metadata
force = False  # forced to exact specified resolution
includeimgp = False  # process _IMGP files
keep = False  # ignore images with matching hres or vres
enlarge = False  # enlarge images with smaller hres or vres
optimize = False  # apply PIL library optimization when saving
quiet = False  # suppress additional information
recurse = False  # recursively process subdirectories
overwrite = False  # remove source images
debug = False  # enable debug information
no_res_opt = False  # --res is not specified
pad = '_IMGP'  # output file suffix when --overwrite is unused
init_time = time.time()  # profile the total time taken
_VERSION_ = '2.2'  # current program version

# Globals for multiprocessing
pool = None
lock = Lock()
count = 0  # successful conversion count
count_lowres = 0  # no-op count due to low resolution


def getres(res):
    '''Retrieve hres and vres as a tuple

    :param res: resolution in hresxvres format
    '''

    hxv = res.split('x')
    if len(hxv) != 2:
        return (0, 0)

    try:
        return(int(hxv[0]), int(hxv[1]))
    except:
        return (0, 0)


def lock_print(msg):
    '''Acquire lock and print a message'''

    global lock

    with lock:
        print(msg)


def cb(data):
    global count, count_lowres

    if data:
        count += data[0]
        count_lowres += data[1]


def traverse_dir(path):
    '''Traverse a directory recursively'''

    global lock, pool

    try:
        for entry in os.scandir(path):
            # Add check for hidden here
            if not dot and entry.name.startswith('.'):
                continue

            if entry.is_dir(follow_symlinks=False):
                if recurse:
                    if not os.access(entry.path, os.W_OK | os.X_OK):
                        lock_print(entry.path + ': no permission')
                        continue
                    traverse_dir(entry.path)
            elif rotate is not None:
                pool.apply_async(rotate_image, args=(entry.path,), callback=cb)
            else:
                pool.apply_async(resize_image, args=(entry.path,), callback=cb)
    except OSError as e:
        with lock:
            print(e)


def rotate_image(src):
    '''Rotate a source image'''

    global lock

    converted = False
    name, ext = os.path.splitext(src)
    if name.endswith(pad) and not includeimgp:
        return

    try:
        img = Image.open(src)
        _format = img.format
        if not (_format == 'JPEG' or _format == 'PNG'):
            if debug:
                lock_print(src + ': not JPEG or PNG image')
            return
        try:
            if eraseexif:
                exifdata = b''
            else:
                exifdata = img.info['exif']
        except:
            exifdata = b''
    except OSError as e:
        if not str(e).startswith('cannot identify image file'):
            with lock:
                print(e)
        return

    # Image.rotate(angle, resample=0, expand=0) expand = True means
    # the image will be enlarged to hold the entire rotated image
    img = img.rotate(-rotate, 0, True)

    if not quiet:
        stats = src + '\n%ld bytes ->' % os.path.getsize(src)

    if overwrite:
        suffix = ''
    else:
        suffix = pad

    if convert and _format == 'PNG':
        dest = name + suffix + '.jpg'
        _format = 'JPEG'
        converted = True
    else:
        dest = name + suffix + ext

    try:
        if not optimize:
            img.save(dest, _format, exif=exifdata)
        else:
            img.save(dest, _format, optimize=True, exif=exifdata)

        if not quiet:
            lock_print('%s %ld bytes\n' % (stats, os.path.getsize(dest)))

        if overwrite and converted:
            os.remove(src)

        return [1, 0]
    except OSError as e:
        with lock:
            print(e)


def resize_image(src):
    '''Resize a source image'''

    global HRES, VRES, lock

    converted = False
    name, ext = os.path.splitext(src)
    if name.endswith(pad) and not includeimgp:
        return

    try:
        img = Image.open(src)
        _format = img.format

        if not (_format == 'JPEG' or _format == 'PNG'):
            if debug:
                lock_print(src + ': not JPEG or PNG image')
            return

        # if 'only' convert option is specified,
        # nothing to be done for non-PNG images
        if no_res_opt and not optimize and _format != 'PNG':
            if debug:
                lock_print('resize_image: cannot convert non-PNG')
            return

        try:
            if eraseexif:
                exifdata = b''
            else:
                exifdata = img.info['exif']
        except:
            exifdata = b''
    except OSError as e:
        if not str(e).startswith('cannot identify image file'):
            with lock:
                print(e)
        return

    hres = img.size[0]
    vres = img.size[1]
    ImageFile.MAXBLOCK = hres * vres

    if keep and (HRES == hres or VRES == vres or scale == 100):
        if not (_format == 'PNG' and convert):
            return
    elif scale:
        if scale == 100:
            _hres = hres
            _vres = vres
        else:
            _hres = hres * scale/float(100)
            # use custom round()
            if _hres >= int(_hres) + .5:
                _hres = _hres + 1
            _hres = int(_hres)

            _vres = vres * scale/float(100)
            if _vres >= int(_vres) + .5:
                _vres = _vres + 1
            _vres = int(_vres)

        img = img.resize((_hres, _vres), PIL.Image.ANTIALIAS)
        if debug:
            lock_print('resize_image: scaled to [%dx%d]' % (HRES, VRES))
    elif HRES == 0 or VRES == 0:
        if HRES:
            if not enlarge and HRES > hres:
                if not quiet:
                    lock_print('%s has lesser hres\n' % src)
                return [0, 1]

            hratio = HRES/float(hres)
            target_vres = vres * hratio
            if target_vres >= int(target_vres) + .5:
                target_vres = target_vres + 1
            target_vres = int(target_vres)

            img = img.resize((HRES, target_vres), PIL.Image.ANTIALIAS)
            if debug:
                lock_print('resize_image: 0 vres %dx%d' % (HRES, target_vres))

        if VRES:
            if not enlarge and VRES > vres:
                if not quiet:
                    lock_print('%s has lesser vres\n' % src)
                return [0, 1]

            vratio = VRES/float(vres)
            target_hres = hres * vratio
            if target_hres >= int(target_hres) + .5:
                target_hres = target_hres + 1
            target_hres = int(target_hres)

            img = img.resize((target_hres, VRES), PIL.Image.ANTIALIAS)
            if debug:
                lock_print('resize_image: 0 hres %dx%d' % (target_hres, VRES))
    elif HRES == hres and VRES == vres:
        if debug:
            lock_print('resize_image: same res [%dx%d]' % (HRES, VRES))
    elif force:
        img = img.resize((HRES, VRES), PIL.Image.ANTIALIAS)
        if debug:
            lock_print('resize_image: brute force res to [%dx%d]'
                       % (HRES, VRES))
    elif adapt:
        ratio_img = float(hres/vres)
        ratio_target = float(HRES/VRES)

        if (ratio_target >= 1 and ratio_img > 1) \
                or (ratio_target < 1 and ratio_img >= 1):
            # same orientation (H >= V and h > v) or
            # cross orientation (H < V and h >= v)
            if not enlarge and HRES > hres:
                if not quiet:
                    lock_print('%s has lesser hres\n' % src)
                return [0, 1]

            if HRES != hres:
                hratio = HRES/float(hres)
                target_vres = vres * hratio
                if target_vres >= int(target_vres) + .5:
                    target_vres = target_vres + 1
                target_vres = int(target_vres)

                img = img.resize((HRES, target_vres), PIL.Image.ANTIALIAS)
                if debug:
                    lock_print('resize_image: vres adapted %dx%d'
                               % (HRES, target_vres))
        elif (ratio_target < 1 and ratio_img < 1) \
                or (ratio_target >= 1 and ratio_img <= 1):
            # same orientation (H < V and h < v) or
            # cross orientation (H >= V and h <= v)
            if not enlarge and VRES > vres:
                if not quiet:
                    lock_print('%s has lesser vres\n' % src)
                return [0, 1]

            if VRES != vres:
                vratio = VRES/float(vres)
                target_hres = hres * vratio
                if target_hres >= int(target_hres) + .5:
                    target_hres = target_hres + 1
                target_hres = int(target_hres)

                img = img.resize((target_hres, VRES), PIL.Image.ANTIALIAS)
                if debug:
                    lock_print('resize_image: hres adapted %dx%d'
                               % (target_hres, VRES))
    else:
        ratio_img = float(hres/vres)
        ratio_target = float(HRES/VRES)

        if ratio_img >= ratio_target:
            if not enlarge and HRES > hres:
                if not quiet:
                    lock_print('%s has lesser hres\n' % src)
                return [0, 1]

            # re-sample as per target horizontal resolution
            hratio = HRES/float(hres)
            target_vres = vres * hratio
            if target_vres >= int(target_vres) + .5:
                target_vres = target_vres + 1
            target_vres = int(target_vres)

            img = img.resize((HRES, target_vres), PIL.Image.ANTIALIAS)
            if debug:
                lock_print('resize_image: vres calculated %dx%d'
                           % (HRES, target_vres))
        else:
            if not enlarge and VRES > vres:
                if not quiet:
                    lock_print('%s has lesser vres\n' % src)
                return [0, 1]

            # re-sample as per target vertical resolution
            vratio = VRES/float(vres)
            target_hres = hres * vratio
            if target_hres >= int(target_hres) + .5:
                target_hres = target_hres + 1
            target_hres = int(target_hres)

            img = img.resize((target_hres, VRES), PIL.Image.ANTIALIAS)
            if debug:
                lock_print('resize_image: hres calculated %dx%d'
                           % (target_hres, VRES))

    if not quiet:
        stats = src + '\n%dx%d -> %dx%d\n%ld bytes ->' \
            % (hres, vres, img.size[0], img.size[1], os.path.getsize(src))

    if overwrite:
        suffix = ''
    else:
        suffix = pad

    if convert and _format == 'PNG':
        dest = name + suffix + '.jpg'
        _format = 'JPEG'
        converted = True
    else:
        dest = name + suffix + ext

    try:
        if not optimize:
            img.save(dest, _format, exif=exifdata)
        else:
            img.save(dest, _format, optimize=True, exif=exifdata)

        if not quiet:
            lock_print('%s %ld bytes\n' % (stats, os.path.getsize(dest)))

        if overwrite and converted:
            os.remove(src)

        return [1, 0]
    except OSError as e:
        with lock:
            print(e)


class ExtendedArgumentParser(argparse.ArgumentParser):
    '''Extend classic argument parser'''

    # Print additional help and info
    @staticmethod
    def print_extended_help(file=None):
        if not file:
            file = sys.stderr

        file.write('''
Version %s
Copyright © 2016-2017 Arun Prakash Jana <engineerarun@gmail.com>
License: GPLv3
Webpage: https://github.com/jarun/imgp
''' % _VERSION_)

    # Help
    def print_help(self, file=None):
        super(ExtendedArgumentParser, self).print_help(file)
        self.print_extended_help(file)


def parse_args(args=None, namespace=None):
    '''Parse imgp arguments/options.
    Parameters
    ----------
    args : list, optional
        Arguments to parse. Default is ``sys.argv``.
    namespace : argparse.Namespace
        Namespace to write to. Default is a new namespace.
    Returns
    -------
    argparse.Namespace
        Namespace with parsed arguments / options.
    '''

    argparser = ExtendedArgumentParser(
                    description='Resize, rotate JPEG and PNG images.')
    addarg = argparser.add_argument
    addarg('-x', '--res', nargs=1, metavar='res',
           help='output resolution in HxV or percentage')
    addarg('-o', '--rotate', type=int, metavar='deg',
           help='rotate clockwise by angle (in degrees)')
    addarg('-a', '--adapt', action='store_true',
           help='adapt to resolution by orientation [default: off]')
    addarg('-c', '--convert', action='store_true',
           help='convert PNG to JPG format [default: off]')
    addarg('-d', '--dot', action='store_true',
           help='include hidden files (on Linux) [default: off]')
    addarg('-e', '--eraseexif', action='store_true',
           help='erase exif metadata [default: off]')
    addarg('-f', '--force', action='store_true',
           help='force to exact specified resolution [default: off]')
    addarg('-i', '--includeimgp', action='store_true',
           help='re-process _IMGP files. * RISKY: refer to docs')
    addarg('-k', '--keep', action='store_true',
           help='skip (but convert, if opted) images with matching specified \
                 hres or vres or --res=100 [default: off]')
    addarg('-n', '--enlarge', action='store_true',
           help='enlarge smaller images [default: off]')
    addarg('-p', '--optimize', action='store_true',
           help='optimize the output images [default: off]')
    addarg('-q', '--quiet', action='store_true',
           help='operate silently [default: informative]')
    addarg('-r', '--recurse', action='store_true',
           help='process directories recursively [default: off] \
                 symbolic links to directories are ignored')
    addarg('-w', '--overwrite', action='store_true',
           help='overwrite source images [default: off]')
    addarg('-z', '--debug', action='store_true',
           help='enable debug logs [default: off]')
    addarg('keywords', nargs='*', metavar='PATH',
           help='source file or dir [default: current dir]')

    # Show help and exit if no arguments
    if len(sys.argv) < 2:
        argparser.print_help(sys.stderr)
        sys.exit(1)

    return argparser.parse_args(args, namespace)


def main():
    global HRES, VRES, scale, rotate, adapt, convert, dot, eraseexif, \
           force, includeimgp, keep, enlarge, optimize, quiet, recurse, \
           overwrite, debug, no_res_opt, pool, init_time, count

    args = parse_args()

    if args.rotate is not None and args.res is not None:
        print('options --rotate and --res cannot be used together')
        return

    if args.res is not None:
        if 'x' in args.res[0]:
            HRES, VRES = getres(args.res[0])
            if HRES == 0 and VRES == 0:
                print('-ve values not allowed, \
                       hres and vres cannot be 0 together')
                return
        else:
            try:
                if args.res[0].endswith('%'):
                    scale = int(args.res[0][:-1])
                else:
                    scale = int(args.res[0])

                if (scale <= 0):
                    print('scale should be > 0%')
                    return
                if scale > 100 and not args.enlarge:
                    print('use --enlarge is to scale > 100%')
                    return
            except Exception:
                print('invalid value for --res')
                return
    elif args.rotate is not None:
        rotate = args.rotate % 360
        if not rotate:
            print('cannot rotate by 0 degree')
            return
    elif args.convert or args.optimize:
        no_res_opt = True
        scale = 100
    else:
        print('missing image transformation')
        return

    adapt = args.adapt
    convert = args.convert
    dot = args.dot
    eraseexif = args.eraseexif
    force = args.force
    includeimgp = args.includeimgp
    keep = args.keep
    enlarge = args.enlarge
    optimize = args.optimize
    quiet = args.quiet
    recurse = args.recurse
    overwrite = args.overwrite
    debug = args.debug

    pool = Pool()

    for path in args.keywords:
        if not os.path.exists(path):
            print('%s does not exist' % path)
        elif os.path.isdir(path):
            traverse_dir(path)
        elif args.rotate is not None:
            if rotate_image(path):
                count += 1
        else:
            if resize_image(path):
                count += 1

    # Process current directory if PATH is unspecified
    if len(args.keywords) == 0:
        if not quiet:
            print('Source omitted. Processing current directory...\n')
        traverse_dir('.')

    pool.close()
    pool.join()

    if not quiet:
        # Show number of images converted
        if count:
            print('%s processed in %.4f seconds.'
                  % (count, time.time() - init_time))

        # Show number of images not ceonverted due to low resolution
        if count_lowres:
            print('%s not processed due to low resolution.' % count_lowres)

if __name__ == '__main__':
    main()