Skip to content

Commit 9d70b98

Browse files
Kazuhito00Kazuhito00
committed
Initial commit
1 parent a1484ad commit 9d70b98

14 files changed

+2095
-0
lines changed

‎.gitignore

+3
Original file line numberDiff line numberDiff line change
@@ -127,3 +127,6 @@ dmypy.json
127127

128128
# Pyre type checker
129129
.pyre/
130+
131+
# bat
132+
*.bat

‎MoveNet_tf2onnx.ipynb

+474
Large diffs are not rendered by default.

‎demo_multipose.py

+336
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,336 @@
1+
#!/usr/bin/env python
2+
# -*- coding: utf-8 -*-
3+
4+
import copy
5+
import time
6+
import argparse
7+
8+
import cv2 as cv
9+
import numpy as np
10+
import tensorflow as tf
11+
import tensorflow_hub as tfhub
12+
13+
14+
def get_args():
15+
parser = argparse.ArgumentParser()
16+
17+
parser.add_argument("--device", type=int, default=0)
18+
parser.add_argument("--file", type=str, default=None)
19+
parser.add_argument("--width", help='cap width', type=int, default=960)
20+
parser.add_argument("--height", help='cap height', type=int, default=540)
21+
22+
parser.add_argument('--mirror', action='store_true')
23+
24+
parser.add_argument("--keypoint_score", type=float, default=0.4)
25+
parser.add_argument("--bbox_score", type=float, default=0.2)
26+
27+
args = parser.parse_args()
28+
29+
return args
30+
31+
32+
def run_inference(model, input_size, image):
33+
image_width, image_height = image.shape[1], image.shape[0]
34+
35+
# 前処理
36+
input_image = cv.resize(image, dsize=(input_size, input_size)) # リサイズ
37+
input_image = cv.cvtColor(input_image, cv.COLOR_BGR2RGB) # BGR→RGB変換
38+
input_image = input_image.reshape(-1, input_size, input_size, 3) # リシェイプ
39+
input_image = tf.cast(input_image, dtype=tf.int32) # int32へキャスト
40+
41+
# 推論
42+
outputs = model(input_image)
43+
44+
keypoints_with_scores = outputs['output_0'].numpy()
45+
keypoints_with_scores = np.squeeze(keypoints_with_scores)
46+
47+
# キーポイント、バウンディングボックス、スコア取り出し
48+
keypoints_list, scores_list = [], []
49+
bbox_list = []
50+
for keypoints_with_score in keypoints_with_scores:
51+
keypoints = []
52+
scores = []
53+
# キーポイント
54+
for index in range(17):
55+
keypoint_x = int(image_width *
56+
keypoints_with_score[(index * 3) + 1])
57+
keypoint_y = int(image_height *
58+
keypoints_with_score[(index * 3) + 0])
59+
score = keypoints_with_score[(index * 3) + 2]
60+
61+
keypoints.append([keypoint_x, keypoint_y])
62+
scores.append(score)
63+
64+
# バウンディングボックス
65+
bbox_ymin = int(image_height * keypoints_with_score[51])
66+
bbox_xmin = int(image_width * keypoints_with_score[52])
67+
bbox_ymax = int(image_height * keypoints_with_score[53])
68+
bbox_xmax = int(image_width * keypoints_with_score[54])
69+
bbox_score = keypoints_with_score[55]
70+
71+
# 6人分のデータ格納用のリストに追加
72+
keypoints_list.append(keypoints)
73+
scores_list.append(scores)
74+
bbox_list.append(
75+
[bbox_xmin, bbox_ymin, bbox_xmax, bbox_ymax, bbox_score])
76+
77+
return keypoints_list, scores_list, bbox_list
78+
79+
80+
def main():
81+
# 引数解析 #################################################################
82+
args = get_args()
83+
cap_device = args.device
84+
cap_width = args.width
85+
cap_height = args.height
86+
87+
if args.file is not None:
88+
cap_device = args.file
89+
90+
mirror = args.mirror
91+
keypoint_score_th = args.keypoint_score
92+
bbox_score_th = args.bbox_score
93+
94+
# カメラ準備 ###############################################################
95+
cap = cv.VideoCapture(cap_device)
96+
cap.set(cv.CAP_PROP_FRAME_WIDTH, cap_width)
97+
cap.set(cv.CAP_PROP_FRAME_HEIGHT, cap_height)
98+
99+
# モデルロード #############################################################
100+
model_url = "https://tfhub.dev/google/movenet/multipose/lightning/1"
101+
input_size = 256
102+
103+
module = tfhub.load(model_url)
104+
model = module.signatures['serving_default']
105+
106+
while True:
107+
start_time = time.time()
108+
109+
# カメラキャプチャ #####################################################
110+
ret, frame = cap.read()
111+
if not ret:
112+
break
113+
if mirror:
114+
frame = cv.flip(frame, 1) # ミラー表示
115+
debug_image = copy.deepcopy(frame)
116+
117+
# 検出実施 ##############################################################
118+
keypoints_list, scores_list, bbox_list = run_inference(
119+
model,
120+
input_size,
121+
frame,
122+
)
123+
124+
elapsed_time = time.time() - start_time
125+
126+
# デバッグ描画
127+
debug_image = draw_debug(
128+
debug_image,
129+
elapsed_time,
130+
keypoint_score_th,
131+
keypoints_list,
132+
scores_list,
133+
bbox_score_th,
134+
bbox_list,
135+
)
136+
137+
# キー処理(ESC：終了) ##################################################
138+
key = cv.waitKey(1)
139+
if key == 27: # ESC
140+
break
141+
142+
# 画面反映 #############################################################
143+
cv.imshow('MoveNet(multipose) Demo', debug_image)
144+
145+
cap.release()
146+
cv.destroyAllWindows()
147+
148+
149+
def draw_debug(
150+
image,
151+
elapsed_time,
152+
keypoint_score_th,
153+
keypoints_list,
154+
scores_list,
155+
bbox_score_th,
156+
bbox_list,
157+
):
158+
debug_image = copy.deepcopy(image)
159+
160+
# 0:鼻 1:左目 2:右目 3:左耳 4:右耳 5:左肩 6:右肩 7:左肘 8:右肘 # 9:左手首
161+
# 10:右手首 11:左股関節 12:右股関節 13:左ひざ 14:右ひざ 15:左足首 16:右足首
162+
for keypoints, scores in zip(keypoints_list, scores_list):
163+
# Line：鼻 → 左目
164+
index01, index02 = 0, 1
165+
if scores[index01] > keypoint_score_th and scores[
166+
index02] > keypoint_score_th:
167+
point01 = keypoints[index01]
168+
point02 = keypoints[index02]
169+
cv.line(debug_image, point01, point02, (255, 255, 255), 4)
170+
cv.line(debug_image, point01, point02, (0, 0, 0), 2)
171+
# Line：鼻 → 右目
172+
index01, index02 = 0, 2
173+
if scores[index01] > keypoint_score_th and scores[
174+
index02] > keypoint_score_th:
175+
point01 = keypoints[index01]
176+
point02 = keypoints[index02]
177+
cv.line(debug_image, point01, point02, (255, 255, 255), 4)
178+
cv.line(debug_image, point01, point02, (0, 0, 0), 2)
179+
# Line：左目 → 左耳
180+
index01, index02 = 1, 3
181+
if scores[index01] > keypoint_score_th and scores[
182+
index02] > keypoint_score_th:
183+
point01 = keypoints[index01]
184+
point02 = keypoints[index02]
185+
cv.line(debug_image, point01, point02, (255, 255, 255), 4)
186+
cv.line(debug_image, point01, point02, (0, 0, 0), 2)
187+
# Line：右目 → 右耳
188+
index01, index02 = 2, 4
189+
if scores[index01] > keypoint_score_th and scores[
190+
index02] > keypoint_score_th:
191+
point01 = keypoints[index01]
192+
point02 = keypoints[index02]
193+
cv.line(debug_image, point01, point02, (255, 255, 255), 4)
194+
cv.line(debug_image, point01, point02, (0, 0, 0), 2)
195+
# Line：鼻 → 左肩
196+
index01, index02 = 0, 5
197+
if scores[index01] > keypoint_score_th and scores[
198+
index02] > keypoint_score_th:
199+
point01 = keypoints[index01]
200+
point02 = keypoints[index02]
201+
cv.line(debug_image, point01, point02, (255, 255, 255), 4)
202+
cv.line(debug_image, point01, point02, (0, 0, 0), 2)
203+
# Line：鼻 → 右肩
204+
index01, index02 = 0, 6
205+
if scores[index01] > keypoint_score_th and scores[
206+
index02] > keypoint_score_th:
207+
point01 = keypoints[index01]
208+
point02 = keypoints[index02]
209+
cv.line(debug_image, point01, point02, (255, 255, 255), 4)
210+
cv.line(debug_image, point01, point02, (0, 0, 0), 2)
211+
# Line：左肩 → 右肩
212+
index01, index02 = 5, 6
213+
if scores[index01] > keypoint_score_th and scores[
214+
index02] > keypoint_score_th:
215+
point01 = keypoints[index01]
216+
point02 = keypoints[index02]
217+
cv.line(debug_image, point01, point02, (255, 255, 255), 4)
218+
cv.line(debug_image, point01, point02, (0, 0, 0), 2)
219+
# Line：左肩 → 左肘
220+
index01, index02 = 5, 7
221+
if scores[index01] > keypoint_score_th and scores[
222+
index02] > keypoint_score_th:
223+
point01 = keypoints[index01]
224+
point02 = keypoints[index02]
225+
cv.line(debug_image, point01, point02, (255, 255, 255), 4)
226+
cv.line(debug_image, point01, point02, (0, 0, 0), 2)
227+
# Line：左肘 → 左手首
228+
index01, index02 = 7, 9
229+
if scores[index01] > keypoint_score_th and scores[
230+
index02] > keypoint_score_th:
231+
point01 = keypoints[index01]
232+
point02 = keypoints[index02]
233+
cv.line(debug_image, point01, point02, (255, 255, 255), 4)
234+
cv.line(debug_image, point01, point02, (0, 0, 0), 2)
235+
# Line：右肩 → 右肘
236+
index01, index02 = 6, 8
237+
if scores[index01] > keypoint_score_th and scores[
238+
index02] > keypoint_score_th:
239+
point01 = keypoints[index01]
240+
point02 = keypoints[index02]
241+
cv.line(debug_image, point01, point02, (255, 255, 255), 4)
242+
cv.line(debug_image, point01, point02, (0, 0, 0), 2)
243+
# Line：右肘 → 右手首
244+
index01, index02 = 8, 10
245+
if scores[index01] > keypoint_score_th and scores[
246+
index02] > keypoint_score_th:
247+
point01 = keypoints[index01]
248+
point02 = keypoints[index02]
249+
cv.line(debug_image, point01, point02, (255, 255, 255), 4)
250+
cv.line(debug_image, point01, point02, (0, 0, 0), 2)
251+
# Line：左股関節 → 右股関節
252+
index01, index02 = 11, 12
253+
if scores[index01] > keypoint_score_th and scores[
254+
index02] > keypoint_score_th:
255+
point01 = keypoints[index01]
256+
point02 = keypoints[index02]
257+
cv.line(debug_image, point01, point02, (255, 255, 255), 4)
258+
cv.line(debug_image, point01, point02, (0, 0, 0), 2)
259+
# Line：左肩 → 左股関節
260+
index01, index02 = 5, 11
261+
if scores[index01] > keypoint_score_th and scores[
262+
index02] > keypoint_score_th:
263+
point01 = keypoints[index01]
264+
point02 = keypoints[index02]
265+
cv.line(debug_image, point01, point02, (255, 255, 255), 4)
266+
cv.line(debug_image, point01, point02, (0, 0, 0), 2)
267+
# Line：左股関節 → 左ひざ
268+
index01, index02 = 11, 13
269+
if scores[index01] > keypoint_score_th and scores[
270+
index02] > keypoint_score_th:
271+
point01 = keypoints[index01]
272+
point02 = keypoints[index02]
273+
cv.line(debug_image, point01, point02, (255, 255, 255), 4)
274+
cv.line(debug_image, point01, point02, (0, 0, 0), 2)
275+
# Line：左ひざ → 左足首
276+
index01, index02 = 13, 15
277+
if scores[index01] > keypoint_score_th and scores[
278+
index02] > keypoint_score_th:
279+
point01 = keypoints[index01]
280+
point02 = keypoints[index02]
281+
cv.line(debug_image, point01, point02, (255, 255, 255), 4)
282+
cv.line(debug_image, point01, point02, (0, 0, 0), 2)
283+
# Line：右肩 → 右股関節
284+
index01, index02 = 6, 12
285+
if scores[index01] > keypoint_score_th and scores[
286+
index02] > keypoint_score_th:
287+
point01 = keypoints[index01]
288+
point02 = keypoints[index02]
289+
cv.line(debug_image, point01, point02, (255, 255, 255), 4)
290+
cv.line(debug_image, point01, point02, (0, 0, 0), 2)
291+
# Line：右股関節 → 右ひざ
292+
index01, index02 = 12, 14
293+
if scores[index01] > keypoint_score_th and scores[
294+
index02] > keypoint_score_th:
295+
point01 = keypoints[index01]
296+
point02 = keypoints[index02]
297+
cv.line(debug_image, point01, point02, (255, 255, 255), 4)
298+
cv.line(debug_image, point01, point02, (0, 0, 0), 2)
299+
# Line：右ひざ → 右足首
300+
index01, index02 = 14, 16
301+
if scores[index01] > keypoint_score_th and scores[
302+
index02] > keypoint_score_th:
303+
point01 = keypoints[index01]
304+
point02 = keypoints[index02]
305+
cv.line(debug_image, point01, point02, (255, 255, 255), 4)
306+
cv.line(debug_image, point01, point02, (0, 0, 0), 2)
307+
308+
# Circle：各点
309+
for keypoint, score in zip(keypoints, scores):
310+
if score > keypoint_score_th:
311+
cv.circle(debug_image, keypoint, 6, (255, 255, 255), -1)
312+
cv.circle(debug_image, keypoint, 3, (0, 0, 0), -1)
313+
314+
# バウンディングボックス
315+
for bbox in bbox_list:
316+
if bbox[4] > bbox_score_th:
317+
cv.rectangle(debug_image, (bbox[0], bbox[1]), (bbox[2], bbox[3]),
318+
(255, 255, 255), 4)
319+
cv.rectangle(debug_image, (bbox[0], bbox[1]), (bbox[2], bbox[3]),
320+
(0, 0, 0), 2)
321+
322+
# 処理時間
323+
cv.putText(debug_image,
324+
"Elapsed Time : " + '{:.1f}'.format(elapsed_time * 1000) + "ms",
325+
(10, 30), cv.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 4,
326+
cv.LINE_AA)
327+
cv.putText(debug_image,
328+
"Elapsed Time : " + '{:.1f}'.format(elapsed_time * 1000) + "ms",
329+
(10, 30), cv.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 2,
330+
cv.LINE_AA)
331+
332+
return debug_image
333+
334+
335+
if __name__ == '__main__':
336+
main()

0 commit comments

Comments
 (0)