Detectron 2 with custom dataset

2022.07.05 - Detectron2 공부하기

Detectron2 공부하기

 

이전에 쓰던 글을 바탕으로 추가적으로 작성하려고 합니다. 이 글 또한 공부를 하며 정리하며 작성하는 거라 정보가 부정확할수도 있습니다. 참고하며 봐주시면 감사하겠습니다. (_ _)

Preview

Detectron2 란?

Facebook(현 Meta) 인공지능 연구자들이 만든 pytorch 기반 object detection와 sementic segemanation을 위한 training inferecne 플랫폼입니다.

Detectron의 특징

Detectron2는 학습 루프를 pytorch로 짜인 engine을 통해 학습합니다. 기존 pytorch style이 아닌 caffe2 학습 스타일을 pytorch로 구현한 것입니다. 또한 python 최적화를 위해 연산량이 많이 드는 부분을 python이 아닌 Cuda와 C로 구현하였습니다.

초기 세팅

먼저 gpu를 사용하기 위해 저는 colab환경에서 코드를 작성하였습니다.

colab 환경에서 gdrive 불러오기

from google.colab import drive
drive.mount('/content/gdrive/')

해당 gdrvie로 이동한 후 test를 위한 directory를 생성합니다.

%cd /content/gdrive/MyDrive
!mkdir custom_detecton
%cd custom_detecton

이후 detectron2 git을 불러옵니다.

!git clone https://github.com/facebookresearch/detectron2.git
!python3 -m pip install -e detectron2
!python -m pip install pyyaml==5.1

이후 modeling을 하기위한 module들을 불러옵니다.

import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()

# import some common libraries
import numpy as np
import cv2
import random
from google.colab.patches import cv2_imshow

# import some common detectron2 utilities
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog
from detectron2.data.catalog import DatasetCatalog

이때 import module들이 불러올 때 오류가 나실 수도 있습니다. 이때 런타임을 재시작한번 해주시면 정상적으로 진행됩니다.
이후 train할 dataset을 불러와야합니다. 저는 roboflow에서 cocodataset형태로 데이터를 불러왔습니다.

!curl -L "https://app.roboflow.com/ds/vDaAxs9QK1?key=k7iUeK3Joc" > roboflow.zip; unzip roboflow.zip; rm roboflow.zip

train,val,test 파일들을 압축을 푸셨다면 이를 coco_instace에 등록해야 합니다.

from detectron2.data.datasets import register_coco_instances
register_coco_instances("my_dataset_train", {}, "/content/gdrive/MyDrive/custom_detecton/train/_annotations.coco.json", "/content/gdrive/MyDrive/custom_detecton/train")
register_coco_instances("my_dataset_val", {}, "/content/gdrive/MyDrive/custom_detecton/valid/_annotations.coco.json", "/content/gdrive/MyDrive/custom_detecton/valid")
register_coco_instances("my_dataset_test", {}, "/content/gdrive/MyDrive/custom_detecton/test/_annotations.coco.json", "/content/gdrive/MyDrive/custom_detecton/test")

인스턴스 감지, 인스턴스 분할 및 키포인트 감지를 위해 COCO의 json 주석 형식으로 데이터 세트를 등록려면 해당 작업을 진행해야합니다.

cocodataset이란?

 

2022.06.27 - COCO dataset

COCO dataset

Train and Test

#visualize training data
my_dataset_train_metadata = MetadataCatalog.get("my_dataset_train")
dataset_dicts = DatasetCatalog.get("my_dataset_train")

import random
from detectron2.utils.visualizer import Visualizer

for d in random.sample(dataset_dicts, 3):
    img = cv2.imread(d["file_name"])
    print(img)
    visualizer = Visualizer(img[:, :, ::-1], metadata=my_dataset_train_metadata, scale=0.5)
    vis = visualizer.draw_dataset_dict(d)
    cv2_imshow(vis.get_image()[:, :, ::-1])

이후 데이터셋에 대한 정보를 저장하는 dictionary를 만듭니다. 이후 해당 정보를 바탕으로 Visualizer를 통해 dataset을 이미지화 하며, MeatadataCatalog을 통해 해당 이미지에 맞는 annotation 정보를 가져옵니다. 

이후 modeling을 진행 한후  model file들을 저장할 directory를 하나 생성합니다.

#We are importing our own Trainer Module here to use the COCO validation evaluation during training. Otherwise no validation eval occurs.

from detectron2.engine import DefaultTrainer
from detectron2.evaluation import COCOEvaluator

class CocoTrainer(DefaultTrainer):

  @classmethod
  def build_evaluator(cls, cfg, dataset_name, output_folder=None):

    if output_folder is None:
        os.makedirs("coco_eval", exist_ok=True)
        output_folder = "coco_eval"

    return COCOEvaluator(dataset_name, cfg, False, output_folder)

 

이후 train 작업을 진행합니다.

 

#from .detectron2.tools.train_net import Trainer
#from detectron2.engine import DefaultTrainer
# select from modelzoo here: https://github.com/facebookresearch/detectron2/blob/master/MODEL_ZOO.md#coco-object-detection-baselines

from detectron2.config import get_cfg
#from detectron2.evaluation.coco_evaluation import COCOEvaluator
import os

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("my_dataset_train",)
cfg.DATASETS.TEST = ("my_dataset_val",)

cfg.DATALOADER.NUM_WORKERS = 4
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml")  # Let training initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = 4
cfg.SOLVER.BASE_LR = 0.01


cfg.SOLVER.WARMUP_ITERS = 1600
cfg.SOLVER.MAX_ITER = 1500 #adjust up if val mAP is still rising, adjust down if overfit
cfg.SOLVER.STEPS = (1000, 1500)
cfg.SOLVER.GAMMA = 0.05




cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 64
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 7 #your number of classes + 1

cfg.TEST.EVAL_PERIOD = 500


os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = CocoTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()

CoCo-detection에 있는 faseter_rcnn 모델을 config 파일에 합칩니다. 이후 훈련시킬 데이터셋을 설정한후 

num_worker를 통해 학습 도중 몇개의 cpu를 사용할건지 설정합니다. 이후 model zoo에 있는 모델을 실행 시킵니다.

• Base_Lr라는 매개변수는 초기 학습률을 설정하는 변수입니다.
• Gamma라는 매개 변수는 학습률의 매번 변화 정도를 나타내고 값은 실수이다.
• stepvalue라는 매개 변수는 언제 훈련을 해야 하는지를 나타내는 다음 과정으로 값은 정수이다.주로 lr_policy이 multistep인 경우에 쓰인다.
• max_iter 이 매개 변수는 신경 네트워크 교체를 훈련하는 최대 횟수를 나타냅니다. 값은 정수입니다.

이후 train이 다끝났으면 훈련된 mode file을 바탕으로 test를 진행합니다.

 

#test evaluation
from detectron2.config import get_cfg
#from detectron2.evaluation.coco_evaluation import COCOEvaluator
import os
from detectron2.data import DatasetCatalog, MetadataCatalog, build_detection_test_loader
from detectron2.evaluation import COCOEvaluator, inference_on_dataset

cfg.MODEL.WEIGHTS = os.path.join("/content/gdrive/MyDrive/custom_detecton/output", "model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7
predictor = DefaultPredictor(cfg)
evaluator = COCOEvaluator("my_dataset_test", cfg, False, output_dir="./output/")
val_loader = build_detection_test_loader(cfg, "my_dataset_test")
inference_on_dataset(trainer.model, val_loader, evaluator)

#test evaluation
from detectron2.data import DatasetCatalog, MetadataCatalog, build_detection_test_loader
from detectron2.evaluation import COCOEvaluator, inference_on_dataset

cfg.MODEL.WEIGHTS = os.path.join("/content/gdrive/MyDrive/custom_detecton/output/", "model_final.pth")
cfg.DATASETS.TEST = ("my_dataset_test", )
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7   # set the testing threshold for this model
predictor = DefaultPredictor(cfg)
test_metadata = MetadataCatalog.get("my_dataset_test")

from detectron2.utils.visualizer import ColorMode
import glob

for imageName in glob.glob('/content/gdrive/MyDrive/custom_detecton/test/*jpg'):
  im = cv2.imread(imageName)
  outputs = predictor(im)
  v = Visualizer(im[:, :, ::-1],
                metadata=test_metadata, 
                scale=0.8
                 )
  out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
  cv2_imshow(out.get_image()[:, :, ::-1])

 

 

해당 테스트 데이터들이 잘 훈련된 모습을 볼 수 있습니다.

이상으로 custom dataset을 이용하여 detectron2를 진행하신 과정을 보실수 있었습니다.

틀린 부분과 잘못생각한 부분을 보셨으면 말해주시면 감사하겠습니다. 

 

 

참조: 

https://intrepidgeeks.com/tutorial/introduction-to-solver-parameter-in-cafe

Caffe의 Solver 매개 변수 소개