【Python-Module】 SSD 수행하기 - OpenCV DNN 모듈

이전 Post를 통해서 SSD의 이론에 대해서 공부했다. OpenCV DNN 모델을 사용해서 SSD를 수행해보자.

OpenCV DNN 모듈을 사용해서 SSD 수행하기
/DLCV/Detection/SSD/OpenCV_SSD_이미지와_영상_Detection.ipynb 참조
OpenCV 사용과정 요약은 이전 Post 참조

1. OpenCV DNN 패키지를 이용하여 SSD 수행하기

  • Tensorflow 에서 Pretrained 된 모델 파일을 OpenCV에서 로드하여 이미지와 영상에 대한 Object Detection 수행.
### 입력 이미지로 사용될 이미지 보기
import cv2
import matplotlib.pyplot as plt
%matplotlib inline

img = cv2.imread('../../data/image/avengers.jpg')
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

print('image shape:', img.shape)
plt.figure(figsize=(12, 12))
plt.imshow(img_rgb)
image shape: (330, 499, 3)

drawing

1-1 Tensorflow에서 Pretrained 된 Inference모델(Frozen graph)와 환경파일을 다운로드 받은 후 이를 이용해 OpenCV에서 Inference 모델 생성

  • https://github.com/opencv/opencv/wiki/TensorFlow-Object-Detection-API 에 다운로드 URL 있음.
  • Inception-SSD v2 2017_11_17 이것 사용할 예정
  • pretrained 모델은 http://download.tensorflow.org/models/object_detection/ssd_inception_v2_coco_2017_11_17.tar.gz 에서 다운로드 후 압축 해제
  • pretrained 모델을 위한 환경 파일은 https://github.com/opencv/opencv_extra/blob/master/testdata/dnn/ssd_inception_v2_coco_2017_11_17.pbtxt 에서 다운로드
  • download된 모델 파일과 config 파일을 인자로 하여 inference 모델을 DNN에서 로딩함.
# %cd /home/sb020518/DLCV/Detection/ssd
# !mkdir pretrained 
# %cd ./pretrained
# !wget http://download.tensorflow.org/models/object_detection/ssd_inception_v2_coco_2017_11_17.tar.gz
# !wget https://raw.githubusercontent.com/opencv/opencv_extra/master/testdata/dnn/ssd_inception_v2_coco_2017_11_17.pbtxt
# !tar -xvf ssd*.tar.gz
# !mv ssd_inception_v2_coco_2017_11_17.pbtxt graph.pbtxt
# cd ../
!pwd
!ls pretrained/ssd_inception_v2_coco_2017_11_17
/home/sb020518/DLCV/Detection/ssd
checkpoint		   model.ckpt.data-00000-of-00001  model.ckpt.meta
frozen_inference_graph.pb  model.ckpt.index		   saved_model
cv_net = cv2.dnn.readNetFromTensorflow('./pretrained/ssd_inception_v2_coco_2017_11_17/frozen_inference_graph.pb', 
                                     './pretrained/graph.pbtxt')
  • coco 데이터 세트의 클래스id별 클래스명 지정.
  • 왜 90번 까지 하는지는 여기 참조
  • 90번까지 해야할지 91번까지 해야할지는 실험적으로 확인하는 수 밖에 없다.
labels_to_names = {1:'person',2:'bicycle',3:'car',4:'motorcycle',5:'airplane',6:'bus',7:'train',8:'truck',9:'boat',10:'traffic light',
                    11:'fire hydrant',12:'street sign',13:'stop sign',14:'parking meter',15:'bench',16:'bird',17:'cat',18:'dog',19:'horse',20:'sheep',
                    21:'cow',22:'elephant',23:'bear',24:'zebra',25:'giraffe',26:'hat',27:'backpack',28:'umbrella',29:'shoe',30:'eye glasses',
                    31:'handbag',32:'tie',33:'suitcase',34:'frisbee',35:'skis',36:'snowboard',37:'sports ball',38:'kite',39:'baseball bat',40:'baseball glove',
                    41:'skateboard',42:'surfboard',43:'tennis racket',44:'bottle',45:'plate',46:'wine glass',47:'cup',48:'fork',49:'knife',50:'spoon',
                    51:'bowl',52:'banana',53:'apple',54:'sandwich',55:'orange',56:'broccoli',57:'carrot',58:'hot dog',59:'pizza',60:'donut',
                    61:'cake',62:'chair',63:'couch',64:'potted plant',65:'bed',66:'mirror',67:'dining table',68:'window',69:'desk',70:'toilet',
                    71:'door',72:'tv',73:'laptop',74:'mouse',75:'remote',76:'keyboard',77:'cell phone',78:'microwave',79:'oven',80:'toaster',
                    81:'sink',82:'refrigerator',83:'blender',84:'book',85:'clock',86:'vase',87:'scissors',88:'teddy bear',89:'hair drier',90:'toothbrush',
                    91:'hair brush'}

1-2 이미지를 preprocessing 수행하여 Network에 입력하고 Object Detection 수행 후 결과를 이미지에 시각화

# 원본 이미지 (633, 806)를 네트웍에 입력시에는 (300, 300)로 resize 함. 
# 이후 결과가 출력되면 resize된 이미지 기반으로 bounding box 위치가 예측 되므로 이를 다시 원복하기 위해 원본 이미지 shape정보 필요
rows = img.shape[0]
cols = img.shape[1]
# cv2의 rectangle()은 인자로 들어온 이미지 배열에 직접 사각형을 업데이트 하므로 그림 표현을 위한 별도의 이미지 배열 생성. 
draw_img = img.copy()

# 원본 이미지 배열을 사이즈 (300, 300)으로, BGR을 RGB로 변환하여 배열 입력
cv_net.setInput(cv2.dnn.blobFromImage(img,  size=(300, 300), swapRB=True, crop=False))
# Object Detection 수행하여 결과를 cv_out으로 반환 
cv_out = cv_net.forward()
print(cv_out.shape)

# bounding box의 테두리와 caption 글자색 지정
green_color=(0, 255, 0)
red_color=(0, 0, 255)

# detected 된 object들을 iteration 하면서 정보 추출
for detection in cv_out[0,0,:,:]:
    score = float(detection[2])
    class_id = int(detection[1])
    # detected된 object들의 score가 0.3 이상만 추출
    if score > 0.3:
        # detected된 object들은 image 크기가 (300, 300)으로 scale된 기준으로 예측되었으므로 다시 원본 이미지 비율로 계산
        left = detection[3] * cols
        top = detection[4] * rows
        right = detection[5] * cols
        bottom = detection[6] * rows
        # labels_to_names 딕셔너리로 class_id값을 클래스명으로 변경. opencv에서는 class_id + 1로 매핑해야함.
        caption = "{}: {:.4f}".format(labels_to_names[class_id], score)
        
        #cv2.rectangle()은 인자로 들어온 draw_img에 사각형을 그림. 위치 인자는 반드시 정수형.
        cv2.rectangle(draw_img, (int(left), int(top)), (int(right), int(bottom)), color=green_color, thickness=2)
        cv2.putText(draw_img, caption, (int(left), int(top - 5)), cv2.FONT_HERSHEY_SIMPLEX, 0.7, red_color, 2)
        print(caption, class_id)

img_rgb = cv2.cvtColor(draw_img, cv2.COLOR_BGR2RGB)

plt.figure(figsize=(12, 12))
plt.imshow(img_rgb)
(1, 1, 100, 7)

drawing

1-3 위에서 한 작업 그대로 함수로 def

import time

def get_detected_img(cv_net, img_array, score_threshold, use_copied_array=True, is_print=True):
    
    rows = img_array.shape[0]
    cols = img_array.shape[1]
    
    draw_img = None
    if use_copied_array:
        draw_img = img_array.copy()
        #draw_img = cv2.cvtColor(draw_img, cv2.COLOR_BGR2RGB)
    else:
        draw_img = img_array
    
    cv_net.setInput(cv2.dnn.blobFromImage(img_array, size=(300, 300), swapRB=True, crop=False))
    
    start = time.time()
    cv_out = cv_net.forward()
    
    green_color=(0, 255, 0)
    red_color=(0, 0, 255)

    # detected 된 object들을 iteration 하면서 정보 추출
    for detection in cv_out[0,0,:,:]:
        score = float(detection[2])
        class_id = int(detection[1])
        # detected된 object들의 score가 0.4 이상만 추출
        if score > score_threshold:
            # detected된 object들은 image 크기가 (300, 300)으로 scale된 기준으로 예측되었으므로 다시 원본 이미지 비율로 계산
            left = detection[3] * cols
            top = detection[4] * rows
            right = detection[5] * cols
            bottom = detection[6] * rows
            # labels_to_names 딕셔너리로 class_id값을 클래스명으로 변경. opencv에서는 class_id + 1로 매핑해야함.
            caption = "{}: {:.4f}".format(labels_to_names[class_id], score)

            #cv2.rectangle()은 인자로 들어온 draw_img에 사각형을 그림. 위치 인자는 반드시 정수형.
            cv2.rectangle(draw_img, (int(left), int(top)), (int(right), int(bottom)), color=green_color, thickness=2)
            cv2.putText(draw_img, caption, (int(left), int(top - 5)), cv2.FONT_HERSHEY_SIMPLEX, 0.7, red_color, 2)
    if is_print:
        print('Detection 수행시간:',round(time.time() - start, 2),"초")

    return draw_img

drawing

# 위에서 def 한 함수 사용해서 Detection 수행하기.
# 0.23초 지린다. 
img = cv2.imread('../../data/image/EPL01.jpg')

#coco dataset 클래스명 매핑

# tensorflow inference 모델 로딩
cv_net = cv2.dnn.readNetFromTensorflow('./pretrained/ssd_inception_v2_coco_2017_11_17/frozen_inference_graph.pb', 
                                     './pretrained/graph.pbtxt')
# Object Detetion 수행 후 시각화 
draw_img = get_detected_img(cv_net, img, score_threshold=0.4, use_copied_array=True, is_print=True)

img_rgb = cv2.cvtColor(draw_img, cv2.COLOR_BGR2RGB)

plt.figure(figsize=(12, 12))
plt.imshow(img_rgb)
Detection 수행시간: 0.27 초

2. Video Object Detection 수행

원본 영상 보기

from IPython.display import clear_output, Image, display, Video, HTML
Video('../../data/video/John_Wick_small.mp4')
  • VideoCapture와 VideoWriter 설정하기
  • VideoCapture를 이용하여 Video를 frame별로 capture 할 수 있도록 설정
  • VideoCapture의 속성을 이용하여 Video Frame의 크기 및 FPS 설정.
  • VideoWriter를 위한 인코딩 코덱 설정 및 영상 write를 위한 설정
  • 총 Frame 별로 iteration 하면서 Object Detection 수행. 개별 frame별로 단일 이미지 Object Detection과 유사
  • video detection 전용 함수 생성.
def do_detected_video(cv_net, input_path, output_path, score_threshold, is_print):
    
    cap = cv2.VideoCapture(input_path)

    codec = cv2.VideoWriter_fourcc(*'XVID')

    vid_size = (round(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),round(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
    vid_fps = cap.get(cv2.CAP_PROP_FPS)

    vid_writer = cv2.VideoWriter(output_path, codec, vid_fps, vid_size) 

    frame_cnt = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    print('총 Frame 갯수:', frame_cnt, )

    green_color=(0, 255, 0)
    red_color=(0, 0, 255)
    while True:
        hasFrame, img_frame = cap.read()
        if not hasFrame:
            print('더 이상 처리할 frame이 없습니다.')
            break
        
        returned_frame = get_detected_img(cv_net, img_frame, score_threshold=score_threshold, use_copied_array=True, is_print=True)
        vid_writer.write(returned_frame)
    # end of while loop

    vid_writer.release()
    cap.release()
do_detected_video(cv_net, '../../data/video/John_Wick_small.mp4', '../../data/output/John_Wick_small_ssd01.avi', 0.4, True)
!gsutil cp ../../data/output/John_Wick_small_ssd01.avi gs://my_bucket_dlcv/data/output/John_Wick_small_ssd01.avi

3. SSD+Mobilenet으로 Object Detection 수행.

  • https://github.com/opencv/opencv/wiki/TensorFlow-Object-Detection-API 에 다운로드 URL 있음.
  • weight파일은 http://download.tensorflow.org/models/object_detection/ssd_mobilenet_v2_coco_2018_03_29.tar.gz 에서 다운로드

  • 수행결과 : 위에서는 SSD+inceptionV2를 사용했다. 확실히 MobileNet의 속도가 훨씬 빠른것을 확인할 수 있다.
# !mkdir pretrained; cd pretrained
# !wget  http://download.tensorflow.org/models/object_detection/ssd_mobilenet_v2_coco_2018_03_29.tar.gz
# !wget https://raw.githubusercontent.com/opencv/opencv_extra/master/testdata/dnn/ssd_mobilenet_v2_coco_2018_03_29.pbtxt
# cd ssd_mobilenet_v2_coco_2018_03_29; mv ssd_mobilenet_v2_coco_2018_03_29.pbtxt graph.pbtxt
cv_net_mobile = cv2.dnn.readNetFromTensorflow('./pretrained/ssd_mobilenet_v2_coco_2018_03_29/frozen_inference_graph.pb', 
                                     './pretrained/ssd_mobilenet_v2_coco_2018_03_29/graph.pbtxt')
#### 영상 Detection
do_detected_video(cv_net_mobile, '../../data/video/John_Wick_small.mp4', '../../data/output/John_Wick_small_ssd_mobile01.avi', 0.2, True)
!gsutil cp ../../data/output/John_Wick_small_ssd_mobile01.avi gs://my_bucket_dlcv/data/output/John_Wick_small_ssd_mobile01.avi
#### Image Detection
img = cv2.imread('../../data/image/avengers.jpg')

#coco dataset 클래스명 매핑

cv_net_mobile = cv2.dnn.readNetFromTensorflow('./pretrained/ssd_mobilenet_v2_coco_2018_03_29/frozen_inference_graph.pb', 
                                     './pretrained/ssd_mobilenet_v2_coco_2018_03_29/graph.pbtxt')
# Object Detetion 수행 후 시각화 
draw_img = get_detected_img(cv_net_mobile, img, score_threshold=0.4, use_copied_array=True, is_print=True)

img_rgb = cv2.cvtColor(draw_img, cv2.COLOR_BGR2RGB)

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

© All rights reserved By Junha Song.