Project

General

Profile

Download (6.51 KB) Statistics
| Branch: | Tag: | Revision:

root/software/lepton_wrapper/lepton_wrapper.py @ 5e0d193b

e2a82fa9 David Sorber
#!/opt/local/bin/python3

a29fbdb0 David Sorber
#!/usr/bin/python3

import argparse
import glob
2aa55fa2 David Sorber
import math
from multiprocessing import Process
from multiprocessing import Lock
import multiprocessing
a29fbdb0 David Sorber
import os
import re
import subprocess
import sys

e2a82fa9 David Sorber
LEPTON_PATH = '/Users/dsorber/Documents/lepton/lepton'
a29fbdb0 David Sorber
# Bash text formatting codes
BOLD = '\033[1m'
ENDC = '\033[0m'
RED = '\033[31m'
YELLOW = '\033[33m'
PURPLE = '\033[35m'

e2a82fa9 David Sorber
def file_walker(top_path, regex):
a29fbdb0 David Sorber
e2a82fa9 David Sorber
files = []
a29fbdb0 David Sorber
e2a82fa9 David Sorber
for dirpath, dirnames, filenames in os.walk(top_path):
a29fbdb0 David Sorber
for filename in filenames:
fullpath = os.path.join(dirpath, filename)
e2a82fa9 David Sorber
if re.search(regex, fullpath, re.IGNORECASE):
a29fbdb0 David Sorber
#~ print(fullpath)
e2a82fa9 David Sorber
files.append(fullpath)
a29fbdb0 David Sorber
for dirname in dirnames:
e2a82fa9 David Sorber
file_walker(os.path.join(dirpath, dirname), regex)
return files
a29fbdb0 David Sorber
2aa55fa2 David Sorber
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
a29fbdb0 David Sorber
2aa55fa2 David Sorber
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))
e2a82fa9 David Sorber
2aa55fa2 David Sorber
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
a29fbdb0 David Sorber
e2a82fa9 David Sorber
def do_decompress(top_path, remove_orig):
# Find all *.lep files (ignoring case)
images = file_walker(sys.argv[1], '\.lep$')
num_images = len(images)
print('Found {:d} *.lep files'.format(num_images))
# Iterate over each image and decompress 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))
cmd = '{:s} "{:s}" "{:s}"'.format(LEPTON_PATH, image, image[:-4])
#~ 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))
# Only remove original file if specified
if remove_orig:
os.remove(image)
else:
print('{:s}{:s}ERROR{:s}'.format(BOLD, RED, ENDC))
print('\n')

2aa55fa2 David Sorber
def chunks(l, n):
return (l[i:i+n] for i in xrange(0, len(l), n))

e2a82fa9 David Sorber
def main():
# Make sure the specified path to lepton is correct
if not (os.path.isfile(LEPTON_PATH) and os.access(LEPTON_PATH, os.X_OK)):
print('{:s}{:s}ERROR:{:s} unable to find lepton executable'
.format(BOLD, RED, ENDC))
return -1

# Parse command line arguments
parser = argparse.ArgumentParser(description='Lepton Wrapper')
parser.add_argument('path',
help='path in which to look for images')
parser.add_argument('-d', '--decompress', action='store_true',
default=False, help='decompress *.lep files')
parser.add_argument('-r', '--remove', action='store_true', default=False,
help='remove orignal files after compress/decompress')
args = parser.parse_args()

# Either compress or decompress depending on mode parameter
if args.decompress:
do_decompress(args.path, args.remove==True)
else:
2aa55fa2 David Sorber
# 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))
e2a82fa9 David Sorber
2aa55fa2 David Sorber
# 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')
e2a82fa9 David Sorber
a29fbdb0 David Sorber
if __name__ == '__main__':
sys.exit(main())