온정완

AI/플랫폼 개발을 공부합니다.

촬영 및 가공과정

모델이 마무리되고, 이제 테스트셋을 구축하는 데 집중했다. 배경이나 화각을 달리 하기 위해 테스트셋 촬영 장소를 회사 밖으로 옮겼다. 현장 cctv와 화각을 비슷하게 하기 위해서 벤치 위에 올라가거나 계단 위에 올라가서 촬영을 진행하였고, 다양한 case 확보를 위해 여러 각도에서 촬영을 진행하였다. 자연스럽게 걸어가거나 쓰레기를 투기하는 행위를 동영상으로 찍고, 해당 동영상을 OpenCV를 통해 프레임별로 뽑아주었다. 그 후 사람이 나오지 않는 프레임은 제거하고 사람이 나온 프레임만 모델에 돌리는 과정으로 사전 테스트를 진행하였다.



테스트셋 가공과정 자동화

테스트셋이 많아질 경우 위에서 언급한 가공 과정이 매우 귀찮고 지루한 반복 작업이 될 수 있기 때문에, 간단한 자동화 코드를 만들어보았다.

0) 사용 Library 및 상수선언

사람 검출을 위해서 모델이 들어있는 run_test_수정본.py 파일을 import해서 사용해주었다. 해당 파일에 있는 SSD 네트워크 모델을 사용하였으며, TensorRT를 사용하여 연산 속도를 높였다.

# image handling
import cv2
import matplotlib.pyplot as plt
import run_test_수정본

# system handle
import os
import sys

# visualization
from tqdm import tqdm

VIDEO_PATH에 저장되어 있는 video들을 전부 불러와서, FPS마다 프레임을 추출하여 FRAME_SAVE_PATH에 저장한다. FRAME_SAVE_PATH에 저장되어 있는 프레임을 다시 모델에 넣고, 사람이 감지된 프레임을 다시 HUMAN_SAVE_PATH에 저장하는 구조이다. 이 때, FRAME_SAVE_PATH를 거치지 않고 바로 모델에 프레임을 넣을 수도 있었으나, 원본 이미지 확보를 위하여 step을 하나 더 추가해 주었다. MODE에 remove를 넣으면 모든 파일이 리셋되게 하였다.

VIDEO_PATH = '/home/cctv_lab/cctv_dumping_verification/ojw/테스트셋/video/video'
FRAME_SAVE_PATH = '/home/cctv_lab/cctv_dumping_verification/ojw/테스트셋/video/frame'
HUMAN_SAVE_PATH = '/home/cctv_lab/cctv_dumping_verification/ojw/테스트셋/test_image/human_detect'
MODE = ''
FPS = 1

if MODE == 'remove':
    for (root, dirs, files) in os.walk(FRAME_SAVE_PATH):
        for file in files:
            file_name, file_ext = os.path.splitext(file)   
            os.remove(os.path.join(FRAME_SAVE_PATH,file))   

    for (root, dirs, files) in os.walk(HUMAN_SAVE_PATH):
        for file in files:
            os.remove(os.path.join(HUMAN_SAVE_PATH,file))                   
    sys.exit()


1) 동영상에서 프레임 추출

동영상에서 프레임을 추출하는 과정을 함수로 만들어서 관리하였다. 동영상, 저장 파일명, 프레임 단위를 input값으로 받으며 기본적으로 10프레임씩 자른다. 가령, 1초에 30프레임이 찍히는 동영상을 해당 함수에 넣으면, file_name_n.jpg라는 이름으로 0.3초마다 프레임이 저장되는 형태이다.

def save_frame(video, file_name, fps=10):
    count = 0
    break_count = 0

    while True:
        ret, image = video.read()

        if(int(video.get(1)) % fps == 0):
            break_count = 0
            tmp = os.path.join(FRAME_SAVE_PATH, file_name + '_{}.jpg'.format(count))
            print(tmp)

            try:
                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

                plt.imsave(tmp, image)
                count += 1

            except:
                break

        break_count += 1
        if break_count > fps:
            break

    video.release()

for (root, dirs, files) in os.walk(VIDEO_PATH):
    for file in files:
        file_name, file_ext = os.path.splitext(file) 

        video = cv2.VideoCapture(os.path.join(VIDEO_PATH, file))        
        save_frame(video, file_name.replace('_',''), FPS)


2) 사람이 검출된 이미지만 저장

run_test_수정본의 running_process() 함수를 사용하여 model을 돌린다. 기존 running_process() 함수는 객체 탐지 모델 외에도 다른 모델도 돌리기 때문에 __name__이 __main__인 경우와 그렇지 않은 경우에 대해 다르게 선언해주었다. 자세한 코드는 보안 문제때문에 생략한다. model에서 계산된 confidence 중, threshold를 넘는 값이 하나라도 있으면 obj_yn에 True를 리턴하도록 하였고, obj_yn이 True이면 사람이 검출된 것으로 간주하여 HUMAN_SAVE_PATH에 저장해주었다.

for (root, dirs, files) in os.walk(FRAME_SAVE_PATH):
    for file in tqdm(files):
        test_image = cv2.imread(os.path.join(FRAME_SAVE_PATH, file), cv2.IMREAD_COLOR)
        img, obj_yn = run_test_수정본.running_process(test_image)

        if obj_yn:
            plt.imsave( os.path.join(HUMAN_SAVE_PATH, file), img )

image


3) 사람 검출 이미지에 대하여 투기 여부 계산

최종적으로 사람이 검출된 이미지에 대해서 모델을 돌려주면 모든 과정이 끝이난다. image

4) 전체 스크립트

import cv2
import os
import matplotlib.pyplot as plt
from tqdm import tqdm
import sys
import run_test_수정본

VIDEO_PATH = '비디오_경로'
FRAME_SAVE_PATH = 'frame_저장_경로'
HUMAN_SAVE_PATH = '사람_검출_이미지_저장_'
MODE = ''
FPS = 1

if MODE == 'remove':
    for (root, dirs, files) in os.walk(FRAME_SAVE_PATH):
        for file in files:
            file_name, file_ext = os.path.splitext(file)   
            os.remove(os.path.join(FRAME_SAVE_PATH,file))   

    for (root, dirs, files) in os.walk(HUMAN_SAVE_PATH):
        for file in files:
            os.remove(os.path.join(HUMAN_SAVE_PATH,file))                   
    sys.exit()

def save_frame(video, success_path, file_name, fps=10):
    count = 0
    break_count = 0

    while True:
        ret, image = video.read()

        if(int(video.get(1)) % fps == 0):
            break_count = 0

            tmp = os.path.join(success_path, file_name + '_{}.jpg'.format(count))
            print(tmp)

            try:
                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

                plt.imsave(tmp, image)
                count += 1

            except:
                break

        break_count += 1
        if break_count > fps:
            break

    video.release()

if __name__ == '__main__':
    ## 프레임 저장
    for (root, dirs, files) in os.walk(VIDEO_PATH):
        for file in files:
            file_name, file_ext = os.path.splitext(file) 
            file_name = file_name.replace('_','')

            video = cv2.VideoCapture(os.path.join(VIDEO_PATH, file))        
            
            success_path = os.path.join(FRAME_SAVE_PATH, file_name)

            if not os.path.isdir(success_path):
                os.mkdir(success_path)

            save_frame(video, success_path, file_name.replace('_',''), FPS)
            
    ## 사람 검출 이미지만 저장
    for dir in tqdm(os.listdir(FRAME_SAVE_PATH)):
        PATH = os.path.join(FRAME_SAVE_PATH, dir)

        for (root, dirs, files) in os.walk(PATH):
            for file in tqdm(files):
                file_name, file_ext = os.path.splitext(file)
                test_image = cv2.imread(os.path.join(PATH, file), cv2.IMREAD_COLOR)
                img, obj_yn = run_test_수정본.running_process(test_image)

                if obj_yn:
                    success_path = os.path.join(HUMAN_SAVE_PATH, dir)

                    if not os.path.isdir(success_path):
                        os.mkdir(success_path)

                    plt.imsave( os.path.join(success_path, file), img )

댓글남기기

참고