• 3D视觉——3.人体姿态估计(Pose Estimation) 算法对比 即 效果展示——MediaPipe与OpenPose

    上一话

    3D视觉——2.人体姿态估计(Pose Estimation)入门——OpenPose含安装、编译、使用（单帧、实时视频）

    1.MediaPipe

    代码

    1. import cv2
    2. import time
    3. import numpy as np
    4. from tqdm import tqdm
    5. import mediapipe as mp
    6.  
    7. mp_pose = mp.solutions.pose
    8. pose = mp_pose.Pose(static_image_mode=True,
    9.                     model_complexity=2,
    10.                     smooth_landmarks=True,
    11.                     min_detection_confidence=0.5,
    12.                     min_tracking_confidence=0.5)
    13. drawing = mp.solutions.drawing_utils
    14.  
    15.  
    16. def process_frame(img):
    17.     height, width, channels = img.shape
    18.     start = time.time()
    19.     results = pose.process(img)
    20.     if results.pose_landmarks:
    21.         drawing.draw_landmarks(img, results.pose_landmarks, mp_pose.POSE_CONNECTIONS)
    22.         coords = np.array(results.pose_landmarks.landmark)
    23.         for index, each in enumerate(coords):
    24.             cx = int(each.x * width)
    25.             cy = int(each.y * height)
    26.             cz = each.z
    27.             radius = 5
    28.             # nose
    29.             if index == 0:
    30.                 img = cv2.circle(img, (cx, cy), radius, (0, 0, 255), -1)
    31.             # shoulder
    32.             elif index in [11, 12]:
    33.                 img = cv2.circle(img, (cx, cy), radius, (193, 182, 255), -1)
    34.             # hip joint
    35.             elif index in [23, 24]:
    36.                 img = cv2.circle(img, (cx, cy), radius, (16, 144, 247), -1)
    37.             # elbow
    38.             elif index in [13, 14]:
    39.                 img = cv2.circle(img, (cx, cy), radius, (1, 240, 255), -1)
    40.             # lap
    41.             elif index in [25, 26]:
    42.                 img = cv2.circle(img, (cx, cy), radius, (140, 47, 240), -1)
    43.             # wrist and ankle
    44.             elif index in [15, 16, 27, 28]:
    45.                 img = cv2.circle(img, (cx, cy), radius, (223, 155, 60), -1)
    46.             # left hand
    47.             elif index in [17, 19, 21]:
    48.                 img = cv2.circle(img, (cx, cy), radius, (16, 144, 247), -1)
    49.             # right hand
    50.             elif index in [18, 20, 22]:
    51.                 img = cv2.circle(img, (cx, cy), radius, (1, 240, 255), -1)
    52.             # left feet
    53.             elif index in [27, 29, 31]:
    54.                 img = cv2.circle(img, (cx, cy), radius, (140, 47, 240), -1)
    55.             # right feet
    56.             elif index in [28, 30, 32]:
    57.                 img = cv2.circle(img, (cx, cy), radius, (223, 155, 6), -1)
    58.             # mouth
    59.             elif index in [9, 10]:
    60.                 img = cv2.circle(img, (cx, cy), radius, (16, 144, 247), -1)
    61.             # face and eye
    62.             elif index in [1, 2, 3, 4, 5, 6, 7, 8]:
    63.                 img = cv2.circle(img, (cx, cy), radius, (1, 240, 255), -1)
    64.             # other
    65.             else:
    66.                 img = cv2.circle(img, (cx, cy), radius, (140, 47, 240), -1)
    67.     else:
    68.         fail = "fail detection"
    69.         img = cv2.putText(img, fail, (25, 100), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 3)
    70.  
    71.     FPS = 1 / (time.time() - start)
    72.     img = cv2.putText(img, "FPS" + str(int(FPS)), (25, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 3)
    73.     return img
    74.  
    75.  
    76. def out_video(input):
    77.     file = input.split("/")[-1]
    78.     output = "out-optim-" + file
    79.     print("It will start processing video: {}".format(input))
    80.     cap = cv2.VideoCapture(input)
    81.     frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    82.     frame_size = (cap.get(cv2.CAP_PROP_FRAME_WIDTH), cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    83.     # create VideoWriter,VideoWriter_fourcc is video decode
    84.     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    85.     fps = cap.get(cv2.CAP_PROP_FPS)
    86.     out = cv2.VideoWriter(output, fourcc, fps, (int(frame_size[0]), int(frame_size[1])))
    87.     # the progress bar
    88.     with tqdm(range(frame_count)) as pbar:
    89.  
    90.         while cap.isOpened():
    91.             success, frame = cap.read()
    92.             if not success:
    93.                 break
    94.             try:
    95.                 frame = process_frame(frame)
    96.                 out.write(frame)
    97.                 pbar.update(1)
    98.             except:
    99.                 print("ERROR")
    100.                 pass
    101.     pbar.close()
    102.     cv2.destroyAllWindows()
    103.     out.release()
    104.     cap.release()
    105.     print("{} finished!".format(output))
    106.  
    107.  
    108. if __name__ == '__main__':
    109.     video_dirs = "1.mp4"
    110.     out_video(video_dirs)

    运行结果

    MediaPipe人体姿态估计效果展示

    2.OpenPose

    使用ffmpeg写入(Pytorch深度学习框架只使用OpenPose的models)代码

    demo_video.py

    1. import copy
    2. import math
    3. import time
    4.  
    5. import numpy as np
    6. import cv2
    7. from glob import glob
    8. import os
    9. import argparse
    10. import json
    11.  
    12. # video file processing setup
    13. # from: https://stackoverflow.com/a/61927951
    14. import argparse
    15. import subprocess
    16. import sys
    17. from pathlib import Path
    18. from typing import NamedTuple
    19.  
    20.  
    21. class FFProbeResult(NamedTuple):
    22.     return_code: int
    23.     json: str
    24.     error: str
    25.  
    26.  
    27. def ffprobe(file_path) -> FFProbeResult:
    28.     command_array = ["ffprobe",
    29.                      "-v", "quiet",
    30.                      "-print_format", "json",
    31.                      "-show_format",
    32.                      "-show_streams",
    33.                      file_path]
    34.     result = subprocess.run(command_array, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
    35.     return FFProbeResult(return_code=result.returncode,
    36.                          json=result.stdout,
    37.                          error=result.stderr)
    38.  
    39.  
    40. # openpose setup
    41. from src import model
    42. from src import util
    43. from src.body import Body
    44. from src.hand import Hand
    45.  
    46. body_estimation = Body('model/body_pose_model.pth')
    47. hand_estimation = Hand('model/hand_pose_model.pth')
    48.  
    49.  
    50. def process_frame(frame, body=True, hands=True):
    51.     start = time.time()
    52.     canvas = copy.deepcopy(frame)
    53.     if body:
    54.         candidate, subset = body_estimation(frame)
    55.         canvas = util.draw_bodypose(canvas, candidate, subset)
    56.     if hands:
    57.         hands_list = util.handDetect(candidate, subset, frame)
    58.         all_hand_peaks = []
    59.         for x, y, w, is_left in hands_list:
    60.             peaks = hand_estimation(frame[y:y + w, x:x + w, :])
    61.             peaks[:, 0] = np.where(peaks[:, 0] == 0, peaks[:, 0], peaks[:, 0] + x)
    62.             peaks[:, 1] = np.where(peaks[:, 1] == 0, peaks[:, 1], peaks[:, 1] + y)
    63.             all_hand_peaks.append(peaks)
    64.         canvas = util.draw_handpose(canvas, all_hand_peaks)
    65.     FPS = math.ceil(1 / (time.time() - start))
    66.     canvas = cv2.putText(canvas, "FPS" + str(int(FPS)), (25, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 3)
    67.     return canvas
    68.  
    69.  
    70. # writing video with ffmpeg because cv2 writer failed
    71. # https://stackoverflow.com/questions/61036822/opencv-videowriter-produces-cant-find-starting-number-error
    72. import ffmpeg
    73.  
    74. # open specified video
    75. parser = argparse.ArgumentParser(
    76.     description="Process a video annotating poses detected.")
    77. parser.add_argument('--file', type=str, default='video/1.avi', help='Video file location to process.')
    78. parser.add_argument('--no_hands', action='store_true', help='No hand pose')
    79. parser.add_argument('--no_body', action='store_true', help='No body pose')
    80. args = parser.parse_args()
    81. video_file = args.file
    82. cap = cv2.VideoCapture(video_file)
    83.  
    84. # get video file info
    85. ffprobe_result = ffprobe(args.file)
    86. info = json.loads(ffprobe_result.json)
    87. videoinfo = [i for i in info["streams"] if i["codec_type"] == "video"][0]
    88. input_fps = videoinfo["avg_frame_rate"]
    89. # input_fps = float(input_fps[0])/float(input_fps[1])
    90. input_pix_fmt = videoinfo["pix_fmt"]
    91. input_vcodec = videoinfo["codec_name"]
    92.  
    93. # define a writer object to write to a movidified file
    94. postfix = info["format"]["format_name"].split(",")[0]
    95. output_file = ".".join(video_file.split(".")[:-1]) + ".processed." + postfix
    96.  
    97.  
    98. class Writer():
    99.     def __init__(self, output_file, input_fps, input_framesize, input_pix_fmt,
    100.                  input_vcodec):
    101.         if os.path.exists(output_file):
    102.             os.remove(output_file)
    103.         self.ff_proc = (
    104.             ffmpeg
    105.                 .input('pipe:',
    106.                        format='rawvideo',
    107.                        pix_fmt="bgr24",
    108.                        s='%sx%s' % (input_framesize[1], input_framesize[0]),
    109.                        r=input_fps)
    110.                 .output(output_file, pix_fmt=input_pix_fmt, vcodec=input_vcodec)
    111.                 .overwrite_output()
    112.                 .run_async(pipe_stdin=True)
    113.         )
    114.  
    115.     def __call__(self, frame):
    116.         self.ff_proc.stdin.write(frame.tobytes())
    117.  
    118.     def close(self):
    119.         self.ff_proc.stdin.close()
    120.         self.ff_proc.wait()
    121.  
    122.  
    123. writer = None
    124. iteration = 0
    125. while (cap.isOpened()):
    126.     ret, frame = cap.read()
    127.     if frame is None:
    128.         break
    129.  
    130.     posed_frame = process_frame(frame, body=not args.no_body,
    131.                                 hands=not args.no_hands)
    132.  
    133.     if writer is None:
    134.         input_framesize = posed_frame.shape[:2]
    135.         writer = Writer(output_file, input_fps, input_framesize, input_pix_fmt,
    136.                         input_vcodec)
    137.  
    138.     # cv2.imshow('frame', posed_frame)
    139.  
    140.     # write the frame
    141.     writer(posed_frame)
    142.     iteration += 1
    143.     print("iteration: {}".format(iteration))
    144.  
    145.     if cv2.waitKey(1) & 0xFF == ord('q'):
    146.         break
    147.  
    148. cap.release()
    149. writer.close()
    150. cv2.destroyAllWindows()

    src/util.py

    1. import numpy as np
    2. import math
    3. import cv2
    4. import matplotlib
    5. from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
    6. from matplotlib.figure import Figure
    7. import numpy as np
    8. import matplotlib.pyplot as plt
    9. import cv2
    10.  
    11.  
    12. def padRightDownCorner(img, stride, padValue):
    13.     h = img.shape[0]
    14.     w = img.shape[1]
    15.  
    16.     pad = 4 * [None]
    17.     pad[0] = 0 # up
    18.     pad[1] = 0 # left
    19.     pad[2] = 0 if (h % stride == 0) else stride - (h % stride) # down
    20.     pad[3] = 0 if (w % stride == 0) else stride - (w % stride) # right
    21.  
    22.     img_padded = img
    23.     pad_up = np.tile(img_padded[0:1, :, :]*0 + padValue, (pad[0], 1, 1))
    24.     img_padded = np.concatenate((pad_up, img_padded), axis=0)
    25.     pad_left = np.tile(img_padded[:, 0:1, :]*0 + padValue, (1, pad[1], 1))
    26.     img_padded = np.concatenate((pad_left, img_padded), axis=1)
    27.     pad_down = np.tile(img_padded[-2:-1, :, :]*0 + padValue, (pad[2], 1, 1))
    28.     img_padded = np.concatenate((img_padded, pad_down), axis=0)
    29.     pad_right = np.tile(img_padded[:, -2:-1, :]*0 + padValue, (1, pad[3], 1))
    30.     img_padded = np.concatenate((img_padded, pad_right), axis=1)
    31.  
    32.     return img_padded, pad
    33.  
    34. # transfer caffe model to pytorch which will match the layer name
    35. def transfer(model, model_weights):
    36.     transfered_model_weights = {}
    37.     for weights_name in model.state_dict().keys():
    38.         transfered_model_weights[weights_name] = model_weights['.'.join(weights_name.split('.')[1:])]
    39.     return transfered_model_weights
    40.  
    41. # draw the body keypoint and lims
    42. def draw_bodypose(canvas, candidate, subset):
    43.     stickwidth = 4
    44.     limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], \
    45.                [10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17], \
    46.                [1, 16], [16, 18], [3, 17], [6, 18]]
    47.  
    48.     colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], \
    49.               [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], \
    50.               [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]]
    51.     for i in range(18):
    52.         for n in range(len(subset)):
    53.             index = int(subset[n][i])
    54.             if index == -1:
    55.                 continue
    56.             x, y = candidate[index][0:2]
    57.             cv2.circle(canvas, (int(x), int(y)), 4, colors[i], thickness=-1)
    58.     for i in range(17):
    59.         for n in range(len(subset)):
    60.             index = subset[n][np.array(limbSeq[i]) - 1]
    61.             if -1 in index:
    62.                 continue
    63.             cur_canvas = canvas.copy()
    64.             Y = candidate[index.astype(int), 0]
    65.             X = candidate[index.astype(int), 1]
    66.             mX = np.mean(X)
    67.             mY = np.mean(Y)
    68.             length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
    69.             angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
    70.             polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
    71.             cv2.fillConvexPoly(cur_canvas, polygon, colors[i])
    72.             canvas = cv2.addWeighted(canvas, 0.4, cur_canvas, 0.6, 0)
    73.     # plt.imsave("preview.jpg", canvas[:, :, [2, 1, 0]])
    74.     # plt.imshow(canvas[:, :, [2, 1, 0]])
    75.     return canvas
    76.  
    77. def draw_handpose(canvas, all_hand_peaks, show_number=False):
    78.     edges = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10], \
    79.              [10, 11], [11, 12], [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20]]
    80.     fig = Figure(figsize=plt.figaspect(canvas))
    81.  
    82.     fig.subplots_adjust(0, 0, 1, 1)
    83.     fig.subplots_adjust(bottom=0, top=1, left=0, right=1)
    84.     bg = FigureCanvas(fig)
    85.     ax = fig.subplots()
    86.     ax.axis('off')
    87.     ax.imshow(canvas)
    88.  
    89.     width, height = ax.figure.get_size_inches() * ax.figure.get_dpi()
    90.  
    91.     for peaks in all_hand_peaks:
    92.         for ie, e in enumerate(edges):
    93.             if np.sum(np.all(peaks[e], axis=1)==0)==0:
    94.                 x1, y1 = peaks[e[0]]
    95.                 x2, y2 = peaks[e[1]]
    96.                 ax.plot([x1, x2], [y1, y2], color=matplotlib.colors.hsv_to_rgb([ie/float(len(edges)), 1.0, 1.0]))
    97.  
    98.         for i, keyponit in enumerate(peaks):
    99.             x, y = keyponit
    100.             ax.plot(x, y, 'r.')
    101.             if show_number:
    102.                 ax.text(x, y, str(i))
    103.     bg.draw()
    104.     canvas = np.fromstring(bg.tostring_rgb(), dtype='uint8').reshape(int(height), int(width), 3)
    105.     return canvas
    106.  
    107. # image drawed by opencv is not good.
    108. def draw_handpose_by_opencv(canvas, peaks, show_number=False):
    109.     edges = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10], \
    110.              [10, 11], [11, 12], [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20]]
    111.     # cv2.rectangle(canvas, (x, y), (x+w, y+w), (0, 255, 0), 2, lineType=cv2.LINE_AA)
    112.     # cv2.putText(canvas, 'left' if is_left else 'right', (x, y), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
    113.     for ie, e in enumerate(edges):
    114.         if np.sum(np.all(peaks[e], axis=1)==0)==0:
    115.             x1, y1 = peaks[e[0]]
    116.             x2, y2 = peaks[e[1]]
    117.             cv2.line(canvas, (x1, y1), (x2, y2), matplotlib.colors.hsv_to_rgb([ie/float(len(edges)), 1.0, 1.0])*255, thickness=2)
    118.  
    119.     for i, keyponit in enumerate(peaks):
    120.         x, y = keyponit
    121.         cv2.circle(canvas, (x, y), 4, (0, 0, 255), thickness=-1)
    122.         if show_number:
    123.             cv2.putText(canvas, str(i), (x, y), cv2.FONT_HERSHEY_SIMPLEX, 0.3, (0, 0, 0), lineType=cv2.LINE_AA)
    124.     return canvas
    125.  
    126. # detect hand according to body pose keypoints
    127. # please refer to https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/src/openpose/hand/handDetector.cpp
    128. def handDetect(candidate, subset, oriImg):
    129.     # right hand: wrist 4, elbow 3, shoulder 2
    130.     # left hand: wrist 7, elbow 6, shoulder 5
    131.     ratioWristElbow = 0.33
    132.     detect_result = []
    133.     image_height, image_width = oriImg.shape[0:2]
    134.     for person in subset.astype(int):
    135.         # if any of three not detected
    136.         has_left = np.sum(person[[5, 6, 7]] == -1) == 0
    137.         has_right = np.sum(person[[2, 3, 4]] == -1) == 0
    138.         if not (has_left or has_right):
    139.             continue
    140.         hands = []
    141.         #left hand
    142.         if has_left:
    143.             left_shoulder_index, left_elbow_index, left_wrist_index = person[[5, 6, 7]]
    144.             x1, y1 = candidate[left_shoulder_index][:2]
    145.             x2, y2 = candidate[left_elbow_index][:2]
    146.             x3, y3 = candidate[left_wrist_index][:2]
    147.             hands.append([x1, y1, x2, y2, x3, y3, True])
    148.         # right hand
    149.         if has_right:
    150.             right_shoulder_index, right_elbow_index, right_wrist_index = person[[2, 3, 4]]
    151.             x1, y1 = candidate[right_shoulder_index][:2]
    152.             x2, y2 = candidate[right_elbow_index][:2]
    153.             x3, y3 = candidate[right_wrist_index][:2]
    154.             hands.append([x1, y1, x2, y2, x3, y3, False])
    155.  
    156.         for x1, y1, x2, y2, x3, y3, is_left in hands:
    157.             # pos_hand = pos_wrist + ratio * (pos_wrist - pos_elbox) = (1 + ratio) * pos_wrist - ratio * pos_elbox
    158.             # handRectangle.x = posePtr[wrist*3] + ratioWristElbow * (posePtr[wrist*3] - posePtr[elbow*3]);
    159.             # handRectangle.y = posePtr[wrist*3+1] + ratioWristElbow * (posePtr[wrist*3+1] - posePtr[elbow*3+1]);
    160.             # const auto distanceWristElbow = getDistance(poseKeypoints, person, wrist, elbow);
    161.             # const auto distanceElbowShoulder = getDistance(poseKeypoints, person, elbow, shoulder);
    162.             # handRectangle.width = 1.5f * fastMax(distanceWristElbow, 0.9f * distanceElbowShoulder);
    163.             x = x3 + ratioWristElbow * (x3 - x2)
    164.             y = y3 + ratioWristElbow * (y3 - y2)
    165.             distanceWristElbow = math.sqrt((x3 - x2) ** 2 + (y3 - y2) ** 2)
    166.             distanceElbowShoulder = math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
    167.             width = 1.5 * max(distanceWristElbow, 0.9 * distanceElbowShoulder)
    168.             # x-y refers to the center --> offset to topLeft point
    169.             # handRectangle.x -= handRectangle.width / 2.f;
    170.             # handRectangle.y -= handRectangle.height / 2.f;
    171.             x -= width / 2
    172.             y -= width / 2  # width = height
    173.             # overflow the image
    174.             if x < 0: x = 0
    175.             if y < 0: y = 0
    176.             width1 = width
    177.             width2 = width
    178.             if x + width > image_width: width1 = image_width - x
    179.             if y + width > image_height: width2 = image_height - y
    180.             width = min(width1, width2)
    181.             # the max hand box value is 20 pixels
    182.             if width >= 20:
    183.                 detect_result.append([int(x), int(y), int(width), is_left])
    184.  
    185.     '''
    186.     return value: [[x, y, w, True if left hand else False]].
    187.     width=height since the network require squared input.
    188.     x, y is the coordinate of top left 
    189.     '''
    190.     return detect_result
    191.  
    192. # get max index of 2d array
    193. def npmax(array):
    194.     arrayindex = array.argmax(1)
    195.     arrayvalue = array.max(1)
    196.     i = arrayvalue.argmax()
    197.     j = arrayindex[i]
    198.     return i, j
    src/hand.py

    1. import cv2
    2. import json
    3. import numpy as np
    4. import math
    5. import time
    6. from scipy.ndimage.filters import gaussian_filter
    7. import matplotlib.pyplot as plt
    8. import matplotlib
    9. import torch
    10. from skimage.measure import label
    11.  
    12. from src.model import handpose_model
    13. from src import util
    14.  
    15. class Hand(object):
    16.     def __init__(self, model_path):
    17.         self.model = handpose_model()
    18.         if torch.cuda.is_available():
    19.             self.model = self.model.cuda()
    20.         model_dict = util.transfer(self.model, torch.load(model_path))
    21.         self.model.load_state_dict(model_dict)
    22.         self.model.eval()
    23.  
    24.     def __call__(self, oriImg):
    25.         scale_search = [0.5, 1.0, 1.5, 2.0]
    26.         # scale_search = [0.5]
    27.         boxsize = 368
    28.         stride = 8
    29.         padValue = 128
    30.         thre = 0.05
    31.         multiplier = [x * boxsize / oriImg.shape[0] for x in scale_search]
    32.         heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 22))
    33.         # paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
    34.  
    35.         for m in range(len(multiplier)):
    36.             scale = multiplier[m]
    37.             imageToTest = cv2.resize(oriImg, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
    38.             imageToTest_padded, pad = util.padRightDownCorner(imageToTest, stride, padValue)
    39.             im = np.transpose(np.float32(imageToTest_padded[:, :, :, np.newaxis]), (3, 2, 0, 1)) / 256 - 0.5
    40.             im = np.ascontiguousarray(im)
    41.  
    42.             data = torch.from_numpy(im).float()
    43.             if torch.cuda.is_available():
    44.                 data = data.cuda()
    45.             # data = data.permute([2, 0, 1]).unsqueeze(0).float()
    46.             with torch.no_grad():
    47.                 output = self.model(data).cpu().numpy()
    48.                 # output = self.model(data).numpy()q
    49.  
    50.             # extract outputs, resize, and remove padding
    51.             heatmap = np.transpose(np.squeeze(output), (1, 2, 0))  # output 1 is heatmaps
    52.             heatmap = cv2.resize(heatmap, (0, 0), fx=stride, fy=stride, interpolation=cv2.INTER_CUBIC)
    53.             heatmap = heatmap[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
    54.             heatmap = cv2.resize(heatmap, (oriImg.shape[1], oriImg.shape[0]), interpolation=cv2.INTER_CUBIC)
    55.  
    56.             heatmap_avg += heatmap / len(multiplier)
    57.  
    58.         all_peaks = []
    59.         for part in range(21):
    60.             map_ori = heatmap_avg[:, :, part]
    61.             one_heatmap = gaussian_filter(map_ori, sigma=3)
    62.             binary = np.ascontiguousarray(one_heatmap > thre, dtype=np.uint8)
    63.             # 全部小于阈值
    64.             if np.sum(binary) == 0:
    65.                 all_peaks.append([0, 0])
    66.                 continue
    67.             label_img, label_numbers = label(binary, return_num=True, connectivity=binary.ndim)
    68.             max_index = np.argmax([np.sum(map_ori[label_img == i]) for i in range(1, label_numbers + 1)]) + 1
    69.             label_img[label_img != max_index] = 0
    70.             map_ori[label_img == 0] = 0
    71.  
    72.             y, x = util.npmax(map_ori)
    73.             all_peaks.append([x, y])
    74.         return np.array(all_peaks)
    75.  
    76. if __name__ == "__main__":
    77.     hand_estimation = Hand('../model/hand_pose_model.pth')
    78.  
    79.     # test_image = '../images/hand.jpg'
    80.     test_image = '../images/hand.jpg'
    81.     oriImg = cv2.imread(test_image)  # B,G,R order
    82.     peaks = hand_estimation(oriImg)
    83.     canvas = util.draw_handpose(oriImg, peaks, True)
    84.     cv2.imshow('', canvas)
    85.     cv2.waitKey(0)

    src/body.py

    1. import cv2
    2. import numpy as np
    3. import math
    4. import time
    5. from scipy.ndimage.filters import gaussian_filter
    6. import matplotlib.pyplot as plt
    7. import matplotlib
    8. import torch
    9. from torchvision import transforms
    10.  
    11. from src import util
    12. from src.model import bodypose_model
    13.  
    14. class Body(object):
    15.     def __init__(self, model_path):
    16.         self.model = bodypose_model()
    17.         if torch.cuda.is_available():
    18.             self.model = self.model.cuda()
    19.         model_dict = util.transfer(self.model, torch.load(model_path))
    20.         self.model.load_state_dict(model_dict)
    21.         self.model.eval()
    22.  
    23.     def __call__(self, oriImg):
    24.         # scale_search = [0.5, 1.0, 1.5, 2.0]
    25.         scale_search = [0.5]
    26.         boxsize = 368
    27.         stride = 8
    28.         padValue = 128
    29.         thre1 = 0.1
    30.         thre2 = 0.05
    31.         multiplier = [x * boxsize / oriImg.shape[0] for x in scale_search]
    32.         heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 19))
    33.         paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
    34.  
    35.         for m in range(len(multiplier)):
    36.             scale = multiplier[m]
    37.             imageToTest = cv2.resize(oriImg, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC)
    38.             imageToTest_padded, pad = util.padRightDownCorner(imageToTest, stride, padValue)
    39.             im = np.transpose(np.float32(imageToTest_padded[:, :, :, np.newaxis]), (3, 2, 0, 1)) / 256 - 0.5
    40.             im = np.ascontiguousarray(im)
    41.  
    42.             data = torch.from_numpy(im).float()
    43.             if torch.cuda.is_available():
    44.                 data = data.cuda()
    45.             # data = data.permute([2, 0, 1]).unsqueeze(0).float()
    46.             with torch.no_grad():
    47.                 Mconv7_stage6_L1, Mconv7_stage6_L2 = self.model(data)
    48.             Mconv7_stage6_L1 = Mconv7_stage6_L1.cpu().numpy()
    49.             Mconv7_stage6_L2 = Mconv7_stage6_L2.cpu().numpy()
    50.  
    51.             # extract outputs, resize, and remove padding
    52.             # heatmap = np.transpose(np.squeeze(net.blobs[output_blobs.keys()[1]].data), (1, 2, 0))  # output 1 is heatmaps
    53.             heatmap = np.transpose(np.squeeze(Mconv7_stage6_L2), (1, 2, 0))  # output 1 is heatmaps
    54.             heatmap = cv2.resize(heatmap, (0, 0), fx=stride, fy=stride, interpolation=cv2.INTER_CUBIC)
    55.             heatmap = heatmap[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
    56.             heatmap = cv2.resize(heatmap, (oriImg.shape[1], oriImg.shape[0]), interpolation=cv2.INTER_CUBIC)
    57.  
    58.             # paf = np.transpose(np.squeeze(net.blobs[output_blobs.keys()[0]].data), (1, 2, 0))  # output 0 is PAFs
    59.             paf = np.transpose(np.squeeze(Mconv7_stage6_L1), (1, 2, 0))  # output 0 is PAFs
    60.             paf = cv2.resize(paf, (0, 0), fx=stride, fy=stride, interpolation=cv2.INTER_CUBIC)
    61.             paf = paf[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
    62.             paf = cv2.resize(paf, (oriImg.shape[1], oriImg.shape[0]), interpolation=cv2.INTER_CUBIC)
    63.  
    64.             heatmap_avg += heatmap_avg + heatmap / len(multiplier)
    65.             paf_avg += + paf / len(multiplier)
    66.  
    67.         all_peaks = []
    68.         peak_counter = 0
    69.  
    70.         for part in range(18):
    71.             map_ori = heatmap_avg[:, :, part]
    72.             one_heatmap = gaussian_filter(map_ori, sigma=3)
    73.  
    74.             map_left = np.zeros(one_heatmap.shape)
    75.             map_left[1:, :] = one_heatmap[:-1, :]
    76.             map_right = np.zeros(one_heatmap.shape)
    77.             map_right[:-1, :] = one_heatmap[1:, :]
    78.             map_up = np.zeros(one_heatmap.shape)
    79.             map_up[:, 1:] = one_heatmap[:, :-1]
    80.             map_down = np.zeros(one_heatmap.shape)
    81.             map_down[:, :-1] = one_heatmap[:, 1:]
    82.  
    83.             peaks_binary = np.logical_and.reduce(
    84.                 (one_heatmap >= map_left, one_heatmap >= map_right, one_heatmap >= map_up, one_heatmap >= map_down, one_heatmap > thre1))
    85.             peaks = list(zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0]))  # note reverse
    86.             peaks_with_score = [x + (map_ori[x[1], x[0]],) for x in peaks]
    87.             peak_id = range(peak_counter, peak_counter + len(peaks))
    88.             peaks_with_score_and_id = [peaks_with_score[i] + (peak_id[i],) for i in range(len(peak_id))]
    89.  
    90.             all_peaks.append(peaks_with_score_and_id)
    91.             peak_counter += len(peaks)
    92.  
    93.         # find connection in the specified sequence, center 29 is in the position 15
    94.         limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], \
    95.                    [10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17], \
    96.                    [1, 16], [16, 18], [3, 17], [6, 18]]
    97.         # the middle joints heatmap correpondence
    98.         mapIdx = [[31, 32], [39, 40], [33, 34], [35, 36], [41, 42], [43, 44], [19, 20], [21, 22], \
    99.                   [23, 24], [25, 26], [27, 28], [29, 30], [47, 48], [49, 50], [53, 54], [51, 52], \
    100.                   [55, 56], [37, 38], [45, 46]]
    101.  
    102.         connection_all = []
    103.         special_k = []
    104.         mid_num = 10
    105.  
    106.         for k in range(len(mapIdx)):
    107.             score_mid = paf_avg[:, :, [x - 19 for x in mapIdx[k]]]
    108.             candA = all_peaks[limbSeq[k][0] - 1]
    109.             candB = all_peaks[limbSeq[k][1] - 1]
    110.             nA = len(candA)
    111.             nB = len(candB)
    112.             indexA, indexB = limbSeq[k]
    113.             if (nA != 0 and nB != 0):
    114.                 connection_candidate = []
    115.                 for i in range(nA):
    116.                     for j in range(nB):
    117.                         vec = np.subtract(candB[j][:2], candA[i][:2])
    118.                         norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
    119.                         norm = max(0.001, norm)
    120.                         vec = np.divide(vec, norm)
    121.  
    122.                         startend = list(zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \
    123.                                             np.linspace(candA[i][1], candB[j][1], num=mid_num)))
    124.  
    125.                         vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0] \
    126.                                           for I in range(len(startend))])
    127.                         vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1] \
    128.                                           for I in range(len(startend))])
    129.  
    130.                         score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(vec_y, vec[1])
    131.                         score_with_dist_prior = sum(score_midpts) / len(score_midpts) + min(
    132.                             0.5 * oriImg.shape[0] / norm - 1, 0)
    133.                         criterion1 = len(np.nonzero(score_midpts > thre2)[0]) > 0.8 * len(score_midpts)
    134.                         criterion2 = score_with_dist_prior > 0
    135.                         if criterion1 and criterion2:
    136.                             connection_candidate.append(
    137.                                 [i, j, score_with_dist_prior, score_with_dist_prior + candA[i][2] + candB[j][2]])
    138.  
    139.                 connection_candidate = sorted(connection_candidate, key=lambda x: x[2], reverse=True)
    140.                 connection = np.zeros((0, 5))
    141.                 for c in range(len(connection_candidate)):
    142.                     i, j, s = connection_candidate[c][0:3]
    143.                     if (i not in connection[:, 3] and j not in connection[:, 4]):
    144.                         connection = np.vstack([connection, [candA[i][3], candB[j][3], s, i, j]])
    145.                         if (len(connection) >= min(nA, nB)):
    146.                             break
    147.  
    148.                 connection_all.append(connection)
    149.             else:
    150.                 special_k.append(k)
    151.                 connection_all.append([])
    152.  
    153.         # last number in each row is the total parts number of that person
    154.         # the second last number in each row is the score of the overall configuration
    155.         subset = -1 * np.ones((0, 20))
    156.         candidate = np.array([item for sublist in all_peaks for item in sublist])
    157.  
    158.         for k in range(len(mapIdx)):
    159.             if k not in special_k:
    160.                 partAs = connection_all[k][:, 0]
    161.                 partBs = connection_all[k][:, 1]
    162.                 indexA, indexB = np.array(limbSeq[k]) - 1
    163.  
    164.                 for i in range(len(connection_all[k])):  # = 1:size(temp,1)
    165.                     found = 0
    166.                     subset_idx = [-1, -1]
    167.                     for j in range(len(subset)):  # 1:size(subset,1):
    168.                         if subset[j][indexA] == partAs[i] or subset[j][indexB] == partBs[i]:
    169.                             subset_idx[found] = j
    170.                             found += 1
    171.  
    172.                     if found == 1:
    173.                         j = subset_idx[0]
    174.                         if subset[j][indexB] != partBs[i]:
    175.                             subset[j][indexB] = partBs[i]
    176.                             subset[j][-1] += 1
    177.                             subset[j][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
    178.                     elif found == 2:  # if found 2 and disjoint, merge them
    179.                         j1, j2 = subset_idx
    180.                         membership = ((subset[j1] >= 0).astype(int) + (subset[j2] >= 0).astype(int))[:-2]
    181.                         if len(np.nonzero(membership == 2)[0]) == 0:  # merge
    182.                             subset[j1][:-2] += (subset[j2][:-2] + 1)
    183.                             subset[j1][-2:] += subset[j2][-2:]
    184.                             subset[j1][-2] += connection_all[k][i][2]
    185.                             subset = np.delete(subset, j2, 0)
    186.                         else:  # as like found == 1
    187.                             subset[j1][indexB] = partBs[i]
    188.                             subset[j1][-1] += 1
    189.                             subset[j1][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
    190.  
    191.                     # if find no partA in the subset, create a new subset
    192.                     elif not found and k < 17:
    193.                         row = -1 * np.ones(20)
    194.                         row[indexA] = partAs[i]
    195.                         row[indexB] = partBs[i]
    196.                         row[-1] = 2
    197.                         row[-2] = sum(candidate[connection_all[k][i, :2].astype(int), 2]) + connection_all[k][i][2]
    198.                         subset = np.vstack([subset, row])
    199.         # delete some rows of subset which has few parts occur
    200.         deleteIdx = []
    201.         for i in range(len(subset)):
    202.             if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:
    203.                 deleteIdx.append(i)
    204.         subset = np.delete(subset, deleteIdx, axis=0)
    205.  
    206.         # subset: n*20 array, 0-17 is the index in candidate, 18 is the total score, 19 is the total parts
    207.         # candidate: x, y, score, id
    208.         return candidate, subset
    209.  
    210. if __name__ == "__main__":
    211.     body_estimation = Body('../model/body_pose_model.pth')
    212.  
    213.     test_image = '../images/ski.jpg'
    214.     oriImg = cv2.imread(test_image)  # B,G,R order
    215.     candidate, subset = body_estimation(oriImg)
    216.     canvas = util.draw_bodypose(oriImg, candidate, subset)
    217.     plt.imshow(canvas[:, :, [2, 1, 0]])
    218.     plt.show()

    运行结果

    OpenPose效果展示

    使用cv2写入代码

    需要安装OpenPose，请看我的另一篇博客：

    3D视觉——2.人体姿态估计(Pose Estimation)入门——OpenPose含安装、编译、使用（单帧、实时视频）

    1. import os
    2. import time
    3.  
    4. import cv2
    5. import sys
    6. from tqdm import tqdm
    7. from sys import platform
    8.  
    9. BASE_DIR = os.path.dirname(os.path.realpath(__file__))
    10. if platform == 'win32':
    11.     lib_dir = 'Release'
    12.     bin_dir = 'bin'
    13.     x64_dir = 'x64'
    14.     lib_path = os.path.join(BASE_DIR, lib_dir)
    15.     bin_path = os.path.join(BASE_DIR, bin_dir)
    16.     x64_path = os.path.join(BASE_DIR, x64_dir)
    17.     sys.path.append(lib_path)
    18.     os.environ['PATH'] += ';' + bin_path + ';' + x64_path + '\Release;'
    19.     try:
    20.         import pyopenpose as op
    21.  
    22.         print("successful, import pyopenpose!")
    23.     except ImportError as e:
    24.         print("fail to import pyopenpose!")
    25.         raise e
    26. else:
    27.     print(f"当前电脑环境:\n{platform}\n")
    28.     sys.exit(-1)
    29.  
    30.  
    31. def out_video(input):
    32.     datum = op.Datum()
    33.     opWrapper = op.WrapperPython()
    34.     params = dict()
    35.     params["model_folder"] = BASE_DIR + "\models"
    36.     params["model_pose"] = "BODY_25"
    37.     params["number_people_max"] = 3
    38.     params["disable_blending"] = False
    39.     opWrapper.configure(params)
    40.     opWrapper.start()
    41.     file = input.split("/")[-1]
    42.     output = "video/out-optim-" + file
    43.     print("It will start processing video: {}".format(input))
    44.     cap = cv2.VideoCapture(input)
    45.     frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    46.     frame_size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
    47.     # create VideoWriter,VideoWriter_fourcc is video decode
    48.     fourcc = cv2.VideoWriter_fourcc('D', 'I', 'V', 'X')
    49.     fps = cap.get(cv2.CAP_PROP_FPS)
    50.     out = cv2.VideoWriter(output, fourcc, fps, frame_size)
    51.     # the progress bar
    52.     with tqdm(range(frame_count)) as pbar:
    53.  
    54.         while cap.isOpened():
    55.             start = time.time()
    56.             success, frame = cap.read()
    57.             if success:
    58.                 datum.cvInputData = frame
    59.                 opWrapper.emplaceAndPop(op.VectorDatum([datum]))
    60.                 opframe = datum.cvOutputData
    61.                 FPS = 1 / (time.time() - start)
    62.                 opframe = cv2.putText(opframe, "FPS" + str(int(FPS)), (25, 50), cv2.FONT_HERSHEY_SIMPLEX, 1,
    63.                                       (0, 255, 0), 3)
    64.                 out.write(opframe)
    65.                 pbar.update(1)
    66.             else:
    67.                 break
    68.  
    69.     pbar.close()
    70.     cv2.destroyAllWindows()
    71.     out.release()
    72.     cap.release()
    73.     print("{} finished!".format(output))
    74.  
    75.  
    76. if __name__ == "__main__":
    77.     video_dir = "video/2.avi"
    78.     out_video(video_dir)

    运行结果

    OpenPose运行结果

    评论

    其实：

    1.MediaPipe比较成熟，并且不需要下载或者配置一些乱七八糟的东西，而OpenPose就得配置一大堆乱七八糟的东西，很乱，除非是把caffe的模型转化成pytorch的模型，本人更熟悉pytorch。

    2.MediaPipe推理速度比OpenPose快，快了不止一点点！

    3.OpenPose的效果比MediaPipe的效果更好，并且有多人在画框时OpenPose能检测到所有人，而MediaPipe却只能检测到一个。

    下一话

    MediaPipe手势识别

  • 相关阅读:
    【初试404分】杭电843学长经验分享
    能源区块链实验室同俄罗斯碳基金签署重要战略合作协议
    从0搭建前端脚手架详解（小白也可以搭建）
    UGUI自动布局Layout
    pytest自动化测试-element not interactable、xpath元素定位不到、多页签下子表定位、pip使用问题
    卡方检验-python代码
    迅为iTOP3568开发板Android11获取root权限关闭selinux
    virtio-blk简易驱动
    es笔记六之聚合操作之指标聚合
    UEFI之DXE阶段
  • 原文地址：https://blog.csdn.net/XiaoyYidiaodiao/article/details/125571632