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