본문 바로가기
ML&DL

selective search 코드 리뷰 - 파이썬, 컴퓨터 비전

by Back2Baek 2024. 2. 11.

 

selective search 란,

object detection 분야에서 객체가 있을만한 위치를 제안해주는 알고리즘이다.

 

비슷한 목적의 알고리즘인 sliding window 방식보다 효율적으로 처리할 수 있다.

 

( 구글 코랩의 환경에서 제작되었습니다.)

import selectivesearch
import cv2
import matplotlib.pyplot as plt
import os
from google.colab.patches import cv2_imshow

 

 

img_path = '/content/NMMP_dolphin_with_locator.jpeg'
img = cv2.imread(img_path)

 

img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
plt.imshow(img_rgb)

 

cv2 라이브러리로 이미지를 읽어들일 경우, 흔히 아는 RGB 형태가 아닌 BGR 형태로 읽어들이게 된다.

따라서 cv2.cvtColor를 통해 RGB로 변환해주어야 한다.

 

변환 후 출력한 이미지는 다음과 같다.

돌고래 - 위키백과, 우리 모두의 백과사전 (wikipedia.org)

def display_selective_search(min_size=100):
    _, regions = selectivesearch.selective_search(img_rgb, min_size=min_size)
    rects = [cand['rect'] for cand in regions]
    border = (0, 0, 200)
    img_copy = img_rgb.copy()

    for rect in rects:
        img_copy = cv2.rectangle(img_copy, (rect[0], rect[1]), (rect[0] + rect[2], rect[1] + rect[3]), border, 2)

    plt.imshow(img_copy)

 

 

selective search 함수 파라미터로 제안 영역의 최소 사이즈를 지정할 수 있다. 최소 사이즈가 작을 수록 실행 시간은 늘어나게 된다.

 

selective_search 실행 후 반환 한 regions 값은 다음과 같이 딕셔너리 형식이다.

[{'rect': (0, 0, 372, 348), 'size': 61689, 'labels': [0.0]}, {'rect': (294, 0, 6, 4), 'size': 28, 'labels': [1.0]}, ...]

 

이 중 제안 영역에 대한 박스의 정보(x, y, w, h)를 담고 있는 rect 정보만 rects 변수에 할당.

 

 

각 최소 사이즈별로 selective를 실행한 결과 및 제안 영역 박스 시각화

display_selective_search(100)

display_selective_search(1000)

display_selective_search(5000)

 

 

 

다음은 해당 코드를 캐글 편집기에 맞게 수정한 코드입니다.

https://www.kaggle.com/code/backtobaek/selective-search