[02 Object Detection] - MMdetection train/val 학습 방법

지난 블로그에서는 MMDetection의 기본적인 설치와 사용 방법에 대해 다뤘습니다. 이번 포스트에서는 train/val 데이터를 나누고 학습을 진행하는 방법에 대해 알아보겠습니다. 특히, COCO 형식의 데이터셋을 사용해 객체 탐지 모델을 학습하는 방법을 중점적으로 설명하겠습니다.

[02 Object Detection] - MMdetection 설치 및 기본사용법

---

1. 데이터셋 준비 및 Train/Val 분할

 

객체 탐지 모델을 학습할 때, 보통 학습 데이터와 검증 데이터를 8:2의 비율로 나누어 학습을 진행합니다. 이번 블로그에서는 COCO 형식의 데이터셋을 사용하여 train/val 데이터를 8:2로 나누는 코드를 소개하겠습니다. 이 코드는 데이터셋을 준비한 후, 간단한 파이썬 코드를 통해 데이터를 분리할 수 있도록 합니다.

 

# %%
import json
import os
import random
import shutil
from pycocotools.coco import COCO

# 설정
data_dir = '../dataset'  # 데이터가 저장된 디렉토리
coco_file = os.path.join(data_dir, 'train.json')  # COCO 형식의 JSON 파일
output_dir = os.path.join(data_dir, 'split')  # 결과를 저장할 디렉토리
train_ratio = 0.8  # 훈련 데이터 비율

# COCO 데이터셋 로드
with open(coco_file, 'r') as f:
    coco_data = json.load(f)

# 이미지 목록과 갯수
images = coco_data['images']
num_images = len(images)

# 랜덤
random.seed(42)  
random.shuffle(images)

# 훈련 데이터와 검증 데이터 split
train_count = int(num_images * train_ratio)
train_images = images[:train_count]
val_images = images[train_count:]

# 훈련 및 검증 데이터를 저장할 디렉토리 생성
os.makedirs(os.path.join(output_dir, 'train'), exist_ok=True)
os.makedirs(os.path.join(output_dir, 'val'), exist_ok=True)

# %%

# 훈련 데이터 저장
train_data = {
    'info': coco_data['info'],
    'licenses': coco_data['licenses'],
    'categories': coco_data['categories'], 
    'images': train_images,
    'annotations': []
}
for img in train_images:
    img_id = img['id']
    train_data['annotations'].extend([ann for ann in coco_data['annotations'] if ann['image_id'] == img_id])

with open(os.path.join(output_dir,  'train.json'), 'w') as f:
    json.dump(train_data, f)

# %%
# 검증 데이터 저장
val_data = {
    'info': coco_data['info'],
    'licenses': coco_data['licenses'],
    'categories': coco_data['categories'],  
    'images': val_images,
    'annotations': []
}
for img in val_images:
    img_id = img['id']
    val_data['annotations'].extend([ann for ann in coco_data['annotations'] if ann['image_id'] == img_id])

with open(os.path.join(output_dir, 'val.json'), 'w') as f:
    json.dump(val_data, f)

# %%

# train 이미지 파일 복사
for img in train_images:
    img_filename = img['file_name'] 
    src_path = os.path.join(data_dir, img_filename)  # 원본 경로
    dst_path = os.path.join(data_dir, 'split', 'train', img_filename)  # 복사할 경로 (train 폴더)

    # 원본 파일 확인 후 복사
    if os.path.exists(src_path):
        # 대상 디렉토리가 없으면 생성
        os.makedirs(os.path.dirname(dst_path), exist_ok=True)
        shutil.copy(src_path, dst_path)  # 파일 복사
        # print(f"복사 완료: {src_path} -> {dst_path}")
    else:
        print(f"파일을 찾을 수 없습니다: {src_path}")

print("Train 이미지 복사가 완료되었습니다!")

# %%
# val 이미지 파일 복사
for img in val_images:
    img_filename = img['file_name'] 
    src_path = os.path.join(data_dir, img_filename)  # 원본 경로
    dst_path = os.path.join(data_dir, 'split', 'val', os.path.basename(img_filename))  # 복사할 경로 (val 폴더)

    # 원본 파일 확인 후 복사
    if os.path.exists(src_path):
        # 대상 디렉토리가 없으면 생성
        os.makedirs(os.path.dirname(dst_path), exist_ok=True)
        shutil.copy(src_path, dst_path)  # 파일 복사
        # print(f"복사 완료: {src_path} -> {dst_path}")
    else:
        print(f"파일을 찾을 수 없습니다: {src_path}")

print("Val 이미지 복사가 완료되었습니다!")

print("Train/Val 데이터셋이 성공적으로 분할되었습니다!")

 

디렉토리 구조는 다음과 같이 됩니다.

dataset/
├── split/
│   ├── train.json
│   ├── val.json
│   ├── train/  # 학습 이미지
│   └── val/    # 검증 이미지

 

---

 

2. Yolo 모델로 train/val 학습 방법

YOLO 모델을 학습하기 전에, 먼저 MMDetection에서 YOLO 모델을 사용할 수 있도록 설정 파일을 준비해야 합니다. MMDetection은 다양한 모델을 쉽게 설정하고 학습할 수 있는 프레임워크로, YOLO 모델도 기본적으로 지원합니다.

우선, YOLO 모델의 기본 설정 파일을 불러와서 학습에 필요한 클래스 정보, 데이터 경로 등을 수정할 수 있습니다.

# 기본 설정 파일 로드
_base_ = './configs/yolo/yolov3_d53_8xb8-320-273e_coco.py'


# 모델 정의
model = dict(
    data_preprocessor=dict(
        # 객체 탐지 사용을 위한 설정
        pad_mask=False,  # mask 사용 false
        batch_augments=None  # 배치 단위 증강 none
    )
)

# class 변경
metainfo = {
    'classes': ("General trash", "Paper", "Paper pack", "Metal", "Glass",
                "Plastic", "Styrofoam", "Plastic bag", "Battery", "Clothing"),
}

train_pipeline = [
    dict(type='LoadImageFromFile', backend_args=_base_.backend_args),  # 파일로부터 이미지를 로드
    dict(type='LoadAnnotations', with_bbox=True, with_mask=False),  # 바운딩 박스 정보를 포함한 어노테이션 로드, mask 는 false
    dict(type='Resize', scale=(640, 640), keep_ratio=True),
    dict(type='PackDetInputs')  # 모델에 입력하기 위해 데이터를 포장
]

val_pipeline = [
    dict(type='LoadImageFromFile', backend_args=_base_.backend_args),  # 파일로부터 이미지를 로드
    dict(type='LoadAnnotations', with_bbox=True, with_mask=False),  # 바운딩 박스 정보를 포함한 어노테이션 로드, mask 는 false
    dict(type='Resize', scale=(640, 640), keep_ratio=True),
    dict(type='PackDetInputs')  # 모델에 입력하기 위해 데이터를 포장
]

data_root = '../dataset'  # dataset root

# train data
train_dataloader = dict(
    dataset=dict(
        _delete_=True,
        metainfo=metainfo,  # 클래스 정보와 팔레트 설정
        type=_base_.dataset_type,  # 사용할 데이터셋의 유형, 기본 설정에서 정의된 dataset_type 사용
        data_root=data_root,  # dataset root
        ann_file='./split/train.json',  # 학습에 사용할 어노테이션
        data_prefix=dict(img=data_root),  # 이미지 파일이 위치한 경로
        filter_cfg=dict(filter_empty_gt=False, min_size=32),  # 필터 설정
        pipeline=train_pipeline,  # augmentation
        backend_args=_base_.backend_args),  # 데이터 로드 방식 정의
    batch_size=4  # 배치사이즈 설정
)

# training configuration 추가
train_cfg = dict(
    type='EpochBasedTrainLoop',
    max_epochs=50,
    val_interval=1  # 매 1 에폭마다 검증을 실행하도록 설정
)


# validation data 설정
val_dataloader = dict(
    dataset=dict(
        metainfo=metainfo,  # 클래스 정보
        type='CocoDataset',  # COCO 형식 데이터셋 사용
        data_root=data_root,  # 검증 데이터 루트
        ann_file='./split/val.json',  # 검증 데이터 어노테이션 파일
        data_prefix=dict(img=data_root),  # 이미지 경로 설정
        filter_cfg=dict(filter_empty_gt=False, min_size=32),  # 필터 설정
        pipeline=val_pipeline,  # validation augmentation
    ),
    batch_size=4,
    num_workers=2,  # 데이터 로딩에 사용할 프로세스의 개수
    persistent_workers=True,
    drop_last=False,  # 배치 크기가 작아도 포함하여 학습
    sampler=dict(type='DefaultSampler', shuffle=False),  # 데이터셋 로딩 순서 설정
)

# validation configuration 설정
val_cfg = dict()  # 빈 dict로 설정해도 문제가 없음

# validation evaluator 설정
val_evaluator = dict(
    type='CocoMetric',
    ann_file='../dataset/split/val.json',
    metric=['bbox'],  # 필요한 metric 설정
    format_only=False,
    iou_thrs=[0.50, 0.75],  # Iou 성능 평가 기준
)


# test 없음
test_dataloader = None
test_cfg = None
test_evaluator = None

work_dir = './work_dirs/yolo'  # 저장할 workdir

 

위의 코드를 yolo_split_train.py 로 해서 만들어주고,

아래 명령어로 실행 합니다.

 

python tools/train.py yolo_split.py

 

 

mmdection 에 포함된 yolo 로 학습한 방법 입니다.

 

---

 

3. ATSS 모델로 학습

 

위의 YOLO 처럼 ATSS 모델로 학습을 시키면 아래 코드와 같습니다.

달라진 점은 모델에 numclass 수가 추가되었고, workdir 부분도 atss 로 수정했습니다.

 

# 기본 설정 파일 로드
_base_ = './configs/yolo/atss_r101_fpn_8xb8-amp-lsj-200e_coco.py'

num_classes = 10

# 모델 정의
model = dict(
    bbox_head=dict(
        num_classes=num_classes  # 10개 클래스에 맞게 설정
    ),
    data_preprocessor=dict(
        # 객체 탐지 사용을 위한 설정
        pad_mask=False,  # mask 사용 false
        batch_augments=None  # 배치 단위 증강 none
    )
)

# class 변경
metainfo = {
    'classes': ("General trash", "Paper", "Paper pack", "Metal", "Glass",
                "Plastic", "Styrofoam", "Plastic bag", "Battery", "Clothing"),
}

train_pipeline = [
    dict(type='LoadImageFromFile', backend_args=_base_.backend_args),  # 파일로부터 이미지를 로드
    dict(type='LoadAnnotations', with_bbox=True, with_mask=False),  # 바운딩 박스 정보를 포함한 어노테이션 로드, mask 는 false
    dict(type='Resize', scale=(640, 640), keep_ratio=True),
    dict(type='PackDetInputs')  # 모델에 입력하기 위해 데이터를 포장
]

val_pipeline = [
    dict(type='LoadImageFromFile', backend_args=_base_.backend_args),  # 파일로부터 이미지를 로드
    dict(type='LoadAnnotations', with_bbox=True, with_mask=False),  # 바운딩 박스 정보를 포함한 어노테이션 로드, mask 는 false
    dict(type='Resize', scale=(640, 640), keep_ratio=True),
    dict(type='PackDetInputs')  # 모델에 입력하기 위해 데이터를 포장
]

data_root = '../dataset'  # dataset root

# train data
train_dataloader = dict(
    dataset=dict(
        _delete_=True,
        metainfo=metainfo,  # 클래스 정보와 팔레트 설정
        type=_base_.dataset_type,  # 사용할 데이터셋의 유형, 기본 설정에서 정의된 dataset_type 사용
        data_root=data_root,  # dataset root
        ann_file='./split/train.json',  # 학습에 사용할 어노테이션
        data_prefix=dict(img=data_root),  # 이미지 파일이 위치한 경로
        filter_cfg=dict(filter_empty_gt=False, min_size=32),  # 필터 설정
        pipeline=train_pipeline,  # augmentation
        backend_args=_base_.backend_args),  # 데이터 로드 방식 정의
    batch_size=4  # 배치사이즈 설정
)

# training configuration 추가
train_cfg = dict(
    type='EpochBasedTrainLoop',
    max_epochs=50,
    val_interval=1  # 매 1 에폭마다 검증을 실행하도록 설정
)


# validation data 설정
val_dataloader = dict(
    dataset=dict(
        metainfo=metainfo,  # 클래스 정보
        type='CocoDataset',  # COCO 형식 데이터셋 사용
        data_root=data_root,  # 검증 데이터 루트
        ann_file='./split/val.json',  # 검증 데이터 어노테이션 파일
        data_prefix=dict(img=data_root),  # 이미지 경로 설정
        filter_cfg=dict(filter_empty_gt=False, min_size=32),  # 필터 설정
        pipeline=val_pipeline,  # validation augmentation
    ),
    batch_size=4,
    num_workers=2,  # 데이터 로딩에 사용할 프로세스의 개수
    persistent_workers=True,
    drop_last=False,  # 배치 크기가 작아도 포함하여 학습
    sampler=dict(type='DefaultSampler', shuffle=False),  # 데이터셋 로딩 순서 설정
)

# validation configuration 설정
val_cfg = dict()  # 빈 dict로 설정해도 문제가 없음

# validation evaluator 설정
val_evaluator = dict(
    type='CocoMetric',
    ann_file='../dataset/split/val.json',
    metric=['bbox'],  # 필요한 metric 설정
    format_only=False,
    iou_thrs=[0.50, 0.75],  # Iou 성능 평가 기준
)


# test 없음
test_dataloader = None
test_cfg = None
test_evaluator = None

work_dir = './work_dirs/atss'  # 저장할 workdir

 

 

 

이상으로 train/val 학습 방법에 대해서 공유 드렸습니다.

다음 블로그는 학습된 파일을 tensorborad 에 적용하는 방법에 대해서 공유 하겠습니다.

 

감사합니다.