Refactor CLI implementation and move it to it's own file
This commit is contained in:
Kristóf Tóth
parent
8e69f70cf3
commit
20225770cb
|
|
@ -0,0 +1,72 @@
|
|||
#!/usr/bin/env python3
|
||||
from argparse import ArgumentParser
|
||||
from os import remove
|
||||
from sys import exit
|
||||
from signal import signal, SIGINT
|
||||
from multiprocessing import Pool, cpu_count
|
||||
|
||||
from imgrate import imgrate
|
||||
|
||||
|
||||
class Cli:
|
||||
def __init__(self):
|
||||
self.args = None
|
||||
self.image_ratings = None
|
||||
self.best_image = None
|
||||
|
||||
def __call__(self):
|
||||
signal(SIGINT, lambda a, b: exit('\nExiting!'))
|
||||
|
||||
self.parse_arguments()
|
||||
|
||||
if not hasattr(imgrate, self.args.method):
|
||||
exit('Invalid --method option!')
|
||||
|
||||
self.calculate_ratings()
|
||||
self.best_image = max(self.image_ratings, key=self.image_ratings.get)
|
||||
self.print_ratings()
|
||||
|
||||
if self.args.delete:
|
||||
self.keep_best_quality_image_only()
|
||||
|
||||
def parse_arguments(self):
|
||||
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.')
|
||||
self.args = ap.parse_args()
|
||||
|
||||
def calculate_ratings(self):
|
||||
ratings = []
|
||||
with Pool(processes=cpu_count()) as pool:
|
||||
for image in self.args.images:
|
||||
ratings.append((image, pool.apply_async(self.calculate_rating,
|
||||
(image, self.args.method))))
|
||||
self.image_ratings = {image: promise.get() for image, promise in ratings}
|
||||
|
||||
@staticmethod
|
||||
def calculate_rating(image, method):
|
||||
return getattr(imgrate(image), method)
|
||||
|
||||
def print_ratings(self):
|
||||
for image in self.image_ratings.keys():
|
||||
if self.args.quiet:
|
||||
print(self.image_ratings[image])
|
||||
else:
|
||||
max_marker = '*' if image == self.best_image and len(self.image_ratings) > 1 else ' '
|
||||
print(f'{max_marker}imgrate("{image}") = {self.image_ratings[image]}')
|
||||
|
||||
def keep_best_quality_image_only(self):
|
||||
for image in self.image_ratings.keys():
|
||||
if image != self.best_image:
|
||||
remove(image)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
Cli()()
|
||||
59
imgrate.py
59
imgrate.py
|
|
@ -1,5 +1,3 @@
|
|||
#!/usr/bin/env python3
|
||||
from argparse import ArgumentParser
|
||||
from functools import wraps
|
||||
|
||||
import cv2
|
||||
|
|
@ -45,60 +43,3 @@ class imgrate:
|
|||
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
|
||||
from multiprocessing import Pool, cpu_count
|
||||
|
||||
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):
|
||||
if not hasattr(imgrate, args.method):
|
||||
exit('Invalid --method option!')
|
||||
|
||||
ratings = calculate_all_ratings(args.images, args.method)
|
||||
maximum = max(ratings, key=ratings.get)
|
||||
|
||||
fancy_print_ratings(ratings, args.quiet, maximum)
|
||||
|
||||
if args.delete:
|
||||
[remove(image) for image in ratings if image != maximum]
|
||||
|
||||
|
||||
def calculate_all_ratings(images, method):
|
||||
ratings = []
|
||||
with Pool(processes=cpu_count()) as pool:
|
||||
for image in images:
|
||||
ratings.append((image, pool.apply_async(calculate_rating, (image, method))))
|
||||
|
||||
return {image: promise.get() for image, promise in ratings}
|
||||
|
||||
|
||||
def calculate_rating(image, method):
|
||||
return getattr(imgrate(image), method)
|
||||
|
||||
|
||||
def fancy_print_ratings(ratings, quiet, maximum):
|
||||
for rating in ratings.keys():
|
||||
if quiet:
|
||||
print(ratings[rating])
|
||||
else:
|
||||
maxmark = '*' if rating == maximum and len(ratings) > 1 else ' '
|
||||
print('{}imgrate("{}") = {}'.format(maxmark, rating, ratings[rating]))
|
||||
|
||||
|
||||
run(parse_arguments())
|
||||
|
|
|
|||
Loading…
Reference in New Issue