Permalink
Switch branches/tags
Nothing to show
Find file
Fetching contributors…
Cannot retrieve contributors at this time
executable file 500 lines (417 sloc) 14.7 KB
#!/usr/bin/env python
import sys
import os
import operator
import subprocess
import logging
import tempfile
from PIL import Image, ImageOps, ImageEnhance
import numpy
class SafeError(RuntimeError): pass
class Stereo(object):
suffix = 'out'
dest = None
def __init__(self, left, right, basename):
self.left = left
self.right = right
self.basename = basename
self.default_action = None
self._outputs = None
@classmethod
def from_single(cls, filename):
img = Image.open(filename)
width, height = img.size
single_width = width // 2
middle_bit = width - (2*single_width)
left = img.copy().crop((0, 0, single_width, height))
right = img.copy().crop((single_width+middle_bit, 0, width, height))
instance = cls(left, right, os.path.basename(filename))
instance.default_action = instance.merge
return instance
@classmethod
def from_pair(cls, left_filename, right_filename):
left = Image.open(left_filename)
right = Image.open(right_filename)
instance = cls(left, right, os.path.basename(left_filename))
instance.default_action = instance.split
return instance
def copy(self):
return type(self)(self.left, self.right, self.basename)
def swap(self, opts=None):
self.left, self.right = self.right, self.left
@property
def outputs(self):
if self._outputs is None:
self._outputs = []
return self._outputs
@property
def has_outputs(self):
return self._outputs is not None
def add_outputs(self, *outputs):
logging.debug("adding %s outputs: %r" % (len(outputs),outputs))
self.outputs.extend(outputs)
def merge(self, opts=None):
def render():
left = self.left.copy()
right = self.right.copy()
lw, lh = left.size
rw, rh = right.size
assert lh == rh
logging.debug("created image %rx%r to house %rx%r and %rx%r images" % (
lw + rw, lh, lw, lh, rw, rh))
final = Image.new(left.mode, (lw + rw, lh))
final.paste(left, (0,0))
final.paste(right, (lw, 0))
return final
self.add_outputs(Output(path=self.add_suffix(), render=render))
def add_suffix(self, suffixes=[], ext=None):
base, existing_ext = os.path.splitext(self.basename)
if ext is not None:
ext = os.path.extsep + ext
else:
ext = existing_ext
if self.dest:
base = os.path.join(self.dest, base)
return "%s-%s%s"% (base, "-".join([self.suffix] + list(suffixes)), ext)
def crop_axis(self, points, percent, sign):
a, b = points
size = b - a
assert size > 0
diff = int((size * (percent / 100.0)))
diff *= sign
if diff == 0: return points
sign = -1 if diff < 0 else 1
if sign < 0:
b += diff
else:
a += diff
return (a,b)
def slice(self,percent):
def slice(sign, img):
w,h = img.size
left, top, right, bottom = 0, 0, w, h
left, right = self.crop_axis(points=(0, w), percent=percent, sign=sign)
return img.crop((left, top, right, bottom))
self.left = slice(-1, self.left)
self.right = slice(1, self.right)
def align(self, percent):
def slice(sign, img):
w,h = img.size
left, top, right, bottom = 0, 0, w, h
top, bottom = self.crop_axis(points=(0, h), percent=percent, sign=sign)
return img.crop((left, top, right, bottom))
self.left = slice(-1, self.left)
self.right = slice(1, self.right)
def split(self, opts=None):
left = Output(path=self.add_suffix(['left']), render=lambda: self.left)
right = Output(path=self.add_suffix(['right']), render=lambda: self.right)
self.add_outputs(left, right)
def squash(self, opts=None):
def resize(img):
(w, h) = img.size
new_size = (w, h / 2)
return img.resize(new_size, Image.ANTIALIAS)
self.map(resize)
def match(self, opts=None):
# http://en.wikipedia.org/wiki/Histogram_matching
size = 256
def chunk(l):
for i in xrange(0, len(l), size):
endpoint = i+size
yield l[i:i+size]
assert len(l) == endpoint # make sure chunking was complete and exhaustive
def find_closest_idx(target, items):
# items MUST be a monotonically increasing sequence
last_value = None
for i, value in enumerate(items):
if value > target:
if last_value is None or abs(target - value) < abs(target - last_value):
return i
else:
return i - 1
last_value = value
return len(items)-1
both_histograms = zip(self.left.histogram(), self.right.histogram())
average = lambda tup: operator.add(*tup) / 2.0
merged_histograms = list(chunk(map(average, both_histograms)))
def match_merged(img):
current_histograms = chunk(img.histogram())
new_levels = []
for current, merged in zip(current_histograms, merged_histograms):
current = numpy.cumsum(current)
merged = numpy.cumsum(merged)
def new_level(level):
existing_value = current[level]
return find_closest_idx(existing_value, merged)
new_levels.extend(map(new_level, range(0, size)))
return new_levels
self.left = self.left.point(match_merged(self.left))
self.right = self.right.point(match_merged(self.right))
def equalize(self, opts=None):
self.map(ImageOps.equalize)
def map(self, fn):
self.left, self.right = map(fn, (self.left, self.right))
def brightness(self, amount):
self.map(lambda img: ImageEnhance.Brightness(img).enhance(amount))
def contrast(self, amount):
self.map(lambda img: ImageEnhance.Contrast(img).enhance(amount))
def apply_adjustments(self, opts):
for attr in ('brightness', 'contrast', 'slice', 'align'):
fn = getattr(self, attr)
fn(getattr(opts, attr))
def animate(self, opts):
resize = getattr(opts, 'resize', None)
delay = getattr(opts, 'delay', None)
gif = AnimatedOuptut(stereo=self, path=self.add_suffix(ext='gif'), resize=resize, delay=delay)
self.add_outputs(gif)
def save(self, conflict_resolver):
if not self.has_outputs:
logging.debug("no outputs specififed; performing default action")
self.default_action()
assert self.outputs is not None
for output in self.outputs:
output.save(conflict_resolver)
def __repr__(self):
return "<Image from %s with left=%r and right=%r>" % (self.basename, self.left, self.right)
class ConflictResolver(object):
REPLACE = object()
ABORT = object()
IGNORE = object()
@classmethod
def always(cls, response):
return cls(response_fn = lambda *a: response)
@classmethod
def prompt(cls):
def handle(path):
if handle.always:
return handle.always
response = raw_input("\n%s already exists.\nOverwrite it? [(Y)es / (n)o / replace (a)ll / (s)kip existing ] " % (path,))
print
response = response.lower().strip() or 'y'
if response == 'y':
return ConflictResolver.REPLACE
if response == 'a':
handle.always = ConflictResolver.REPLACE
return handle.always
if response == 's':
handle.always = ConflictResolver.IGNORE
return handle.always
else:
return ConflictResolver.ABORT
handle.always = None
return cls(handle)
def __init__(self, response_fn):
self.response_fn = response_fn
def resolve(self, path):
if os.path.exists(path):
response = self.response_fn(path)
if response is ConflictResolver.REPLACE:
os.remove(path)
elif response is ConflictResolver.ABORT:
raise RuntimeError("path %r already exists!" % (path,))
elif response is ConflictResolver.IGNORE:
return
else:
raise RuntimeError("Unknown response!")
class Output(object):
def __init__(self, path, render):
self.path = path
self.render = render
def override_path(self, path):
#TODO: copy ext!
self.path = path
def save(self, conflict_resolver):
img = self.render()
conflict_resolver.resolve(self.path)
img.save(self.path)
class AnimatedOuptut(object):
def __init__(self, stereo, path, resize=None, delay=None):
self.path = path
self.stereo = stereo
self.resize = resize
self.delay = delay
def save(self, conflict_resolver):
stereo = self.stereo.copy()
stereo.split()
def output_to_temp(output):
base, ext = os.path.splitext(output.path)
temp_path = tempfile.NamedTemporaryFile(suffix=ext, delete=False).name
output.override_path(temp_path)
map(output_to_temp, stereo.outputs)
stereo.save(conflict_resolver = ConflictResolver.always(ConflictResolver.REPLACE))
paths = [output.path for output in stereo.outputs]
try:
cmd = ['convert', '-delay', self.delay or '15', '-loop', '0']
if self.resize:
cmd.extend(['-resize', self.resize])
cmd.extend(paths)
cmd.append(self.path)
conflict_resolver.resolve(self.path)
subprocess.check_call(cmd)
finally:
try:
map(os.remove, paths)
except StandardError, e:
print >> sys.stderr, "Warning: could not remove files (%s): %r" % (e, paths)
class Gui(object):
def __init__(self, image, opts):
self.image = image
self.approved = False
self.opts = opts
self.slice = opts.slice
self.align = opts.align
self.brightness = opts.brightness
self.contrast = opts.contrast
def display(self):
from PIL import ImageTk
image = self.update(self.image.copy())
left, right = image.left, image.right
blended = Image.blend(left, right, 0.5)
blended.thumbnail((600,400))
tk_img = ImageTk.PhotoImage(blended)
self.label.configure(image = tk_img)
self.label.image = tk_img
self.label.update()
def update(self, image):
image.apply_adjustments(self)
return image
def opts_description(self):
keys = "slice", "align", "brightness", "contrast"
return " ".join(["--%s=%s" % (key, getattr(self, key)) for key in keys])
def run(self):
#TODO: allow for re-ordering / adding / removing operations
import Tkinter
def quit(event):
event.widget.quit()
def accept(event):
logging.debug("applying GUI changes to real image...")
logging.info("accepted values:\n%s" % (self.opts_description()))
self.update(self.image)
self.approved = True
quit(event)
root = Tkinter.Tk()
root.resizable(width=Tkinter.FALSE, height=Tkinter.FALSE)
def make_callback(fn):
def generated(*a):
def callback(*ignored):
fn(*a)
logging.info("updated settings to: %s" % (self.opts_description(),))
self.display()
return callback
return generated
@make_callback
def align(amount):
self.align += amount
@make_callback
def slice(amount):
self.slice += amount
@make_callback
def brightness(amount):
self.brightness += amount
@make_callback
def contrast(amount):
self.contrast += amount
@make_callback
def reset():
self.brightness = 1.0
self.contrast = 1.0
root.bind('<Escape>', quit)
root.bind('<Return>', accept)
def bind_arrows(up, down, left, right):
just = lambda x: "<%s>" % (x,)
shift = lambda x: just(x.upper()) if len(x) == 1 else "<Shift-%s>" % (x,)
root.bind(just(up), align(1))
root.bind(just(down), align(-1))
root.bind(shift(up), align(0.1))
root.bind(shift(down), align(-0.1))
root.bind(just(left), slice(1))
root.bind(just(right), slice(-1))
root.bind(shift(left), slice(0.1))
root.bind(shift(right), slice(-0.1))
bind_arrows('Up','Down','Left','Right')
bind_arrows('j','k','l','h')
#def debug(evt):
# print evt.keysym
#root.bind('<Key>', debug)
root.bind('<equal>', brightness(0.05))
root.bind('<minus>', brightness(-0.05))
root.bind('<plus>', contrast(0.05))
root.bind('<underscore>', contrast(-0.05))
root.bind('<0>', reset())
self.label = Tkinter.Label(root)
self.label.pack()
self.display()
root.title(self.image.basename)
root.mainloop()
if not self.approved:
raise SafeError("Cancelled.")
def main():
from optparse import OptionParser, OptionGroup
p = OptionParser("usage: %prog [OPTIONS] input1 [input2 [input3]]")
p.add_option('-v', '--verbose', action='store_true')
p.add_option('-f', '--force', help='replace existing files without asking', action='store_true')
input_opts = OptionGroup(p, "Input / Output")
input_opts.add_option('-l', '--left', help='inptut image (left)')
input_opts.add_option('-r', '--right', help='inptut image (right)')
input_opts.add_option('-o', '--suffix', help='output suffix')
input_opts.add_option('-d', '--dest', help='destination directory')
actions = OptionGroup(p, title="Image Actions")
actions.add_option('--squash', action='append_const', const='squash', dest='actions', default=[], help='squash half vertically')
actions.add_option('--split', action='append_const', const='split', dest='actions', help='split into left & right')
actions.add_option('--merge', action='append_const', const='merge', dest='actions', help='merge left & right images')
actions.add_option('--swap', action='append_const', const='swap', dest='actions', help='swap left & right')
actions.add_option('--equalize', action='append_const', const='equalize', dest='actions', help='equalize brightness (auto brightness / contrast)')
actions.add_option('--match', action='append_const', const='match', dest='actions', help='match left & right colours')
actions.add_option('--animate', action='append_const', const='animate', dest='actions', help='make animated gif')
gui_opts = OptionGroup(p, title="GUI Actions")
gui_opts.add_option('-g', '--gui', help='modify slice / align interactively', action='store_true')
gui_opts.add_option('--resize', help='(animation only) resize image to fit in rect (e.g "200x120")')
gui_opts.add_option('--size', dest='resize', help='alias for --resize')
gui_opts.add_option('--slice', default=0, type='float')
gui_opts.add_option('--align', default=0, type='float')
gui_opts.add_option('--brightness', default=1.0, type='float')
gui_opts.add_option('--contrast', default=1.0, type='float')
p.add_option_group(input_opts)
p.add_option_group(actions)
p.add_option_group(gui_opts)
opts, args = p.parse_args()
logging.basicConfig(level=logging.DEBUG if opts.verbose else logging.INFO, format="%(message)s")
if opts.suffix:
Stereo.suffix = opts.suffix
if opts.dest:
if not os.path.exists(opts.dest):
logging.debug("creating %s" % (opts.dest,))
os.makedirs(opts.dest)
Stereo.dest = opts.dest
conflict_resolver = ConflictResolver.prompt()
if opts.force:
conflict_resolver = ConflictResolver.always(ConflictResolver.REPLACE)
#TODO: use a graphical conflict resolver for --gui
def process(image):
logging.info("Processing %s" % image.basename)
for action in opts.actions:
logging.debug("performing action: %s" % (action,))
getattr(image, action)(opts)
logging.debug("after %s, image = %r" % (action, image))
if opts.gui:
Gui(image, opts).run()
else:
image.apply_adjustments(opts)
image.save(conflict_resolver = conflict_resolver)
if opts.left and opts.right:
assert opts.left and opts.right
process(Stereo.from_pair(opts.left, opts.right))
else:
assert not (opts.left or opts.right)
assert len(args) > 0, "Please provide an input file! (or try --help)"
for file in args:
process(Stereo.from_single(file))
assert len(opts.actions) > 0, "Please provide at least one action!"
if __name__ == '__main__':
try:
sys.exit(main())
except (SafeError, AssertionError), e:
print >> sys.stderr, e
sys.exit(1)
except (KeyboardInterrupt, EOFError):
sys.exit(1)