Google Colaboratory에서 OpenPose를 이용하여 영상 속 인물의 자세 추정하기

제가 다니는 대학교에서 인공지능 관련 교과목을 들으며 진행한 프로젝트의 과정입니다.

 

우선 내 드라이브에 파일을 하나 만듭니다. 폴더명은 자기가 하고싶은걸로 

저는 만들기 귀찮아서 그냥 Colab Notebooks 폴더를 사용했습니다.

 

파일 만들기

OpenPose에 있는 CNN네트워크는 BODY-25(출력관절이 25개), COCO(출력관절이 18개), MPII(출력관절이15개)가 있습니다. 오늘은 COCO를 이용하겠습니다.

COCO를 이용하기 위해서 관련 파일을 Colab Notebooks안에 옮겨놓아야 합니다. 아래 사이트에서 파일을 다운 받습니다.

https://github.com/CMU-Perceptual-Computing-Lab/openpose

GitHub - CMU-Perceptual-Computing-Lab/openpose: OpenPose: Real-time multi-person keypoint detection library for body, face, hand

파일을 다운 받고 압축 해제 후 models에서getModels.bat을 실행하면 각 모델을 다운로드 해줍니다. (getModels.sh는 Linux용 입니다)

 

모델 다운로드가 끝나면 우리는 COCO를 사용할 것이기 때문에 models/pose/mpii 폴더 안에서 pose_deploy_linevec.prototxt 파일과 pose_iter_440000.caffemodel 파일들을 자신이 만든 폴더 안에 옮겨 놓습니다. (저는 Colab Notebooks에 옮겨놓았습니다)

 

COCO외에도 MPII와 BODY-25를 사용하실 분들은 필요한 파일들이 각각 다르니 유의해서 파일을 옮겨주세요.

(MPII:  pose_iter_160000.caffemodel, pose_deploy_linevec_faster_4_stages.prototxt

BODY-25 :  pose_iter_584000.caffemodel, pose_deploy.prototxt)

 

옮기는 작입이 다 끝났으면 Colab Notebooks 폴더 안에 새로운 파일을 하나 만듭니다. 파일명도 자유

새로 만든 파일을 열고 아래의 코드를 입력해줍니다.

from google.colab import drive
drive.mount('/content/gdrive')

 

 

위의 코드는 코랩과 자신의 드라이브를 연결시키겠다는 내용입니다. 위에 코드가 없으면 드라이브내 옮겨놓았던 OpenPose 관련 파일들과 영상들을 사용할 수없으니 반드시 입력해주세요.

 

 

다음 아래의 코드를 차례 차례 입력해줍니다. 코드를 한번에 돌리면 나중에 오류 잡기가 어려우니 나눠서 돌리는걸 추천합니다. 

 

import cv2
import numpy as np
from tqdm import tqdm
import pandas as pd
import csv

protoFile_COCO = "/home/hir5/intelligence/pose_deploy_linevec.prototxt"
weightsFile_COCO = "/home/hir5/intelligence/pose_iter_440000.caffemodel"

BODY_PARTS_COCO =  {0: "Nose", 1: "Neck", 2: "RShoulder", 3: "RElbow", 4: "RWrist",
                   5: "LShoulder", 6: "LElbow", 7: "LWrist", 8: "RHip", 9: "RKnee",
                   10: "RAnkle", 11: "LHip", 12: "LKnee", 13: "LAnkle", 14: "REye",
                   15: "LEye", 16: "REar", 17: "LEar", 18: "Background"}

POSE_PAIRS_COCO = [[0, 1], [0, 14], [0, 15], [1, 2], [1, 5], [1, 8], [1, 11], [2, 3], [3, 4],
                   [5, 6], [6, 7], [8, 9], [9, 10], [12, 13], [11, 12], [14, 16], [15, 17]]

 

 

만약 MPII나 BODY-25 사용하실 분들은 변수 값을 인터넷에서 검색해 바꿔주셔야 합니다. 

 

def output_keypoints(threshold, BODY_PARTS):
    global points

    # 네트워크 지정
    net = net_openpose

    # 입력 이미지의 사이즈 정의
    image_height = 368
    image_width = 368

    # 네트워크에 넣기 위한 전처리
    input_blob = cv2.dnn.blobFromImage(frame, 1.0 / 255, (image_width, image_height), (0, 0, 0),
                                       swapRB=False, crop=False)

    # 전처리된 blob 네트워크에 입력
    net.setInput(input_blob)

    # 결과 받아오기
    out = net.forward()
    # The output is a 4D matrix :
    # The first dimension being the image ID ( in case you pass more than one image to the network ).
    # The second dimension indicates the index of a keypoint.
    # The model produces Confidence Maps and Part Affinity maps which are all concatenated.
    # For COCO model it consists of 57 parts – 18 keypoint confidence Maps + 1 background + 19*2 Part Affinity Maps. Similarly, for MPI, it produces 44 points.
    # We will be using only the first few points which correspond to Keypoints.
    # The third dimension is the height of the output map.
    out_height = out.shape[2]
    # The fourth dimension is the width of the output map.
    out_width = out.shape[3]

    # 원본 이미지의 높이, 너비를 받아오기
    frame_height, frame_width = frame.shape[:2]

   # 포인트 리스트 초기화
    points = []

    for i in range(len(BODY_PARTS)):
        # 신체 부위의 confidence map
        prob_map = out[0, i, :, :]

        # 최소값, 최대값, 최소값 위치, 최대값 위치
        min_val, prob, min_loc, point = cv2.minMaxLoc(prob_map)

        # 원본 이미지에 맞게 포인트 위치 조정
        x = (frame_width * point[0]) / out_width
        x = int(x)
        y = (frame_height * point[1]) / out_height
        y = int(y)

        if prob > threshold:  # [pointed]
            cv2.circle(frame_openpose, (x, y), 5, (0, 255, 255), thickness=-1, lineType=cv2.FILLED)
            cv2.circle(frame_zeros, (x, y), 5, (0, 255, 255), thickness=-1, lineType=cv2.FILLED)
            cv2.putText(frame_openpose, str(i), (x, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 1,
                        lineType=cv2.LINE_AA)
            points.append((x, y))
            # csv file
            with open(r"/home/hir5/intelligence/footmove1.csv",'a') as f:
                wr=csv.writer(f)
                wr.writerow([{BODY_PARTS[i]},prob,x,y])

        else:  # [not pointed]
            cv2.circle(frame_openpose, (x, y), 5, (0, 255, 255), thickness=-1, lineType=cv2.FILLED)
            cv2.putText(frame_openpose, str(i), (x, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 1,
                        lineType=cv2.LINE_AA)
            points.append(None)
            # csv file
            with open(r"/home/hir5/intelligence/footmove1.csv",'a') as f:
                wr=csv.writer(f)
                wr.writerow([{BODY_PARTS[i]},0,0,0])

    return frame_openpose

 

 

함수 정의 해 주고

 

 

def output_keypoints_with_lines(POSE_PAIRS):
    # POSE_PAIRS 갯수만큼 RGB 배열을 생성 (frame_zeros에 쓰일 색)
    #colors = []
    #for i in range(3):
    #    value_r = value_g = value_b = 50
    #    for j in range(int(len(POSE_PAIRS_BODY_25B) / 3)):
    #        if i == 0:
    #            value_b += 31
    #        if i == 1:
    #            value_g += 31
    #        if i == 2:
    #            value_r += 31
    #        colors.append((value_b, value_g, value_r))

    for index, pair in enumerate(POSE_PAIRS):
        part_a = pair[0]  # 0 (Head)
        part_b = pair[1]  # 1 (Neck)
        if points[part_a] and points[part_b]:
            cv2.line(frame_openpose, points[part_a], points[part_b], (0, 255, 0), 3)
            #cv2.line(frame_zeros, points[part_a], points[part_b], colors[index], 2)

    return frame_openpose

또 함수 정의해주고

def in_box():
    dic = {}
    for key, value in nms_boxes.items():
        if key.split(' ')[0] != "person" and key.split(' ')[0] not in classes_custom:  # person, classes_custom 객체 제외
            points_list = []
            for index, point in enumerate(points):
                if point is not None:  # point 가 없을 경우 제외
                    x = value[0]
                    y = value[1]
                    width = value[0] + value[2]
                    height = value[1] + value[3]
                    if (point[0] > x) and (point[0] < width) and (point[1] > y) and (point[1] < height):
                        points_list.append(index)
            if len(points_list) > 0:
                dic[key] = points_list
                
    return dic

 

또 함수 정의합니다.

 

여기까지 코드가 문제 없이 돌아간다면 다음 단계로 넘어가기 전에 작업을 하나 더 해줘야 합니다.

오늘의 목표는 영상을 입력 Data로 받아 영상 속 인물의 자세를 추정하는 것 입니다. 그러기 위해서는 영상을 프레임 별로 나눠서 input data로 만들어 주고 학습을 돌려야 하는데 그러기 위해서는 우리가 만든 폴더안에 학습에 필요한 영상을 넣어줘야 합니다. (저는 Colab Notebooks폴더안에 영상을 넣었습니다.)

 

현재까지 자신이 만든 폴더 안에는 Openpose 에 필요한 파일 두개와 사용할 영상 총 3개의 파일이 들어있어야 합니다 

 

영상을 넣어준 후 자신이 만든 폴더 안에 또 다른 폴더를 하나 만들어 줘야 합니다. 이 폴더 안에는 프레임별로 학습한 자세 추정이 완료된 새로운 프레임들이 들어갈 폴더 입니다.

 

이후 아래 코드를 입력하면 학습을 시작합니다 .

# 키포인트를 저장할 빈 리스트
points = []

# NMS 를 거친 객체 box 들의 정보를 저장할 딕셔너리 생성 (in_box() 함수에서 사용됨)
nms_boxes = {}

# 네트워크 불러오기
net_openpose = cv2.dnn.readNetFromCaffe(protoFile_body_25b, weightsFile_body_25b)

net_openpose.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
net_openpose.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)

#동영상 경로입니다. 파일명 틀리지 않게 조심하세요
video = "/content/gdrive/My Drive/Colab Notebooks/foot_move.mp4"

capture = cv2.VideoCapture(video)
total_frame = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))

for now_frame in tqdm(range(total_frame)):
    ret, frame = capture.read()

    # 프레임 복사
    frame_openpose = frame.copy()
    frame_yolo = frame.copy()

    # 이미지의 크기와 같은 검은색 프레임 생성
    frame_zeros = np.zeros((frame.shape[0], frame.shape[1], 3), np.uint8)

    # openpose (joint 포인팅)
    output_keypoints(threshold=0.1, BODY_PARTS=BODY_PARTS_BODY_25B)

    # 각 joint 를 선으로 연결한 프레임
    output_keypoints_with_lines(POSE_PAIRS=POSE_PAIRS_BODY_25B)

    # 프레임 세로로 이어붙이기
    #frame_horizontal = cv2.vconcat([frame_openpose, frame_yolo])

    # 이미지 파일 생성
    # 저는 내 드라이브에 바로 폴더를 만들었지만 자신이 만든 폴더 안에 폴더를 만들어도 됩니다 그러면 MyDrive 뒤에 자신이 만든 폴더 명이 들어가야겠죠?
    # 폴더 안에 만든 새로운 폴더 명을 입력해주세요 ( 저는 footmove라는 폴더를 만들어 이 폴더에 새로운 프레임을 저장하도록 했습니다)
    cv2.imwrite(f"/content/gdrive/My Drive/footmove/{now_frame:.0f}.jpg",frame_openpose)

capture.release()

저는 약 40초짜리 영상이라 엄청 오래걸리지 않았으나 영상 길이가 길면 오래 걸리고 또한 컴퓨터 사양도 중요합니다. 컴퓨터 사양에 비해 데이터가 너무 오버스펙이면 중간에 뻑 나니 틈틈히 학습이 잘 돌아가는지 봐줘야 합니다.

 

학습이 끝난 후 새로운 프레임을 저장한 폴더에 들어가면 학습된 프레임들이 저장되어 있을 겁니다.

아래의 코드는 새로 저장한 프레임을 다시 영상으로 만드는 코드 입니다. 참고하세요

import glob
import cv2
import numpy as np
from tqdm import tqdm

img_array = []

for file in tqdm(glob.glob('/content/gdrive/My Drive/footmove/*.jpg')):
    img = cv2.imread(file)
    img_array.append(img)

height, width = img.shape[:2]
size = (width, height)
#새로 만들 video 파일명과 저장 경로 입니다. 저는 내 드라이브에 result_footmove1이라는 영상을 저장하도록 했습니다.
video = cv2.VideoWriter('/content/gdrive/My Drive/result_footmove1.mp4',cv2.VideoWriter_fourcc(*'DIVX'), 20, size)

for img in img_array:
    video.write(img)

video.release()

 

 

동영상으로 잘 저장된것을 볼 수 있습니다.

 

 

제가 오늘 진행한 과정의 출처는 아래 블로그 입니다. 

https://hanryang1125.tistory.com/19

Google Colaboratory에서 OpenCV, YOLOv4, OpenPose를 이용하여 영상 속 객체 탐지 및 인물의 자세 추정 (Colab)