Поиск по образцу. Как сделать быстрый поиск вхождения картинок-символов на большом чертеже?

Question

[imageman]

Добрый день.
Прибежала ко мне одна задача - смотреть чертежи.
Есть легенда (иногда с мусором по типу каких-то мелких символов или стрелочек) и символы с этой легенды нужно найти на чертеже. Для простоты считаем, что масштаб почти не меняется. Основных проблем две - все символы на этапе обучения мы не знаем (открытый мир; YOLO не подходит); символы могут быть с перекрытием; символы могут быть похожими.

Самый плохой случай символ "кольцо" и символ "ёлочка" которые в одном и том же месте чертежа (друг на друге, но хорошо видны). Ну и всякий мусор по типу фонов, надписей, рамок и т.п.

Нужно на большом чертеже относительно быстро найти все вхождения. Я пробовал через ViT патчи, но сообразил, что косинусное сходство симметрично и при поиске по патчам совпадения перекрывающейся "кольца" или "елочки" на чертеже у меня или одно даст хорошее сходство или другое (а скорее всего ни то, ни другое). И не очень могу понять куда идти дальше. Ещё раз - на момент обучения мы не знаем всех символов.
Может подскажете?

Синтетически сгенерированный пример, нужно найти все ёлочки и все круги


#Python, #pytorch, #timm, #transformers

Answer 1

[dim5x]

import cv2
import numpy as np

# Загрузка изображений:
dr_image = cv2.imread('original.png')

# Преобразуем в серое:
gray = cv2.cvtColor(dr_image, cv2.COLOR_BGR2GRAY)

# Бинаризация: черные линии → белые (на чёрном фоне).
_, binary = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY_INV)

# Шаблоны:
template1 = cv2.imread('temp2.png', cv2.IMREAD_GRAYSCALE)  # ёлочка
template2 = cv2.imread('temp1.png', cv2.IMREAD_GRAYSCALE)  # круг

_, template1 = cv2.threshold(template1, 127, 255, cv2.THRESH_BINARY_INV)
_, template2 = cv2.threshold(template2, 127, 255, cv2.THRESH_BINARY_INV)

# Параметры:
threshold1 = 0.7  # для ёлочки
threshold2 = 0.5  # для круга

# Поиск "ёлочки":
result1 = cv2.matchTemplate(binary, template1, cv2.TM_CCOEFF_NORMED)
loc1 = np.where(result1 >= threshold1)

# Поиск "круга":
result2 = cv2.matchTemplate(binary, template2, cv2.TM_CCOEFF_NORMED)
loc2 = np.where(result2 >= threshold2)

# Результат:
result = dr_image.copy()
detected_boxes = []

# Обработка "ёлочки":
for pt in zip(*loc1[::-1]):
    x, y = pt
    w, h = template1.shape[1], template1.shape[0]
    box = (x, y, x + w, y + h)

    if not any(np.allclose(box, b, atol=10) for b in detected_boxes):
        detected_boxes.append(box)
        cv2.rectangle(result, pt, (pt[0] + w, pt[1] + h), (0, 0, 255), 2)
        cv2.putText(result, "Tree", (pt[0], pt[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)

# Обработка "круга":
for pt in zip(*loc2[::-1]):
    x, y = pt
    w, h = template2.shape[1], template2.shape[0]
    box = (x, y, x + w, y + h)

    if not any(np.allclose(box, b, atol=10) for b in detected_boxes):
        detected_boxes.append(box)
        cv2.rectangle(result, pt, (pt[0] + w, pt[1] + h), (0, 255, 0), 2)
        cv2.putText(result, "Circle", (pt[0], pt[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)

# Сохраняем результат:
cv2.imwrite('result_with_labels.png', result)
print(f"Найдено объектов: {len(detected_boxes)}")

# Показываем:
cv2.imshow('Detected', result)
cv2.waitKey(0)
cv2.destroyAllWindows()

>>> Найдено объектов: 12

Comments

[imageman]

Спасибо. В принципе - да, для простых случаев подходит (при условии, что масштаб в легенде совпадает). Правда даже тут мы видим, что пороги подобраны индивидуально. В реальности чертежи более шумные и порог хотелось бы один универсальный.

Из интересных алгоритмов нашелся Chamfer distance (правда на русском описание не могу найти, только так https://chatgpt.com/share/68d3fe50-1b04-8000-b18b-...)

[dim5x]

imageman, пожалуйста. Если помогло - отметьте решением.

[imageman]

dim5x, к сожалению только помогло, это для меня не решение.... (слишком много ложных срабатываний)

В реальности нечто подобное нужно
https://goodproject.by/wp-content/uploads/2023/02/...
https://54land.ru/wp-content/uploads/17%D0%94%D0%B...