Minor updates and bug fixes.
[OpenColorIO-Configs.git] / aces_1.0.0 / python / aces_ocio / generate_lut.py
index 6059636..ac42d58 100755 (executable)
@@ -10,7 +10,6 @@ from __future__ import division
 
 import array
 import os
-import sys
 
 import OpenImageIO as oiio
 
@@ -25,6 +24,9 @@ __status__ = 'Production'
 
 __all__ = ['generate_1d_LUT_image',
            'write_SPI_1d',
+           'write_CSP_1d',
+           'write_CTL_1d',
+           'write_1d',
            'generate_1d_LUT_from_image',
            'generate_3d_LUT_image',
            'generate_3d_LUT_from_image',
@@ -78,12 +80,12 @@ def generate_1d_LUT_image(ramp_1d_path,
     ramp.close()
 
 
-def write_SPI_1d(filename, 
-                 from_min, 
-                 from_max, 
-                 data, 
-                 entries, 
-                 channels, 
+def write_SPI_1d(filename,
+                 from_min,
+                 from_max,
+                 data,
+                 entries,
+                 channels,
                  components=3):
     """
     Object description.
@@ -120,12 +122,12 @@ def write_SPI_1d(filename,
         fp.write('}\n')
 
 
-def write_CSP_1d(filename, 
-                 from_min, 
-                 from_max, 
-                 data, 
-                 entries, 
-                 channels, 
+def write_CSP_1d(filename,
+                 from_min,
+                 from_max,
+                 data,
+                 entries,
+                 channels,
                  components=3):
     """
     Object description.
@@ -149,8 +151,8 @@ def write_CSP_1d(filename,
         fp.write('CSPLUTV100\n')
         fp.write('1D\n')
         fp.write('\n')
-        fp.write('BEGIN METADATA')
-        fp.write('END METADATA')
+        fp.write('BEGIN METADATA\n')
+        fp.write('END METADATA\n')
 
         fp.write('\n')
 
@@ -168,25 +170,113 @@ def write_CSP_1d(filename,
 
         fp.write('%d\n' % entries)
         if components == 1:
-          for i in range(0, entries):
-              entry = ''
-              for j in range(3):
-                  entry = '%s %s' % (entry, data[i * channels])
-              fp.write('%s\n' % entry)
+            for i in range(0, entries):
+                entry = ''
+                for j in range(3):
+                    entry = '%s %s' % (entry, data[i * channels])
+                fp.write('%s\n' % entry)
         else:
-          for i in range(entries):
-              entry = ''
-              for j in range(components):
-                  entry = '%s %s' % (entry, data[i * channels + j])
-              fp.write('%s\n' % entry)
+            for i in range(entries):
+                entry = ''
+                for j in range(components):
+                    entry = '%s %s' % (entry, data[i * channels + j])
+                fp.write('%s\n' % entry)
         fp.write('\n')
 
-def write_1d(filename, 
-             from_min, 
-             from_max, 
-             data, 
-             data_entries, 
-             data_channels, 
+
+def write_CTL_1d(filename,
+                 from_min,
+                 from_max,
+                 data,
+                 entries,
+                 channels,
+                 components=3):
+    """
+    Object description.
+
+    Parameters
+    ----------
+    parameter : type
+        Parameter description.
+
+    Returns
+    -------
+    type
+         Return value description.
+    """
+
+    # May want to use fewer components than there are channels in the data
+    # Most commonly used for single channel LUTs
+    components = min(3, components, channels)
+
+    with open(filename, 'w') as fp:
+        fp.write('// %d x %d LUT generated by "generate_lut"\n' % (
+            entries, components))
+        fp.write('\n')
+        fp.write('const float min1d = %3.9f;\n' % from_min)
+        fp.write('const float max1d = %3.9f;\n' % from_max)
+        fp.write('\n')
+
+        # Write LUT
+        if components == 1:
+            fp.write('const float lut[] = {\n')
+            for i in range(0, entries):
+                fp.write('%s' % data[i * channels])
+                if i != (entries - 1):
+                    fp.write(',')
+                fp.write('\n')
+            fp.write('};\n')
+            fp.write('\n')
+        else:
+            for j in range(components):
+                fp.write('const float lut%d[] = {\n' % j)
+                for i in range(0, entries):
+                    fp.write('%s' % data[i * channels])
+                    if i != (entries - 1):
+                        fp.write(',')
+                    fp.write('\n')
+                fp.write('};\n')
+                fp.write('\n')
+
+        fp.write('void main\n')
+        fp.write('(\n')
+        fp.write('  input varying float rIn,\n')
+        fp.write('  input varying float gIn,\n')
+        fp.write('  input varying float bIn,\n')
+        fp.write('  input varying float aIn,\n')
+        fp.write('  output varying float rOut,\n')
+        fp.write('  output varying float gOut,\n')
+        fp.write('  output varying float bOut,\n')
+        fp.write('  output varying float aOut\n')
+        fp.write(')\n')
+        fp.write('{\n')
+        fp.write('  float r = rIn;\n')
+        fp.write('  float g = gIn;\n')
+        fp.write('  float b = bIn;\n')
+        fp.write('\n')
+        fp.write('  // Apply LUT\n')
+        if components == 1:
+            fp.write('  r = lookup1D(lut, min1d, max1d, r);\n')
+            fp.write('  g = lookup1D(lut, min1d, max1d, g);\n')
+            fp.write('  b = lookup1D(lut, min1d, max1d, b);\n')
+        elif components == 3:
+            fp.write('  r = lookup1D(lut0, min1d, max1d, r);\n')
+            fp.write('  g = lookup1D(lut1, min1d, max1d, g);\n')
+            fp.write('  b = lookup1D(lut2, min1d, max1d, b);\n')
+        fp.write('\n')
+        fp.write('  rOut = r;\n')
+        fp.write('  gOut = g;\n')
+        fp.write('  bOut = b;\n')
+        fp.write('  aOut = aIn;\n')
+        fp.write('}\n')
+
+
+def write_1d(filename,
+             from_min,
+             from_max,
+             data,
+             data_entries,
+             data_channels,
              lut_components=3,
              format='spi1d'):
     """
@@ -203,29 +293,39 @@ def write_1d(filename,
          Return value description.
     """
 
-    ocioFormatsToExtensions = {'cinespace' : 'csp',
-                               'flame'     : '3dl',
-                               'icc'       : 'icc',
-                               'houdini'   : 'lut',
-                               'lustre'    : '3dl'}
-
-    if format in ocioFormatsToExtensions:
-      if ocioFormatsToExtensions[format] == 'csp':
-        write_CSP_1d(filename,
+    ocio_formats_to_extensions = {'cinespace': 'csp',
+                                  'flame': '3dl',
+                                  'icc': 'icc',
+                                  'houdini': 'lut',
+                                  'lustre': '3dl',
+                                  'ctl': 'ctl'}
+
+    if format in ocio_formats_to_extensions:
+        if ocio_formats_to_extensions[format] == 'csp':
+            write_CSP_1d(filename,
+                         from_min,
+                         from_max,
+                         data,
+                         data_entries,
+                         data_channels,
+                         lut_components)
+        elif ocio_formats_to_extensions[format] == 'ctl':
+            write_CTL_1d(filename,
+                         from_min,
+                         from_max,
+                         data,
+                         data_entries,
+                         data_channels,
+                         lut_components)
+    else:
+        write_SPI_1d(filename,
                      from_min,
                      from_max,
                      data,
                      data_entries,
                      data_channels,
                      lut_components)
-    else:
-      write_SPI_1d(filename,
-                   from_min,
-                   from_max,
-                   data,
-                   data_entries,
-                   data_channels,
-                   lut_components)
+
 
 def generate_1d_LUT_from_image(ramp_1d_path,
                                output_path=None,
@@ -261,8 +361,8 @@ def generate_1d_LUT_from_image(ramp_1d_path,
     type = oiio.FLOAT
     ramp_data = ramp.read_image(type)
 
-    write_1d(output_path, min_value, max_value, 
-      ramp_data, ramp_width, ramp_channels, channels, format)
+    write_1d(output_path, min_value, max_value,
+             ramp_data, ramp_width, ramp_channels, channels, format)
 
 
 def generate_3d_LUT_image(ramp_3d_path, resolution=32):
@@ -293,8 +393,8 @@ def generate_3d_LUT_image(ramp_3d_path, resolution=32):
     lut_extract.execute()
 
 
-def generate_3d_LUT_from_image(ramp_3d_path, 
-                               output_path=None, 
+def generate_3d_LUT_from_image(ramp_3d_path,
+                               output_path=None,
                                resolution=32,
                                format='spi3d'):
     """
@@ -314,56 +414,56 @@ def generate_3d_LUT_from_image(ramp_3d_path,
     if output_path is None:
         output_path = '%s.%s' % (ramp_3d_path, 'spi3d')
 
-    ocioFormatsToExtensions = {'cinespace' : 'csp',
-                               'flame'     : '3dl',
-                               'icc'       : 'icc',
-                               'houdini'   : 'lut',
-                               'lustre'    : '3dl'}
-
-    if format == 'spi3d' or not (format in ocioFormatsToExtensions):
-      # Extract a spi3d LUT
-      args = ['--extract',
-              '--cubesize',
-              str(resolution),
-              '--maxwidth',
-              str(resolution * resolution),
-              '--input',
-              ramp_3d_path,
-              '--output',
-              output_path]
-      lut_extract = Process(description='extract a 3d LUT',
-                            cmd='ociolutimage',
-                            args=args)
-      lut_extract.execute()
+    ocio_formats_to_extensions = {'cinespace': 'csp',
+                                  'flame': '3dl',
+                                  'icc': 'icc',
+                                  'houdini': 'lut',
+                                  'lustre': '3dl'}
+
+    if format == 'spi3d' or not (format in ocio_formats_to_extensions):
+        # Extract a spi3d LUT
+        args = ['--extract',
+                '--cubesize',
+                str(resolution),
+                '--maxwidth',
+                str(resolution * resolution),
+                '--input',
+                ramp_3d_path,
+                '--output',
+                output_path]
+        lut_extract = Process(description='extract a 3d LUT',
+                              cmd='ociolutimage',
+                              args=args)
+        lut_extract.execute()
 
     else:
-      output_path_spi3d = '%s.%s' % (output_path, 'spi3d')
-
-      # Extract a spi3d LUT
-      args = ['--extract',
-              '--cubesize',
-              str(resolution),
-              '--maxwidth',
-              str(resolution * resolution),
-              '--input',
-              ramp_3d_path,
-              '--output',
-              output_path_spi3d]
-      lut_extract = Process(description='extract a 3d LUT',
-                            cmd='ociolutimage',
-                            args=args)
-      lut_extract.execute()
-
-      # Convert to a different format
-      args = ['--lut',
-              output_path_spi3d,
-              '--format',
-              format,
-              output_path]
-      lut_convert = Process(description='convert a 3d LUT',
-                            cmd='ociobakelut',
-                            args=args)
-      lut_convert.execute()
+        output_path_spi3d = '%s.%s' % (output_path, 'spi3d')
+
+        # Extract a spi3d LUT
+        args = ['--extract',
+                '--cubesize',
+                str(resolution),
+                '--maxwidth',
+                str(resolution * resolution),
+                '--input',
+                ramp_3d_path,
+                '--output',
+                output_path_spi3d]
+        lut_extract = Process(description='extract a 3d LUT',
+                              cmd='ociolutimage',
+                              args=args)
+        lut_extract.execute()
+
+        # Convert to a different format
+        args = ['--lut',
+                output_path_spi3d,
+                '--format',
+                format,
+                output_path]
+        lut_convert = Process(description='convert a 3d LUT',
+                              cmd='ociobakelut',
+                              args=args)
+        lut_convert.execute()
 
 
 def apply_CTL_to_image(input_image,
@@ -394,6 +494,10 @@ def apply_CTL_to_image(input_image,
 
     if len(ctl_paths) > 0:
         ctlenv = os.environ
+
+        if "/usr/local/bin" not in ctlenv['PATH'].split(':'):
+            ctlenv['PATH'] = "%s:/usr/local/bin" % ctlenv['PATH']
+
         if aces_ctl_directory is not None:
             if os.path.split(aces_ctl_directory)[1] != 'utilities':
                 ctl_module_path = os.path.join(aces_ctl_directory, 'utilities')
@@ -449,7 +553,7 @@ def convert_bit_depth(input_image, output_image, depth):
 def generate_1d_LUT_from_CTL(lut_path,
                              ctl_paths,
                              lut_resolution=1024,
-                             identity_LUT_bit_depth='half',
+                             identity_lut_bit_depth='half',
                              input_scale=1,
                              output_scale=1,
                              global_params=None,
@@ -478,30 +582,30 @@ def generate_1d_LUT_from_CTL(lut_path,
 
     lut_path_base = os.path.splitext(lut_path)[0]
 
-    identity_LUT_image_float = '%s.%s.%s' % (lut_path_base, 'float', 'tiff')
-    generate_1d_LUT_image(identity_LUT_image_float,
+    identity_lut_image_float = '%s.%s.%s' % (lut_path_base, 'float', 'tiff')
+    generate_1d_LUT_image(identity_lut_image_float,
                           lut_resolution,
                           min_value,
                           max_value)
 
-    if identity_LUT_bit_depth not in ['half', 'float']:
-        identity_LUT_image = '%s.%s.%s' % (lut_path_base, 'uint16', 'tiff')
-        convert_bit_depth(identity_LUT_image_float,
-                          identity_LUT_image,
-                          identity_LUT_bit_depth)
+    if identity_lut_bit_depth not in ['half', 'float']:
+        identity_lut_image = '%s.%s.%s' % (lut_path_base, 'uint16', 'tiff')
+        convert_bit_depth(identity_lut_image_float,
+                          identity_lut_image,
+                          identity_lut_bit_depth)
     else:
-        identity_LUT_image = identity_LUT_image_float
+        identity_lut_image = identity_lut_image_float
 
-    transformed_LUT_image = '%s.%s.%s' % (lut_path_base, 'transformed', 'exr')
-    apply_CTL_to_image(identity_LUT_image,
-                       transformed_LUT_image,
+    transformed_lut_image = '%s.%s.%s' % (lut_path_base, 'transformed', 'exr')
+    apply_CTL_to_image(identity_lut_image,
+                       transformed_lut_image,
                        ctl_paths,
                        input_scale,
                        output_scale,
                        global_params,
                        aces_ctl_directory)
 
-    generate_1d_LUT_from_image(transformed_LUT_image,
+    generate_1d_LUT_from_image(transformed_lut_image,
                                lut_path,
                                min_value,
                                max_value,
@@ -509,14 +613,14 @@ def generate_1d_LUT_from_CTL(lut_path,
                                format)
 
     if cleanup:
-        os.remove(identity_LUT_image)
-        if identity_LUT_image != identity_LUT_image_float:
-            os.remove(identity_LUT_image_float)
-        os.remove(transformed_LUT_image)
+        os.remove(identity_lut_image)
+        if identity_lut_image != identity_lut_image_float:
+            os.remove(identity_lut_image_float)
+        os.remove(transformed_lut_image)
 
 
-def correct_LUT_image(transformed_LUT_image,
-                      corrected_LUT_image,
+def correct_LUT_image(transformed_lut_image,
+                      corrected_lut_image,
                       lut_resolution):
     """
     Object description.
@@ -532,7 +636,7 @@ def correct_LUT_image(transformed_LUT_image,
          Return value description.
     """
 
-    transformed = oiio.ImageInput.open(transformed_LUT_image)
+    transformed = oiio.ImageInput.open(transformed_lut_image)
 
     transformed_spec = transformed.spec()
     width = transformed_spec.width
@@ -546,14 +650,14 @@ def correct_LUT_image(transformed_LUT_image,
                   height,
                   lut_resolution * lut_resolution,
                   lut_resolution))
-        print('Generating %s' % corrected_LUT_image)
+        print('Generating %s' % corrected_lut_image)
 
         # Forcibly read data as float, the Python API doesn't handle half-float
         # well yet.
         type = oiio.FLOAT
         source_data = transformed.read_image(type)
 
-        correct = oiio.ImageOutput.create(corrected_LUT_image)
+        correct = oiio.ImageOutput.create(corrected_lut_image)
 
         correct_spec = oiio.ImageSpec()
         correct_spec.set_format(oiio.FLOAT)
@@ -561,7 +665,7 @@ def correct_LUT_image(transformed_LUT_image,
         correct_spec.height = width
         correct_spec.nchannels = channels
 
-        correct.open(corrected_LUT_image, correct_spec, oiio.Create)
+        correct.open(corrected_lut_image, correct_spec, oiio.Create)
 
         dest_data = array.array('f',
                                 ('\0' * correct_spec.width *
@@ -581,17 +685,17 @@ def correct_LUT_image(transformed_LUT_image,
         correct.close()
     else:
         # shutil.copy(transformedLUTImage, correctedLUTImage)
-        corrected_LUT_image = transformed_LUT_image
+        corrected_lut_image = transformed_lut_image
 
     transformed.close()
 
-    return corrected_LUT_image
+    return corrected_lut_image
 
 
 def generate_3d_LUT_from_CTL(lut_path,
                              ctl_paths,
                              lut_resolution=64,
-                             identity_LUT_bit_depth='half',
+                             identity_lut_bit_depth='half',
                              input_scale=1,
                              output_scale=1,
                              global_params=None,
@@ -617,49 +721,50 @@ def generate_3d_LUT_from_CTL(lut_path,
 
     lut_path_base = os.path.splitext(lut_path)[0]
 
-    identity_LUT_image_float = '%s.%s.%s' % (lut_path_base, 'float', 'tiff')
-    generate_3d_LUT_image(identity_LUT_image_float, lut_resolution)
+    identity_lut_image_float = '%s.%s.%s' % (lut_path_base, 'float', 'tiff')
+    generate_3d_LUT_image(identity_lut_image_float, lut_resolution)
 
-    if identity_LUT_bit_depth not in ['half', 'float']:
-        identity_LUT_image = '%s.%s.%s' % (lut_path_base,
-                                           identity_LUT_bit_depth,
+    if identity_lut_bit_depth not in ['half', 'float']:
+        identity_lut_image = '%s.%s.%s' % (lut_path_base,
+                                           identity_lut_bit_depth,
                                            'tiff')
-        convert_bit_depth(identity_LUT_image_float,
-                          identity_LUT_image,
-                          identity_LUT_bit_depth)
+        convert_bit_depth(identity_lut_image_float,
+                          identity_lut_image,
+                          identity_lut_bit_depth)
     else:
-        identity_LUT_image = identity_LUT_image_float
+        identity_lut_image = identity_lut_image_float
 
-    transformed_LUT_image = '%s.%s.%s' % (lut_path_base, 'transformed', 'exr')
-    apply_CTL_to_image(identity_LUT_image,
-                       transformed_LUT_image,
+    transformed_lut_image = '%s.%s.%s' % (lut_path_base, 'transformed', 'exr')
+    apply_CTL_to_image(identity_lut_image,
+                       transformed_lut_image,
                        ctl_paths,
                        input_scale,
                        output_scale,
                        global_params,
                        aces_ctl_directory)
 
-    corrected_LUT_image = '%s.%s.%s' % (lut_path_base, 'correct', 'exr')
-    corrected_LUT_image = correct_LUT_image(transformed_LUT_image,
-                                            corrected_LUT_image,
+    corrected_lut_image = '%s.%s.%s' % (lut_path_base, 'correct', 'exr')
+    corrected_lut_image = correct_LUT_image(transformed_lut_image,
+                                            corrected_lut_image,
                                             lut_resolution)
 
-    generate_3d_LUT_from_image(corrected_LUT_image, 
-                               lut_path, 
-                               lut_resolution, 
+    generate_3d_LUT_from_image(corrected_lut_image,
+                               lut_path,
+                               lut_resolution,
                                format)
 
     if cleanup:
-        os.remove(identity_LUT_image)
-        if identity_LUT_image != identity_LUT_image_float:
-            os.remove(identity_LUT_image_float)
-        os.remove(transformed_LUT_image)
-        if corrected_LUT_image != transformed_LUT_image:
-            os.remove(corrected_LUT_image)
+        os.remove(identity_lut_image)
+        if identity_lut_image != identity_lut_image_float:
+            os.remove(identity_lut_image_float)
+        os.remove(transformed_lut_image)
+        if corrected_lut_image != transformed_lut_image:
+            os.remove(corrected_lut_image)
         if format != 'spi3d':
             lut_path_spi3d = '%s.%s' % (lut_path, 'spi3d')
             os.remove(lut_path_spi3d)
 
+
 def main():
     """
     Object description.
@@ -723,31 +828,24 @@ def main():
         for param in options.ctlRenderParam:
             params[param[0]] = float(param[1])
 
-    try:
-        args_start = sys.argv.index('--') + 1
-        args = sys.argv[args_start:]
-    except:
-        args_start = len(sys.argv) + 1
-        args = []
-
     if generate_1d:
         print('1D LUT generation options')
     else:
         print('3D LUT generation options')
 
-    print('lut                 : %s' % lut)
-    print('format              : %s' % format)
-    print('ctls                : %s' % ctls)
-    print('lut res 1d          : %s' % lut_resolution_1d)
-    print('lut res 3d          : %s' % lut_resolution_3d)
-    print('min value           : %s' % min_value)
-    print('max value           : %s' % max_value)
-    print('input scale         : %s' % input_scale)
-    print('output scale        : %s' % output_scale)
-    print('ctl render params   : %s' % params)
-    print('ctl release path    : %s' % ctl_release_path)
-    print('bit depth of input  : %s' % bit_depth)
-    print('cleanup temp images : %s' % cleanup)
+    print('LUT                 : %s' % lut)
+    print('Format              : %s' % format)
+    print('CTLs                : %s' % ctls)
+    print('LUT Res 1d          : %s' % lut_resolution_1d)
+    print('LUT Res 3d          : %s' % lut_resolution_3d)
+    print('Min Value           : %s' % min_value)
+    print('Max Value           : %s' % max_value)
+    print('Input Scale         : %s' % input_scale)
+    print('Output Scale        : %s' % output_scale)
+    print('CTL Render Params   : %s' % params)
+    print('CTL Release Path    : %s' % ctl_release_path)
+    print('Input Bit Depth     : %s' % bit_depth)
+    print('Cleanup Temp Images : %s' % cleanup)
 
     if generate_1d:
         generate_1d_LUT_from_CTL(lut,
@@ -775,10 +873,9 @@ def main():
                                  ctl_release_path,
                                  format=format)
     else:
-        print(('\n\nNo LUT generated. '
+        print(('\n\nNo LUT generated! '
                'You must choose either 1D or 3D LUT generation\n\n'))
 
 
 if __name__ == '__main__':
     main()
-