본문 바로가기

컴퓨터비전

Faster R-CNN 실습해보기

지금의 COCO dataset으로 pretrained된 inference 모델은 tensorflow만으로 사용가능하다.

하지만 향후 사용해보게 될 tensorflow object detection API는 tensorflow에 들어있지 않고 따로 install 해야한다. 이는 custom data를 학습 시키는 데에 사용된다.

이를 통해 만들어진 모델은 그냥 tensorflow에서 object detection을 수행할 수 있다.

 

#inference graph를 읽음
with tf.gfile.FastGFile(os.path.join(default_rcnn_dir, 'pretrained/faster_rcnn_resnet50_coco_2018_01_28/frozen_inference_graph.pb'), 'rb') as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
    
with tf.Session() as sess:
    # Session 시작하고 inference graph 모델 로딩
    sess.graph.as_default()
    tf.import_graph_def(graph_def, name='')
    
    # 입력 이미지 생성 및 BGR을 RGB로 변경 
    img = cv2.imread(os.path.join(default_dir, 'data/image/beatles01.jpg'))
    draw_img = img.copy()

    rows = img.shape[0]
    cols = img.shape[1]
    
    #수동으로 BGR2RGB
    inp = img[:, :, [2, 1, 0]]
    start = time.time()
    # Object Detection 수행. 
    # 네 개의 텐서를 가져오게 됨 박스에는 0~1 사이의 좌표값
    # input image가 4차원인 이유는 한 이미지가 아닌 여.러. 이미지에 대해서도 수행할 수 있어서
    # out에는 get한 파라미터 값이 들어간다. 
    out = sess.run([sess.graph.get_tensor_by_name('num_detections:0'),
                    sess.graph.get_tensor_by_name('detection_scores:0'),
                    sess.graph.get_tensor_by_name('detection_boxes:0'),
                    sess.graph.get_tensor_by_name('detection_classes:0')],
                   feed_dict={'image_tensor:0': inp.reshape(1, inp.shape[0], inp.shape[1], 3)})
    print('type of out:', type(out), 'length of out:',len(out))
    print(out)
    green_color=(0, 255, 0)
    red_color=(0, 0, 255)
    
    # Bounding Box 시각화 
    num_detections = int(out[0][0])
    for i in range(num_detections):
        classId = int(out[3][0][i])
        score = float(out[1][0][i])
        bbox = [float(v) for v in out[2][0][i]]
        if score > 0.5:
        #순서가 y1, x1, y2, x2임 
        #scale 값으로 bbox에 들어가져 있어서 이미지 사이즈를 곱해야 함
            left = bbox[1] * cols
            top = bbox[0] * rows
            right = bbox[3] * cols
            bottom = bbox[2] * rows
            cv2.rectangle(draw_img, (int(left), int(top)), (int(right), int(bottom)), green_color, thickness=2)
            caption = "{}: {:.4f}".format(labels_to_names[classId], score)
            print(caption)
            cv2.putText(draw_img, caption, (int(left), int(top - 5)), cv2.FONT_HERSHEY_SIMPLEX, 0.4, red_color, 1)
    
    print('Detection 수행시간:',round(time.time() - start, 2),"초")
    
img_rgb = cv2.cvtColor(draw_img, cv2.COLOR_BGR2RGB)

plt.figure(figsize=(12, 12))
plt.imshow(img_rgb)

 

  • OpenCV로 돌릴 때에는 SwapRB가 있었으나, tensorflow로 돌릴 시 수동으로 BGR을 RGB로 바꿔줘야 한다.
  • out의 type은 list로 각 요소에는 numpy array가 존재한다.
  • 위의 각 요소는 sess.run의 파라미터로 들어있는 것을 get 한다.
  • 3번째에 들어있는 box에는 scale된 0~1 사이의 값이 들어있는데 [y1, x1, y2, x2]의 순이다.
  • 그러므로 scale된 값 x 이미지 사이즈를 곱해야 한다(left, top, right, bottom)

'컴퓨터비전' 카테고리의 다른 글

Fast R-CNN과 Faster R-CNN  (0) 2020.10.03
R-CNN  (0) 2020.09.21