Detectron2(Centernet) 이미지 인식기 만들기(Linux)(2)

 

저번 글에 이어 이번에는 직접 이미지에 대해서 학습하고 인식하는 것 까지 시도 해보도록 하겠습니다!

 

Detectron2(Centernet) 이미지 인식기 만들기(Linux)(1)

 

 

먼저 학습을 돌릴 데이터들을 만들어 줍시다. 

 

저는 간단하게 제 카드를 가지고 실험 해 보았습니다. 

 

 

 

먼저 라벨 툴로 라벨링 해주고 

 

 

 

다음으로는 이 데이터들을 train, test 셋으로 분리 하면서 coco형태로 변형해 주어야 합니다.

 

train셋은 말 그대로 학습에 사용하는 데이터들이고

한번의 학습이 끝나면 test셋을 이용해 검증, 보완해나가는 구조입니다. 

 

 

 

방법은 직접 코드로 분리해도 되지만 labelme를 통해 라벨링 데이터를 생성하였다면 

 

 

아래와 같은 명령어로 분리 가능합니다. (명령어가 실행된 위치에 run파일이 생김)

labelme2coco path/to/labelme/dir --train_split_rate 0.85

이미지 인식 라벨링을 위한 Labelme Tool 사용 및 수정(1)

 

이렇게 분리하고 나면 

이렇게 두 파일이 생성되는데 이것이 각각 셋의 역할을 합니다 .

 

이제 학습, 테스트 코드를 구성하러 가봅시다.

 

 

 

 

학습데이터 폴더는 아래같은 형태로 구성했으니 참고 하시면 될 것 같습니다. 

 

 

학습기

 

학습은 아래의 코드로 이루어 집니다.

from detectron2.utils.logger import setup_logger
from detectron2.data.datasets import register_coco_instances
from detectron2.engine import DefaultTrainer

import os
import pickle
from utils import *


setup_logger()

config_file_path="COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml"  ## 사용모델
checkpoint_url="COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml"

outpu_dir="./DetectronCard/output" ## 결과값들 저장 위치

device="cuda"

train_dataset_name="Fake_train"   ## train 데이터셋 /주소
train_images_path="./DetectronCard/CardTest2_10fps_json"
train_json_annot_path="./DetectronCard/train.json"

test_dataset_name="Fake_test"     ## test 데이터셋 /주소
test_images_path="./DetectronCard/CardTest2_10fps_json"
test_json_annot_path="./DetectronCard/val.json"

cfg_save_path="./DetectronCard/TestCard.pickle"    ## 학습 가중치 파일 // 저장할곳

현재 데이터들이 모여있는 파일과 coco로 생성한 파일의 위치, 결과값 저장위치들을 지정해 줍시다.

그리고 가장 중요한 학습 가중치 파일을 따로 저장해줍시다. (학습 방법, 횟수 등을 컨트롤하는 파일!)

 

 

 

 

다음은 라벨링의 카테고리를 위한 메타데이터를 생성해줍니다. (현재는 카드가 1장이라..)

num_classes=1   ## 카테고리 갯수 ******************************************************

##############################################  메타 데이터 생성

register_coco_instances(name=train_dataset_name,metadata={},
                        json_file=train_json_annot_path,image_root=train_images_path)


register_coco_instances(name=test_dataset_name,metadata={},
                        json_file=test_json_annot_path,image_root=test_images_path)



#################################################

 

 

 

 

마지막으로 main에서 가중치 파일을 열어 이미 지정된 학습 루틴으로 학습을 진행합니다

def main():
    cfg=get_train_cfg(config_file_path,checkpoint_url,train_dataset_name,test_dataset_name,num_classes,device,outpu_dir)

    with open(cfg_save_path,'wb')as f:
        pickle.dump(cfg,f,protocol=pickle.HIGHEST_PROTOCOL)


    os.makedirs(cfg.OUTPUT_DIR,exist_ok=True)

    trainer=DefaultTrainer(cfg)
    trainer.resume_or_load(resume=False)

    trainer.train()


   
if __name__ =='__main__':
    main()

 

여기까지가 학습 끝! 하지만 아직 가중치 파일에 대해서 만들어 주지 않아 실행 되지 않을 것 입니다. 

 

 

다음은

 

 

가중치 파일 (utils) 

 

생각보다 많은 라이브러리를 받아오고..

from detectron2.data import DatasetCatalog,MetadataCatalog
from detectron2.utils.visualizer import Visualizer
from detectron2.config import get_cfg
from detectron2 import model_zoo

from detectron2.utils.visualizer import ColorMode

import random
import cv2
import matplotlib.pyplot as plt


import cv2
import torch
import os
import json

from detectron2.data import MetadataCatalog

from detectron2.utils.visualizer import ColorMode, Visualizer

 

 

가중치를 설정하는 부분!

내부 요소들은 학습의 정도에 따라 설정해주시면 될 것 같습니다.

def get_train_cfg(config_file_path,checkpoint_url,train_dataset_name,test_dataset_name,num_classes,device,output_dir):
    cfg=get_cfg()

    cfg.merge_from_file(model_zoo.get_config_file(config_file_path))
    cfg.MODEL.WEIGHTS=model_zoo.get_checkpoint_url(checkpoint_url)
    cfg.DATASETS.TRAIN=(train_dataset_name,)
    cfg.DATASETS.TEST=(test_dataset_name,)

    cfg.DATALOADER.NUM_WORKERS=4 # 작업자의 수  >> 적당한 값으로 설정
    cfg.SOLVER.IMS_PER_BATCH=4   # 한 번 학습에 사용할 이미지 수 
    cfg.SOLVER.BASE_LR=0.00025  # 가중치 업데이트 속도
    cfg.SOLVER.MAX_ITER=1000    # 반복 횟수
    cfg.SOLVER.STEPS=[]

    cfg.MODEL.ROI_HEADS.NUM_CLASSES=num_classes # 클래스수
    cfg.MODEL.DEVICE=device # 어떤 디바이스로 실행할지 cpu gpu
    cfg.OUTPUT_DIR=output_dir # 출력값 저장위치

    return cfg

 

 

그리고 추가적으로 테스트 결과를 출력해줄 on_image 함수도 만들어 주었습니다. 

이미지와 json을 같이 생성해줍니다.

def on_image(image_path,save_path,predictor,metadata):
    im=cv2.imread(image_path)
    outputs=predictor(im)
    
    V=Visualizer(im[:,:,::-1],metadata=metadata,scale=1,instance_mode=ColorMode.SEGMENTATION)
    predictions=outputs["instances"].to("cpu")

    v=V.draw_instance_predictions(predictions)

    data = {}
    data["shapes"]=[]

    points=predictions.pred_boxes.tensor.tolist()
    label=predictions.pred_classes.tolist()

    for i in range(len(points)):
        shape = {} 
        shape["label"]=label[i]
        shape["points"]=[[points[i][0],points[i][1]],[points[i][2],points[i][3]]]

        data["shapes"].append(shape)

    file_name = os.path.basename(image_path)
    imageHeight,imageWidth ,__=im.shape
    
    
    file_path = file_name
    data["imagePath"]=file_name

    with open(save_path+str(file_name)+".json", "w") as json_file:
        json.dump(data, json_file)


    # plt.figure(figsize=(14,10))
    # plt.imshow(v.get_image())
    # plt.show()
    
    img_rgb = v.get_image()
    img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
    cv2.imwrite(save_path+str(file_name)+".jpg",img_bgr)

 

 

 

여기까지 코드를  생성하면 이제 train이 가능합니다. 

학습 진행시 아래처럼 로그가 찍히고

 

학습이 끝나면 아까 설정해준 output저장위치에

이런식의 모델이 만들어집니다.

 

 

이제 이걸 가지고 테스트해봅시다.

 

테스트

 

from detectron2.engine import DefaultPredictor

import os
import pickle
from glob import glob
import json

from utils import*

cfg_save_path="./DetectronCard/TestCard.pickle"  ## 테스트 가중치 


metadata={0: "Card"}

metadata = {
    "thing_classes": [
        "Card"
    ]
}

학습으로 생성된 가중치를 지정해주고 

클래스 명들을 설정해줍시다. 

 

 

 

다음은 가중치파일을 불러와 모델을 실행하여 테스트 결과를 저장해줍시다.

with open(cfg_save_path,'rb') as f:
    cfg=pickle.load(f)

cfg.MODEL.WEIGHTS=os.path.join(cfg.OUTPUT_DIR,"model_final.pth")  ## 테스트에 사용할 모델
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST=0.5

predictor=DefaultPredictor(cfg)


images = glob('./images/Cardtest/*.jpg')  ## 테스트할 이미지폴더

for i in images:## 모든 이미지 테스트후 저장
    image_path=i   ## 테스트 이미지
    save_path="./Result/"
    on_image(image_path,save_path,predictor,metadata)

 

 

 

1000번 정도 학습시켰는데 배경이 단조로워서 생각보다 잘되는것 같습니다!

 

 

 

Detectron으로 학습하기 끝!

 

 

 

<참고 유튜브>

 

 

 

틀린점이 있다면 댓 달아주세요!