summaryrefslogtreecommitdiff
path: root/utils/raspberrypi/ctt/ctt_pretty_print_json.py
blob: a4cae62d824fa907c5980742ef5d07ff08bc4fde (plain)
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
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
#!/usr/bin/env python3
#
# SPDX-License-Identifier: BSD-2-Clause
#
# Copyright 2022 Raspberry Pi Ltd
#
# Script to pretty print a Raspberry Pi tuning config JSON structure in
# version 2.0 and later formats.

import argparse
import json
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 = {
            'weights': 15,
            'table': 16,
            'luminance_lut': 16,
            'ct_curve': 3,
            'ccm': 3,
            'lut_rx': 9,
            'lut_bx': 9,
            'lut_by': 9,
            'lut_ry': 9,
            'gamma_curve': 2,
            'y_target': 2,
            'prior': 2,
            'tonemap': 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)


def pretty_print(in_json: dict, custom_elems={}) -> 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')

    encoder = Encoder(indent=4, sort_keys=False)
    encoder.custom_elems |= custom_elems
    return encoder.encode(in_json) #json.dumps(in_json, cls=Encoder, indent=4, sort_keys=False)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter, description=
                    'Prettify a version 2.0 camera tuning config JSON file.')
    parser.add_argument('-t', '--target', type=str, help='Target platform', choices=['pisp', 'vc4'], default='vc4')
    parser.add_argument('input', type=str, help='Input tuning file.')
    parser.add_argument('output', type=str, nargs='?',
                        help='Output converted tuning file. If not provided, the input file will be updated in-place.',
                        default=None)
    args = parser.parse_args()

    with open(args.input, 'r') as f:
        in_json = json.load(f)

    if args.target == 'pisp':
        from ctt_pisp import grid_size
    elif args.target == 'vc4':
        from ctt_vc4 import grid_size

    out_json = pretty_print(in_json, custom_elems={'table': grid_size[0], 'luminance_lut': grid_size[0]})

    with open(args.output if args.output is not None else args.input, 'w') as f:
        f.write(out_json)