Python script using OpenImageIO to automatically expose scene linear aces ap0 openexr source files and render into jpeg

photo_autoexpose.py

import OpenImageIO as oiio
from OpenImageIO import ImageInput, ImageOutput
from OpenImageIO import ImageBuf, ImageSpec, ImageBufAlgo


import os, sys
import shlex, subprocess
import glob

#--------------------------------------------
# CONFIG


OCIO_CONFIG = os.getenv("OCIO")
if not OCIO_CONFIG:
	OCIO_CONFIG = "/mnt/cave/dev/ocio/aces/config.ocio"

CENTER_PERCENT = 0.5
AUTOEXPOSE_SLICES = 12
EXPOSURE_TARGET = 0.3
RESIZE_WIDTH = 3414
IMG_CLIP = 10.0

# def run_command(command):
# 	process = subprocess.Popen(shlex.split(command), stdout=subprocess.PIPE)
# 	while True:
# 		output = process.stdout.readline()
# 		if output == '' and process.poll() is not None:
# 			break
# 		if output:
# 			print output.strip()
# 	rc = process.poll()
# 	return rc

def run_command(args, wait=False):

    try:
        if (wait):
            p = subprocess.Popen(
                args, 
                stdout = subprocess.PIPE)
            p.wait()
        else:
            p = subprocess.Popen(
                args, 
                stdin = None, stdout = None, stderr = None, close_fds = True)

        (result, error) = p.communicate()
        
    except subprocess.CalledProcessError as e:
        sys.stderr.write(
            "common::run_command() : [ERROR]: output = %s, error code = %s\n" 
            % (e.output, e.returncode))

    return result 



def write_jpeg(imgpath, exp):
	"""
		Write out a jpeg from the exr with auto exposure adjustment 
	"""
	# Output file path
	img_output_path = os.path.join(os.getcwd(), os.path.splitext(imgpath)[0] + ".jpg")

	# imgspec = imginput.spec()
	# imginput.close()

	# # Create ImageBuf from ImageInput
	# imgbuf = ImageBuf(imgpath)

	# # Auto exposure adjustment
	# ImageBufAlgo.mul(imgbuf, imgbuf, (exp, exp, exp, 1.0))

	# ImageBufAlgo.ociodisplay(imgbuf, imgbuf, "ACES", "RRT", colorconfig=OCIO_CONFIG)

	
	
	# # --quality 92 --attrib "jpeg:subsampling" "4:4:4" --resize 50% \
	# halfres_roi = oiio.ROI(0, RESIZE_WIDTH, 0, RESIZE_WIDTH)

	# imgbuf = ImageBufAlgo.fit(imgbuf, "lanczos3", 6.0, roi=halfres_roi)

	# # imgbuf.specmod().attribute("quality", 50) # doesn't work for some reason
	# # imgbuf.specmod().attribute("jpeg:subsampling", "4:4:4")
	# # imgbuf.write(img_output_path)
	
	# imgspec_out = imgbuf.spec()
	# imgspec_out.attribute("quality", 90)
	# imgspec_out.attribute("jpeg:subsampling", "4:4:4")
	# imgoutput = ImageOutput.create(img_output_path)
	# imgoutput.open(img_output_path, imgspec_out)
	# if not imgoutput:
	# 	print "Error:", oiio.geterror()
	# outspec = imgoutput.spec()

	# # for i in range(len(outspec.extra_attribs)):
	# # 	print i, outspec.extra_attribs[i].name, str(outspec.extra_attribs[i].type), " :"
	# # 	print "\t", outspec.extra_attribs[i].value

	# imgoutput.close()


	### This is massively complicated and doesn't work to set jpeg compression... so I'll use subprocess

	# oiiotool -v -i ../${dir}_exr/${filename}.exr --ociodisplay ACES RRT --quality 92 --attrib "jpeg:subsampling" "4:4:4" --resize 50% -o ../${dir}_jpg/${filename}.jpg
	

	# Get aspect ratio
	imginput = ImageInput.open(imgpath)
	imgspec = imginput.spec()
	imginput.close()
	input_ar = float(imgspec.width) / float(imgspec.height)
	
	if input_ar > 1.0:
		vertical = True
	else:
		vertical = False

	# Handle vertical format images
	if vertical:
		resize = "{0}x0".format(RESIZE_WIDTH)
	else:
		resize = "0x{0}".format(RESIZE_WIDTH)

	command = 'oiiotool -i {0} --mulc {1},{1},{1},1.0 --ociodisplay ACES RRT --quality 90 --attrib "jpeg:subsampling" "4:4:4" --fit:filter=lanczos3 {2} -o {3}'.format(imgpath, exp, resize, img_output_path)
	args = shlex.split(command)
	run_command(args)

def get_center_average(imgpath):
	""" Get a center average from the image
	"""

	# Create ImageBuf from ImageInput
	imgbuf = ImageBuf(imgpath)
	imgspec = imgbuf.spec()

	# Construct ROI from CENTER_PERCENT
	center = (int(imgspec.width/2), int(imgspec.height/2))
	size = (int(imgspec.width*CENTER_PERCENT), int(imgspec.height*CENTER_PERCENT))
	print "Size of box:", size
	box_roi = oiio.ROI(center[0]-size[0]/2, center[0]+size[0]/2, center[1]-size[1]/2, center[1]+size[1]/2)

	pixels = imgbuf.get_pixels(format=oiio.FLOAT, roi=box_roi)
	# print pixels.size, pixels.shape
	# print pixels[0].size, pixels[0].shape
	# print pixels[1].size, pixels[1].shape
	

	# img_min = pixels.min()
	# img_max = pixels.max()
	# print "Min: \t{0}\nMax:\t{1}".format(img_min, img_max)
	
	img_average = pixels.mean()
	print "Average:", img_average

	# # scale average by max value
	# clip_ratio = IMG_CLIP / img_max
	# mix_back = 0.8
	# img_average = img_average * (clip_ratio * mix_back + 1 * (1 - mix_back))
	# print "Scaling average by", clip_ratio, "\tNEW Average:", img_average

	return img_average
	

def get_average(imgpath):
	"""
		gather the pixel values of every scanline slice and return the average
	"""
	
	imginput = ImageInput.open(imgpath)
	if imginput == None:
		print "Error:", oiio.geterror()
		return

	imgspec = imginput.spec()
	
	if imgspec.tile_width != 0:
		print "Skipping Tiled image", imgpath
		return None



	# pixels = imginput.read_scanlines(imgspec.y, imgspec.y + imgspec.height, 0, 0, imgspec.nchannels)
	scanline_sample_width = int(round(imgspec.height / AUTOEXPOSE_SLICES))

	# imgmin = []
	# imgmax = []
	imgmean = []
	for scanline_slice in range(imgspec.y, imgspec.y+imgspec.height, scanline_sample_width):		
		pixels = imginput.read_scanline(scanline_slice, imgspec.z, "float")
		img_average = pixels.mean()
		imgmean.append(img_average)
		# imgmax.append(pixels.max())
		# imgmin.append(pixels.img_average
	# max_pixel = max(imgmax)
	# min_pixel = min(imgmin)

	img_average = reduce(lambda x, y: x + y, imgmean) / len(imgmean)

	return img_average


def gather_images():
	ARGS = sys.argv

	if len(ARGS) > 1:
		input_dir = ARGS[1]
	else:
		input_dir = os.getcwd()

	exr_images = glob.glob(os.path.join(input_dir, "*.exr"))
	exr_images.sort()
	for imgpath in exr_images:
		print "Processing... \t", imgpath

		# Calc exposure adjustment
		# img_average = get_average(imgpath)
		img_average = get_center_average(imgpath)

		exp = EXPOSURE_TARGET / img_average
		print exp



		write_jpeg(imgpath, exp)


if __name__=="__main__":
	gather_images()