Метрика Mutial information с многоканальными изображениями в батче

Подробнее о метрике MI: https://matthew-brett.github.io/teaching/mutual_information.html

Реализация: https://scikit-learn.org/stable/modules/generated/sklearn.metrics.mutual_info_score.html?highlight=mutual_info_score

В качестве метрики хотела использовать реализацию: https://scikit-learn.org/stable/modules/generated/sklearn.metrics.normalized_mutual_info_score.html?highlight=mutual_info#sklearn.metrics.normalized_mutual_info_score

На вход подаются одномерные массивы labels_true и labels_pred.

Хотела использовать эту метрику в свёрточной нейронной сети для набора изображений (batch_size, height, width, channels=3), поэтому изначально бездумно привела к весь батч к одномерному представлению посредством np_images_true.ravel(). А потом подумала:

Cейчас изображения в батчах никак структурно не связаны по измерению batch_size, а приводя весь батч к одномерному представлению я строю совместную 2D гистограмму по всем значениям интенсивностей пикселей со всех изображений. Полученная оценка сильно отличается от усредненного по батчу значения метрики, посчитанного для каждого изображения отдельно. Что вообще отображает в этом случае метрика?

В случае же одноканальных изображений, считая MI для каждого снимка, я получу интерпретируемое усредненное по батчу значение метрики. Но что делать с многоканальными изображениями? Считать для каждого канала каждого изображения значение MI, затем усреднять по каналам, а потом по batch_size? Или же строить совместную 2D гистограмму по всем каналам?

Update: В scikit-image (0.19.0.dev0) есть реализация, которая принимает на вход 2 многомерных представления изображений: https://scikit-image.org/docs/dev/api/skimage.metrics.html#skimage.metrics.normalized_mutual_information, но она довольно странная: берут входное изображение и считают максимальный размер по всем измерениям max_shape, потом делают паддинг, дополняя нулями представление по каждому измерению до max_shape. Только после этого они применяют np.reshape() до одного измерения.

def _pad_to(arr, shape):
    """Pad an array with trailing zeros to a given target shape.
    Parameters
    ----------
    arr : ndarray
        The input array.
    shape : tuple
        The target shape.
    Returns
    -------
    padded : ndarray
        The padded array.
    Examples
    --------
    >>> _pad_to(np.ones((1, 1), dtype=int), (1, 3))
    array([[1, 0, 0]])
    """
    if not all(s >= i for s, i in zip(shape, arr.shape)):
        raise ValueError(f"Target shape {shape} cannot be smaller than input"
                         f"shape {arr.shape} along any axis.")
    padding = [(0, s-i) for s, i in zip(shape, arr.shape)]
    return np.pad(arr, pad_width=padding, mode='constant', constant_values=0)


def normalized_mutual_information(image0, image1, *, bins=100):
    if image0.ndim != image1.ndim:
        raise ValueError('NMI requires images of same number of dimensions. '
                         'Got {}D for `image0` and {}D for `image1`.'
                         .format(image0.ndim, image1.ndim))
    if image0.shape != image1.shape:
        max_shape = np.maximum(image0.shape, image1.shape)
        padded0 = _pad_to(image0, max_shape)
        padded1 = _pad_to(image1, max_shape)
    else:
        padded0, padded1 = image0, image1

    hist, bin_edges = np.histogramdd(
            [np.reshape(padded0, -1), np.reshape(padded1, -1)],
            bins=bins,
            density=True,
            )

    H0 = entropy(np.sum(hist, axis=0))
    H1 = entropy(np.sum(hist, axis=1))
    H01 = entropy(np.reshape(hist, -1))

    return (H0 + H1) / H01

https://github.com/scikit-image/scikit-image/blob/main/skimage/metrics/simple_metrics.py

Update 2: Другая метрика Normalized Cross Correlation также не определяет работу с многоканальными изображениями, обсуждение:
https://www.researchgate.net/post/How-to-calculate-NCC-normalized-cross-correlation-of-RGB-image-Is-there-is-any-inbuilt-function-for-that