4. YOLO v8을 이용한 이상행동 탐지

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. 컴퓨터비전/Abnormal'
file_root = f'{data_root}/data'

먼저 데이터 저

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

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 = 0  # Initialize num before the loop

    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()

3개 폴더 생성

# 프레임이 있는 파일 리스트만 추출
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(2024)
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

# 라벨링 정보를 재구성
# 예) 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)

print(len(frame_list))

frame_list

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()

csv 파일 생성됨

# train.csv
# valid.csv
# test.csv
# 데이터셋 만들기