summaryrefslogtreecommitdiff
path: root/utils/tuning/libtuning/utils.py
blob: c70dfae049bc00959ed53410b136dba3282e914a (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
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
# SPDX-License-Identifier: BSD-2-Clause
#
# Copyright (C) 2019, Raspberry Pi Ltd
# Copyright (C) 2022, Paul Elder <paul.elder@ideasonboard.com>
#
# Utilities for libtuning

import cv2
import decimal
import math
import numpy as np
import os
from pathlib import Path
import re
import sys
import logging

import libtuning as lt
from libtuning.image import Image
from libtuning.macbeth import locate_macbeth

logger = logging.getLogger(__name__)

# Utility functions


def get_module_by_type_name(modules, name):
    for module in modules:
        if module.type == name:
            return module
    return None


# Private utility functions


def _list_image_files(directory):
    d = Path(directory)
    files = [d.joinpath(f) for f in os.listdir(d)
             if re.search(r'\.(jp[e]g$)|(dng$)', f)]
    files.sort()
    return files


def _parse_image_filename(fn: Path):
    lsc_only = False
    color_temperature = None
    lux = None

    parts = fn.stem.split('_')
    for part in parts:
        if part == 'alsc':
            lsc_only = True
            continue
        r = re.match(r'(\d+)[kK]', part)
        if r:
            color_temperature = int(r.group(1))
            continue
        r = re.match(r'(\d+)[lLuU]', part)
        if r:
            lux = int(r.group(1))

    if color_temperature is None:
        logger.error(f'The file name of "{fn.name}" does not contain a color temperature')

    if lux is None and lsc_only is False:
        logger.error(f'The file name of "{fn.name}" must either contain alsc or a lux level')

    return color_temperature, lux, lsc_only


# \todo Implement this from check_imgs() in ctt.py
def _validate_images(images):
    return True


# Public utility functions


# @brief Load images into a single list of Image instances
# @param input_dir Directory from which to load image files
# @param config Configuration dictionary
# @param load_nonlsc Whether or not to load non-lsc images
# @param load_lsc Whether or not to load lsc-only images
# @return A list of Image instances
def load_images(input_dir: str, config: dict, load_nonlsc: bool, load_lsc: bool) -> list:
    files = _list_image_files(input_dir)
    if len(files) == 0:
        logger.error(f'No images found in {input_dir}')
        return None

    images = []
    for f in files:
        color, lux, lsc_only = _parse_image_filename(f)
        if color is None:
            continue

        # Skip lsc image if we don't need it
        if lsc_only and not load_lsc:
            logger.warning(f'Skipping {f.name} as this tuner has no LSC module')
            continue

        # Skip non-lsc image if we don't need it
        if not lsc_only and not load_nonlsc:
            logger.warning(f'Skipping {f.name} as this tuner only has an LSC module')
            continue

        # Load image
        try:
            image = Image(f)
        except Exception as e:
            logger.error(f'Failed to load image {f.name}: {e}')
            continue

        # Populate simple fields
        image.lsc_only = lsc_only
        image.color = color
        image.lux = lux

        # Black level comes from the TIFF tags, but they are overridable by the
        # config file.
        if 'blacklevel' in config['general']:
            image.blacklevel_16 = config['general']['blacklevel']

        if lsc_only:
            images.append(image)
            continue

        # Handle macbeth
        macbeth = locate_macbeth(config)
        if macbeth is None:
            continue

        images.append(image)

    if not _validate_images(images):
        return None

    return images



"""
Some code that will save virtual macbeth charts that show the difference between optimised matrices and non optimised matrices

The function creates an image that is 1550 by 1050 pixels wide, and fills it with patches which are 200x200 pixels in size
Each patch contains the ideal color, the color from the original matrix, and the color from the final matrix
_________________
|               |
|  Ideal Color  |
|_______________|
|  Old  |  new  |
| Color | Color |
|_______|_______|

Nice way of showing how the optimisation helps change the colors and the color matricies
"""
def visualise_macbeth_chart(macbeth_rgb, original_rgb, new_rgb, output_filename):
    image = np.zeros((1050, 1550, 3), dtype=np.uint8)
    colorindex = -1
    for y in range(6):
        for x in range(4):  # Creates 6 x 4 grid of macbeth chart
            colorindex += 1
            xlocation = 50 + 250 * x  # Means there is 50px of black gap between each square, more like the real macbeth chart.
            ylocation = 50 + 250 * y
            for g in range(200):
                for i in range(100):
                    image[xlocation + i, ylocation + g] = macbeth_rgb[colorindex]
            xlocation = 150 + 250 * x
            ylocation = 50 + 250 * y
            for i in range(100):
                for g in range(100):
                    image[xlocation + i, ylocation + g] = original_rgb[colorindex]  # Smaller squares below to compare the old colors with the new ones
            xlocation = 150 + 250 * x
            ylocation = 150 + 250 * y
            for i in range(100):
                for g in range(100):
                    image[xlocation + i, ylocation + g] = new_rgb[colorindex]

    im_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    cv2.imwrite(f'{output_filename} Generated Macbeth Chart.png', im_bgr)