骨架对比整合到类中

Author: yhu-ctrl

angle.py

@@ -1,45 +1,70 @@
 import numpy as np
 from mxnet import ndarray
+from sklearn.metrics import r2_score
+
 
 # 需要测量角度的部位，每个部位需要用它本身和与之连接的两个节点来计算角度
 # 第一点是关节点
 KeyPoints = [
     (5, 6, 7),       # 左肩
     (6, 7, 8),       # 右肩
-    (7, 5, 9),      # 左臂
+    (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)))
+class AngeleCal():
+
+    def __init__(self, filename):
+        self.stdAngles = np.loadtxt(filename, delimiter='\t')
+        self.pos = 0
+
+    # 计算角度
+    @staticmethod
+    def cal(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)
+        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()
+                    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 ])
+                    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)
+                    angles[i][j] = np.dot(v1, v2) / np.linalg.norm(v1) / np.linalg.norm(v2)
 
+                else:
+                    angles[i][j] = np.nan
+    
+        return angles
+
+    # 角度对比
+    def compare(self, angles):
+        # 每次读取一行标准角度
+        stdAngle = self.stdAngles[self.pos]
+        scores = []
+        visibles = ~np.isnan(stdAngle)          # 样本中没有缺失值的
+        for angle in angles:
+            angle_v = angle[visibles]           # 过滤样本中也有缺失值的点
+            if np.isnan(angle_v).any():         # 还有缺失值
+                scores.append('NaN')
             else:
-                angles[i][j] = np.nan
-   
-    return angles
+                scores.append('{:.4f}'.format(r2_score(angle_v, stdAngle[visibles])))
+        self.pos += 1
+
+        return scores

data.py

@@ -8,7 +8,7 @@
 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 angle import AngeleCal
 
 # 读取参数
 parser = argparse.ArgumentParser()
@@ -53,7 +53,7 @@
         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]
+        X = AngeleCal.cal(pred_coords, confidence)[0]
         print(X)
         features.append(X)
     else:

run.py

@@ -8,8 +8,7 @@
 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
+from angle import AngeleCal
 
 # 读取参数
 parser = argparse.ArgumentParser()
@@ -39,7 +38,7 @@
 cap2 = cv.VideoCapture(args.demo)
 
 # 标准特征
-stdAngle = np.loadtxt(args.data, delimiter='\t')
+angeleCal = AngeleCal(args.data)
 pos = 0
 
 ret1, frame1 = cap1.read()
@@ -62,24 +61,17 @@
         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
+        angles = AngeleCal.cal(pred_coords, confidence)
+        results = angeleCal.compare(angles)
+    else:
+        results = ['NaN']
 
+    print('result', results)
     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')
+        left_txt_list=results, canvas_name='pose')
     fps_time = time.time()
-    # cv.imshow('demo', frame2)
+    cv.imshow('demo', frame2)
 
     # ESC键退出
     if cv.waitKey(1) == 27: