[파이썬] Rotate-invariant realtime template matching using SAM model

 

회전왜곡에 강건한 템플릿매칭 예시1

 

회전왜곡에 강건한 템플릿매칭 예시2

 

- 메타의 SAM(세그먼트 애니띵) 모델로 실시간에 이미지내 물체의 템플릿을 만든후 만들어진 템플릿으로 이미지내에서 비슷한 물체를 매칭(OpenCV로 수학적으로)하는 방식입니다.

- 실시간 템플릿은 이미지에서 마우스 포인터로 찾고자하는 물체내부를 클릭하면(Visual Prompt ) SAM 모델로 물체의 윤곽선을 세그먼테이션하며 그 윤곽선으로 Best-fit된 bbox를 수학적으로 만들고 그 내부 이미지를 템플릿으로 저장합니다.

- 그리고 저장된 템플릿으로 템플릿 매칭을 하는데 회전왜곡에 불변하도록 매칭하는 알고리즘으로 회전된 물체도 같은 모양이면 찾아주게 됩니다.

- 단,  90도 배수에 해당하는 각도로 회전된 물체는 매칭율이 잘 않나오는 문제가 있습니다. 나머지 모든각도에 대해선 잘 찾아줍니다.

---

import cv2
import tkinter as tk
from PIL import Image, ImageTk
import numpy as np
import torch

MASK_COLOR = (255, 0, 0)

# Load the Lena image
image_path = "/Users/user1/Downloads/PL테크 Sensing Tab 촬상 이미지(20231215)/Vision Align Ⅱ/1_2.bmp"

image = cv2.imread(image_path)
# 새로운 크기 설정
width = 512
height = 512
dim = (width, height)

# 이미지 크기 조정
image = cv2.resize(image, dim, interpolation = cv2.INTER_AREA)


original_image = image.copy()  # Keep a copy of the original image

-  이미지를 로딩하고 크기를 조정합니다.
- 실시간 반응성을 고려해 GUI패키지로 tkInter 를 사용합니다.

import sys
sys.path.append("..")
from segment_anything import sam_model_registry, SamPredictor

sam_checkpoint = "/Users/user1/Downloads/sam_vit_b_01ec64.pth"
model_type = "vit_b"
device = "cpu"

sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)
predictor = SamPredictor(sam)
predictor.set_image(image)

 - 파이토치로 SAM 모델을 로딩하고 프리딕터를 설정합니다. 맥에선 cpu 모드 말고도 mps 모드로 로딩하면 추론속도가 빨라집니다.

template = None  # Initialize the template

def make_mask_2_img(mask):
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * np.array(MASK_COLOR).reshape(1, 1, -1)
    mask_image = mask_image.astype(np.uint8)
    return mask_image

def on_mouse_move(event):
    global image, template

    input_point = np.array([[event.x, event.y]])
    input_label = np.array([1])

    mask, _, _ = predictor.predict(
        point_coords=input_point,
        point_labels=input_label,
        multimask_output=False,
    )

    mask_img = make_mask_2_img(mask)
    gray_mask = cv2.cvtColor(mask_img, cv2.COLOR_BGR2GRAY)
    _, binary_img = cv2.threshold(gray_mask, 1, 255, cv2.THRESH_BINARY)

    contours, _ = cv2.findContours(binary_img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    for cnt in contours:
        x, y, w, h = cv2.boundingRect(cnt)
        template = image[y:y+h, x:x+w]  # Create the template from the original image within the bounding rectangle

    image_rgb = cv2.cvtColor(mask_img, cv2.COLOR_BGR2RGB)

    im = Image.fromarray(image_rgb)
    img = ImageTk.PhotoImage(im)

    img_label_proc.img = img
    img_label_proc.config(image=img)

def rotate_image(image, angle):
    h, w = image.shape[:2]
    center = (w // 2, h // 2)
    M = cv2.getRotationMatrix2D(center, angle, 1.0)
    rotated_image = cv2.warpAffine(image, M, (w, h))
    return rotated_image

- 필요한 함수들과 마우스 무브 이벤트핸드러를 정의 합니다.

- 마우스 이동시에 좌측버튼이 눌리면 SAM 마스킹이 일어나고 바운딩박스까지 구해지게 됩니다. 

def on_key_press(event):
    global image, template, original_image

    if event.char == 'm' and template is not None:
        # 매칭된 위치를 저장할 리스트
        locations = []

        # 원본 이미지에 템플릿 매칭을 위한 회전된 템플릿 생성
        for angle in range(0, 180, 1):  # 0도에서 360도까지 10도씩 회전
            # 템플릿 회전
            (h, w) = template.shape[:2]
            center = (w // 2, h // 2)
            M = cv2.getRotationMatrix2D(center, angle, 1.0)
            rotated_template = cv2.warpAffine(template, M, (w, h))

            # 회전된 템플릿에 대해 템플릿 매칭 실행
            res = cv2.matchTemplate(image, rotated_template, cv2.TM_CCOEFF_NORMED)
            threshold = 0.6  # 매칭 임계값 설정
            loc = np.where(res >= threshold)
            for pt in zip(*loc[::-1]):  # 임계값을 넘는 위치 찾기
                locations.append((pt[0], pt[1], pt[0] + w, pt[1] + h))

        # 겹치는 사각형을 처리하기 위해 cv2.groupRectangles 사용
        locations, _ = cv2.groupRectangles(locations, groupThreshold=1, eps=0.5)
        for (x1, y1, x2, y2) in locations:
            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)

        # 결과 이미지를 RGB로 변환하고 화면에 표시
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        im = Image.fromarray(image_rgb)
        img = ImageTk.PhotoImage(im)
        img_label_proc.img = img
        img_label_proc.config(image=img)

    if event.char == 'r':  # Reload the original image
        image = original_image.copy()
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        im = Image.fromarray(image_rgb)
        img = ImageTk.PhotoImage(im)

        img_label_orig.img = img
        img_label_orig.config(image=img)

        img_label_proc.img = img
        img_label_proc.config(image=img)

- 하지만 키보드에서 "m"키가 눌려야지만 비로소 마스크가 템플릿으로 저장되게 됩니다.

root = tk.Tk()

image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
im = Image.fromarray(image_rgb)
img = ImageTk.PhotoImage(im)

img_label_orig = tk.Label(root, image=img)
img_label_orig.grid(row=0, column=0)

img_label_proc = tk.Label(root)
img_label_proc.grid(row=0, column=1)

root.bind("<Button-1>", on_mouse_move)
root.bind("<Key>", on_key_press)

root.mainloop()

- 나머지 tkinter 화면 정의 및 이벤트 연결 코드입니다.

 

- SAM과 같은 딥러닝 모델은 강력한 마스킹 성능을 제공합니다. OpenCV에선 매우 어려웠던 기능이었죠. 전처리를 어떻게 하느냐에 따라 매번 걀과가 달라졌고 조명이나 초점, 그림자나 배경에 영향을 많이 받았거든요~

- 반대로  윤곽선 후처리나 측정  또는 고등통계적 윤곽선 피쳐추출은 딥러닝이 하지 못하는 영역입니다. 이것은 OpenCV나 skimage 같은 수학적 영상처리가 떠 뛰어나고 속도도 빠르죠.  

- 이런식으로 데이터-드리븐 방식인 딥러닝 모델과 수학적 영상처리의 만남은 더욱더 강력한 결과물을 만들어 줍니다.

씨너지라는 말은 이럴때 쓰는 것이죠~