[파이썬] 이미지 객체 탐지 - CCTV 이미지를 활용한 교통량 검지

CCTV 이미지 객체 탐지

---

이미지 객체 탐지  - CCTV 이미지를 활용한 교통량 검지

---

들어가며

---

yolov5와 OpenCV를 활용하여 이미지 안의 객체를 탐지하는 방법을 소개합니다. yolov5는 pytorch를 기반으로 구현된 모델이므로 사전에 pytorch 설치가 필요합니다. 공식 github를 참고하세요. 이 글에서는 모든 패키지가 설치되어 있다는 가정 하에 패키지 로드부터 시작합니다.

이 글에서는 CCTV 이미지를 활용하여 고속도로의 교통량을 검지하는 방법을 소개합니다. 이미지 안에 차량의 개수를 산출하고 더 나아가 차종을 구분하여 교통량을 산출해 보겠습니다.

---

• CCTV 이미지
• 이미지 객체 탐지
• 전체 코드

---

CCTV 이미지

---

CCTV 이미지는 'AI Hub'의 교통문제 해결을 위한 CCTV 교통 영상(고속도로)에서 다운로드했습니다. sample 이미지는 아래와 같습니다.

CCTV 이미지

---

이미지 객체 탐지

---

먼저 필요한 패키지인 OpenCV와 pytorch를 import 합니다.

import cv2
import torch

---

이 글에서 활용하는 yolov5는 사전에 훈련된 모델입니다. 따라서 pytorch 허브를 통해 yolov5 모델을 불러옵니다.

model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

---

다음은 이미지입니다. 이미지 경로와 함께 OpenCV(cv2)를 사용하여 이미지를 불러옵니다. 

OpenCV는 이미지를 numpy 배열로 반환하기 때문에 img를 출력하면 숫자 배열을 확인할 수 있습니다. 각 숫자는 해당 픽셀의 색상 값을 나타냅니다. 이미지의 shape을 출력했을 때 결과는 (1080, 1920, 3)이 나오므로 이 이미지는 가로 1080 세로 1920의 이미지이며 RGB 채널 3은 컬러를 의미합니다.

image_path = r'C:\Users\USER\Desktop\Gangneung_hyangho2brdg_20201019_1730_MON_15m_NH_highway_TW2_sunny_FHD_038.png'
img = cv2.imread(image_path)

print('이미지 크기: ', img.shape)
img

img 출력 결과

---

OpenCV에서 이미지를 불러오면 기본적으로 BGR 형식으로 저장되기 때문에 RGB 형식으로 변환해야 합니다. 다음과 같이 cvtColor로 이미지 색상 공간을 변환합니다. img_rgb는 RGB 형식으로 변환된 이미지입니다.

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

---

이제 이미지 객체 탐지를 진행합니다. 다음과 같이 yolov5 모델을 사용하여 img_rgb의 객체를 탐지합니다. 출력한 results에서 이미지의 크기와 함께 2개의 차량과 1개의 트럭이 탐지된 것을 확인할 수 있습니다. 

results = model(img_rgb)
results

 

results 출력 결과

---

results 출력 결과를 편리하게 처리할 수 있도록 데이터프레임 형태로 변환하겠습니다. 탐지한 객체에 대한 상세 정보는 다음과 같습니다. x_min, y_min(바운딩 박스의 왼쪽 위 모서리 좌표), x_max, y_max(바운딩 박스의 오른쪽 아래 모서리 좌표), confidence(신뢰도 점수), class(감지된 객체의 클래스), name(감지된 객체의 이름)입니다.

detections = results.pandas().xyxy[0]
print('교통량: ', len(detections))
detections

detections 출력 결과

---

차종을 구분한 결과는 다음과 같습니다. 샘플 이미지에서는 자동차 2대, 트럭 1대를 탐지했습니다.

num_cars = len(detections[detections['name'] == 'car'])
print('num_cars: ', num_cars)

num_trucks = len(detections[detections['name'] == 'truck'])
print('num_trucks: ', num_trucks)

차종 구분 결과

반응형

---

전체 코드

---

## 패키지, 모델

import cv2
import torch

model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

# --------------------------------------------------------------------------------- */

## sample 이미지

image_path = r'C:\Users\USER\Desktop\Gangneung_hyangho2brdg_20201019_1730_MON_15m_NH_highway_TW2_sunny_FHD_038.png'
img = cv2.imread(image_path)
print('이미지 크기: ', img.shape)

img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
results = model(img_rgb)

# --------------------------------------------------------------------------------- */

## 객체 탐지

detections = results.pandas().xyxy[0]
print('교통량: ', len(detections))

num_cars = len(detections[detections['name'] == 'car'])
print('num_cars: ', num_cars)

num_trucks = len(detections[detections['name'] == 'truck'])
print('num_trucks: ', num_trucks)

---