TFOD( Tensorflow Object Detection )

TFOD는 하나의 Object가 아닌 여러개의 Object를 Detection할 수 있다.

 

 

TFOD 환경 만들기

https://durian9s-coding-tree.tistory.com/151

TensorFlow Object Detection 환경구축

 

 

Protobuff란 

jeong-pro.tistory.com/190

Protocol Buffer 원리로 배우는 고성능 직렬화, 역직렬화 전략! Protocol Buffer 예제 테스트(구글이 쓰는

Coco API란

cocodataset.org/#explore

COCO - Common Objects in Context

 

 

코딩 순서

1. 모델을 불러온다.
2. 하나의 이미지를 모델에 입력해 물체를 검출하고 그에 대한 데이터들을 뽑아내는 함수를 만든다.
3. 2의 함수의 출력을 이용해 이미지와 검출된 카테고리의 이름과 위치를 출력한다.
4. 불러온 데이터와 이미지 데이터가 있는 경로를 3의 함수에 입력한다.

import numpy as np
import os
import six.moves.urllib as urllib
import sys
import tarfile
import tensorflow as tf
import zipfile

from collections import defaultdict
from io import StringIO
from matplotlib import pyplot as plt
from PIL import Image

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 vis_util

import pathlib
import cv2

utils_ops.tf = tf.compat.v1
tf.gfile = tf.io.gfile

def load_model(model_name):
    base_url = 'http://download.tensorflow.org/models/object_detection/'
    model_file = model_name + '.tar.gz'
    model_dir = tf.keras.utils.get_file(
                                        fname=model_name, 
                                        origin=base_url + model_file,
                                        untar=True)

    model_dir = pathlib.Path(model_dir)/"saved_model"
    model = tf.saved_model.load(str(model_dir))

    return model

# 여기 이름을 바꿔서 다른 모델을 사용합니다. base_url도 바꿔야하는 경우가 있습니다.
model_name = 'ssd_mobilenet_v1_coco_2017_11_17'

# 모델 가져옵니다.
detection_model = load_model(model_name)

# 모델이 저장된 위치입니다.
PATH_TO_LABELS = 'C:\\Users\\bear0\\Documents\\Tensorllow\\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)


# 모델과 이미지를 입력받아 딕셔너리 형태로 인식한 카테고리와 해당 물체에 박스치는 정보를 출력합니다.
def run_inference_for_single_image(model, image):
    # 학습한 데이터 형식이 tensor이기 떄문에 형식을 바꿔줍니다.
    image = np.asarray(image)
    input_tensor = tf.convert_to_tensor(image)
    input_tensor = input_tensor[tf.newaxis,...]

    # 모델을 가져와 데이터를 입력합니다.
    model_fn = model.signatures['serving_default']
    output_dict = model_fn(input_tensor)

    # 결과로 나온 값들을 딕셔너리 형태로 바꿔줍니다.
    num_detections = int(output_dict.pop('num_detections'))
    output_dict = {key:value[0, :num_detections].numpy() 
                    for key,value in output_dict.items()}
    output_dict['num_detections'] = num_detections

    # 클래스들의 타입을 숫자형식으로 바꿔줍니다.
    output_dict['detection_classes'] = output_dict['detection_classes'].astype(np.int64)
   
    # 데이터 형식들을 알맞게 변환시켜줍니다.
    if 'detection_masks' in output_dict:
        output_dict['detection_masks'] = tf.convert_to_tensor(output_dict['detection_masks'], dtype=tf.float32)
        output_dict['detection_boxes'] = tf.convert_to_tensor(output_dict['detection_boxes'], dtype=tf.float32)
        detection_masks_reframed = utils_ops.reframe_box_masks_to_image_masks(
                output_dict['detection_masks'], output_dict['detection_boxes'],
                image.shape[0], image.shape[1])  
        detection_masks_reframed = tf.cast(detection_masks_reframed > 0.5,
                                        tf.uint8)
        output_dict['detection_masks_reframed'] = detection_masks_reframed.numpy()
        
    return output_dict

# 현재까지의 함수가 잘 되었는지 하나의 경로를 입력해 확인해봅니다.
# imgae_np = cv2.imread('data/images/image2.jpg')
# output_dict = run_inference_for_single_image(detection_model, imgae_np)
# print(output_dict)
# print(output_dict.keys())


# 모델과 이미지 데이터를 입력받아 물체를 인식하고 인식한 물체에 네모 박스와 이름을 적어 출력합니다.
def show_inference(model, image_path):
    image_np = np.array(Image.open(image_path))

    output_dict = run_inference_for_single_image(model, image_np)

    vis_util.visualize_boxes_and_labels_on_image_array(
                    image_np,
                    np.array(output_dict['detection_boxes']),
                    output_dict['detection_classes'],
                    output_dict['detection_scores'],
                    category_index,
                    instance_masks=output_dict.get('detection_masks_reframed',None),
                    use_normalized_coordinates=True,
                    line_thickness=8)
    
    # OpenCV는 BGR이기 때문에 변환합니다. 만약 모델에 입력하기 전에 변환한다면 예측값이 이상하게 나올 수 있습니다.
    image_np = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
                
    cv2.imshow(str(image_path), image_np)


# 이미지가 3개가 들어있는 폴더의 경로 밑의 .jpg 파일들을 모두 모델에 입력할 것입니다.
PATH_TO_TEST_IMAGES_DIR = pathlib.Path('data\\images')
TEST_IMAGE_PATHS = sorted(list(PATH_TO_TEST_IMAGES_DIR.glob("*.jpg")))

# 하나씩 모델에 입력해 Object Detection 후 결과를 출력합니다.
for image_path in TEST_IMAGE_PATHS:
    show_inference(detection_model, image_path)

cv2.waitKey()
cv2.destroyAllWindows()