[OpenCV] 동영상에서 이미지 추출하기

YOLO를 이용한 커스텀 학습을 위해서는 이미지 데이터를 만들어야 한다.

 

우선 파이썬과 opencv가 설치되어 있어야 한다.

---

import cv2

count = 1
test = 1
train = 1

for i in range(1,4):
    vidcap = cv2.VideoCapture("D:/video/%d.mp4" %i)

    print("%d번째 영상 시작" %i)

    while(vidcap.isOpened()):
        ret, image = vidcap.read()

        if(ret==False):
            print("%d번째 영상 끝" %i)
            print("train=%d" %train)
            print("test=%d" %test)
            break

        if(int(vidcap.get(1)) % 5 == 0):

            num=count % 10

            if num==0 or num==5: # 20%는 test data, 80%는 train data
                cv2.imwrite("D:/labeling/test1/test%s.jpg" %str(test).zfill(5), image)
                test +=1
            else:
                cv2.imwrite("D:/labeling/train1/train%s.jpg" %str(train).zfill(5), image)
                train += 1

            count += 1

    vidcap.release()

print("devide end")

※ 주의할 점

1. 영상의 확장자가 mp4인지 확인할 것. 아니라면 바꿔야 함.

2. 영상 제목을 1, 2, 3.. 처럼 숫자로 수정해야 함. 아니면 코드를 바꾸는 것도 방법임.

3. test, train 데이터의 비율을 1:4로 나누기 위해 count 변수를 두고 10으로 나눈 나머지 값을 사용해서 분류함.

4. 나중에 YOLO 학습할 때 이미지의 순서가 뒤죽박죽되지 않게 만들기 위해 zfill 함수를 이용해서 자릿수를 통일시킴.