动作对比功能已完成

Author: yhu-ctrl

.gitignore

@@ -127,3 +127,6 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+.vscode
+data/*

angle.py

@@ -0,0 +1,45 @@
+import numpy as np
+from mxnet import ndarray
+
+# 需要测量角度的部位，每个部位需要用它本身和与之连接的两个节点来计算角度
+# 第一点是关节点
+KeyPoints = [
+    (5, 6, 7),       # 左肩
+    (6, 7, 8),       # 右肩
+    (7, 5, 9),      # 左臂
+    (8, 6, 10),      # 右臂
+    (11, 5, 13),     # 左胯
+    (12, 6, 14),     # 右胯
+    (13, 11, 15),    # 左膝
+    (14, 12, 16),    # 右膝
+]
+
+# 计算所有人关键部位的夹角的余弦值
+def CalAngle(coords, confidence, keypoint_thresh=0.2):
+    joint_visible = confidence[:, :, 0] > keypoint_thresh
+    angles = np.empty((coords.shape[0], len(KeyPoints)))
+
+    for i, pts in enumerate(coords):
+        # 某个人
+        for j, keyPoint in enumerate(KeyPoints):
+            # 是否识别到这个关节
+            if joint_visible[i, keyPoint[0]] and joint_visible[i, keyPoint[1]] and joint_visible[i, keyPoint[2]]:
+                # 计算
+                # print(pts)
+
+                p0x = pts[keyPoint[0], 0].asscalar()
+                p0y = pts[keyPoint[0], 1].asscalar()
+                p1x = pts[keyPoint[1], 0].asscalar()
+                p1y = pts[keyPoint[1], 1].asscalar()
+                p2x = pts[keyPoint[2], 0].asscalar()
+                p2y = pts[keyPoint[2], 1].asscalar()
+
+                v1 = np.array([ p1x - p0x, p1y - p0y ])
+                v2 = np.array([ p2x - p0x, p2y - p0y ])
+
+                angles[i][j] = np.dot(v1, v2) / np.linalg.norm(v1) / np.linalg.norm(v2)
+
+            else:
+                angles[i][j] = np.nan
+   
+    return angles

data.py

@@ -0,0 +1,69 @@
+import argparse, time, os
+import cv2 as cv
+import mxnet as mx
+import numpy as np
+import gluoncv as gcv
+from gluoncv import data
+from gluoncv.model_zoo import get_model
+from gluoncv.data.transforms.pose import detector_to_simple_pose, heatmap_to_coord
+from gluoncv.utils.viz import cv_plot_image, cv_plot_keypoints
+from mxnet.gluon.data.vision import transforms
+from angle import CalAngle
+
+# 读取参数
+parser = argparse.ArgumentParser()
+parser.add_argument('--video')
+args = parser.parse_args()
+
+fps_time = 0
+
+# 设置模型
+ctx = mx.gpu()
+
+detector_name = "ssd_512_mobilenet1.0_coco"
+detector = get_model(detector_name, pretrained=True, ctx=ctx)
+
+estimator_name = "simple_pose_resnet18_v1b"
+estimator = get_model(estimator_name, pretrained='ccd24037', ctx=ctx)
+
+detector.reset_class(classes=['person'], reuse_weights={'person':'person'})
+
+detector.hybridize()
+estimator.hybridize()
+
+# 视频读取
+cap = cv.VideoCapture(args.video)
+
+ret, frame = cap.read()
+features = []
+while ret:
+
+    # 目标检测
+    frame = mx.nd.array(cv.cvtColor(frame, cv.COLOR_BGR2RGB)).astype('uint8')
+
+    x, img = gcv.data.transforms.presets.ssd.transform_test(frame, short=512)
+    x = x.as_in_context(ctx)
+    class_IDs, scores, bounding_boxs = detector(x)
+
+    pose_input, upscale_bbox = detector_to_simple_pose(img, class_IDs, scores, bounding_boxs, output_shape=(128, 96), ctx=ctx)
+
+    # 只识别一个人的姿态
+    if len(upscale_bbox) == 1:
+        predicted_heatmap = estimator(pose_input)
+        pred_coords, confidence = heatmap_to_coord(predicted_heatmap, upscale_bbox)
+        img = cv_plot_keypoints(img, pred_coords, confidence, class_IDs, bounding_boxs, scores)
+
+        X = CalAngle(pred_coords, confidence)[0]
+        print(X)
+        features.append(X)
+    else:
+        # 人数不对就插入nan
+        print(np.nan)
+        features.append(np.nan)
+
+    ret, frame = cap.read()
+
+cap.release()
+
+# 将一个视频的特征保存到文件
+np.savetxt(os.path.join('data', 'demo.tsv'), np.array(features), delimiter='\t', fmt='%4f')

run.py

@@ -0,0 +1,94 @@
+import argparse, time, os
+import cv2 as cv
+import mxnet as mx
+import numpy as np
+import gluoncv as gcv
+from gluoncv import data
+from gluoncv.model_zoo import get_model
+from gluoncv.data.transforms.pose import detector_to_simple_pose, heatmap_to_coord
+from gluoncv.utils.viz import cv_plot_image, cv_plot_keypoints
+from mxnet.gluon.data.vision import transforms
+from angle import CalAngle
+from sklearn.metrics import r2_score
+
+# 读取参数
+parser = argparse.ArgumentParser()
+parser.add_argument('--input', default=0)
+parser.add_argument('--demo', required=True)
+parser.add_argument('--data', required=True)
+args = parser.parse_args()
+
+fps_time = 0
+
+# 设置模型
+ctx = mx.gpu()
+
+detector_name = "ssd_512_mobilenet1.0_coco"
+detector = get_model(detector_name, pretrained=True, ctx=ctx)
+
+estimator_name = "simple_pose_resnet18_v1b"
+estimator = get_model(estimator_name, pretrained='ccd24037', ctx=ctx)
+
+detector.reset_class(classes=['person'], reuse_weights={'person':'person'})
+
+detector.hybridize()
+estimator.hybridize()
+
+# 视频读取
+cap1 = cv.VideoCapture(args.input)
+cap2 = cv.VideoCapture(args.demo)
+
+# 标准特征
+stdAngle = np.loadtxt(args.data, delimiter='\t')
+pos = 0
+
+ret1, frame1 = cap1.read()
+ret2, frame2 = cap2.read()
+while ret1 and ret2:
+
+    # 目标检测
+    frame = mx.nd.array(cv.cvtColor(frame1, cv.COLOR_BGR2RGB)).astype('uint8')
+
+    x, img = gcv.data.transforms.presets.ssd.transform_test(frame, short=512)
+    x = x.as_in_context(ctx)
+    class_IDs, scores, bounding_boxs = detector(x)
+
+    pose_input, upscale_bbox = detector_to_simple_pose(img, class_IDs, scores, bounding_boxs, output_shape=(128, 96), ctx=ctx)
+
+    # 姿态识别
+    if len(upscale_bbox) > 0:
+        predicted_heatmap = estimator(pose_input)
+        pred_coords, confidence = heatmap_to_coord(predicted_heatmap, upscale_bbox)
+        img = cv_plot_keypoints(img, pred_coords, confidence, class_IDs, bounding_boxs, scores)
+
+        # 动作对比
+        scores = []
+        # print(stdAngle[pos])
+        visibles = ~np.isnan(stdAngle[pos])     # 样本中没有缺失值的点
+        angles = CalAngle(pred_coords, confidence)
+        for angle in angles:
+            angle_v = angle[visibles]           # 过滤样本中也有缺失值的点
+            print(angle_v)
+            if np.isnan(angle_v).any():         # 还有缺失值
+                scores.append('NaN')
+            else:
+                scores.append('{:.4f}'.format(r2_score(angle_v, stdAngle[pos][visibles])))
+        pos += 1
+
+    cv_plot_image(img, 
+        upperleft_txt=f"FPS:{(1.0 / (time.time() - fps_time)):.2f}", upperleft_txt_corner=(10,25),
+        left_txt_list=scores, canvas_name='pose')
+    fps_time = time.time()
+    # cv.imshow('demo', frame2)
+    
+    # ESC键退出
+    if cv.waitKey(1) == 27:
+        break
+
+    ret1, frame1 = cap1.read()
+    ret2, frame2 = cap2.read()
+
+cv.destroyAllWindows()
+
+cap1.release()
+cap2.release()