[Python] YOLOv5를 활용한 객체 인식(Object Detection) (1)

============= 목표 =============

YOLOv5 모델을 통한 선박 객체 인식

==============================

 

1. 데이터 셋 준비하기

2. 데이터 셋 형식 변환

3. 데이터 셋을 통한 모델 학습

4. 훈련 결과를 통한 객체 인식 검증

 

 

《 데이터 셋 준비하기 》

모델 학습을 위한 데이터 셋 준비가 필수적이다.

YOLO를 학습하기 위한 데이터 셋 형태는 아래와 같다

'''

학습을 위한 이미지 파일

각 이미지와 1:1로 매핑되는 YOLO 형식의 annotaion 텍스트 파일

'''

 

YOLOv5의 어노테이션 형식은 .txt의 텍스트파일으로

해당 텍스트 파일들은 각 이미지와 매핑되어 이미지의 바운딩 박스에

대한 설명을 위한 자료로 활용된다.

 

예를 들어 아래와 같은 이미지와 1대1로 매핑되는 어노테이션

파일의 경우 아래와 같은 형식으로 명시된다.

 

한 줄마다 한 객체의 정보를 담고 있으며

바운딩 박스의 각 좌표들은 해당 이미지를 0-1로

정규화 시켰을 때의 값을 의미한다.

 

각 줄에는 클래스(인덱스), x, y, w, h의 정보를 의미하고 있다.

0(Person) x1, y1, w1, h1

27(Tie) x2, y2, w2, h2 

 

이제 본인이 어떤 객체 인식을 요구하느냐에 따라서

기존에 공개되어 있는 데이터 셋을 활용해서 해당 모델을 사용할 수도 있고,

본인이 직접 이미지 파일들을 준비하고, 라벨링 하여 데이터 셋을 구축하여

YOLO 모델을 사용할 수도 있다.

 

 

1. 기존 데이터셋 활용

여러 활용되는 데이터셋들이 존재하지만

그 중에서 보편적으로 사용되는 COCO 데이터 셋의 경우

공식 홈페이지에서 직접 다운로드 받을 수 있다.

 

코코데이터셋

COCO - Common Objects in Context

공식 홈페이지의 데이터셋 다운로드 페이지로 가게 되면

1대1 매핑되는 이미지와 어노테이션 파일들이 있습니다.

클릭해서 다운로드 되지 않을 경우 마우스 우클릭으로 링크 주소 복사를 한 후

URL에 직접 붙여넣는 식으로 다운로드 가능합니다.

 

 

2. 사용자 커스텀 데이터 셋 생성

우선 저는 기존의 데이터 셋을 활용해 모델을 훈련시키는 것이 아닌

직접 이미지를 구해서 데이터 셋을 생성하는 작업을 해보았습니다.

 

[ 이미지 수집 ]

우선 이미지를 수집해야 합니다.

무료 이미지 제공 사이트를 이용하거나 직접 일일히

다운로드 하는 방법도 존재하지만,

저는 우선 간단한 파이썬 Selenium 패키지를 이용한 코드를 통해

구글 크롬에서 이미지를 크롤링 해왔습니다.

 

『구글 이미지 크롤링』

import ssl
ssl._create_default_https_context = ssl._create_unverified_context

from selenium import webdriver
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.common.by import By
from selenium.webdriver.chrome.service import Service as ChromeService
from webdriver_manager.chrome import ChromeDriverManager
import time
import urllib.request

searchKey = input('검색 키워드 입력:')

# 파일 경로에 크롬 드라이버(크롬 드라이버.exe) 파일이 없는경우
driver = webdriver.Chrome(service=ChromeService(ChromeDriverManager().install()))
# 파일 경로에 크롬 드라이버 파일이 있는 경우
# driver = webdriver.Chrome()

driver.get("https://www.google.co.kr/imghp?hl=ko&tab=wi&authuser=0&ogbl")
elem = driver.find_element("name", "q")

elem.send_keys(searchKey)
elem.send_keys(Keys.RETURN)
SCROLL_PAUSE_TIME = 1

# Get scroll height
last_height = driver.execute_script("return document.body.scrollHeight")
while True:
    # Scroll down to bottom
    driver.execute_script("window.scrollTo(0, document.body.scrollHeight);")
    # Wait to load page
    time.sleep(SCROLL_PAUSE_TIME)
    # Calculate new scroll height and compare with last scroll height
    new_height = driver.execute_script("return document.body.scrollHeight")
    if new_height == last_height:
        try:
            driver.find_element(By.CSS_SELECTOR, ".mye4qd").click()
        except:
            break
    last_height = new_height

    images = driver.find_elements(By.CSS_SELECTOR, ".rg_i.Q4LuWd")
    count = 1
    for image in images:
        try:
            image.click()
            time.sleep(0.5)
            imgUrl = driver.find_element(
                By.XPATH,
                '//*[@id="Sva75c"]/div[2]/div[2]/div[2]/div[2]/c-wiz/div/div/div/div[3]/div[1]/a/img[1]'
            ).get_attribute("src")
            opener = urllib.request.build_opener()
            opener.addheaders = [
                ('User-Agent',
                 'Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/36.0.1941.0 Safari/537.36')
            ]
            urllib.request.install_opener(opener)
            urllib.request.urlretrieve(imgUrl, f'./images/{searchKey}{str(count)}.jpg')
            count = count + 1
        except KeyboardInterrupt:
            pass
        except Exception as e:
            print('e : ', e)
            pass
    driver.close()

 

해당 파이썬 코드를 통해 선박 관련 이미지 100개를 크롤링했습니다. 

데이터셋을 만들기 이전에 이미지 어노테이션 작업을 해주어야 합니다.

이미지 어노테이션(Image Annotation)이란 위의 YOLOv5의 어노테이션 형식의

텍스트 파일을 위해 학습 이미지 전체에 태그를 지정, 또는 분류하고

바운딩 박스를 지정하는 등의 행위를 의미합니다.

이는 컴퓨터 비전 모델 훈련에 필수적인 작업이기도 합니다.

 

 

해당 작업은 CVAT을 통해 수행하였습니다.

회원가입 후 무료 이용 가능하며,

유료 회원의 경우 AI를 통한 자동 라벨링, 원본 이미지 지원 등

다양한 추가 기능을 지원하고 있습니다.

 

CVAT

Computer Vision Annotation Tool

 

 

[ 이미지 어노테이션 ]

해당 페이지에서 Create new task 버튼을 통해

새로운 작업을 생성, 라벨명을 설정 한 뒤

이미지를 하나하나 수동으로 지정하여 박스 바운딩 작업을 완료해줍니다.

 

완료 된 작업의 경우 해당 작업의 우측 Actions 옆의 메뉴를 클릭하여

Export task dataset란을 선택해줍니다.

 

이 때 annotation 변환 타입은 COCO.json 형식을 지원하는

COCO 1.0을 이용했습니다.

 

여기까지 수행했으면 YOLOv5 데이터셋 변환을 위한

이미지 준비까지 완료입니다.

 

다음 작업에서 라벨링 데이터를 YOLOv5 형식으로 변환하고

이를 통해 데이터 셋을 준비하고 모델 훈련하는 과정까지 해보겠습니다.