modules/py_image: Upgrade ccm to work well with ulab.

This commit is contained in:
Kwabena W. Agyeman 2024-01-03 22:00:23 -05:00
parent 4c659b9511
commit a99d9e09ea
3 changed files with 153 additions and 19 deletions

View File

@ -0,0 +1,111 @@
# This work is licensed under the MIT license.
# Copyright (c) 2013-2024 OpenMV LLC. All rights reserved.
# https://github.com/openmv/openmv/blob/master/LICENSE
#
# Color Correction Example
#
# This example shows off using the color correction matrix multiplication
# method to apply generic matrix multiplications to images.
#
# In the example below we are going to:
# 1. Convert the image to YUV from RGB.
# 2. Apply a rotation matrix to the UV components to cause a Hue Shift.
# 3. Apply a scaling value to the UV components to cause a Saturation Shift.
# 4. Convert the image back from YUV to RGB.
#
# However, instead of applying these 4 steps to the image one a time we can apply
# them to each other before hand to produce a 3x3 matrix suitable for feeding to
# the color correction matrix method.
#
# By converting the color space to YUV this separates the "value" of pixels from
# their "hue" and "saturation". The "hue" then is just the rotation angle given
# the U/V components and the "saturation" is their magnitude.
#
# |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| = INVERSE|-0.168736, -0.331264, 0.5| * |U_rot_scaled|
# |R_rot_scaled| | 0.5, -0.418688, -0.081312| |V_rot_scaled|
#
# Note: The ccm() method can accept 3x3 and 3x4 matrices. 3x4 matrices are for
# if you want an offset to by applied. In this case things look like:
#
# |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|
#
# Keep in mind that the CCM method is just doing:
#
# |R'| |R| |R'| |R|
# |G'| = 3x3 Matrix * |G| or |G'| = 3x4 Matrix * |G|
# |B'| |B| |B'| |B|
# |1|
#
# If you are creating intermediate values using matrix math you need to end
# up back at RGB values for the final matrix you pass to CCM().
#
# Finally, you are free to do the matrix formation using 4x4 matrices. The last
# row will be ignored if you pass a 4x4 matrix (e.g. it will be treated as 3x4).
from ulab import numpy as np
import sensor
import time
import math
# Set to 0 for a grayscale image. Set above 1.0 to pump-up the saturation.
UV_SCALE = 1.0
sensor.reset()
sensor.set_pixformat(sensor.RGB565)
sensor.set_framesize(sensor.QVGA)
sensor.skip_frames(time=2000)
# These are the standard coefficents for converting RGB to 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
# Now get the inverse so we can get back to RGB from YUV.
yuv2rgb = np.linalg.inv(rgb2yuv)
clock = time.clock()
# r will be the angle by which we rotate the colors on the UV plane by.
r = 0
while True:
clock.tick()
# Increment in a loop.
r = (r + 1) % 360
a = math.radians(r)
# This is a rotation matrix which we will apply on the UV components of YUV values.
# 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
# This is the scale matrix
scale = np.array([[1, 0, 0], # noqa
[0, UV_SCALE, 0], # noqa
[0, 0, UV_SCALE]], dtype=np.float) # noqa
# Now compute the final matrix using matrix multiplication.
m = np.dot(yuv2rgb, np.dot(scale, np.dot(rot, rgb2yuv)))
# Apply the color transformation (m.flatten().tolist() also works)
img = sensor.snapshot().ccm(m.tolist())
print(clock.fps())

View File

@ -963,15 +963,9 @@ void imlib_awb(image_t *img, bool max) {
}
void imlib_ccm(image_t *img, float *ccm, bool offset) {
float rr = ccm[0], rg = ccm[3], rb = ccm[6], ro = 0.f;
float gr = ccm[1], gg = ccm[4], gb = ccm[7], go = 0.f;
float br = ccm[2], bg = ccm[5], bb = ccm[8], bo = 0.f;
if (offset) {
ro = ccm[9];
go = ccm[10];
bo = ccm[11];
}
float rr = ccm[0], rg = ccm[1], rb = ccm[2], ro = ccm[3];
float gr = ccm[4], gg = ccm[5], gb = ccm[6], go = ccm[7];
float br = ccm[8], bg = ccm[9], bb = ccm[10], bo = ccm[11];
int i_rr = IM_MIN(fast_roundf(rr * 64), 1024);
int i_rg = IM_MIN(fast_roundf(rg * 32), 512);

View File

@ -1861,23 +1861,52 @@ STATIC mp_obj_t py_awb(uint n_args, const mp_obj_t *args, mp_map_t *kw_args) {
STATIC MP_DEFINE_CONST_FUN_OBJ_KW(py_awb_obj, 1, py_awb);
STATIC mp_obj_t py_ccm(mp_obj_t img_obj, mp_obj_t ccm_obj) {
image_t *arg_img =
py_helper_arg_to_image(img_obj, ARG_IMAGE_MUTABLE);
image_t *image = py_helper_arg_to_image(img_obj, ARG_IMAGE_MUTABLE);
float ccm[12] = {};
bool offset = false;
size_t len;
mp_obj_t *items;
mp_obj_get_array(ccm_obj, &len, &items);
if ((len != 9) && (len != 12)) {
mp_raise_msg(&mp_type_ValueError, MP_ERROR_TEXT("Expected a 3x3 or 4x3 matrix!"));
// Form [[rr, rg, rb], [gr, gg, gb], [br, bg, bb]]
// Form [[rr, rg, rb], [gr, gg, gb], [br, bg, bb], [xx, xx, xx]]
// Form [[rr, rg, rb, ro], [gr, gg, gb, go], [br, bg, bb, bo]]
// Form [[rr, rg, rb, ro], [gr, gg, gb, go], [br, bg, bb, bo], [xx, xx, xx, xx]]
if ((len == 3) || (len == 4)) {
for (size_t i = 0; i < 3; i++) {
size_t row_len;
mp_obj_t *row_items;
mp_obj_get_array(items[i], &row_len, &row_items);
offset = offset || (row_len == 4);
if ((row_len == 3) || (row_len == 4)) {
for (size_t j = 0; j < row_len; j++) {
ccm[(i * 4) + j] = mp_obj_get_float(row_items[j]);
}
} else {
mp_raise_msg(&mp_type_ValueError, MP_ERROR_TEXT("Unexpected matrix dimensions!"));
}
}
// Form [rr, rg, rb, gr, gg, gb, br, bg, bb]
} else if (len == 9) {
for (size_t i = 0; i < 3; i++) {
for (size_t j = 0; j < 3; j++) {
ccm[(i * 4) + j] = mp_obj_get_float(items[(i * 3) + j]);
}
}
// Form [rr, rg, rb, ro, gr, gg, gb, go, br, bg, bb, bo]
// Form [rr, rg, rb, ro, gr, gg, gb, go, br, bg, bb, bo, xx, xx, xx, xx]
} else if (len == 12 || len == 16) {
offset = true;
for (size_t i = 0; i < 12; i++) {
ccm[i] = mp_obj_get_float(items[i]);
}
} else {
mp_raise_msg(&mp_type_ValueError, MP_ERROR_TEXT("Unexpected matrix dimensions!"));
}
float ccm[12] = {};
for (size_t i = 0; i < len; i++) {
ccm[i] = mp_obj_get_float(items[i]);
}
imlib_ccm(arg_img, ccm, len == 12);
imlib_ccm(image, ccm, offset);
return img_obj;
}
STATIC MP_DEFINE_CONST_FUN_OBJ_2(py_ccm_obj, py_ccm);