Поиск субимажа внутри изображения Numpy

У меня есть два массива Numpy (3-мерный uint8), преобразованный из изображений PIL.

Я хочу найти, если первое изображение содержит второе изображение, и если да, найдите координаты верхнего левого пикселя внутри первого изображения, в котором находится совпадение.

Есть ли способ сделать это чисто в Numpy, достаточно быстро, вместо того, чтобы использовать (4! очень медленные) чистые контуры Python?

Пример 2D:

a = numpy.array([
    [0, 1,  2,  3],
    [4, 5,  6,  7],
    [8, 9, 10, 11]
])
b = numpy.array([
    [2, 3],
    [6, 7]
])

Как сделать что-то подобное?

position = a.find(b)

position будет тогда (0, 2).

Ответ 1

Это можно сделать, используя scipy correlate2d, а затем используя argmax, чтобы найти пик в кросс-корреляции.

Ниже приведено более полное объяснение математики и идей и некоторые примеры.

Если вы хотите остаться в чистом Numpy и даже не использовать scipy, или если изображения большие, вам, вероятно, лучше всего использовать подход на основе FFT для кросс-корреляций.

Изменить: вопрос специально задан для решения чистого Numpy. Но если вы можете использовать OpenCV или другие инструменты обработки изображений, очевидно, что проще использовать один из них. Примером такого является приведенный ниже PiQuer, который я бы рекомендовал, если вы можете его использовать.

Ответ 2

Я делаю это с помощью OpenCV matchTemplate. Существует превосходная привязка к python для OpenCV, которая использует numpy внутри, поэтому изображения представляют собой просто массивы с несколькими числами. Например, допустим, что у вас есть файл BGR testx.bmp размером 100x100 пикселей. Мы берем суб-изображение 10x10 в позиции (30,30) и находим его в оригинале.

import cv2
import numpy as np

image = cv2.imread("testimage.bmp")
template = image[30:40,30:40,:]

result = cv2.matchTemplate(image,template,cv2.TM_CCOEFF_NORMED)
print np.unravel_index(result.argmax(),result.shape)

Вывод:

(30, 30)

Вы можете выбрать один из нескольких алгоритмов для соответствия шаблону оригиналу, cv2.TM_CCOEFF_NORMED - только один из них. Более подробную информацию см. В документации, некоторые алгоритмы указывают совпадения как минимумы, другие - как максимумы в массиве результатов. Предупреждение: OpenCV по умолчанию использует порядок каналов BGR, поэтому будьте осторожны, например. когда вы сравниваете изображение, загруженное с помощью cv2.imread, на изображение, которое вы преобразовали из PIL в numpy. Вы всегда можете использовать cv2.cvtColor для преобразования между форматами.

Чтобы найти все совпадения выше заданного порогового значения confidence, я использую что-то по строкам этого, чтобы извлечь согласованные координаты из массива результатов:

match_indices = np.arange(result.size)[(result>confidence).flatten()]
np.unravel_index(match_indices,result.shape)

Это дает кортеж массивов длиной 2, каждый из которых является совпадающей координатой.

Ответ 3

Я только что закончил писать автономную реализацию нормализованной кросс-корреляции для N-мерных массивов. Вы можете получить его из здесь.

Кросс-корреляция рассчитывается либо напрямую, используя scipy.ndimage.correlate, либо в частотной области, используя scipy.fftpack.fftn/ifftn в зависимости от того, что будет быстрее для заданных размеров ввода.

Ответ 4

Фактически вы можете уменьшить эту проблему до простого строкового поиска, используя regex, например, следующую реализацию: принимает два объекта PIL.Image и находит координаты needle в пределах haystack. Это примерно на 127 раз быстрее, чем использование поэтапного поиска.

def subimg_location(haystack, needle):
    haystack = haystack.convert('RGB')
    needle   = needle.convert('RGB')

    haystack_str = haystack.tostring()
    needle_str   = needle.tostring()

    gap_size = (haystack.size[0] - needle.size[0]) * 3
    gap_regex = '.{' + str(gap_size) + '}'

    # Split b into needle.size[0] chunks
    chunk_size = needle.size[0] * 3
    split = [needle_str[i:i+chunk_size] for i in range(0, len(needle_str), chunk_size)]

    # Build regex
    regex = re.escape(split[0])
    for i in xrange(1, len(split)):
        regex += gap_regex + re.escape(split[i])

    p = re.compile(regex)
    m = p.search(haystack_str)

    if not m:
        return None

    x, _ = m.span()

    left = x % (haystack.size[0] * 3) / 3
    top  = x / haystack.size[0] / 3

    return (left, top)