Detectron2 - Example 따라하기 1

https://colab.research.google.com/drive/16jcaJoc6bCFAQ96jDe2HwtXj7BMD_-m5#scrollTo=4Qg7zSVOulkb

Detectron2 Tutorial.ipynb

 

구글 코랩에 올라와있는 샘플 코드를 기반으로 로컬에서 실행해본다.

 

# check pytorch installation:
import torch, torchvision
print(torch.__version__, torch.cuda.is_available())
assert torch.__version__.startswith("1.9")   # please manually install torch 1.9 if Colab changes its default version

# Some basic setup:
# Setup detectron2 logger
import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()

# import some common libraries
import numpy as np
import os, json, cv2, random

# 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, DatasetCatalog

im = cv2.imread("./input.jpg")


cfg = get_cfg()

# without cuda
cfg.MODEL.DEVICE = "cpu"

# add project-specific config (e.g., TensorMask) here if you're not running a model in detectron2's core library
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5  # set threshold for this model
# Find a model from detectron2's model zoo. You can use the https://dl.fbaipublicfiles... url as well
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
predictor = DefaultPredictor(cfg)
outputs = predictor(im)

print(outputs["instances"].pred_classes)
print(outputs["instances"].pred_boxes)

v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
cv2.imshow("image", out.get_image()[:, :, ::-1])

cv2.waitKey(0)
cv2.destroyAllWindows()

 

같은 폴더에 우선 테스트 이미지를 다운로드 받고 위의 코드를 실행 해야 한다.

wget http://images.cocodataset.org/val2017/000000439715.jpg -q -O input.jpg

 

torch 기본이 cuda를 사용하기 때문에 지원하는 GPU 가 없을 경우 반드시 cfg 옵션에 다음과 같이 옵션을 추가해주어야 한다.

cfg.MODEL.DEVICE = "cpu"

 

코드를 실행하면, 샘플과 동일한 결과가 나온것을 확인할수 있다.