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