generator.py (4579B)
1 #!/usr/bin/env python3 2 # Copyright (c) the JPEG XL Project Authors. All rights reserved. 3 # 4 # Use of this source code is governed by a BSD-style 5 # license that can be found in the LICENSE file. 6 """Tool for generating a conformance testing corpus from a set of .jxl files. 7 8 This is not the JPEG XL conformance test runner. This is a tool to generate a 9 conformance testing corpus from a set of .jxl files. 10 """ 11 12 import argparse 13 import itertools 14 import json 15 import os 16 import shutil 17 import subprocess 18 import sys 19 20 21 def GenerateConformanceCorpus(args): 22 """Generate the conformance test corpus for the given arguments.""" 23 files = [] 24 for jxl in args.inputs: 25 if os.path.isdir(jxl): 26 # Add all the .jxl files recursively. 27 for root, _, dir_files in os.walk(jxl): 28 files.extend( 29 os.path.join(root, filename) for filename in dir_files 30 if filename.lower().endswith('.jxl')) 31 else: 32 files.append(jxl) 33 34 os.makedirs(args.output, 0o755, exist_ok=True) 35 36 test_ids = [] 37 for jxl in files: 38 # Generate a unique test_id for this file based on the filename. 39 test_id = os.path.basename(jxl).lower() 40 if test_id.endswith('.jxl'): 41 test_id = test_id[:-4] 42 if test_id in test_ids: 43 for i in itertools.count(2): 44 candidate = test_id + '%02d' % i 45 if candidate not in test_ids: 46 test_id = candidate 47 break 48 test_ids.append(test_id) 49 50 test_dir = os.path.join(args.output, test_id) 51 os.makedirs(test_dir, 0o755, exist_ok=True) 52 print('Generating %s' % (test_id, )) 53 input_file = os.path.join(test_dir, 'input.jxl') 54 shutil.copy(jxl, input_file) 55 56 # The test descriptor file. 57 descriptor = {} 58 descriptor['jxl'] = 'input.jxl' 59 60 original_icc_filename = os.path.join(test_dir, 'original.icc') 61 reconstructed_filename = os.path.join(test_dir, 'reconstructed.jpg') 62 pixel_prefix = os.path.join(test_dir, 'reference') 63 output_file = pixel_prefix + '_image.npy' 64 cmd = [args.decoder, input_file, output_file] 65 metadata_filename = os.path.join(test_dir, 'test.json') 66 cmd.extend(['--metadata_out', metadata_filename]) 67 cmd.extend(['--icc_out', pixel_prefix + '.icc']) 68 69 # Decode and generate the reference files. 70 subprocess.check_call(cmd) 71 72 with open(metadata_filename, 'r') as f: 73 metadata = json.load(f) 74 75 if os.path.exists(original_icc_filename): 76 metadata['original_icc'] = "original.icc" 77 78 if os.path.exists(reconstructed_filename): 79 metadata['reconstructed_jpeg'] = "reconstructed.jpg" 80 81 for frame in metadata['frames']: 82 frame['rms_error'] = args.rmse 83 frame['peak_error'] = args.peak_error 84 85 if 'preview' in metadata: 86 metadata['preview']['rms_error'] = args.rmse 87 metadata['preview']['peak_error'] = args.peak_error 88 89 # Create the test descriptor file. 90 with open(metadata_filename, 'w') as f: 91 json.dump(metadata, f, indent=2) 92 93 # Generate a corpus descriptor with the list of the all the test_id names, 94 # one per line. 95 with open(os.path.join(args.output, 'corpus.txt'), 'w') as f: 96 f.write(''.join(line + '\n' for line in test_ids)) 97 98 99 def main(): 100 parser = argparse.ArgumentParser(description=__doc__) 101 parser.add_argument('--decoder', 102 metavar='DECODER', 103 required=True, 104 help='path to the decoder binary under test.') 105 parser.add_argument('--output', 106 metavar='DIR', 107 required=True, 108 help='path to the output directory') 109 parser.add_argument('--peak_error', 110 metavar='PEAK_ERROR', 111 type=float, 112 required=True, 113 help='peak error for each testcase') 114 parser.add_argument('--rmse', 115 metavar='RMSE', 116 type=float, 117 required=True, 118 help='max RMSE for each testcase') 119 parser.add_argument('inputs', 120 metavar='JXL', 121 nargs='+', 122 help='path to input .jxl file(s)') 123 args = parser.parse_args() 124 GenerateConformanceCorpus(args) 125 126 127 if __name__ == '__main__': 128 main()