color_correction.py

# 本作品采用MIT许可证授权。
# 版权所有 (c) 2013-2024 OpenMV LLC。保留所有权利。
# https://github.com/openmv/openmv/blob/master/LICENSE
#
# 色彩校正示例
#
# 此示例展示了如何使用色彩校正矩阵乘法
# 方法对图像应用通用矩阵乘法。
#
# 在下面的示例中，我们将：
# 1. 将图像从RGB转换为YUV。
# 2. 对UV分量应用旋转矩阵以引起色调偏移。
# 3. 对UV分量应用缩放值以引起饱和度偏移。
# 4. 将图像从YUV转换回RGB。
#
# 然而，我们不必依次对图像应用这4个步骤，而是可以
# 预先将它们相互结合，生成一个适合输入到
# 色彩校正矩阵方法的3x3矩阵。
#
# 通过转换色彩空间至YUV，可将像素"值"从
# 它们的“色调”和“饱和度”。“色调”就是U/V分量给出的旋转角度，
# 而“饱和度”是它们的幅值。
#
# |Y|   |    0.299,     0.587,     0.114|   |R|
# |U| = |-0.168736, -0.331264,       0.5| * |G|
# |V|   |      0.5, -0.418688, -0.081312|   |B|
#
# |Y_rot|   |1,           0,            0|   |Y|
# |U_rot| = |0, math.cos(a), -math.sin(a)| * |U|
# |V_rot|   |0, math.sin(a),  math.cos(a)|   |V|
#
# |Y_rot_scaled|   |1, 0,  0|   |Y|
# |U_rot_scaled| = |0, s,  0| * |U|
# |V_rot_scaled|   |0, 0,  s|   |V|
#
# |R_rot_scaled|          |    0.299,     0.587,     0.114|   |Y_rot_scaled|
# |R_rot_scaled| = 逆矩阵|-0.168736, -0.331264,       0.5| * |U_rot_scaled|
# |R_rot_scaled|          |      0.5, -0.418688, -0.081312|   |V_rot_scaled|
#
# 注意：ccm()方法可以接受3x3和3x4矩阵。3x4矩阵用于
# 需要应用偏移量的情况。此时公式如下：
#
# |Y|   |    0.299,     0.587,     0.114,  y_offset|   |R|
# |U| = |-0.168736, -0.331264,       0.5,  u_offset| * |G|
# |V|   |      0.5, -0.418688, -0.081312,  v_offset|   |B|
#                                                      |1|
#
# 请记住，CCM方法只是执行以下操作：
#
# |R'|                |R|      |R'|                |R|
# |G'| = 3x3 矩阵 * |G|  或  |G'| = 3x4 矩阵 * |G|
# |B'|                |B|      |B'|                |B|
#                                                  |1|
#
# 如果您使用矩阵运算创建中间值，最终传给CCM()的矩阵
# 必须回到RGB值。
#
# 最后，您也可以使用4x4矩阵来构造矩阵。如果传入4x4矩阵，
# 最后一行将被忽略（即按3x4处理）。

from ulab import numpy as np
import csi
import time
import math

# 设为0得到灰度图像。设为大于1.0可提升饱和度。
UV_SCALE = 1.0

csi0 = csi.CSI()
csi0.reset()
csi0.pixformat(csi.RGB565)
csi0.framesize(csi.QVGA)
csi0.snapshot(time=2000)

# 这些是RGB转YUV的标准系数。
rgb2yuv = np.array([[    0.299,     0.587,     0.114], # noqa
                    [-0.168736, -0.331264,       0.5], # noqa
                    [      0.5, -0.418688, -0.081312]], dtype=np.float) # noqa

# 现在求逆矩阵，以便从YUV转回RGB。
yuv2rgb = np.linalg.inv(rgb2yuv)

clock = time.clock()

# r是我们在UV平面上旋转颜色的角度。
r = 0

while True:
    clock.tick()

    # 在循环中递增。
    r = (r + 1) % 360
    a = math.radians(r)

    # 这是我们将应用于YUV值中UV分量的旋转矩阵。
    # https://en.wikipedia.org/wiki/Rotation_matrix
    rot = np.array([[1,           0,            0], # noqa
                    [0, math.cos(a), -math.sin(a)], # noqa
                    [0, math.sin(a),  math.cos(a)]], dtype=np.float) # noqa

    # 这是缩放矩阵
    scale = np.array([[1,        0, 0], # noqa
                      [0, UV_SCALE, 0], # noqa
                      [0,        0, UV_SCALE]], dtype=np.float) # noqa

    # 现在用矩阵乘法计算最终矩阵。
    m = np.dot(yuv2rgb, np.dot(scale, np.dot(rot, rgb2yuv)))

    # 应用颜色变换（m.flatten().tolist()也可以）
    img = csi0.snapshot().ccm(m.tolist())

    print(clock.fps())