한강 공개 cctv 영상을 가지고 custom dataset 만들어서 yolo v5 학습 및 실시간 object detection
이전 자료 링크 <- 한강공원 자전거 사고 방지 알림이 시스템[1] – python, 파이썬,cctv 영상 웹 스크래핑, cctv 영상 웹 크롤링
이전 자료 링크 <- 한강공원 자전거 사고 방지 알림 시스템[2] – python, 파이썬,cctv 영상, yolo v3, object detection
영상에서 image 추출
import cv2
import os
OUTPUT_PATH = ''
THRESHOLD = 0 # 절대 차이의 합이 이 값 이상일 때만 이미지를 저장합니다.
MAX_SAVES = 100000 # 저장할 수 있는 최대 프레임 수
SKIP_FRAMES = 0
if not os.path.exists(OUTPUT_PATH):
os.mkdir(OUTPUT_PATH)
VIDEO_LIST = [
]
for VIDEO_PATH in VIDEO_LIST:
cap = cv2.VideoCapture(VIDEO_PATH)
video_name = os.path.splitext(os.path.basename(VIDEO_PATH))[0]
ret, prev_frame = cap.read() # 첫 번째 프레임 읽기
save_count = 0
frame_count = 0
while ret:
ret, current_frame = cap.read()
if not ret or current_frame is None:
break
# 두 프레임 간의 절대 차이 계산
height = current_frame.shape[0]
start_row = int(height / 3)
diff = cv2.absdiff(prev_frame[start_row:], current_frame[start_row:])
sum_diff = diff.sum()
print(sum_diff, end=', ')
if sum_diff > THRESHOLD:
# 차이가 큰 경우만 이미지 저장
filename = os.path.join(OUTPUT_PATH, f'frame_{video_name}_{frame_count}.jpg')
print(filename)
cv2.imwrite(filename, current_frame)
save_count += 1
if save_count >= MAX_SAVES:
# 최대 저장 수에 도달하면 종료
break
# 현재 프레임을 화면에 표시
cv2.imshow('Current Frame', current_frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
# 현재 프레임을 이전 프레임으로 설정
frame_count += 1
prev_frame = current_frame
# 지정된 프레임 수만큼 뛰어넘기
for _ in range(SKIP_FRAMES):
ret, _ = cap.read()
frame_count += 1
if not ret:
break
cap.release()
cv2.destroyAllWindows()
이 코드를 활용하여 이미지를 저장할 수 있다.
필요에 따라 전 프레임과 특정 임계 값 이상의 차이가 발생할 때만 영상을 캡처 하게 하거나 한 프레임 단위가 아닌 여러 프레임 단위로 캡처하게 할 수 있다.
그다음 이미지가 저장된 폴더를 가지고 다음과 같은 구조를 만든다. labels.cache 는 지금은 없는 것이 정상, 추후 과정을 밟다 보면 생성된다.
images 와 labels 는 폴더 이름도 똑같이 하기!
dataset 만들기
dataset labelimg 프로그램 설치 (window 기준)
labelimg.exe 라고 하는 프로그램 설치 -> 링크
해당 파일을 windows (c:) 위치에 둔다. (한글 경로가 없도록 하기 위한 가장 확실한 위치)
data 안에 predefined_classes 에서 원하는 라벨 추가
프로그램 실행하면 화면이 위와 같은데 1번을 눌러 이미지 폴더
2번을 눌러 라벨이 저장될 폴더
3번을 클릭하여 pascalVOC 가 아닌 YOLO 로 바꾸어 준다.
w 를 눌러 라벨을 선택하고 d 를 눌러 다음 이미지로 간다. a를 눌러 이전 이미지로 갈 수 있다.
빠르게 라벨링 하려면 프로그램 윗쪽에 view 버튼을 눌러 auto save mode 를 키고 use default labtel 를 누른 다음 그안에 원하는 label 을 적어놓으면 자동으로 저장된다.
위와 같이 label 파일과 classes.txt 파일이 생성되고
각각의 label 파일은 위와 같이 class 와 x_center,y_center,w,h 로 이루어져 있다.
yolo v5 학습하기
사전 설정
먼저 yolo v5 git 에서 코드를 다운 받는다. -> 링크
위와 같이 다운되면 먼저 dataset.yml 파일을 만들어서 수정한다.
이전 생성한 프로젝트 경로를 path 로 그리고 images 를 적어준다. 그리고 추가한 label 를 적어 주고 추가한 label 갯수에 따라 nc 도 수정해야 한다.
validation set 을 두실 분은 위에 과정을 따로 한번 더 진행하면 될 것이다.
학습시작, anaconda 가상환경
anaconda 가상환경을 설치했다는 가정하에
$ conda activate vision39
$ python train.py –img 320 –batch 32 –epochs 1500 –data dataset.yml –weights yolov5s.pt –workers 2
이렇게 하면 위와 같이 학습이 된다.
학습하는 동안 위에 폴더 위치에 모델이 저장된다. 이를 이용하여 추론하면 된다.
from matplotlib import scale
import torch
import numpy as np
import cv2
model = torch.hub.load('ultralytics/yolov5', 'custom', path='./best.pt', force_reload=False)
video_list = [
'C:/Users/jinu0/Downloads/data/seoul_hackathon/비디오보관/낮_선별/test1.ts',
'C:/Users/jinu0/Downloads/data/seoul_hackathon/비디오보관/낮_선별/test2.ts',
'C:/Users/jinu0/Downloads/data/seoul_hackathon/비디오보관/낮_선별/test3.ts',
'C:/Users/jinu0/Downloads/data/seoul_hackathon/비디오보관/낮_선별/test4.ts',
'C:/Users/jinu0/Downloads/data/seoul_hackathon/비디오보관/낮_선별/test5.ts'
]
for video_path in video_list:
cap = cv2.VideoCapture(video_path)
ret, frame = cap.read()
frame_count = 0
while ret:
ret, frame = cap.read()
if frame is None:
continue
results = model(frame)
cv2.imshow('YOLO', np.squeeze(results.render()))
# print(results)
# results.print()
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
위 코드를 통해 위와 같이 실시각 객체 인식을 할 수 있다.
다음 할일
자전거 알림이 목적이기 때문에 자전거가 어느 방향으로 갈지를 판단해야 한다.
또한 지금은 모델의 추론이 끝나면 바로 다음 프레임으로 넘어가기 때문에 1배속이 아니다. 이를 보상해줘야 한다.
또한 이 결과를 보일 하드웨어 제작과 아두이노 시리얼 통신을 통한 프로토 타입 제작이 필요하다.
0 Comments