🔥알림🔥
테디노트 유튜브 - 구경하러 가기!

12 분 소요

본 포스팅은 데이콘(dacon.io)에서 2023.02.06 ~ 2023.03.13 기간 동안 진행하는 자동차 충돌 분석 AI경진대회에 제출한 베이스라인 코드 입니다.

코드


대회소개

제 1회 코스포 x 데이콘 자동차 충돌 분석 AI경진대회(채용 연계형)

본 대회는 코스포와 데이콘에서 주관한 대회로 블랙박스 영상을 활용하여 자동차의 충돌 상황을 분석하는 AI 모델을 생성하는 것을 목표로 합니다.

대회의 평가 방식은 Macro F1 Score를 사용합니다.

data 압축 해제

대회 링크의 데이터 탭에서 다운로드 받은 데이터를 압축 해제 합니다.

train / test 폴더로 구분되어 있으며, 각 폴더 안에는 .mp4 영상이 있습니다.

# import urllib.request
# import zipfile
# import os

# def unzip(source_file, dest_path):
#     with zipfile.ZipFile(source_file, 'r') as zf:
#         zipInfo = zf.infolist()
#         for member in zipInfo:
#             try:
#                 member.filename = member.filename.encode('cp437').decode('euc-kr', 'ignore')
#                 zf.extract(member, dest_path)
#             except:
#                 print(source_file)
#                 raise Exception('unzipping error')
                

# unzip('./data/open.zip', './data')

SEED 고정

import random
import torch
import numpy as np
import os

SEED = 123

def seed_everything(seed):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = True

# Seed 고정
seed_everything(SEED) 

하이퍼파라미터 설정

CFG = {
    'VIDEO_LENGTH': 50,      # 10프레임 * 5초
    'IMG_SIZE': 224,         # 224 x 224
    'EPOCHS': 10,            # Epochs
    'LEARNING_RATE': 3e-4,   # Learning Rate
    'BATCH_SIZE': 4,         # batch_size
    'SEED': 123,             # SEED 값
    'MODEL_NAME': 'Base-Model'
}

데이터 로드

import pandas as pd
import os
import warnings

warnings.filterwarnings('ignore')

DATA_DIR = 'data'
SUBMISSION_DIR = 'submit'
MODEL_DIR = 'model'
PRETRAINED_DIR = 'pretrained'

train = pd.read_csv(os.path.join(DATA_DIR, 'train.csv'))
test = pd.read_csv(os.path.join(DATA_DIR, 'test.csv'))
# 5개의 행 출력
train.head()
sample_id video_path label
0 TRAIN_0000 ./train/TRAIN_0000.mp4 7
1 TRAIN_0001 ./train/TRAIN_0001.mp4 7
2 TRAIN_0002 ./train/TRAIN_0002.mp4 0
3 TRAIN_0003 ./train/TRAIN_0003.mp4 0
4 TRAIN_0004 ./train/TRAIN_0004.mp4 1

video_path 디렉토리 수정

비디오 경로를 수정

# train set의 video_path 수정
train['video_path'] = train['video_path'].apply(lambda x: os.path.join(DATA_DIR, x[2:]))
train['video_path'].head()
0    data/train/TRAIN_0000.mp4
1    data/train/TRAIN_0001.mp4
2    data/train/TRAIN_0002.mp4
3    data/train/TRAIN_0003.mp4
4    data/train/TRAIN_0004.mp4
Name: video_path, dtype: object
# test set의 video_path 수정
test['video_path'] = test['video_path'].apply(lambda x: os.path.join(DATA_DIR, x[2:]))
test['video_path'].head()
0    data/test/TEST_0000.mp4
1    data/test/TEST_0001.mp4
2    data/test/TEST_0002.mp4
3    data/test/TEST_0003.mp4
4    data/test/TEST_0004.mp4
Name: video_path, dtype: object

Label 값의 데이터 Imbalance 확인

본 대회는 총 13개로 이루어진 label 분류 대회입니다.

각각의 label 은 다음의 표 형식에 따라 분류 됩니다.

0번 label의 분포가 가장 많습니다. 기타 label은 매우 imbalance 한 분포를 가집니다.

# label의 분포를 출력합니다
train['label'].value_counts()
0     1783
1      318
7      317
3       78
2       51
9       34
11      33
8       30
5       28
4       13
12       6
10       4
6        3
Name: label, dtype: int64

학습/검증 셋 분할

from sklearn.model_selection import train_test_split

# 데이터셋 분할
train_data, valid_data = train_test_split(train, 
                                          test_size=0.2, 
                                          stratify=train['label'], 
                                          random_state=SEED)
# 분할된 데이터셋 shape 출력
train_data.shape, valid_data.shape
((2158, 3), (540, 3))

영상 로드 (opencv2)

train.csv 파일에는 video_path 컬럼이 존재하는데 video 영상을 로드하기 위한 경로가 있습니다.

cv2VideoCapture() 함수로 영상을 로드합니다. 그리고, 로드한 영상에서 Frame을 추출하는 작업을 진행합니다.

# 5개 출력
train_data.head()
sample_id video_path label
2491 TRAIN_2491 data/train/TRAIN_2491.mp4 7
236 TRAIN_0236 data/train/TRAIN_0236.mp4 0
1591 TRAIN_1591 data/train/TRAIN_1591.mp4 0
151 TRAIN_0151 data/train/TRAIN_0151.mp4 0
1683 TRAIN_1683 data/train/TRAIN_1683.mp4 0
# 0번째 (셔플이 되어 있으므로 2491번 index) 데이터의 video_path 를 가져옵니다
sample_idx = 0
sample_video_path = train_data['video_path'].iloc[sample_idx]
sample_label = train_data['label'].iloc[sample_idx]
sample_video_path, sample_label
('data/train/TRAIN_2491.mp4', 7)

영상 로드를 위한 cv2 라이브러리를 import 한 뒤 VideoCapture() 함수로 영상을 로드 합니다.

import cv2

captured_output = cv2.VideoCapture(sample_video_path)

아래 코드는 영상을 열어 Frame을 읽어들여 frames list 변수에 frame 한 장씩 추가하는 코드입니다.

1초10 frame, 그리고 5초짜리 영상이니 총 50 frame이 추출됩니다.

frames = []
IMG_SIZE = 224

# 영상 로드
captured_output = cv2.VideoCapture(sample_video_path)

# 5*10 => 5초 * 10프레임
for _ in range(5*10):
    try:
        # 1장 frame
        _, img = captured_output.read()
        # image resize
        img = cv2.resize(img, (IMG_SIZE, IMG_SIZE))
        # 정규화
        img = img / 255.
        # 프레임에 추가
        frames.append(img)
    except Exception as e:
        print(str(e))

총 50장의 frame, 그리고 1장의 frame(사진)은 (224, 224, 3) 의 형태를 가집니다.

# 50 프레임, 1 프레임 (224 X 224)
len(frames), frames[0].shape
(50, (224, 224, 3))

1개 영상 프레임별 시각화

import matplotlib.pyplot as plt


def visualize_frames(frames, label):
    print(f'Label: {label}')
    fig, axes = plt.subplots(10, 5)
    fig.set_size_inches(10, 20)
    for i in range(50):
        axes[i//5, i%5].imshow(frames[i])
        axes[i//5, i%5].axis('off')
        axes[i//5, i%5].set_title(f'frame {i}')

    plt.tight_layout()
    plt.show()
# 1개 동영상 프레임별 시각화
visualize_frames(frames, sample_label)
Label: 7