본문 바로가기

Deep Learning/TF Object Detection API

1.4. run_inference_for_single_image 함수 설계

320x100
320x100




pretrained model 을 메모리에 올리는 과정까지 알아봤다.


이제 메모리에 올려진 계산 그래프에서 예측을 위해 필요한 텐서 두 개를 가져와 실행하는 데에 필요한 함수들을 알아보자.


관련 글로 [Deep Learning/Tensorflow] - MNIST 분류 모델, 조금 다르게 실행해보자 가 있으니 읽고 오길 바란다.


모델 실행을 위해서는


  1. 계산 그래프에서 텐서를 가져와 object detecting 결과를 뽑아주는 함수

  2. 결과값을 통해 이미지에 bounding box를 그려주는 함수 가 필요하다.


이번 글에서는 '1. 계산 그래프에서 텐서를 가져와 object detection 결과를 뽑아주는 함수' 에 대해 알아보도록 하겠다.


이 과정에서 이전에 다운받지 않은 opencv를 사용할 것이기 때문에 가상환경에 opencv 설치 먼저 하자.



추가로 TF_Object_Detecion_API 폴더에 test_images 폴더를 새로 만든 후 이 폴더 안에 object detection 대상 이미지를 넣도록 하자.



이미지의 이름은 'test_image_0.jpg' 의 형식을 따르게끔 지정하자.








object_detection_tutorial.ipynb 에서 run_inference_for_single_image( ) 함수를 읽어보면 


우리가 다운받은 pretrained model 이 내부적으로 어떻게 구현돼있는 지는 알 순 없어도,


object detection 최종 결과 출력을 담당하는 오퍼레이션들의 이름이


'num_detections', 'detection_boxes', 'detection_scores', 'detection_classes', 'detection_masks' 이란 것을 알 수 있다.


따라서 우리는 default graph 에서 위의 이름을 갖는 오퍼레이션에서 출력되는 텐서를 뽑은 뒤, 


session을 통해 실행한 뒤 결과값을 반환해주는 함수를 정의해 사용해주면 된다.



바로 코드를 살펴보자.


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
import cv2
import numpy as np
 
def run_inference_for_single_image(image, graph):
    with tf.Session(graph = graph) as sess:
 
        input_tensor = graph.get_tensor_by_name('image_tensor:0')
        
        target_operation_names = ['num_detections''detection_boxes',
                                  'detection_scores''detection_classes''detection_masks']
        tensor_dict = {}
        for key in target_operation_names:
            op = None
            try:
                op = graph.get_operation_by_name(key)
                
            except:
                continue
 
            tensor = graph.get_tensor_by_name(op.outputs[0].name)
            tensor_dict[key] = tensor
 
        if 'detection_masks' in tensor_dict:
            detection_boxes = tf.squeeze(tensor_dict['detection_boxes'], [0])
 
        output_dict = sess.run(tensor_dict, feed_dict = {input_tensor : [image]})
            
        output_dict['num_detections'= int(output_dict['num_detections'][0])
        output_dict['detection_classes'= output_dict['detection_classes'][0].astype(np.uint8)
        output_dict['detection_boxes'= output_dict['detection_boxes'][0]
        output_dict['detection_scores'= output_dict['detection_scores'][0]
 
        return output_dictColored by Color Scripter
cs
  • Line 4
    : run_inference_for_single_image( )는 분석 대상 이미지와 계산 그래프가 그려져 있는 graph 객체를 input으로 받는다

  • Line 5
    : tf.Session( graph = graph )로 설정을 안 해주면 다음과 같은 에러가 뜨니 설정해줘야 한다.

    이곳을 보고 에러 해결해줄 수 있었다.

  • Line 7
    : 모델 설계 시 데이터가 최초로 입력되는 placeholder에 해당된다. sess.run( )에서 feed_dict를 통해 사용해줘야하기 때문에 불러온다.

  • Line 15
    : 이미 알고 있는 오퍼레이션의 이름을 통해 계산 그래프에서 해당 오퍼레이션을 불러온다.

  • Line 17
    : 'detection_masks'의 경우 mask_rcnn 모델에서만 추가적으로 쓰는 오퍼레이션 같다. 
      우리가 사용하는 모델에는 없는 오퍼레이션이라서 불러오려고 하면 없다고 에러가 뜨기 때문에
      이 문제를 해결하고자 try / except 를 추가했다.

  • Line 20 ~ 21
    : 사용성을 위해 사용하고자 하는 텐서들을 딕셔너리 형태로 저장했다.
      
      ※ op = 오퍼레이션 / op.outputs[0] = output tensor / op.outputs[0].name = output tensor name

  • Line 23 ~ 24
    : 위에서 말했듯 'detection_masks'는 mask_rcnn에서나 존재하는 오퍼레이션인 듯 하다. 
      그래서 사실 이 부분은 없어도 무방한 코드이긴 하나,
      나중에 mask_rcnn 도 사용하게 될 것 같아 남겨두었다.

  • Line 26
    : 불러온 텐서들을 최종적으로 session을 통해 실행하게 된다. 
      sess.run( )은 fetch를 딕셔너리 형태로 입력받으면 반환도 딕셔너리 형태로 반환해준다.

  • Line 23 ~ 31
    : 결과에 리스트가 한 번 더 씌워서 반환하는 탓에 뒤에 [0]을 붙여 정리해줘야한다.


함수가 오류 없이 잘 설계되었는지 확인을 위해 다음의 코드를 추가해서 실행시켜줘보자.


1
2
3
4
5
6
7
8
9
10
PATH_TO_TEST_IMAGE = 'test_images'
n_images = len(os.listdir(PATH_TO_TEST_IMAGE))
TEST_IMAGE_PATHS = [os.path.join(PATH_TO_TEST_IMAGE, 'test_image_{}.jpg' .format(i)) for i in range(n_images)]
print('분석 대상 이미지 경로 지정 완료...')
 
for image_path in TEST_IMAGE_PATHS:    
    img = cv2.imread(image_path)
    height, width, _ = img.shape
    
    output_dict = run_inference_for_single_image(img, detection_graph)Colored by Color Scripter
cs


test image와 이전 글의 detection_graph를 매개변수로 넘겨준 후 실행하면,


다음과 같이 문제 없이 실행이 잘 되는 것을 확인 할 수 있다.








다음은 코드 전체이다.


[ run_inference_tf.Graph.py ]






지금까지의 글들을 통해 


  1. pretrained model 다운받고 압축풀어서

  2. frozen_inference_graph.pb을 메모리에 올린 뒤

  3. 메모리에서 모델 사용에 필요한 텐서들을 불러와서

  4. session을 통해 실행까지 마쳐 결과를 반환받는 것까지 알아봤다.

여기까지 오는데 오래도 걸렸다. 거의 다 왔다.

이제 object detection 결과를 test image에 잘 그려주기만 하면 된다.

다음 글에서는 결과값을 통해 이미지에 bounding box를 그려주는 함수에 대해 알아보도록 하겠다.





< 참고 사이트 >

https://stackoverflow.com/questions/42593771/session-graph-is-empty