Revision 2aa55fa2
Added by David Sorber almost 10 years ago
| software/lepton_wrapper/lepton_wrapper.py | ||
|---|---|---|
|
|
||
|
import argparse
|
||
|
import glob
|
||
|
import math
|
||
|
from multiprocessing import Process
|
||
|
from multiprocessing import Lock
|
||
|
import multiprocessing
|
||
|
import os
|
||
|
import re
|
||
|
import subprocess
|
||
| ... | ... | |
|
file_walker(os.path.join(dirpath, dirname), regex)
|
||
|
return files
|
||
|
|
||
|
def do_compress(top_path, remove_orig):
|
||
|
|
||
|
# Find all *.jpg (or *.jpeg) images (ignoring case)
|
||
|
images = file_walker(sys.argv[1], '(\.jpg$)|(\.jpeg$)')
|
||
|
num_images = len(images)
|
||
|
|
||
|
# Counters to keep track of total compression ratio
|
||
|
total_uncompressed_bytes = 0
|
||
|
total_compressed_bytes = 0
|
||
|
|
||
|
print('Found {:d} *.jpg images'.format(num_images))
|
||
|
|
||
|
# Iterate over each image and compress it using the lepton binary
|
||
|
for idx, image in enumerate(images):
|
||
|
print('Image {:d} of {:d}:'.format(idx + 1, num_images))
|
||
|
print('\t{:s}'.format(image))
|
||
|
class CompressWrapper(Process):
|
||
|
|
||
|
def __init__(self, index, image_queue, remove_orig, print_lock):
|
||
|
Process.__init__(self)
|
||
|
self.index = index
|
||
|
self.image_queue = image_queue
|
||
|
self.remove_orig = remove_orig
|
||
|
self.print_lock = print_lock
|
||
|
self.errors = []
|
||
|
self.running = False
|
||
|
|
||
|
def run(self):
|
||
|
self.running = True
|
||
|
|
||
|
compressed_image = '{:s}.lep'.format(image)
|
||
|
cmd = '{:s} "{:s}" "{:s}"'.format(LEPTON_PATH, image, compressed_image)
|
||
|
#~ rc = subprocess.run(cmd, shell=True)
|
||
|
#~ if rc.returncode == 0:
|
||
|
rc = subprocess.call(cmd, shell=True)
|
||
|
if rc == 0:
|
||
|
print('{:s}SUCCESS{:s}'.format(BOLD, ENDC))
|
||
|
total_uncompressed_bytes += os.path.getsize(image)
|
||
|
total_compressed_bytes += os.path.getsize(compressed_image)
|
||
|
num_images = len(self.image_queue)
|
||
|
ten_percent = math.ceil(num_images * 0.1)
|
||
|
percent = 0
|
||
|
count = 0
|
||
|
|
||
|
# Open log file for this instance to use
|
||
|
output = open('compressor_{:03d}.log'.format(self.index), 'w')
|
||
|
|
||
|
# Iterate over each image and compress it using the lepton binary
|
||
|
for idx, image in enumerate(self.image_queue):
|
||
|
output.write('[{:2d}] Image {:d} of {:d}:\n'.format(self.index, idx + 1, num_images))
|
||
|
output.write('\t{:s}\n'.format(image))
|
||
|
|
||
|
# Only remove original file if specified
|
||
|
if remove_orig:
|
||
|
os.remove(image)
|
||
|
else:
|
||
|
print('{:s}{:s}ERROR{:s}'.format(BOLD, RED, ENDC))
|
||
|
print('\n')
|
||
|
|
||
|
# Calculate and print total compression ratio
|
||
|
percent = (float(total_compressed_bytes) / total_uncompressed_bytes) * 100
|
||
|
print('Total compression is: {:4.2f}%'.format(percent))
|
||
|
if count == ten_percent:
|
||
|
percent += 10
|
||
|
count = 0
|
||
|
self.print_lock.acquire()
|
||
|
print('Worker {:d} -- {:d}% complete'.format(self.index, percent))
|
||
|
self.print_lock.release()
|
||
|
else:
|
||
|
count += 1
|
||
|
|
||
|
compressed_image = '{:s}.lep'.format(image)
|
||
|
cmd = '{:s} "{:s}" "{:s}"'.format(LEPTON_PATH, image, compressed_image)
|
||
|
#~ rc = subprocess.run(cmd, shell=True)
|
||
|
#~ if rc.returncode == 0:
|
||
|
rc = subprocess.call(cmd, shell=True, stdout=output, stderr=output)
|
||
|
if rc == 0:
|
||
|
output.write('SUCCESS\n')
|
||
|
|
||
|
# Only remove original file if specified
|
||
|
if self.remove_orig:
|
||
|
os.remove(image)
|
||
|
else:
|
||
|
# Append this image to the list of errors
|
||
|
output.write('ERROR\n')
|
||
|
self.errors.append(image)
|
||
|
output.write('\n\n')
|
||
|
|
||
|
output.close()
|
||
|
self.print_lock.acquire()
|
||
|
print('Worker {:d} -- [ DONE ]'.format(self.index))
|
||
|
self.print_lock.release()
|
||
|
|
||
|
def get_errors(self):
|
||
|
return self.errors
|
||
|
|
||
|
|
||
|
def do_decompress(top_path, remove_orig):
|
||
|
|
||
| ... | ... | |
|
print('{:s}{:s}ERROR{:s}'.format(BOLD, RED, ENDC))
|
||
|
print('\n')
|
||
|
|
||
|
def chunks(l, n):
|
||
|
return (l[i:i+n] for i in xrange(0, len(l), n))
|
||
|
|
||
|
def main():
|
||
|
|
||
|
# Make sure the specified path to lepton is correct
|
||
| ... | ... | |
|
if args.decompress:
|
||
|
do_decompress(args.path, args.remove==True)
|
||
|
else:
|
||
|
do_compress(args.path, args.remove==True)
|
||
|
|
||
|
# Find all *.jpg (or *.jpeg) images (ignoring case)
|
||
|
images = file_walker(sys.argv[1], '(\.jpg$)|(\.jpeg$)')
|
||
|
num_images = len(images)
|
||
|
print('Found {:d} *.jpg images\n'.format(num_images))
|
||
|
|
||
|
num_workers = max(1, multiprocessing.cpu_count())
|
||
|
print('Using {:d} workers'.format(num_workers))
|
||
|
|
||
|
# Divide input data into chunks
|
||
|
chunk_size = math.ceil(num_images / num_workers)
|
||
|
chunks = [images[x:x+chunk_size]
|
||
|
for x in range(0, num_images, chunk_size)]
|
||
|
|
||
|
# Start compressor worker threads
|
||
|
the_lock = Lock()
|
||
|
workers = []
|
||
|
for worker_idx in range(num_workers):
|
||
|
worker = CompressWrapper(worker_idx, chunks[worker_idx],
|
||
|
args.remove==True, the_lock)
|
||
|
worker.start()
|
||
|
print('\tWorker {:d} - spawned'.format(worker_idx))
|
||
|
workers.append(worker)
|
||
|
print()
|
||
|
|
||
|
# Join compressor worker threads and look for any errors
|
||
|
all_errors = []
|
||
|
for worker in workers:
|
||
|
worker.join()
|
||
|
all_errors.extend(worker.get_errors())
|
||
|
|
||
|
# Check for errors
|
||
|
if all_errors:
|
||
|
print('Errors were encountered when attempting to compress the '
|
||
|
'following images:')
|
||
|
for error in all_errors:
|
||
|
print('\t{:s}'.format(error))
|
||
|
else:
|
||
|
print('COMPLETE')
|
||
|
|
||
|
|
||
|
if __name__ == '__main__':
|
||
|
sys.exit(main())
|
||
Making the compressor side multithreaded.