Python(60)- 영상데이터를 활용한 이상탐지

*이 글을 읽기전에 작성자 개인의견이 있으니, 다른 블로그와 교차로 읽는것을 권장합니다.*

import os
import random
import shutil
import cv2
import glob
import xml.etree.ElementTree as ET
import csv
from tqdm import tqdm

data_root = '/content/drive/MyDrive/KDT/8. 컴퓨터 비전/8. Abnormal'
file_root = f'{data_root}/data'

file_list = os.listdir(f'{file_root}/images_id')
len(file_list)

# mp4에서 이미지 추출하기
# 각 mp4이름으로 디렉토리 생성 후 저장
for file in tqdm(file_list):
    name = file.split('.')[0]
    if not os.path.isdir(f'{file_root}/images/{name}'):
        os.mkdir(f'{file_root}/images/{name}')
    video_path = f'{file_root}/images_id/{file}'
    cap = cv2.VideoCapture(video_path)

    num = 1
    while(cap.isOpened()):
        ret, frame = cap.read()
        if ret:
            cv2.imwrite(f'{file_root}/images/{name}/{name}_{num:06d}.png', frame)
            num += 1
        else:
            break

    cap.release()

# 프레임이 있는 파일 리스트만 추출
file_list = os.listdir(f'{file_root}/images_id')
real_file_list = []
for file in file_list:
    temp = os.listdir(f"{file_root}/images/{file.split('.')[-2]}")
    if len(temp) > 0:
        real_file_list.append(file)

len(real_file_list)

random.seed(2004)
random.shuffle(real_file_list)
test_ratio = 0.1

num_file = len(real_file_list)

test_list = real_file_list[:int(num_file*test_ratio)]
valid_list = real_file_list[int(num_file*test_ratio):int(num_file*test_ratio)*2]
train_list = real_file_list[int(num_file*test_ratio)*2:]

def read_one_xml(xml_path):
    tree = ET.parse(xml_path)
    root = tree.getroot()
    event_dict = {
        'event': None,
        'start_frame': -1,
        'end_frame': -1
    }

    for child in root.iter('track'):
        if child.attrib['label'] == 'theft_start':
            event_dict['event'] = child.attrib['label'].split('_')[0]
            frame = child.find('box').attrib['frame']
            event_dict['start_frame'] = int(frame)
        elif child.attrib['label'] == 'theft_end':
            event_dict['event'] = child.attrib['label'].split('_')[0]
            try:
                frame = child.find('box').attrib['frame']
            except:
                return {}, {}
            event_dict['end_frame'] = int(frame)
    return event_dict

# XML파일에서 행동정보 추출하기
# {'event': 'theft', 'start_frame': X, 'end_frame': X}
file = real_file_list[0]
xml_path = f"{file_root}/labels_id/{file.replace('mp4', 'xml')}"
event_dict = read_one_xml(xml_path)
event_dict

# 라벨링 정보를 재구성
# 예) frame1, frame2, frame3, Label(정상:0 or 이상:1)
# [a1.png, a2.png, a3.png, 0]
# [a4.png, a5.png, a6.png, 0]
# ...
# [a75.png, a76.png, a77.png, 1]
name = file.split('.')[0]
frame_list = []
if event_dict is not None:
    current_frame = 0
    while(current_frame < event_dict['start_frame']+4):
        tmp = []
        if (current_frame+3) >= event_dict['start_frame']:
            event_num = '1'
        else:
            event_num = '0'
        tmp.append([
            f'images/{name}/{name}_{(current_frame+1):06d}.png',
            f'images/{name}/{name}_{(current_frame+2):06d}.png',
            f'images/{name}/{name}_{(current_frame+3):06d}.png',
            event_num
        ])
        current_frame += 3
        frame_list.append(tmp)
    tmp = []
    tmp.append([
            f'images/{name}/{name}_{(current_frame+1):06d}.png',
            f'images/{name}/{name}_{(current_frame+2):06d}.png',
            f'images/{name}/{name}_{(current_frame+3):06d}.png',
            event_num
    ])
    frame_list.append(tmp)

# frame_list[0][0]
# ['images/폴더명/파일명1.png', ''images/폴더명/파일명2.png', ''images/폴더명/파일명3.png', 0 or 1]
print(len(frame_list))
print(frame_list[0][0])

name_list = f'{file_root}/ex.csv'
f = open(name_list, 'w', newline='')
writer = csv.writer(f)
for i in frame_list:
    writer.writerow(i[0])
f.close()