Google Colab에서 TensorFlow Image Object Detection API와 OpenCV 혼합해 이미지 Detection하기

Google Colab(구글 코랩)의 서버 컴퓨터는 리눅스를 사용하기 때문에 TensorFlow Image Object Detection API는 기존 visual studio code에서 하는것과 조금 다릅니다.

 

먼저 텐서플로우 버전을 낮추거나 하지않습니다.

테스트 결과 visual studio code에서 텐서플로우 2.5.0버전에서는 원활하게 실행이 되었지만,

구글 코랩에서는 2.5.0버전에서 import object_detection 에러가 나는것을 확인 했기 때문에  원래 텐서플로우 2.7.0버전으로 테스트 했습니다.

#버전 확인
tf. __version__

구글 코랩에서는 COCO API를 지원하기 때문에 따로 설치를 안해주셔도 됩니다.

 

 

tensorflow github(깃허브)에 있는 repository(레파지토리)를 파이썬 코드로 clone(클론) 하겠습니다.

import os
import pathlib

if "models" in pathlib.Path.cwd().parts :
  while "models" in pathlib.Path.cwd().part :
    os.chdir('..')
elif not pathlib.Path('models').exists() :
  ! git clone --depth 1 https://github.com/tensorflow/models  #클론할 주소 입력

 

추가 : repository(레파지토리)에 pull(풀)하고 싶을때

os.chdir('git pull할 경로')
! git pull

 

Object Detection API 설치하기

%%bash
cd models/research/
protoc object_detection/protos/*.proto --python_out=.
cp object_detection/packages/tf2/setup.py .
python -m pip install .

리눅스에서는 명령어 앞에 !를 붙여야 합니다

하지만 명령에 실행전에 %%bash를 적고 나면 !없이도 실행할 수 있습니다.

 

이제 TensorFlow Image Object Detection API를 하기 위한 라이브러리들을 import 하겠습니다.

import tensorflow as tf
import os
import pathlib

import numpy as np
import zipfile

import matplotlib.pyplot as plt
from PIL import Image

from IPython.display import display #주피터 노트북에서 이미지 쉽게 표시하는 함수 import

from object_detection.utils import ops as utils_ops
from object_detection.utils import label_map_util
from object_detection.utils import visualization_utils as viz_utils

 

레이블 파일을 인덱스와 연결 시켜줍니다.

PATH_TO_LABELS = '/content/models/research/object_detection/data/mscoco_label_map.pbtxt'
category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS,use_display_name=True)

print(category_index)

 

모델 다운로드를 복사하겠습니다.

https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf2_detection_zoo.md

GitHub - tensorflow/models: Models and examples built with TensorFlow

테스트할 모델을 링크 복사 합니다.

#모델 로드하는 함수

#테스트할 모델들
#/20200711/ssd_mobilenet_v2_fpnlite_640x640_coco17_tpu-8.tar.gz

def download_model(model_name, model_date):
    base_url = 'http://download.tensorflow.org/models/object_detection/tf2/' #경로는 변하지 않음
    model_file = model_name + '.tar.gz'
    model_dir = tf.keras.utils.get_file(fname=model_name,
                                        origin=base_url + model_date + '/' + model_file,
                                        untar=True)
    return str(model_dir)

base_url 주소는 변경하지 않습니다.

 

모델을 받습니다.

MODEL_DATE = '20200711'
MODEL_NAME = 'ssd_mobilenet_v2_fpnlite_640x640_coco17_tpu-8'
PATH_TO_MODEL_DIR = download_model(MODEL_NAME, MODEL_DATE)

 

모델을 호출하는 함수와 모델 호출입니다.

#모델 호출 함수
def load_model(model_dir):
    model_full_dir = model_dir + "/saved_model"

    # Load saved model and build the detection function
    detection_model = tf.saved_model.load(model_full_dir)
    return detection_model
    
#모델 호출
detection_model = load_model(PATH_TO_MODEL_DIR)

 

이제 텐서플로우 모델에 있는 샘플 이미지를 가져오겠습니다.(다른 이미지로 하셔도 됩니다.)

PATH_TO_IMAGE_DIR = pathlib.Path('/content/models/research/object_detection/test_images') #사진경로
IMAGE_PATHS = sorted(list(PATH_TO_IMAGE_DIR.glob('*.jpg'))) #확장자명 확인 *은 모든 이미지를 가져온다는 뜻

 

이제 Object Detection 하겠습니다.

def show_inference(detection_model,image_np) :

    # The input needs to be a tensor, convert it using `tf.convert_to_tensor`.
    input_tensor = tf.convert_to_tensor(image_np)
    # The model expects a batch of images, so add an axis with `tf.newaxis`.
    input_tensor = input_tensor[tf.newaxis, ...]

    # input_tensor = np.expand_dims(image_np, 0)
    detections = detection_model(input_tensor)

    # All outputs are batches tensors.
    # Convert to numpy arrays, and take index [0] to remove the batch dimension.
    # We're only interested in the first num_detections.
    num_detections = int(detections.pop('num_detections'))
    detections = {key: value[0, :num_detections].numpy()
                   for key, value in detections.items()}
    detections['num_detections'] = num_detections

    # detection_classes should be ints.
    detections['detection_classes'] = detections['detection_classes'].astype(np.int64)
    #print(detections)
    image_np_with_detections = image_np.copy()

    viz_utils.visualize_boxes_and_labels_on_image_array(
          image_np_with_detections,
          detections['detection_boxes'],
          detections['detection_classes'],
          detections['detection_scores'],
          category_index,
          use_normalized_coordinates=True,
          max_boxes_to_draw=200, #최대 박스 개수
          min_score_thresh=.40, #확률 표시 숫자+프로 일치하는것만 표기 현재는 40프로
          agnostic_mode=False)

    display(Image.fromarray(image_np_with_detections))
    
  for image_path in IMAGE_PATHS :
  image_np = np.array(Image.open(image_path))

  show_inference(detection_model,image_np)

이렇게 image object detection이 잘 되는걸 볼 수 있습니다.