aces_1.0.0/python/aces_ocio/generate_lut.py

   1 #!/usr/bin/env python
   2 # -*- coding: utf-8 -*-
   3
   4 """
   5 Defines objects to generate various kind of 1d, 2d and 3d LUTs in various file
   6 formats.
   7 """
   8
   9 from __future__ import division
  10
  11 import array
  12 import os
  13 import sys
  14
  15 import OpenImageIO as oiio
  16
  17 from aces_ocio.process import Process
  18
  19 __author__ = 'ACES Developers'
  20 __copyright__ = 'Copyright (C) 2014 - 2015 - ACES Developers'
  21 __license__ = ''
  22 __maintainer__ = 'ACES Developers'
  23 __email__ = 'aces@oscars.org'
  24 __status__ = 'Production'
  25
  26 __all__ = ['generate_1d_LUT_image',
  27            'write_SPI_1d',
  28            'generate_1d_LUT_from_image',
  29            'generate_3d_LUT_image',
  30            'generate_3d_LUT_from_image',
  31            'apply_CTL_to_image',
  32            'convert_bit_depth',
  33            'generate_1d_LUT_from_CTL',
  34            'correct_LUT_image',
  35            'generate_3d_LUT_from_CTL',
  36            'main']
  37
  38
  39 def generate_1d_LUT_image(ramp_1d_path,
  40                           resolution=1024,
  41                           min_value=0,
  42                           max_value=1):
  43     """
  44     Object description.
  45
  46     Parameters
  47     ----------
  48     parameter : type
  49         Parameter description.
  50
  51     Returns
  52     -------
  53     type
  54          Return value description.
  55     """
  56
  57     ramp = oiio.ImageOutput.create(ramp_1d_path)
  58
  59     spec = oiio.ImageSpec()
  60     spec.set_format(oiio.FLOAT)
  61     # spec.format.basetype = oiio.FLOAT
  62     spec.width = resolution
  63     spec.height = 1
  64     spec.nchannels = 3
  65
  66     ramp.open(ramp_1d_path, spec, oiio.Create)
  67
  68     data = array.array('f',
  69                        '\0' * spec.width * spec.height * spec.nchannels * 4)
  70     for i in range(resolution):
  71         value = float(i) / (resolution - 1) * (
  72             max_value - min_value) + min_value
  73         data[i * spec.nchannels + 0] = value
  74         data[i * spec.nchannels + 1] = value
  75         data[i * spec.nchannels + 2] = value
  76
  77     ramp.write_image(spec.format, data)
  78     ramp.close()
  79
  80
  81 def write_SPI_1d(filename, from_min, from_max, data, entries, channels):
  82     """
  83     Object description.
  84
  85     Credit to *Alex Fry* for the original single channel version of the spi1d
  86     writer.
  87
  88     Parameters
  89     ----------
  90     parameter : type
  91         Parameter description.
  92
  93     Returns
  94     -------
  95     type
  96          Return value description.
  97     """
  98
  99     with open(filename, 'w') as fp:
 100         fp.write('Version 1\n')
 101         fp.write('From %f %f\n' % (from_min, from_max))
 102         fp.write('Length %d\n' % entries)
 103         fp.write('Components %d\n' % (min(3, channels)))
 104         fp.write('{\n')
 105         for i in range(0, entries):
 106             entry = ''
 107             for j in range(0, min(3, channels)):
 108                 entry = '%s %s' % (entry, data[i * channels + j])
 109             fp.write('        %s\n' % entry)
 110         fp.write('}\n')
 111
 112
 113 def generate_1d_LUT_from_image(ramp_1d_path,
 114                                output_path=None,
 115                                min_value=0,
 116                                max_value=1):
 117     """
 118     Object description.
 119
 120     Parameters
 121     ----------
 122     parameter : type
 123         Parameter description.
 124
 125     Returns
 126     -------
 127     type
 128          Return value description.
 129     """
 130
 131     if output_path is None:
 132         output_path = '%s.%s' % (ramp_1d_path, 'spi1d')
 133
 134     ramp = oiio.ImageInput.open(ramp_1d_path)
 135
 136     spec = ramp.spec()
 137     width = spec.width
 138     channels = spec.nchannels
 139
 140     # Forcibly read data as float, the Python API doesn't handle half-float
 141     # well yet.
 142     type = oiio.FLOAT
 143     data = ramp.read_image(type)
 144
 145     write_SPI_1d(output_path, min_value, max_value, data, width, channels)
 146
 147
 148 def generate_3d_LUT_image(ramp_3d_path, resolution=32):
 149     """
 150     Object description.
 151
 152     Parameters
 153     ----------
 154     parameter : type
 155         Parameter description.
 156
 157     Returns
 158     -------
 159     type
 160          Return value description.
 161     """
 162
 163     args = ['--generate',
 164             '--cubesize',
 165             str(resolution),
 166             '--maxwidth',
 167             str(resolution * resolution),
 168             '--output',
 169             ramp_3d_path]
 170     lut_extract = Process(description='generate a 3d LUT image',
 171                           cmd='ociolutimage',
 172                           args=args)
 173     lut_extract.execute()
 174
 175
 176 def generate_3d_LUT_from_image(ramp_3d_path, output_path=None, resolution=32):
 177     """
 178     Object description.
 179
 180     Parameters
 181     ----------
 182     parameter : type
 183         Parameter description.
 184
 185     Returns
 186     -------
 187     type
 188          Return value description.
 189     """
 190
 191     if output_path is None:
 192         output_path = '%s.%s' % (ramp_3d_path, 'spi1d')
 193
 194     args = ['--extract',
 195             '--cubesize',
 196             str(resolution),
 197             '--maxwidth',
 198             str(resolution * resolution),
 199             '--input',
 200             ramp_3d_path,
 201             '--output',
 202             output_path]
 203     lut_extract = Process(description='extract a 3d LUT',
 204                           cmd='ociolutimage',
 205                           args=args)
 206     lut_extract.execute()
 207
 208
 209 def apply_CTL_to_image(input_image,
 210                        output_image,
 211                        ctl_paths=None,
 212                        input_scale=1,
 213                        output_scale=1,
 214                        global_params=None,
 215                        aces_ctl_directory=None):
 216     """
 217     Object description.
 218
 219     Parameters
 220     ----------
 221     parameter : type
 222         Parameter description.
 223
 224     Returns
 225     -------
 226     type
 227          Return value description.
 228     """
 229
 230     if ctl_paths is None:
 231         ctl_paths = []
 232     if global_params is None:
 233         global_params = {}
 234
 235     if len(ctl_paths) > 0:
 236         ctlenv = os.environ
 237         if aces_ctl_directory is not None:
 238             if os.path.split(aces_ctl_directory)[1] != 'utilities':
 239                 ctl_module_path = os.path.join(aces_ctl_directory, 'utilities')
 240             else:
 241                 ctl_module_path = aces_ctl_directory
 242             ctlenv['CTL_MODULE_PATH'] = ctl_module_path
 243
 244         args = []
 245         for ctl in ctl_paths:
 246             args += ['-ctl', ctl]
 247         args += ['-force']
 248         args += ['-input_scale', str(input_scale)]
 249         args += ['-output_scale', str(output_scale)]
 250         args += ['-global_param1', 'aIn', '1.0']
 251         for key, value in global_params.iteritems():
 252             args += ['-global_param1', key, str(value)]
 253         args += [input_image]
 254         args += [output_image]
 255
 256         ctlp = Process(description='a ctlrender process',
 257                        cmd='ctlrender',
 258                        args=args, env=ctlenv)
 259
 260         ctlp.execute()
 261
 262
 263 def convert_bit_depth(input_image, output_image, depth):
 264     """
 265     Object description.
 266
 267     Parameters
 268     ----------
 269     parameter : type
 270         Parameter description.
 271
 272     Returns
 273     -------
 274     type
 275          Return value description.
 276     """
 277
 278     args = [input_image,
 279             '-d',
 280             depth,
 281             '-o',
 282             output_image]
 283     convert = Process(description='convert image bit depth',
 284                       cmd='oiiotool',
 285                       args=args)
 286     convert.execute()
 287
 288
 289 def generate_1d_LUT_from_CTL(lut_path,
 290                              ctl_paths,
 291                              lut_resolution=1024,
 292                              identity_LUT_bit_depth='half',
 293                              input_scale=1,
 294                              output_scale=1,
 295                              global_params=None,
 296                              cleanup=True,
 297                              aces_ctl_directory=None,
 298                              min_value=0,
 299                              max_value=1):
 300     """
 301     Object description.
 302
 303     Parameters
 304     ----------
 305     parameter : type
 306         Parameter description.
 307
 308     Returns
 309     -------
 310     type
 311          Return value description.
 312     """
 313
 314     if global_params is None:
 315         global_params = {}
 316
 317     lut_path_base = os.path.splitext(lut_path)[0]
 318
 319     identity_LUT_image_float = '%s.%s.%s' % (lut_path_base, 'float', 'tiff')
 320     generate_1d_LUT_image(identity_LUT_image_float,
 321                           lut_resolution,
 322                           min_value,
 323                           max_value)
 324
 325     if identity_LUT_bit_depth != 'half':
 326         identity_LUT_image = '%s.%s.%s' % (lut_path_base, 'uint16', 'tiff')
 327         convert_bit_depth(identity_LUT_image_float,
 328                           identity_LUT_image,
 329                           identity_LUT_bit_depth)
 330     else:
 331         identity_LUT_image = identity_LUT_image_float
 332
 333     transformed_LUT_image = '%s.%s.%s' % (lut_path_base, 'transformed', 'exr')
 334     apply_CTL_to_image(identity_LUT_image,
 335                        transformed_LUT_image,
 336                        ctl_paths,
 337                        input_scale,
 338                        output_scale,
 339                        global_params,
 340                        aces_ctl_directory)
 341
 342     generate_1d_LUT_from_image(transformed_LUT_image,
 343                                lut_path,
 344                                min_value,
 345                                max_value)
 346
 347     if cleanup:
 348         os.remove(identity_LUT_image)
 349         if identity_LUT_image != identity_LUT_image_float:
 350             os.remove(identity_LUT_image_float)
 351         os.remove(transformed_LUT_image)
 352
 353
 354 def correct_LUT_image(transformed_LUT_image,
 355                       corrected_LUT_image,
 356                       lut_resolution):
 357     """
 358     Object description.
 359
 360     Parameters
 361     ----------
 362     parameter : type
 363         Parameter description.
 364
 365     Returns
 366     -------
 367     type
 368          Return value description.
 369     """
 370
 371     transformed = oiio.ImageInput.open(transformed_LUT_image)
 372
 373     transformed_spec = transformed.spec()
 374     width = transformed_spec.width
 375     height = transformed_spec.height
 376     channels = transformed_spec.nchannels
 377
 378     if width != lut_resolution * lut_resolution or height != lut_resolution:
 379         print(('Correcting image as resolution is off. '
 380                'Found %d x %d. Expected %d x %d') % (
 381                   width,
 382                   height,
 383                   lut_resolution * lut_resolution,
 384                   lut_resolution))
 385         print('Generating %s' % corrected_LUT_image)
 386
 387         # Forcibly read data as float, the Python API doesn't handle half-float
 388         # well yet.
 389         type = oiio.FLOAT
 390         source_data = transformed.read_image(type)
 391
 392         correct = oiio.ImageOutput.create(corrected_LUT_image)
 393
 394         correct_spec = oiio.ImageSpec()
 395         correct_spec.set_format(oiio.FLOAT)
 396         correct_spec.width = height
 397         correct_spec.height = width
 398         correct_spec.nchannels = channels
 399
 400         correct.open(corrected_LUT_image, correct_spec, oiio.Create)
 401
 402         dest_data = array.array('f',
 403                                 ('\0' * correct_spec.width *
 404                                  correct_spec.height *
 405                                  correct_spec.nchannels * 4))
 406         for j in range(0, correct_spec.height):
 407             for i in range(0, correct_spec.width):
 408                 for c in range(0, correct_spec.nchannels):
 409                     dest_data[(correct_spec.nchannels *
 410                                correct_spec.width * j +
 411                                correct_spec.nchannels * i + c)] = (
 412                         source_data[correct_spec.nchannels *
 413                                     correct_spec.width * j +
 414                                     correct_spec.nchannels * i + c])
 415
 416         correct.write_image(correct_spec.format, dest_data)
 417         correct.close()
 418     else:
 419         # shutil.copy(transformedLUTImage, correctedLUTImage)
 420         corrected_LUT_image = transformed_LUT_image
 421
 422     transformed.close()
 423
 424     return corrected_LUT_image
 425
 426
 427 def generate_3d_LUT_from_CTL(lut_path,
 428                              ctl_paths,
 429                              lut_resolution=64,
 430                              identity_LUT_bit_depth='half',
 431                              input_scale=1,
 432                              output_scale=1,
 433                              global_params=None,
 434                              cleanup=True,
 435                              aces_ctl_directory=None):
 436     """
 437     Object description.
 438
 439     Parameters
 440     ----------
 441     parameter : type
 442         Parameter description.
 443
 444     Returns
 445     -------
 446     type
 447          Return value description.
 448     """
 449
 450     if global_params is None:
 451         global_params = {}
 452
 453     lut_path_base = os.path.splitext(lut_path)[0]
 454
 455     identity_LUT_image_float = '%s.%s.%s' % (lut_path_base, 'float', 'tiff')
 456     generate_3d_LUT_image(identity_LUT_image_float, lut_resolution)
 457
 458     if identity_LUT_bit_depth != 'half':
 459         identity_LUT_image = '%s.%s.%s' % (lut_path_base,
 460                                            identity_LUT_bit_depth,
 461                                            'tiff')
 462         convert_bit_depth(identity_LUT_image_float,
 463                           identity_LUT_image,
 464                           identity_LUT_bit_depth)
 465     else:
 466         identity_LUT_image = identity_LUT_image_float
 467
 468     transformed_LUT_image = '%s.%s.%s' % (lut_path_base, 'transformed', 'exr')
 469     apply_CTL_to_image(identity_LUT_image,
 470                        transformed_LUT_image,
 471                        ctl_paths,
 472                        input_scale,
 473                        output_scale,
 474                        global_params,
 475                        aces_ctl_directory)
 476
 477     corrected_LUT_image = '%s.%s.%s' % (lut_path_base, 'correct', 'exr')
 478     corrected_LUT_image = correct_LUT_image(transformed_LUT_image,
 479                                             corrected_LUT_image,
 480                                             lut_resolution)
 481
 482     generate_3d_LUT_from_image(corrected_LUT_image, lut_path, lut_resolution)
 483
 484     if cleanup:
 485         os.remove(identity_LUT_image)
 486         if identity_LUT_image != identity_LUT_image_float:
 487             os.remove(identity_LUT_image_float)
 488         os.remove(transformed_LUT_image)
 489         if corrected_LUT_image != transformed_LUT_image:
 490             os.remove(corrected_LUT_image)
 491
 492
 493 def main():
 494     """
 495     Object description.
 496
 497     Parameters
 498     ----------
 499     parameter : type
 500         Parameter description.
 501
 502     Returns
 503     -------
 504     type
 505          Return value description.
 506     """
 507
 508     import optparse
 509
 510     p = optparse.OptionParser(
 511         description='A utility to generate LUTs from CTL',
 512         prog='generateLUT',
 513         version='0.01',
 514         usage='%prog [options]')
 515
 516     p.add_option('--lut', '-l', type='string', default='')
 517     p.add_option('--ctl', '-c', type='string', action='append')
 518     p.add_option('--lutResolution1d', '', type='int', default=1024)
 519     p.add_option('--lutResolution3d', '', type='int', default=33)
 520     p.add_option('--ctlReleasePath', '-r', type='string', default='')
 521     p.add_option('--bitDepth', '-b', type='string', default='float')
 522     p.add_option('--keepTempImages', '', action='store_true')
 523     p.add_option('--minValue', '', type='float', default=0)
 524     p.add_option('--maxValue', '', type='float', default=1)
 525     p.add_option('--inputScale', '', type='float', default=1)
 526     p.add_option('--outputScale', '', type='float', default=1)
 527     p.add_option('--ctlRenderParam', '-p', type='string', nargs=2,
 528                  action='append')
 529
 530     p.add_option('--generate1d', '', action='store_true')
 531     p.add_option('--generate3d', '', action='store_true')
 532
 533     options, arguments = p.parse_args()
 534
 535     lut = options.lut
 536     ctls = options.ctl
 537     lut_resolution_1d = options.lut_resolution_1d
 538     lut_resolution_3d = options.lut_resolution_3d
 539     min_value = options.minValue
 540     max_value = options.maxValue
 541     input_scale = options.inputScale
 542     output_scale = options.outputScale
 543     ctl_release_path = options.ctlReleasePath
 544     generate_1d = options.generate1d is True
 545     generate_3d = options.generate3d is True
 546     bit_depth = options.bitDepth
 547     cleanup = not options.keepTempImages
 548
 549     params = {}
 550     if options.ctlRenderParam is not None:
 551         for param in options.ctlRenderParam:
 552             params[param[0]] = float(param[1])
 553
 554     try:
 555         args_start = sys.argv.index('--') + 1
 556         args = sys.argv[args_start:]
 557     except:
 558         args_start = len(sys.argv) + 1
 559         args = []
 560
 561     if generate_1d:
 562         print('1D LUT generation options')
 563     else:
 564         print('3D LUT generation options')
 565
 566     print('lut                 : %s' % lut)
 567     print('ctls                : %s' % ctls)
 568     print('lut res 1d          : %s' % lut_resolution_1d)
 569     print('lut res 3d          : %s' % lut_resolution_3d)
 570     print('min value           : %s' % min_value)
 571     print('max value           : %s' % max_value)
 572     print('input scale         : %s' % input_scale)
 573     print('output scale        : %s' % output_scale)
 574     print('ctl render params   : %s' % params)
 575     print('ctl release path    : %s' % ctl_release_path)
 576     print('bit depth of input  : %s' % bit_depth)
 577     print('cleanup temp images : %s' % cleanup)
 578
 579     if generate_1d:
 580         generate_1d_LUT_from_CTL(lut,
 581                                  ctls,
 582                                  lut_resolution_1d,
 583                                  bit_depth,
 584                                  input_scale,
 585                                  output_scale,
 586                                  params,
 587                                  cleanup,
 588                                  ctl_release_path,
 589                                  min_value,
 590                                  max_value)
 591
 592     elif generate_3d:
 593         generate_3d_LUT_from_CTL(lut,
 594                                  ctls,
 595                                  lut_resolution_3d,
 596                                  bit_depth,
 597                                  input_scale,
 598                                  output_scale,
 599                                  params,
 600                                  cleanup,
 601                                  ctl_release_path)
 602     else:
 603         print(('\n\nNo LUT generated. '
 604                'You must choose either 1D or 3D LUT generation\n\n'))
 605
 606
 607 if __name__ == '__main__':
 608     main()
 609