imgrate/imgrate.py

#!/usr/bin/env python3
import cv2
import numpy as np
from argparse import ArgumentParser
from functools import wraps



def image_required(fun):
    @wraps(fun)
    def wrapper(self, *args, **kwargs):
        if self.image is None:
            raise RuntimeError('Property {}.{} requires an image!'
                               .format(self.__class__.__name__, fun.__name__))
        return fun(self, *args, **kwargs)
    return wrapper



class imgrate:
    def __init__(self, imgfile=None):
        self.image = None
        if imgfile: self.load_image(imgfile)

    def load_image(self, imgfile):
        image = cv2.imread(imgfile)
        self.image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    @property
    @image_required
    def quality(self):
        return self.laplacian * self.fdibm

    @property
    @image_required
    def laplacian(self):
        return cv2.Laplacian(self.image, cv2.CV_64F).var()

    @property
    @image_required
    def fdibm(self):
        F = np.fft.fft2(self.image)
        Fcentered = np.fft.fftshift(F)
        Fabs = np.abs(Fcentered)
        Fmax = Fabs.max()
        Th = np.count_nonzero(Fabs > Fmax/1000)
        return Th / (F.shape[0] * F.shape[1])



if __name__ == '__main__':
    from os import remove
    from sys import exit
    from signal import signal, SIGINT

    signal(SIGINT, lambda a, b: exit('\nExiting!'))

    def parse_arguments():
        ap = ArgumentParser(description='Find the best quality one among similar images.'
                                        'Note: image dimensions should match!')
        ap.add_argument('images', type=str, nargs='+', help='')
        ap.add_argument('-d', '--delete', action='store_true', help='Delete all but the best quality image.')
        ap.add_argument('-q', '--quiet', action='store_true', help='Print quality measurements only.')
        ap.add_argument('-m', '--method', type=str, default='quality', help='Select quality measurement to use.')
        return ap.parse_args()


    def run(args):
        try: ratings = {image: getattr(imgrate(image), args.method) for image in args.images}
        except AttributeError: exit('Invalid --method option!')
        maximum = max(ratings, key=ratings.get)

        for rating in ratings.keys():
            if args.quiet:
                print(ratings[rating])
            else:
                maxmark = '*' if rating == maximum and len(ratings) > 1 else ' '
                print('{}imgrate("{}") = {}'.format(maxmark, rating, ratings[rating]))

        if args.delete:
            [remove(image) for image in ratings if image != maximum]


    run(parse_arguments())