YOLOv5
Yolov5는 이미지 객체 검출을 위한 딥러닝 모델 중 하나입니다. You Only Look Once (YOLO) 알고리즘을 기반으로 한 모델이며, 좀 더 높은 정확도와 빠른 처리 속도를 제공합니다. Yolov5는 PyTorch 라이브러리를 사용하여 구현되었으며, 다양한 이미지 객체 검출 작업에 사용될 수 있습니다.
Yolov5로 객체를 검출하려면 이미 학습된 모델을 사용하거나, 직접 데이터 셋을 이용해 모델을 학습시켜야 한다. 나는 멧돼지를 검출해야 해서 멧돼지 데이터 셋이 필요했다.
무료로 데이터셋을 제공해주는 여러 사이트가 존재했지만, 이미지를 수집하고 라벨링 해보는 것도 좋은 경험이 될 것 같아서 직접 데이터셋을 만들어보려고 한다.
이미지 크롤링
먼저 데이터 셋을 만들기 위해서는 이미지가 필요하다.
이미지 크롤링을 google과 daum, pixabay에서 시도해 보았다.
pixabay
pixabay에서는 API로 간단하게 성공했다.
우선, pixabay의 API를 얻어야 한다.
- pixabay 개발자 페이지로 이동 후 회원가입
https://pixabay.com/ko/service/about/api/ - Key 확인
키를 얻었으면 코드를 작성해보도록 한다.
import os
import requests
from PIL import Image폴더 생성과 이미지 저장을 위한 os 모듈
이미지 URL 처리를 위한 requests 라이브러리
이미지 전처리를 위한 PIL 라이브러리에서 Image 모듈만 호출
API_KEY = 'Your-API-KEY'
SEARCH_TERM = 'wild boar'
PER_PAGE = 100
NUM_PAGES = 5
SAVE_DIR = './images'pixabay에서 받은 API 키 변수 생성, 검색어로 사용할 단어를 저장
한 페이지에서 100장의 이미지씩 5페이지를 수집하여 총 500장의 이미지
SAVE_DIR에 경로를 입력하고 이 경로로 폴더를 생성, 이미지를 저장하도록 한다.
if not os.path.exists(SAVE_DIR):
os.makedirs(SAVE_DIR)경로에 폴더가 없으면 폴더를 생성해주는 코드
count = 0 # 다운로드한 이미지 개수
for i in range(NUM_PAGES):
url = f'https://pixabay.com/api/?key={API_KEY}&q={SEARCH_TERM}&image_type=photo&per_page={PER_PAGE}&page={i+1}'
response = requests.get(url)
data = response.json()url = API 키와 검색어, photo타입 100장, page 1부터 NUM_PAGES번 만큼 페이지를 이동
url로부터 GET( )하여 API의 JSON 데이터를 가져온다.
이 JSON 데이터를 python에서 사용할 수 있도록 파싱해준다.
for image in data['hits']:
try:
img_url = image['webformatURL']
response = requests.get(img_url)data['hits']는 이미지를 담고 있는 python 리스트이다. 딕셔너리가 리스트로 감싸진 형태이다.
image 리스트에서 해상도가 webformatURL인 key 값의 value(이미지 URL)를 가져온다.
이제 requests.get( )를 사용하여 URL로 이미지를 다운로드 받는다.
Yolov5 모델을 학습시킬 때 --img 옵션으로 이미지 크기를 지정할 수 있는데
나는 416으로 지정할 것이기 때문에 이미지의 크기가 640인 이미지로 받았다.
학습에 사용할 이미지의 크기는 작을수록 빠르지만, 성능은 떨어진다.
이미지 URL 옵션
- webformatURL 640
- largeImageURL 1280
- fullHDURL 1920
img_name = f"{SEARCH_TERM}_{count+1}.jpg"
img_path = os.path.join(SAVE_DIR, img_name)
with open(img_path, 'wb') as f:
f.write(response.content)img_name 변수에 이미지의 이름을 지정해준다.
os.path.join( )으로 ./(SAVE_DIR)/(img_name) 경로를 생성한다.
response.content가 바이너리 형식 데이터이므로
img_path 경로의 파일을 바이너리 쓰기 모드로 열어주고 쓴다.
img = Image.open(img_path)
img_resized = img.resize((416, 416))
img_resized.save(img_path)
count += 1
print(f"{count} image saved.")
if count == 500:
breakImage 모듈의 open( ) 함수를 사용하여 이미지를 열어준다.
이 이미지를 416x416으로 리사이즈 해주고, 저장한다.
이미지 저장에 성공할 때마다 count가 1씩 증가하고 현재까지 저장된 이미지의 수를 출력한다.
이미지가 500장 저장되면 for 루프를 종료한다.
except Exception as e:
print(e)
continue이미지 다운로드 중 예외 처리이다.
이 구문이 있어야 이미지 다운로드에 실패해도 다음 이미지를 다운로드할 수 있다.
이렇게 하면 500장을 채우지 못하고 종료가 되는데,
while 문을 이용해서 count가 500일 때까지 반복하도록 하고
hits 리스트에 값이 더 이상 없으면 코드를 멈추도록 작성해야 한다.
전체 코드
import os
import requests
from PIL import Image
API_KEY = 'Your_API_KEY'
SEARCH_TERM = 'wild boar'
PER_PAGE = 100
NUM_PAGES = 5
SAVE_DIR = './image' # 저장할 폴더 경로
# 폴더 생성
if not os.path.exists(SAVE_DIR):
os.makedirs(SAVE_DIR)
count = 0 # 다운로드한 이미지 개수
for i in range(NUM_PAGES):
url = f'https://pixabay.com/api/?key={API_KEY}&q={SEARCH_TERM}&image_type=photo&per_page={PER_PAGE}&page={i+1}'
response = requests.get(url)
data = response.json()
for image in data['hits']:
try:
# 이미지 다운로드
img_url = image['webformatURL'] # 이미지 해상도
response = requests.get(img_url)
# 이미지 저장
img_name = f"{SEARCH_TERM}_{count+1}.jpg" # 저장할 이미지 파일명
img_path = os.path.join(SAVE_DIR, img_name) # 저장할 이미지 경로
with open(img_path, 'wb') as f:
f.write(response.content)
# 이미지 리사이징 (416x416)
img = Image.open(img_path)
img_resized = img.resize((416, 416))
img_resized.save(img_path)
count += 1
print(f"{count} image saved.")
if count == 500:
break
except Exception as e:
print(e)
continue
if count == 500:
breakgoogle에서는 Python selenium 라이브러리로 이미지 크롤링에 성공했다.
먼저 이미지 크롤링을 위해선 Chromedriver가 필요하다.
자신의 크롬 브라우저 버전을 확인하고 더보기 -> 도움말 -> Chrome 정보
아래 URL에서 가장 근접한 버전으로 다운로드 한다.
https://chromedriver.chromium.org/downloads
from selenium import webdriver
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.common.by import By
import time
import urllib.request
from pathlib import Path
from PIL import Image코드에 필요한 패키지 호출
Search_Word = input("검색어를 입력하세요 : ").strip()
img_num = input("원하는 이미지 개수 입력하세요 : ").strip()
Path(Search_Word).mkdir(exist_ok=True) # 검색어명 폴더 생성검색어와 이미지 개수를 입력으로 받았다.
이번에는 파일을 생성할 때 pathlib을 사용해 봤다.
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(Search_Word)
elem.send_keys(Keys.RETURN)구글 이미지 페이지를 열고 검색창 속성에 검색어 키워드를 전송한다.
SCROLL_PAUSE_TIME = 1
last_height = driver.execute_script("return document.body.scrollHeight") # 브라우저 높이 확인
while True:
driver.execute_script("window.scrollTo(0, document.body.scrollHeight);")
time.sleep(SCROLL_PAUSE_TIME)
new_height = driver.execute_script("return document.body.scrollHeight")
if new_height == last_height: # 스크롤 모두 내렸을 때
try:
driver.find_elements(By.CSS_SELECTOR, ".mye4qd").click() # 결과 더보기 클릭
except:
break
last_height = new_height브라우저의 높이를 파악하는 코드이다.
images = driver.find_elements(By.CSS_SELECTOR,".rg_i.Q4LuWd")CSS 선택자로 class가 .rg_i.Q4LuWd인 이미지들의 정보를 find_elements로 가져온다.
count = 1
for image in images:
if count > int(img_num):
break
try:
image.click()
time.sleep(2)
imgUrl = driver.find_element(By.CSS_SELECTOR, '.r48jcc.pT0Scc.iPVvYb').get_attribute("src") # 이미지 src 값 가져오기
urllib.request.urlretrieve(imgUrl, Search_Word + "/" + Search_Word + "_" + str(count) + ".jpg")
count = count + 1
except:
passimages에 담은 이미지 정보에서 이미지의 URL이 있는 src 값을 하나씩 가져온다.
이번에는 urllib.request.urlretrieve()을 사용해서 이미지를 다운받고 저장했다.
처음에 img 태그가 두 개여서 헤맸다.
자꾸 실패해서 구글에서 크롤링을 막은 줄 알았다.
구글에서 이미지를 막 크롤링하게 되면 문제가 있을 수도 있으니
구글에서 이미지 크롤링할 때는 API를 사용하거나 다른 사이트를 이용해야 할 것 같다.
daum
이 블로그를 참조했다.
https://smart-factory-lee-joon-ho.tistory.com/m/346
마무리
이미지 수집에 관해서 글을 작성해 보았다.
하나하나 상세하게 작성했더니 꽤 오래 걸렸다.
글을 계속 작성해 가면서 실력을 길러야겠다.
이번에 기록하면서 selenium으로 이미지 크롤링하는 원리, JSON 데이터 파이썬에서 사용하기, 딕셔너리 안에 리스트 안에 딕셔너리가 들어있는 구조, continue 예외처리, PIL 라이브러리로 Image 처리, URL 다운로드 정도가 주요하게 볼 것이었다.
다음은 수집한 이미지들로 라벨링과 학습을 시킬 차례다.
