Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python | |
| # -*- coding: utf-8 -*- | |
| import copy | |
| import math | |
| import argparse | |
| import cv2 as cv | |
| import numpy as np | |
| import mediapipe as mp | |
| from utils import CvFpsCalc | |
| def get_args(): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--device", type=int, default=0) | |
| parser.add_argument("--width", help='cap width', type=int, default=640) | |
| parser.add_argument("--height", help='cap height', type=int, default=360) | |
| parser.add_argument('--static_image_mode', action='store_true') | |
| parser.add_argument("--model_complexity", | |
| help='model_complexity(0,1(default),2)', | |
| type=int, | |
| default=1) | |
| parser.add_argument("--min_detection_confidence", | |
| help='min_detection_confidence', | |
| type=float, | |
| default=0.5) | |
| parser.add_argument("--min_tracking_confidence", | |
| help='min_tracking_confidence', | |
| type=int, | |
| default=0.5) | |
| parser.add_argument('--rev_color', action='store_true') | |
| args = parser.parse_args() | |
| return args | |
| def main(): | |
| # 引数解析 ################################################################# | |
| args = get_args() | |
| cap_device = args.device | |
| cap_width = args.width | |
| cap_height = args.height | |
| static_image_mode = args.static_image_mode | |
| model_complexity = args.model_complexity | |
| min_detection_confidence = args.min_detection_confidence | |
| min_tracking_confidence = args.min_tracking_confidence | |
| rev_color = args.rev_color | |
| # カメラ準備 ############################################################### | |
| cap = cv.VideoCapture(cap_device) | |
| cap.set(cv.CAP_PROP_FRAME_WIDTH, cap_width) | |
| cap.set(cv.CAP_PROP_FRAME_HEIGHT, cap_height) | |
| # モデルロード ############################################################# | |
| mp_pose = mp.solutions.pose | |
| pose = mp_pose.Pose( | |
| static_image_mode=static_image_mode, | |
| model_complexity=model_complexity, | |
| min_detection_confidence=min_detection_confidence, | |
| min_tracking_confidence=min_tracking_confidence, | |
| ) | |
| # FPS計測モジュール ######################################################## | |
| cvFpsCalc = CvFpsCalc(buffer_len=10) | |
| # 色指定 | |
| if rev_color: | |
| color = (255, 255, 255) | |
| bg_color = (100, 33, 3) | |
| else: | |
| color = (100, 33, 3) | |
| bg_color = (255, 255, 255) | |
| while True: | |
| display_fps = cvFpsCalc.get() | |
| # カメラキャプチャ ##################################################### | |
| ret, image = cap.read() | |
| if not ret: | |
| break | |
| image = cv.flip(image, 1) # ミラー表示 | |
| debug_image01 = copy.deepcopy(image) | |
| debug_image02 = np.zeros((image.shape[0], image.shape[1], 3), np.uint8) | |
| cv.rectangle(debug_image02, (0, 0), (image.shape[1], image.shape[0]), | |
| bg_color, | |
| thickness=-1) | |
| # 検出実施 ############################################################# | |
| image = cv.cvtColor(image, cv.COLOR_BGR2RGB) | |
| results = pose.process(image) | |
| # 描画 ################################################################ | |
| if results.pose_landmarks is not None: | |
| # 描画 | |
| debug_image01 = draw_landmarks( | |
| debug_image01, | |
| results.pose_landmarks, | |
| ) | |
| debug_image02 = draw_stick_figure( | |
| debug_image02, | |
| results.pose_landmarks, | |
| color=color, | |
| bg_color=bg_color, | |
| ) | |
| cv.putText(debug_image01, "FPS:" + str(display_fps), (10, 30), | |
| cv.FONT_HERSHEY_SIMPLEX, 1.0, (0, 255, 0), 2, cv.LINE_AA) | |
| cv.putText(debug_image02, "FPS:" + str(display_fps), (10, 30), | |
| cv.FONT_HERSHEY_SIMPLEX, 1.0, color, 2, cv.LINE_AA) | |
| # キー処理(ESC：終了) ################################################# | |
| key = cv.waitKey(1) | |
| if key == 27: # ESC | |
| break | |
| # 画面反映 ############################################################# | |
| cv.imshow('Tokyo2020 Debug', debug_image01) | |
| cv.imshow('Tokyo2020 Pictogram', debug_image02) | |
| cap.release() | |
| cv.destroyAllWindows() | |
| def draw_stick_figure( | |
| image, | |
| landmarks, | |
| color=(100, 33, 3), | |
| bg_color=(255, 255, 255), | |
| visibility_th=0.5, | |
| ): | |
| image_width, image_height = image.shape[1], image.shape[0] | |
| # 各ランドマーク算出 | |
| landmark_point = [] | |
| for index, landmark in enumerate(landmarks.landmark): | |
| landmark_x = min(int(landmark.x * image_width), image_width - 1) | |
| landmark_y = min(int(landmark.y * image_height), image_height - 1) | |
| landmark_z = landmark.z | |
| landmark_point.append( | |
| [index, landmark.visibility, (landmark_x, landmark_y), landmark_z]) | |
| # 脚の付け根の位置を腰の中点に修正 | |
| right_leg = landmark_point[23] | |
| left_leg = landmark_point[24] | |
| leg_x = int((right_leg[2][0] + left_leg[2][0]) / 2) | |
| leg_y = int((right_leg[2][1] + left_leg[2][1]) / 2) | |
| landmark_point[23][2] = (leg_x, leg_y) | |
| landmark_point[24][2] = (leg_x, leg_y) | |
| # 距離順にソート | |
| sorted_landmark_point = sorted(landmark_point, | |
| reverse=True, | |
| key=lambda x: x[3]) | |
| # 各サイズ算出 | |
| (face_x, face_y), face_radius = min_enclosing_face_circle(landmark_point) | |
| face_x = int(face_x) | |
| face_y = int(face_y) | |
| face_radius = int(face_radius * 1.5) | |
| stick_radius01 = int(face_radius * (4 / 5)) | |
| stick_radius02 = int(stick_radius01 * (3 / 4)) | |
| stick_radius03 = int(stick_radius02 * (3 / 4)) | |
| # 描画対象リスト | |
| draw_list = [ | |
| 11, # 右腕 | |
| 12, # 左腕 | |
| 23, # 右脚 | |
| 24, # 左脚 | |
| ] | |
| # 背景色 | |
| cv.rectangle(image, (0, 0), (image_width, image_height), | |
| bg_color, | |
| thickness=-1) | |
| # 顔 描画 | |
| cv.circle(image, (face_x, face_y), face_radius, color, -1) | |
| # 腕/脚 描画 | |
| for landmark_info in sorted_landmark_point: | |
| index = landmark_info[0] | |
| if index in draw_list: | |
| point01 = [p for p in landmark_point if p[0] == index][0] | |
| point02 = [p for p in landmark_point if p[0] == (index + 2)][0] | |
| point03 = [p for p in landmark_point if p[0] == (index + 4)][0] | |
| if point01[1] > visibility_th and point02[1] > visibility_th: | |
| image = draw_stick( | |
| image, | |
| point01[2], | |
| stick_radius01, | |
| point02[2], | |
| stick_radius02, | |
| color=color, | |
| bg_color=bg_color, | |
| ) | |
| if point02[1] > visibility_th and point03[1] > visibility_th: | |
| image = draw_stick( | |
| image, | |
| point02[2], | |
| stick_radius02, | |
| point03[2], | |
| stick_radius03, | |
| color=color, | |
| bg_color=bg_color, | |
| ) | |
| return image | |
| def min_enclosing_face_circle(landmark_point): | |
| landmark_array = np.empty((0, 2), int) | |
| index_list = [1, 4, 7, 8, 9, 10] | |
| for index in index_list: | |
| np_landmark_point = [ | |
| np.array( | |
| (landmark_point[index][2][0], landmark_point[index][2][1])) | |
| ] | |
| landmark_array = np.append(landmark_array, np_landmark_point, axis=0) | |
| center, radius = cv.minEnclosingCircle(points=landmark_array) | |
| return center, radius | |
| def draw_stick( | |
| image, | |
| point01, | |
| point01_radius, | |
| point02, | |
| point02_radius, | |
| color=(100, 33, 3), | |
| bg_color=(255, 255, 255), | |
| ): | |
| cv.circle(image, point01, point01_radius, color, -1) | |
| cv.circle(image, point02, point02_radius, color, -1) | |
| draw_list = [] | |
| for index in range(2): | |
| rad = math.atan2(point02[1] - point01[1], point02[0] - point01[0]) | |
| rad = rad + (math.pi / 2) + (math.pi * index) | |
| point_x = int(point01_radius * math.cos(rad)) + point01[0] | |
| point_y = int(point01_radius * math.sin(rad)) + point01[1] | |
| draw_list.append([point_x, point_y]) | |
| point_x = int(point02_radius * math.cos(rad)) + point02[0] | |
| point_y = int(point02_radius * math.sin(rad)) + point02[1] | |
| draw_list.append([point_x, point_y]) | |
| points = np.array((draw_list[0], draw_list[1], draw_list[3], draw_list[2])) | |
| cv.fillConvexPoly(image, points=points, color=color) | |
| return image | |
| def draw_landmarks( | |
| image, | |
| landmarks, | |
| # upper_body_only, | |
| visibility_th=0.5, | |
| ): | |
| image_width, image_height = image.shape[1], image.shape[0] | |
| landmark_point = [] | |
| for index, landmark in enumerate(landmarks.landmark): | |
| landmark_x = min(int(landmark.x * image_width), image_width - 1) | |
| landmark_y = min(int(landmark.y * image_height), image_height - 1) | |
| landmark_z = landmark.z | |
| landmark_point.append([landmark.visibility, (landmark_x, landmark_y)]) | |
| if landmark.visibility < visibility_th: | |
| continue | |
| if index == 0: # 鼻 | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 1: # 右目：目頭 | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 2: # 右目：瞳 | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 3: # 右目：目尻 | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 4: # 左目：目頭 | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 5: # 左目：瞳 | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 6: # 左目：目尻 | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 7: # 右耳 | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 8: # 左耳 | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 9: # 口：左端 | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 10: # 口：左端 | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 11: # 右肩 | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 12: # 左肩 | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 13: # 右肘 | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 14: # 左肘 | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 15: # 右手首 | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 16: # 左手首 | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 17: # 右手1(外側端) | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 18: # 左手1(外側端) | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 19: # 右手2(先端) | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 20: # 左手2(先端) | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 21: # 右手3(内側端) | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 22: # 左手3(内側端) | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 23: # 腰(右側) | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 24: # 腰(左側) | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 25: # 右ひざ | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 26: # 左ひざ | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 27: # 右足首 | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 28: # 左足首 | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 29: # 右かかと | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 30: # 左かかと | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 31: # 右つま先 | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| if index == 32: # 左つま先 | |
| cv.circle(image, (landmark_x, landmark_y), 5, (0, 255, 0), 2) | |
| # if not upper_body_only: | |
| if True: | |
| cv.putText(image, "z:" + str(round(landmark_z, 3)), | |
| (landmark_x - 10, landmark_y - 10), | |
| cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1, | |
| cv.LINE_AA) | |
| # 右目 | |
| if landmark_point[1][0] > visibility_th and landmark_point[2][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[1][1], landmark_point[2][1], | |
| (0, 255, 0), 2) | |
| if landmark_point[2][0] > visibility_th and landmark_point[3][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[2][1], landmark_point[3][1], | |
| (0, 255, 0), 2) | |
| # 左目 | |
| if landmark_point[4][0] > visibility_th and landmark_point[5][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[4][1], landmark_point[5][1], | |
| (0, 255, 0), 2) | |
| if landmark_point[5][0] > visibility_th and landmark_point[6][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[5][1], landmark_point[6][1], | |
| (0, 255, 0), 2) | |
| # 口 | |
| if landmark_point[9][0] > visibility_th and landmark_point[10][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[9][1], landmark_point[10][1], | |
| (0, 255, 0), 2) | |
| # 肩 | |
| if landmark_point[11][0] > visibility_th and landmark_point[12][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[11][1], landmark_point[12][1], | |
| (0, 255, 0), 2) | |
| # 右腕 | |
| if landmark_point[11][0] > visibility_th and landmark_point[13][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[11][1], landmark_point[13][1], | |
| (0, 255, 0), 2) | |
| if landmark_point[13][0] > visibility_th and landmark_point[15][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[13][1], landmark_point[15][1], | |
| (0, 255, 0), 2) | |
| # 左腕 | |
| if landmark_point[12][0] > visibility_th and landmark_point[14][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[12][1], landmark_point[14][1], | |
| (0, 255, 0), 2) | |
| if landmark_point[14][0] > visibility_th and landmark_point[16][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[14][1], landmark_point[16][1], | |
| (0, 255, 0), 2) | |
| # 右手 | |
| if landmark_point[15][0] > visibility_th and landmark_point[17][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[15][1], landmark_point[17][1], | |
| (0, 255, 0), 2) | |
| if landmark_point[17][0] > visibility_th and landmark_point[19][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[17][1], landmark_point[19][1], | |
| (0, 255, 0), 2) | |
| if landmark_point[19][0] > visibility_th and landmark_point[21][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[19][1], landmark_point[21][1], | |
| (0, 255, 0), 2) | |
| if landmark_point[21][0] > visibility_th and landmark_point[15][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[21][1], landmark_point[15][1], | |
| (0, 255, 0), 2) | |
| # 左手 | |
| if landmark_point[16][0] > visibility_th and landmark_point[18][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[16][1], landmark_point[18][1], | |
| (0, 255, 0), 2) | |
| if landmark_point[18][0] > visibility_th and landmark_point[20][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[18][1], landmark_point[20][1], | |
| (0, 255, 0), 2) | |
| if landmark_point[20][0] > visibility_th and landmark_point[22][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[20][1], landmark_point[22][1], | |
| (0, 255, 0), 2) | |
| if landmark_point[22][0] > visibility_th and landmark_point[16][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[22][1], landmark_point[16][1], | |
| (0, 255, 0), 2) | |
| # 胴体 | |
| if landmark_point[11][0] > visibility_th and landmark_point[23][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[11][1], landmark_point[23][1], | |
| (0, 255, 0), 2) | |
| if landmark_point[12][0] > visibility_th and landmark_point[24][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[12][1], landmark_point[24][1], | |
| (0, 255, 0), 2) | |
| if landmark_point[23][0] > visibility_th and landmark_point[24][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[23][1], landmark_point[24][1], | |
| (0, 255, 0), 2) | |
| if len(landmark_point) > 25: | |
| # 右足 | |
| if landmark_point[23][0] > visibility_th and landmark_point[25][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[23][1], landmark_point[25][1], | |
| (0, 255, 0), 2) | |
| if landmark_point[25][0] > visibility_th and landmark_point[27][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[25][1], landmark_point[27][1], | |
| (0, 255, 0), 2) | |
| if landmark_point[27][0] > visibility_th and landmark_point[29][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[27][1], landmark_point[29][1], | |
| (0, 255, 0), 2) | |
| if landmark_point[29][0] > visibility_th and landmark_point[31][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[29][1], landmark_point[31][1], | |
| (0, 255, 0), 2) | |
| # 左足 | |
| if landmark_point[24][0] > visibility_th and landmark_point[26][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[24][1], landmark_point[26][1], | |
| (0, 255, 0), 2) | |
| if landmark_point[26][0] > visibility_th and landmark_point[28][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[26][1], landmark_point[28][1], | |
| (0, 255, 0), 2) | |
| if landmark_point[28][0] > visibility_th and landmark_point[30][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[28][1], landmark_point[30][1], | |
| (0, 255, 0), 2) | |
| if landmark_point[30][0] > visibility_th and landmark_point[32][ | |
| 0] > visibility_th: | |
| cv.line(image, landmark_point[30][1], landmark_point[32][1], | |
| (0, 255, 0), 2) | |
| return image | |
| if __name__ == '__main__': | |
| main() | |