detectron2 colab예제

지난번 포스트에서 detectron2를 windows환경에서 설치해 봤다.
이번 포스트에서는 colab예제를 windows환경에서 실행할 수 있도록 코드를 수정해봤다.

---

설치

!python -m pip install pyyaml==5.1

지난번 포스트에서 detectron2설치는 끝냈으니 필요한 부분만 설치

모듈 import

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

# import some common libraries
import numpy as np
import os, json, cv2, 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, DatasetCatalog

특이한 부분이 있다.

from google.colab.patches import cv2_imshow

colab에서는 cv2.imshow가 충돌을 일으키기 때문에 사용할 수 없다.
그래서 cv2_imshow라는 모듈을 만들어서 사용하는 것 같다.
colab환경이 아니라면 cv2.imshow로 수정해서 사용하자

---

이미지 다운

!wget http://images.cocodataset.org/val2017/000000439715.jpg -q -O input.jpg
im = cv2.imread("./input.jpg")
cv2_imshow(im)

wget은 web get 웹 상에서 필요한 데이터를 다운 받는 명령어이다.
하지만 주피터 노트북에서 사용할 수 없다.
pip install wget으로 설치 후 사용하도록 하자
또한 스크립트 내에 imshow를 하기 위해서 cv2가 아닌 matplotlib 함수로 수정했다.

import wget
from matplotlib import pyplot as plt

# wget.download('http://images.cocodataset.org/val2017/000000439715.jpg') 계속실행하면 (1),(2)로 계속 생김

image = cv2.imread('000000439715.jpg')
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
plt.imshow(image_rgb)
plt.title('Image Window')
plt.show()

예측 결과

결과는 그냥 출력하면 잘나온다.

cfg = get_cfg()
# 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"))

# img_result_1 = out.get_image()[:, :, ::-1]
img_result_1 = out.get_image()[:, :, :]
# image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
plt.imshow(img_result_1)
plt.title('Image Window')
plt.show()

cfg = get_cfg()   # get a fresh new config
cfg.merge_from_file(model_zoo.get_config_file("COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7  # set threshold for this model
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml")

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

img_result_2 = out.get_image()[:, :, :]
# image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
plt.imshow(img_result_2)
plt.title('Image Window')
plt.show()

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-PanopticSegmentation/panoptic_fpn_R_101_3x.yaml"))
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-PanopticSegmentation/panoptic_fpn_R_101_3x.yaml")
predictor = DefaultPredictor(cfg)
panoptic_seg, segments_info = predictor(im)["panoptic_seg"]
v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
out = v.draw_panoptic_seg_predictions(panoptic_seg.to("cpu"), segments_info)

img_result_3 = out.get_image()[:, :, :]
# image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
plt.imshow(img_result_3)
plt.title('Image Window')
plt.show()