path: root/utils/raspberrypi/ctt/ctt_cac.py

# SPDX-License-Identifier: BSD-2-Clause
#
# Copyright (C) 2023, Raspberry Pi Ltd
#
# ctt_cac.py - CAC (Chromatic Aberration Correction) tuning tool

from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm

from ctt_dots_locator import find_dots_locations


# This is the wrapper file that creates a JSON entry for you to append
# to your camera tuning file.
# It calculates the chromatic aberration at different points throughout
# the image and uses that to produce a martix that can then be used
# in the camera tuning files to correct this aberration.


def pprint_array(array):
    # Function to print the array in a tidier format
    array = array
    output = ""
    for i in range(len(array)):
        for j in range(len(array[0])):
            output += str(round(array[i, j], 2)) + ", "
        # Add the necessary indentation to the array
        output += "\n                   "
    # Cut off the end of the array (nicely formats it)
    return output[:-22]


def plot_shifts(red_shifts, blue_shifts):
    # If users want, they can pass a command line option to show the shifts on a graph
    # Can be useful to check that the functions are all working, and that the sample
    # images are doing the right thing
    Xs = np.array(red_shifts)[:, 0]
    Ys = np.array(red_shifts)[:, 1]
    Zs = np.array(red_shifts)[:, 2]
    Zs2 = np.array(red_shifts)[:, 3]
    Zs3 = np.array(blue_shifts)[:, 2]
    Zs4 = np.array(blue_shifts)[:, 3]

    fig, axs = plt.subplots(2, 2)
    ax = fig.add_subplot(2, 2, 1, projection='3d')
    ax.scatter(Xs, Ys, Zs, cmap=cm.jet, linewidth=0)
    ax.set_title('Red X Shift')
    ax = fig.add_subplot(2, 2, 2, projection='3d')
    ax.scatter(Xs, Ys, Zs2, cmap=cm.jet, linewidth=0)
    ax.set_title('Red Y Shift')
    ax = fig.add_subplot(2, 2, 3, projection='3d')
    ax.scatter(Xs, Ys, Zs3, cmap=cm.jet, linewidth=0)
    ax.set_title('Blue X Shift')
    ax = fig.add_subplot(2, 2, 4, projection='3d')
    ax.scatter(Xs, Ys, Zs4, cmap=cm.jet, linewidth=0)
    ax.set_title('Blue Y Shift')
    fig.tight_layout()
    plt.show()


def shifts_to_yaml(red_shift, blue_shift, image_dimensions, output_grid_size=9):
    # Convert the shifts to a numpy array for easier handling and initialise other variables
    red_shifts = np.array(red_shift)
    blue_shifts = np.array(blue_shift)
    # create a grid that's smaller than the output grid, which we then interpolate from to get the output values
    xrgrid = np.zeros((output_grid_size - 1, output_grid_size - 1))
    xbgrid = np.zeros((output_grid_size - 1, output_grid_size - 1))
    yrgrid = np.zeros((output_grid_size - 1, output_grid_size - 1))
    ybgrid = np.zeros((output_grid_size - 1, output_grid_size - 1))

    xrsgrid = []
    xbsgrid = []
    yrsgrid = []
    ybsgrid = []
    xg = np.zeros((output_grid_size - 1, output_grid_size - 1))
    yg = np.zeros((output_grid_size - 1, output_grid_size - 1))

    # Format the grids - numpy doesn't work for this, it wants a
    # nice uniformly spaced grid, which we don't know if we have yet, hence the rather mundane setup
    for x in range(output_grid_size - 1):
        xrsgrid.append([])
        yrsgrid.append([])
        xbsgrid.append([])
        ybsgrid.append([])
        for y in range(output_grid_size - 1):
            xrsgrid[x].append([])
            yrsgrid[x].append([])
            xbsgrid[x].append([])
            ybsgrid[x].append([])

    image_size = (image_dimensions[0], image_dimensions[1])
    gridxsize = image_size[0] / (output_grid_size - 1)
    gridysize = image_size[1] / (output_grid_size - 1)

    # Iterate through each dot, and it's shift values and put these into the correct grid location
    for red_shift in red_shifts:
        xgridloc = int(red_shift[0] / gridxsize)
        ygridloc = int(red_shift[1] / gridysize)
        xrsgrid[xgridloc][ygridloc].append(red_shift[2])
        yrsgrid[xgridloc][ygridloc].append(red_shift[3])

    for blue_shift in blue_shifts:
        xgridloc = int(blue_shift[0] / gridxsize)
        ygridloc = int(blue_shift[1] / gridysize)
        xbsgrid[xgridloc][ygridloc].append(blue_shift[2])
        ybsgrid[xgridloc][ygridloc].append(blue_shift[3])

    # Now calculate the average pixel shift for each square in the grid
    for x in range(output_grid_size - 1):
        for y in range(output_grid_size - 1):
            xrgrid[x, y] = np.mean(xrsgrid[x][y])
            yrgrid[x, y] = np.mean(yrsgrid[x][y])
            xbgrid[x, y] = np.mean(xbsgrid[x][y])
            ybgrid[x, y] = np.mean(ybsgrid[x][y])

    # Next, we start to interpolate the central points of the grid that gets passed to the tuning file
    input_grids = np.array([xrgrid, yrgrid, xbgrid, ybgrid])
    output_grids = np.zeros((4, output_grid_size, output_grid_size))

    # Interpolate the centre of the grid
    output_grids[:, 1:-1, 1:-1] = (input_grids[:, 1:, :-1] + input_grids[:, 1:, 1:] + input_grids[:, :-1, 1:] + input_grids[:, :-1, :-1]) / 4

    # Edge cases:
    output_grids[:, 1:-1, 0] = ((input_grids[:, :-1, 0] + input_grids[:, 1:, 0]) / 2 - output_grids[:, 1:-1, 1]) * 2 + output_grids[:, 1:-1, 1]
    output_grids[:, 1:-1, -1] = ((input_grids[:, :-1, 7] + input_grids[:, 1:, 7]) / 2 - output_grids[:, 1:-1, -2]) * 2 + output_grids[:, 1:-1, -2]
    output_grids[:, 0, 1:-1] = ((input_grids[:, 0, :-1] + input_grids[:, 0, 1:]) / 2 - output_grids[:, 1, 1:-1]) * 2 + output_grids[:, 1, 1:-1]
    output_grids[:, -1, 1:-1] = ((input_grids[:, 7, :-1] + input_grids[:, 7, 1:]) / /**
 * \var ipa_context_ops::map_buffers
 * \brief Map buffers shared between the pipeline handler and the IPA
 * \param[in] ctx The IPA context
 * \param[in] buffers The buffers to map
 * \param[in] num_buffers The number of entries in the \a buffers array
 *
 * The dmabuf file descriptors provided in \a buffers are borrowed from the
 * caller and are only guaranteed to be valid during the map_buffers() call.
 * Should the callee need to store a copy of the file descriptors, it shall
 * duplicate them first with ::%dup().
 *
 * \sa libcamera::IPAInterface::mapBuffers()
 */

/**
 * \var ipa_context_ops::unmap_buffers
 * \brief Unmap buffers shared by the pipeline to the IPA
 * \param[in] ctx The IPA context
 * \param[in] ids The IDs of the buffers to unmap
 * \param[in] num_buffers The number of entries in the \a ids array
 *
 * \sa libcamera::IPAInterface::unmapBuffers()
 */

/**
 * \var ipa_context_ops::process_event
 * \brief Process an event from the pipeline handler
 * \param[in] ctx The IPA context
 *
 * \sa libcamera::IPAInterface::processEvent()
 */

/**
 * \fn ipaCreate()
 * \brief Entry point to the IPA modules
 *
 * This function is the entry point to the IPA modules. It is implemented by
 * every IPA module, and called by libcamera to create a new IPA context.
 *
 * \return A newly created IPA context
 */

namespace libcamera {

/**
 * \struct IPAStream
 * \brief Stream configuration for the IPA interface
 *
 * The IPAStream structure stores stream configuration parameters needed by the
 * IPAInterface::configure() method. It mirrors the StreamConfiguration class
 * that is not suitable for this purpose due to not being serializable.
 */

/**
 * \var IPAStream::pixelFormat
 * \brief The stream pixel format
 */

/**
 * \var IPAStream::size
 * \brief The stream size in pixels
 */

/**
 * \struct IPABuffer
 * \brief Buffer information for the IPA interface
 *
 * The IPABuffer structure associates buffer memory with a unique ID. It is
 * used to map buffers to the IPA with IPAInterface::mapBuffers(), after which
 * buffers will be identified by their ID in the IPA interface.
 */

/**
 * \var IPABuffer::id
 * \brief The buffer unique ID
 *
 * Buffers mapped to the IPA are identified by numerical unique IDs. The IDs
 * are chosen by the pipeline handler to fulfil the following constraints:
 *
 * - IDs shall be positive integers different than zero
 * - IDs shall be unique among all mapped buffers
 *
 * When buffers are unmapped with IPAInterface::unmapBuffers() their IDs are
 * freed and may be reused for new buffer mappings.
 */

/**
 * \var IPABuffer::planes
 * \brief The buffer planes description
 *
 * Stores the dmabuf handle and length for each plane of the buffer.
 */

/**
 * \struct IPAOperationData
 * \brief Parameters for IPA operations
 *
 * The IPAOperationData structure carries parameters for the IPA operations
 * performed through the IPAInterface::processEvent() method and the
 * IPAInterface::queueFrameAction signal.
 */

/**
 * \var IPAOperationData::operation
 * \brief IPA protocol operation
 *
 * The operation field describes which operation the receiver shall perform. It
 * defines, through the IPA protocol, how the other fields of the structure are
 * interpreted. The protocol freely assigns numerical values to operations.
 */

/**
 * \var IPAOperationData::data
 * \brief Operation integer data
 *
 * The interpretation and position of different values in the array are defined
 * by the IPA protocol.
 */

/**
 * \var IPAOperationData::controls
 * \brief Operation controls data
 *
 * The interpretation and position of different values in the array are defined
 * by the IPA protocol.
 */

/**
 * \class IPAInterface
 * \brief C++ Interface for IPA implementation
 *
 * This pure virtual class defines a C++ API corresponding to the ipa_context,
<
    xred_weight = np.sum(np.dot(np.arange(x_num_pixels), red_channel))
    red_sum = np.sum(red_channel)

    green_channel = np.array(dot)[:, :, 1]
    ygreen_weight = np.sum(np.dot(green_channel, np.arange(y_num_pixels)))
    xgreen_weight = np.sum(np.dot(np.arange(x_num_pixels), green_channel * The pipeline handler shall use the IPAManager to locate a compatible
 * IPAInterface. The interface may then be used to interact with the IPA module.
 */

/**
 * \fn IPAInterface::init()
 * \brief Initialise the IPAInterface
 */

/**
 * \fn IPAInterface::configure()
 * \brief Configure the IPA stream and sensor settings
 * \param[in] streamConfig Configuration of all active streams
 * \param[in] entityControls Controls provided by the pipeline entities
 *
 * This method shall be called when the camera is started to inform the IPA of
 * the camera's streams and the sensor settings. The meaning of the numerical
 * keys in the \a streamConfig and \a entityControls maps is defined by the IPA
 * protocol.
 */

/**
 * \fn IPAInterface::mapBuffers()
 * \brief Map buffers shared between the pipeline handler and the IPA
 * \param[in] buffers List of buffers to map
 *
 * This method informs the IPA module of memory buffers set up by the pipeline
 * handler that the IPA needs to access. It provides dmabuf file handles for
 * each buffer, and associates the buffers with unique numerical IDs.
 *
 * IPAs shall map the dmabuf file handles to their address space and keep a
 * cache of the mappings, indexed by the buffer numerical IDs. The IDs are used
 * in all other IPA interface methods to refer to buffers, including the
 * unmapBuffers() method.
 *
 * All buffers that the pipeline handler wishes to share with an IPA shall be
 * mapped with this method. Buffers may be mapped all at once with a single
 * call, or mapped and unmapped dynamically at runtime, depending on the IPA
 * protocol. Regardless of the protocol, all buffers mapped at a given time
 * shall have unique numerical IDs.
 *
 * The numerical IDs have no meaning defined by the IPA interface, and IPA
 * protocols shall not give them any specific meaning either. They should be
 * treated as opaque handles by IPAs, with the only exception that ID zero is
 * invalid.
 *
 * \sa unmapBuffers()
 *
 * \todo Provide a generic implementation of mapBuffers and unmapBuffers for
 * IPAs
 */

/**
 * \fn IPAInterface::unmapBuffers()
 * \brief Unmap buffers shared by the pipeline to the IPA
 * \param[in] ids List of buffer IDs to unmap
 *
 * This method removes mappings set up with mapBuffers(). Buffers may be
 * unmapped all at once with a single call, or selectively at runtime, depending
 * on the IPA protocol. Numerical IDs of unmapped buffers may be reused when
 * mapping new buffers.
 *
 * \sa mapBuffers()
 */

/**
 * \fn IPAInterface::processEvent()
 * \brief Process an event from the pipeline handler
 * \param[in] data IPA operation data
 *
 * This operation is used by pipeline handlers to inform the IPA module of
 * events that occurred during the on-going capture operation.
 *
 * The event notified by the pipeline handler with this method is handled by the
 * IPA, which interprets the operation parameters according to the separately
 * documented IPA protocol.
 */

/**
 * \var IPAInterface::queueFrameAction
 * \brief Queue an action associated with a frame to the pipeline handler
 * \param[in] frame The frame number for the action
 * \param[in] data IPA operation data
 *
 * This signal is emitted when the IPA wishes to queue a FrameAction on the
 * pipeline. The pipeline is still responsible for the scheduling of the action
 * on its timeline.
 *
 * This signal is emitted by the IPA to queue an action to be executed by the
 * pipeline handler on a frame. The type of action is identified by the
 * \a data.operation field, as defined by the IPA protocol, and the rest of the
 * \a data is interpreted accordingly. The pipeline handler shall queue the
 * action and execute it as appropriate.
 */

} /* namespace libcamera */