继续优化代码

Author: yhu-ctrl

data.py

@@ -2,37 +2,21 @@
 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.presets.ssd import transform_test
 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 model import ctx, detector, estimator
 from angle import AngeleCal
 
 # 读取参数
 parser = argparse.ArgumentParser()
-parser.add_argument('--video')
+parser.add_argument('--input')
+parser.add_argument('--output', 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()
-
 # 视频读取
-cap = cv.VideoCapture(args.video)
+cap = cv.VideoCapture(args.input)
 
 ret, frame = cap.read()
 features = []
@@ -41,14 +25,13 @@
     # 目标检测
     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, img = 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:
+    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)
@@ -57,7 +40,7 @@
         print(X)
         features.append(X)
     else:
-        # 人数不对就插入nan
+        # 人数不够就插入nan
         print(np.nan)
         features.append(np.nan)
 
@@ -66,4 +49,4 @@
 cap.release()
 
 # 将一个视频的特征保存到文件
-np.savetxt(os.path.join('data', 'demo.tsv'), np.array(features), delimiter='\t', fmt='%4f')
+np.savetxt(args.output, np.array(features), delimiter='\t', fmt='%4f')

fps.py

@@ -0,0 +1,11 @@
+import time
+
+class FPS():
+    __fps_time = 0
+
+    @staticmethod
+    def fps():
+        fps = f"FPS:{(1.0 / (time.time() - FPS.__fps_time)):.2f}"
+        FPS.__fps_time = time.time()
+
+        return fps

model.py

@@ -0,0 +1,15 @@
+import mxnet as mx
+from gluoncv.model_zoo import get_model
+
+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()

requirements.txt

@@ -1,13 +1,13 @@
-import argparse, time, os
+import argparse, 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.presets.ssd import transform_test
 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 model import ctx, detector, estimator
+from fps import FPS
 from angle import AngeleCal
 
 # 读取参数
@@ -17,38 +17,23 @@
 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()
-
 # 视频读取
+# 1是输入视频，2是示例视频
 cap1 = cv.VideoCapture(args.input)
 cap2 = cv.VideoCapture(args.demo)
 
 # 标准特征
 angeleCal = AngeleCal(args.data)
-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, img = transform_test(frame, short=512)
     x = x.as_in_context(ctx)
     class_IDs, scores, bounding_boxs = detector(x)
 
@@ -66,11 +51,9 @@
     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),
+        upperleft_txt=FPS.fps(), upperleft_txt_corner=(10,25),
         left_txt_list=results, canvas_name='pose')
-    fps_time = time.time()
     cv.imshow('demo', frame2)
 
     # ESC键退出