mmdetection을 이용한 BCCD 데이터셋 학습

BCCD 데이터셋 sample

https://github.com/Shenggan/BCCD_Dataset 에서 BCCD 데이터 셋을 다운받아 보았다.

이번에 해볼 실습에서는 voc형태의 데이터를 coco 형태의 데이터로 변환하는 package를 사용해본다.

https://github.com/yukkyo/voc2coco 에서 package를 다운받았다.

 

package를 다운받은 후

%cd voc2coco
!python voc2coco.py --ann_dir /content/BCCD_Dataset/BCCD/Annotations \
--ann_ids /content/BCCD_Dataset/BCCD/ImageSets/Main/train.txt \
--labels /content/BCCD_Dataset/BCCD/labels.txt \
--output /content/BCCD_Dataset/BCCD/train.json \
--ext xml

!python voc2coco.py --ann_dir /content/BCCD_Dataset/BCCD/Annotations \
--ann_ids /content/BCCD_Dataset/BCCD/ImageSets/Main/val.txt \
--labels /content/BCCD_Dataset/BCCD/labels.txt \
--output /content/BCCD_Dataset/BCCD/val.json \
--ext xml

!python voc2coco.py --ann_dir /content/BCCD_Dataset/BCCD/Annotations \
--ann_ids /content/BCCD_Dataset/BCCD/ImageSets/Main/test.txt \
--labels /content/BCCD_Dataset/BCCD/labels.txt \
--output /content/BCCD_Dataset/BCCD/test.json \
--ext xml

cd 경로를 다운받은 voc2coco로 지정해준 뒤 변환 voc2coco.py를 실행한다. 이때 참조는 다운받은 BCCD데이터 셋의 annotations 폴더를 참조한뒤 각각의 ann_ids와 labels을 파싱한뒤 output으로 원하는 coco데이터셋의 json형태가 나오도록한다.

원하는 coco형태의 json파일이 생성됨을 볼 수 있다.

 

 

 

 

 

 

 

이제 만들어진 coco 형식의 dataset을 mmdetection의 학습 데이터셋으로 변환해줘야 한다. mmdetection은 기본적으로 coco형태의 학습 데이터셋을 요구하기 때문에 비교적 간단한 과정을 통해 변환을 마칠 수 있다.

 

from mmdet.datasets.builder import DATASETS
from mmdet.datasets.coco import CocoDataset

@DATASETS.register_module(force=True)
class BCCDDataset(CocoDataset):
  CLASSES = ('WBC', 'RBC', 'Platelets')

mmdetection에서 CocoDataset 변환 프레임워크를 사용하면 된다.

CLASSES에 원하는 label값인 WBC, RBC, Platelets를 적용해주면 간단하게 끝낼 수 있다.

 

이제 Config파일을 설정하고 checkpoint파일을 받아야된다.

from mmdet.apis import set_random_seed

cfg.dataset_type = 'BCCDDataset'
cfg.data_root = '/content/BCCD_Dataset/BCCD/'

cfg.data.train.type = 'BCCDDataset'
cfg.data.train.data_root = '/content/BCCD_Dataset/BCCD/'
cfg.data.train.ann_file = 'train.json'
cfg.data.train.img_prefix = 'JPEGImages'

cfg.data.val.type = 'BCCDDataset'
cfg.data.val.data_root = '/content/BCCD_Dataset/BCCD/'
cfg.data.val.ann_file = 'val.json'
cfg.data.val.img_prefix = 'JPEGImages'

cfg.data.test.type = 'BCCDDataset'
cfg.data.test.data_root = '/content/BCCD_Dataset/BCCD/'
cfg.data.test.ann_file = 'test.json'
cfg.data.test.img_prefix = 'JPEGImages'

cfg.model.roi_head.bbox_head.num_classes = 3
# pretrained 모델
cfg.load_from = '/content/mmdetection/checkpoints/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth'

# 학습 weight 파일로 로그를 저장하기 위한 디렉토리 설정. 
cfg.work_dir = './tutorial_exps'

# 학습율 변경 환경 파라미터 설정. 
cfg.optimizer.lr = 0.02 / 8
cfg.lr_config.warmup = None
cfg.log_config.interval = 10

cfg.evaluation.metric = 'bbox'
cfg.evaluation.interval = 12
cfg.checkpoint_config.interval = 12

cfg.lr_config.policy='step'

cfg.seed = 0
set_random_seed(0, deterministic=False)
cfg.gpu_ids = range(1)

config파일의 data_root는 아까전에 json파일을 만들었던 경로로 설정해주면 된다.

다음으로 train, val, test의 경로를 알맞게 설정, 총 class의 개수인 3을 설정, pretrained된 모델 경로 설정을 하고 eval_metric을 mAP가 아닌 bbox로 설정해 주어야 mmdetection모듈이 cocodataset에 맞게끔(iou threshold 값을 0.5부터 0.95까지) 학습을 진행한다.

 

이제 학습을 적용할 Dataset을 만들고 본격적으로 모델을 학습하고 inference를 적용하면 된다.

적용할 Dataset을 만들고

from mmdet.datasets import build_dataset
from mmdet.models import build_detector
from mmdet.apis import train_detector

datasets = [build_dataset(cfg.data.train)]

pretrained된 빈 모델을 생성해준다.

model = build_detector(cfg.model, train_cfg=cfg.get('train_cfg'), test_cfg=cfg.get('test_cfg'))

설정이 끝난 후 모델학습을 진행해줬다.

import os.path as osp
mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))

train_detector(model, datasets, cfg, distributed=False, validate=True)

모델 weight저장 위치는 cfg.work_dir로 설정후 기본 epoch인 12회만큼 학습을 진행하였다.

학습 완료 후 precision과 recall

val_dataset에 대해서의 성능에서 small dataset에 대해서는 비교적 아쉬운 성능을 보여주었다.

sample 분류 image