[파이썬] CCTV 이미지를 활용한 이미지 객체 탐지 - 박스 표시, yolov5

객체 탐지 박스 표시

---

CCTV 이미지를 활용한 이미지 객체 탐지 - 박스 표시, yolov5

---

들어가며

---

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

import cv2
import torch

---

• CCTV 이미지
• 이미지 객체 탐지
• 탐지된 객체에 박스 표시

---

CCTV 이미지

---

CCTV 이미지는 'AI Hub'의 교통문제 해결을 위한 CCTV 교통 영상(고속도로)에서 다운로드했습니다. sample 이미지를 matplotlib을 활용하여 확인하겠습니다.

import matplotlib.pyplot as plt
import matplotlib.image as mpimg

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

# 이미지 출력
plt.imshow(img)
plt.axis('off')
plt.show()

CCTV 이미지

---

이미지 객체 탐지

---

이미지가 저장된 경로를 활용하여 이미지를 로드하고 RGB로 변환합니다. YOLO는 RGB 형식을 사용하기 때문에 변환이 필수입니다.

image_path = r'C:\Users\USER\Desktop\Gangneung_hyangho2brdg_20201019_1730_MON_15m_NH_highway_TW2_sunny_FHD_129.png'

# 이미지 로드
img = cv2.imread(image_path)

# 이미지 RGB 변환
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

---

YOLOV5 모델을 로드하고, 객체를 탐지합니다. 탐지된 객체 정보는 detections에 저장합니다. detection 출력 결과는 아래와 같습니다. 2개의 truck과 2개의 car가 탐지됐고, 각 객체의 위치, 신뢰도를 확인할 수 있습니다.

# YOLOv5 모델 로드
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

# YOLOv5로 객체 탐지
results = model(img_rgb)

# 탐지한 객체 정보 저장
detections = results.pandas().xyxy[0]
detections

detections

---

탐지된 객체에 박스 표시

---

위의 detections에서 탐지된 객체 정보를 활용하여 객체에 박스를 표시하겠습니다. xmin, ymin, xmax, ymax 컬럼에 저장된 위치를 사용합니다. 위치를 활용하여 박스를 그리고 객체 이름을 표시합니다.

박스가 표시된 이미지는 detected_objects.jpg 파일로 저장하겠습니다.

for index, row in detections.iterrows():
    # 객체의 위치 정보, 이름
    xmin, ymin, xmax, ymax = int(row['xmin']), int(row['ymin']), int(row['xmax']), int(row['ymax'])
    label = row['name']
    
    # 객체 박스 표시
    cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)  # 녹색 박스
    
    # 객체 이름 표시
    cv2.putText(img, label, (xmin, ymin - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)  # 녹색 글씨

# 박스 표시 결과 저장
cv2.imwrite('detected_objects.jpg', img)

---

박스가 표시된 결과를 확인하면 다음과 같습니다. 위의 detections를 확인하면 2대의 truck과 2대의 car가 탐지됐는데, 아래 이미지에는 1대의 truck과 2대의 car를 확인할 수 있습니다. 첫 번째로 통행하는 차량의 경우, truck과 car로 중복 탐지되고 있습니다.

객체가 흐릿한 경우 다음과 같이 잘못된 정보를 탐지하기도 합니다. 따라서, 기존에 학습된 yolov5가 아닌 새롭게 학습된 모델을 사용하는 것이 좋습니다. 예를 들어, CCTV 이미지를 활용하여 학습 데이터셋을 구성했다면, CCTV 이미지를 학습한 YOLO 모델로 탐지를 하는 것이 정확한 결과를 추출할수 있습니다.

객체 탐지 박스 표시

---

전체 코드

---

import cv2
import torch



## 샘플 이미지 확인 

import matplotlib.pyplot as plt
import matplotlib.image as mpimg

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

# 이미지 출력
plt.imshow(img)
plt.axis('off')
plt.show()



## 이미지 객체 탐지

image_path = r'C:\Users\USER\Desktop\Gangneung_hyangho2brdg_20201019_1730_MON_15m_NH_highway_TW2_sunny_FHD_129.png'

# 이미지 로드
img = cv2.imread(image_path)

# 이미지 RGB 변환
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)



# YOLOv5 모델 로드
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

# YOLOv5로 객체 탐지
results = model(img_rgb)

# 탐지한 객체 정보 저장
detections = results.pandas().xyxy[0]


## 탐지된 객체에 박스 표시

for index, row in detections.iterrows():
    # 객체의 위치 정보, 이름
    xmin, ymin, xmax, ymax = int(row['xmin']), int(row['ymin']), int(row['xmax']), int(row['ymax'])
    label = row['name']
    
    # 객체 박스 표시
    cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)  # 녹색 박스
    
    # 객체 이름 표시
    cv2.putText(img, label, (xmin, ymin - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)  # 녹색 글씨

# 박스 표시 결과 저장
cv2.imwrite('detected_objects.jpg', img)

---