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