update

README.md

 # light pose
 人体关键点检测
 
+## 项目介绍    
+人体关键点检测  
+ ![video](https://codechina.csdn.net/EricLee/light_pose/-/raw/master/samples/s1.gif) 
+## 项目配置   
+* 作者开发环境：   
+* Python 3.7   
+* PyTorch >= 1.5.1  
+
+## 数据集     
+### coco2017 数据集  
+[数据集下载地址(百度网盘 Password:  )]()  
+
+## 预训练模型   
+* [预训练模型下载地址(百度网盘 Password: )]()   
+
+## 项目使用方法  
 * step 1: python prepare_train_labels.py  
 * step 2: python make_val_subset.py  
 * step 3: python train.py   

inference_video.py

@@ -181,12 +181,12 @@ def draw_one_pose(img,keypoints,color_x = [255, 0, 0]):
             cv2.line(img, (int(x_a), int(y_a)), (int(x_b), int(y_b)), color_x, 2)
 
 if __name__ == '__main__':
-    video_path = "./video/rw_7.mp4" # 加载视频
+    video_path = "./video/a.mp4" # 加载视频
     # video_path = 0 # 加载相机
-    model_path = "finetune_model/light_pose.pth"
+    model_path = "light_pose_checkpoints/light_pose-iter_1000.pth"
 
     model_pose = light_pose_model(model_path = model_path,heatmaps_thr = 0.08) # 定义模型推理类
-
+    print("load:{}".format(model_path))
     video_capture = cv2.VideoCapture(video_path)
 
     flag_write_video = True # 是否记录推理 demo 视频
@@ -201,7 +201,7 @@ if __name__ == '__main__':
             if flag_video_start == False  and flag_write_video:
                 loc_time = time.localtime()
                 str_time = time.strftime("%Y-%m-%d_%H-%M-%S", loc_time)
-                video_writer = cv2.VideoWriter("./demo/demo_{}.mp4".format(str_time), cv2.VideoWriter_fourcc(*"mp4v"), fps=25, frameSize=(int(im0.shape[1]), int(im0.shape[0])))
+                video_writer = cv2.VideoWriter("./demo/demo_{}.mp4".format(str_time), cv2.VideoWriter_fourcc(*"mp4v"), fps=24, frameSize=(int(im0.shape[1]), int(im0.shape[0])))
                 flag_video_start = True
 
             pose_dict = model_pose.predict(im0.copy())
@@ -224,6 +224,8 @@ if __name__ == '__main__':
 
             if cv2.waitKey(1) == 27:
                 break
+        else:
+            break
 
     cv2.destroyAllWindows()
     if flag_write_video: