1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
|
# SPDX-License-Identifier: BSD-2-Clause
#
# Copyright 2022 Raspberry Pi Ltd
#
# Generate tuning file in Raspberry Pi's json format
#
# (Copied from ctt_pretty_print_json.py)
from .generator import Generator
import json
from pathlib import Path
import textwrap
class Encoder(json.JSONEncoder):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.indentation_level = 0
self.hard_break = 120
self.custom_elems = {
'table': 16,
'luminance_lut': 16,
'ct_curve': 3,
'ccm': 3,
'gamma_curve': 2,
'y_target': 2,
'prior': 2
}
def encode(self, o, node_key=None):
if isinstance(o, (list, tuple)):
# Check if we are a flat list of numbers.
if not any(isinstance(el, (list, tuple, dict)) for el in o):
s = ', '.join(json.dumps(el) for el in o)
if node_key in self.custom_elems.keys():
# Special case handling to specify number of elements in a row for tables, ccm, etc.
self.indentation_level += 1
sl = s.split(', ')
num = self.custom_elems[node_key]
chunk = [self.indent_str + ', '.join(sl[x:x + num]) for x in range(0, len(sl), num)]
t = ',\n'.join(chunk)
self.indentation_level -= 1
output = f'\n{self.indent_str}[\n{t}\n{self.indent_str}]'
elif len(s) > self.hard_break - len(self.indent_str):
# Break a long list with wraps.
self.indentation_level += 1
t = textwrap.fill(s, self.hard_break, break_long_words=False,
initial_indent=self.indent_str, subsequent_indent=self.indent_str)
self.indentation_level -= 1
output = f'\n{self.indent_str}[\n{t}\n{self.indent_str}]'
else:
# Smaller lists can remain on a single line.
output = f' [ {s} ]'
return output
else:
# Sub-structures in the list case.
self.indentation_level += 1
output = [self.indent_str + self.encode(el) for el in o]
self.indentation_level -= 1
output = ',\n'.join(output)
return f' [\n{output}\n{self.indent_str}]'
elif isinstance(o, dict):
self.indentation_level += 1
output = []
for k, v in o.items():
if isinstance(v, dict) and len(v) == 0:
# Empty config block special case.
output.append(self.indent_str + f'{json.dumps(k)}: {{ }}')
else:
# Only linebreak if the next node is a config block.
sep = f'\n{self.indent_str}' if isinstance(v, dict) else ''
output.append(self.indent_str + f'{json.dumps(k)}:{sep}{self.encode(v, k)}')
output = ',\n'.join(output)
self.indentation_level -= 1
return f'{{\n{output}\n{self.indent_str}}}'
else:
return ' ' + json.dumps(o)
@property
def indent_str(self) -> str:
return ' ' * self.indentation_level * self.indent
def iterencode(self, o, **kwargs):
return self.encode(o)
class RaspberryPiOutput(Generator):
def __init__(self):
super().__init__()
def _pretty_print(self, in_json: dict) -> str:
if 'version' not in in_json or \
'target' not in in_json or \
'algorithms' not in in_json or \
in_json['version'] < 2.0:
raise RuntimeError('Incompatible JSON dictionary has been provided')
return json.dumps(in_json, cls=Encoder, indent=4, sort_keys=False)
def write(self, output_file: Path, output_dict: dict, output_order: list):
# Write json dictionary to file using ctt's version 2 format
out_json = {
"version": 2.0,
'target': 'bcm2835',
"algorithms": [{f'{module.out_name}': output_dict[module]} for module in output_order]
}
with open(output_file, 'w') as f:
f.write(self._pretty_print(out_json))
|
