Yolov8 pose 推理部署笔记
目录
左右手臂设置不同颜色:
相机推理代码:
部署 关于新版本torch的报错
左右手臂设置不同颜色:
import argparse
import copy
import os
import warnings
import cv2
import torch
import numpy as np
import time
import torch
import yaml
from torch.utils import datafrom nets import nn
from utils import util
from utils.dataset import Datasetwarnings.filterwarnings("ignore")# coding=utf-8
import sys
import os
current_dir = os.path.dirname(os.path.abspath(__file__))
os.chdir(current_dir)
print('current_dir', current_dir)
paths = [current_dir, current_dir+'/../']
paths.append(os.path.join(current_dir, 'src'))
for path in paths:sys.path.insert(0, path)os.environ['PYTHONPATH'] = (os.environ.get('PYTHONPATH', '') + ':' + path).strip(':')@torch.no_grad()
def demo(args):# ---------- 硬编码路径 ----------video_path = r"D:\data\huichui\tuiche\err2.mkv"output_path = r"res\result.mp4"weights_path = r"weights\best.pt"weights_path = r"weights\v8_m_pose.pt"input_size = 640 # YOLO 输入尺寸# ---------- 颜色和骨架 ----------palette = np.array([[255, 128, 0], [255, 153, 51], [255, 178, 102], [230, 230, 0], [255, 153, 255],[153, 204, 255], [255, 102, 255], [255, 51, 255], [102, 178, 255], [51, 153, 255],[255, 153, 153], [255, 102, 102], [255, 51, 51], [153, 255, 153], [102, 255, 102],[51, 255, 51], [0, 255, 0], [0, 0, 255], [255, 0, 0], [255, 255, 255]], dtype=np.uint8)skeleton = [[16, 14], [14, 12], [17, 15], [15, 13], [12, 13], [6, 12], [7, 13], [6, 7], [6, 8], [7, 9],[8, 10], [9, 11], [2, 3], [1, 2], [1, 3], [2, 4], [3, 5], [4, 6], [5, 7]]# 定义左右手臂的关键点索引left_arm_kpts = [5, 7, 9] # 左肩、左肘、左手腕right_arm_kpts = [6, 8, 10] # 右肩、右肘、右手腕# 定义左右手臂的骨骼连接索引left_arm_limbs = [8, 10] # 左肩-左肘, 左肘-左手腕right_arm_limbs = [9, 11] # 右肩-右肘, 右肘-右手腕# 设置左右手臂颜色left_arm_color = [0, 0, 255] # 红色 (BGR)right_arm_color = [255, 0, 0] # 绿色 (BGR)# 初始化关键点和骨骼颜色kpt_color = palette[[16, 16, 16, 16, 16, 0, 0, 0, 0, 0, 0, 9, 9, 9, 9, 9, 9]]limb_color = palette[[9, 9, 9, 9, 7, 7, 7, 0, 0, 0, 0, 0, 16, 16, 16, 16, 16, 16, 16]]# 为左右手臂设置不同颜色for kpt_idx in left_arm_kpts:kpt_color[kpt_idx] = left_arm_colorfor kpt_idx in right_arm_kpts:kpt_color[kpt_idx] = right_arm_colorfor limb_idx in left_arm_limbs:limb_color[limb_idx] = left_arm_colorfor limb_idx in right_arm_limbs:limb_color[limb_idx] = right_arm_color# ---------- 加载模型 ----------model = torch.load(weights_path, map_location='cuda')['model'].float()model.half().eval()stride = int(max(model.stride.cpu().numpy()))# ---------- 打开视频 ----------cap = cv2.VideoCapture(video_path)if not cap.isOpened():print(f"Cannot open video: {video_path}")return# ---------- 输出视频 ----------fourcc = cv2.VideoWriter_fourcc(*'mp4v')fps = int(cap.get(cv2.CAP_PROP_FPS))width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))writer = cv2.VideoWriter(output_path, fourcc, fps, (width, height))# ---------- 主循环 ----------while cap.isOpened():ret, frame = cap.read()if not ret:breakif np.prod(frame.shape[:2]) > 1000*1300:x_scale =np.sqrt( 1000*1200 / np.prod(frame.shape[:2]))frame = cv2.resize(frame, None, fx=x_scale, fy=x_scale, interpolation=cv2.INTER_AREA)img = frame.copy()h0, w0 = img.shape[:2]r = min(input_size / h0, input_size / w0)pad_w, pad_h = int(round(w0 * r)), int(round(h0 * r))img_resized = cv2.resize(img, (pad_w, pad_h), interpolation=cv2.INTER_LINEAR)top = (input_size - pad_h) // 2bottom = input_size - pad_h - topleft = (input_size - pad_w) // 2right = input_size - pad_w - leftimg_resized = cv2.copyMakeBorder(img_resized, top, bottom, left, right, cv2.BORDER_CONSTANT)img_resized = img_resized[:, :, ::-1].transpose(2, 0, 1) # BGR->RGB, HWC->CHWimg_resized = np.ascontiguousarray(img_resized)img_tensor = torch.from_numpy(img_resized).unsqueeze(0).cuda().half() / 255start = time.time()# ---------- 推理 ----------outputs = model(img_tensor)outputs = util.non_max_suppression(outputs, 0.5, 0.7, model.head.nc)print("耗时",time.time()-start)for output in outputs:output = output.clone()if len(output):box_output = output[:, :6]kps_output = output[:, 6:].view(len(output), *model.head.kpt_shape)else:continuer0 = min(img_tensor.shape[2] / h0, img_tensor.shape[3] / w0)# 计算填充量pad_x = (img_tensor.shape[3] - w0 * r0) / 2 # 水平填充pad_y = (img_tensor.shape[2] - h0 * r0) / 2 # 垂直填充# 转换边界框坐标box_output[:, [0, 2]] = (box_output[:, [0, 2]] - pad_x) / r0box_output[:, [1, 3]] = (box_output[:, [1, 3]] - pad_y) / r0box_output[:, 0].clamp_(0, w0)box_output[:, 1].clamp_(0, h0)box_output[:, 2].clamp_(0, w0)box_output[:, 3].clamp_(0, h0)# 转换关键点坐标kps_output[..., 0] = (kps_output[..., 0] - pad_x) / r0kps_output[..., 1] = (kps_output[..., 1] - pad_y) / r0kps_output[..., 0].clamp_(0, w0)kps_output[..., 1].clamp_(0, h0)# ---------- 画框 ----------for box in box_output:x1, y1, x2, y2, score, idx = box.cpu().numpy()cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)# ---------- 画关键点和骨架 ----------for kpt in kps_output:for i, k in enumerate(kpt):if len(k) == 3 and k[2] < 0.5:continuecv2.circle(frame, (int(k[0]), int(k[1])), 5, kpt_color[i].tolist(), -1, lineType=cv2.LINE_AA)for i, sk in enumerate(skeleton):pos1 = (int(kpt[sk[0]-1, 0]), int(kpt[sk[0]-1, 1]))pos2 = (int(kpt[sk[1]-1, 0]), int(kpt[sk[1]-1, 1]))if kpt.shape[-1] == 3 and (kpt[sk[0]-1,2]<0.5 or kpt[sk[1]-1,2]<0.5):continuecv2.line(frame, pos1, pos2, limb_color[i].tolist(), 2, lineType=cv2.LINE_AA)# ---------- 写入视频 ----------writer.write(frame)cv2.imshow('Pose Estimation', frame)if cv2.waitKey(10) & 0xFF == ord('q'):breakcap.release()writer.release()cv2.destroyAllWindows()
def profile(args, params):model = nn.yolo_v8_n(len(params['names']))shape = (1, 3, args.input_size, args.input_size)model.eval()model(torch.zeros(shape))params = sum(p.numel() for p in model.parameters())if args.local_rank == 0:print(f'Number of parameters: {int(params)}')def main():parser = argparse.ArgumentParser()parser.add_argument('--input-size', default=640, type=int)parser.add_argument('--local_rank', default=0, type=int)parser.add_argument('--epochs', default=1000, type=int)parser.add_argument('--demo',default=True, action='store_true')args = parser.parse_args()util.setup_seed()util.setup_multi_processes()with open(r'./utils\args.yaml', errors='ignore') as f:params = yaml.safe_load(f)profile(args, params)demo(args)if __name__ == "__main__":main()
相机推理代码:
import argparse
import copy
import os
import warnings
import cv2
import torch
import numpy as np
import time
import torch
import yaml
from torch.utils import datafrom nets import nn
from utils import util
from utils.dataset import Datasetwarnings.filterwarnings("ignore")# coding=utf-8
import sys
import os
current_dir = os.path.dirname(os.path.abspath(__file__))
os.chdir(current_dir)
print('current_dir', current_dir)
paths = [current_dir, current_dir+'/../']
paths.append(os.path.join(current_dir, 'src'))
for path in paths:sys.path.insert(0, path)os.environ['PYTHONPATH'] = (os.environ.get('PYTHONPATH', '') + ':' + path).strip(':')@torch.no_grad()
def demo(args):# ---------- 硬编码路径 ----------video_path = r"D:\data\huichui\tuiche\err2.mkv"output_path = r"res\result.mp4"weights_path = r"weights\best.pt"weights_path = r"weights\v8_m_pose.pt"input_size = 640 # YOLO 输入尺寸# ---------- 颜色和骨架 ----------palette = np.array([[255, 128, 0], [255, 153, 51], [255, 178, 102], [230, 230, 0], [255, 153, 255],[153, 204, 255], [255, 102, 255], [255, 51, 255], [102, 178, 255], [51, 153, 255],[255, 153, 153], [255, 102, 102], [255, 51, 51], [153, 255, 153], [102, 255, 102],[51, 255, 51], [0, 255, 0], [0, 0, 255], [255, 0, 0], [255, 255, 255]], dtype=np.uint8)skeleton = [[16, 14], [14, 12], [17, 15], [15, 13], [12, 13], [6, 12], [7, 13], [6, 7], [6, 8], [7, 9],[8, 10], [9, 11], [2, 3], [1, 2], [1, 3], [2, 4], [3, 5], [4, 6], [5, 7]]kpt_color = palette[[16, 16, 16, 16, 16, 0, 0, 0, 0, 0, 0, 9, 9, 9, 9, 9, 9]]limb_color = palette[[9, 9, 9, 9, 7, 7, 7, 0, 0, 0, 0, 0, 16, 16, 16, 16, 16, 16, 16]]# ---------- 加载模型 ----------model = torch.load(weights_path, map_location='cuda')['model'].float()model.half().eval()stride = int(max(model.stride.cpu().numpy()))# ---------- 打开视频 ----------cap = cv2.VideoCapture(video_path)if not cap.isOpened():print(f"Cannot open video: {video_path}")return# ---------- 输出视频 ----------fourcc = cv2.VideoWriter_fourcc(*'mp4v')fps = int(cap.get(cv2.CAP_PROP_FPS))width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))writer = cv2.VideoWriter(output_path, fourcc, fps, (width, height))# ---------- 主循环 ----------while cap.isOpened():ret, frame = cap.read()if not ret:breakif np.prod(frame.shape[:2]) > 1000*1300:x_scale =np.sqrt( 1000*1200 / np.prod(frame.shape[:2]))frame = cv2.resize(frame, None, fx=x_scale, fy=x_scale, interpolation=cv2.INTER_AREA)img = frame.copy()h0, w0 = img.shape[:2]r = min(input_size / h0, input_size / w0)pad_w, pad_h = int(round(w0 * r)), int(round(h0 * r))img_resized = cv2.resize(img, (pad_w, pad_h), interpolation=cv2.INTER_LINEAR)top = (input_size - pad_h) // 2bottom = input_size - pad_h - topleft = (input_size - pad_w) // 2right = input_size - pad_w - leftimg_resized = cv2.copyMakeBorder(img_resized, top, bottom, left, right, cv2.BORDER_CONSTANT)img_resized = img_resized[:, :, ::-1].transpose(2, 0, 1) # BGR->RGB, HWC->CHWimg_resized = np.ascontiguousarray(img_resized)img_tensor = torch.from_numpy(img_resized).unsqueeze(0).cuda().half() / 255start = time.time()# ---------- 推理 ----------outputs = model(img_tensor)outputs = util.non_max_suppression(outputs, 0.5, 0.7, model.head.nc)print("耗时",time.time()-start)for output in outputs:output = output.clone()if len(output):box_output = output[:, :6]kps_output = output[:, 6:].view(len(output), *model.head.kpt_shape)else:continuer0 = min(img_tensor.shape[2] / h0, img_tensor.shape[3] / w0)box_output[:, [0, 2]] -= (img_tensor.shape[3] - w0 * r0) / 2box_output[:, [1, 3]] -= (img_tensor.shape[2] - h0 * r0) / 2box_output[:, :4] /= r0box_output[:, 0].clamp_(0, w0)box_output[:, 1].clamp_(0, h0)box_output[:, 2].clamp_(0, w0)box_output[:, 3].clamp_(0, h0)kps_output[..., 0] -= (img_tensor.shape[3] - w0 * r0) / 2kps_output[..., 1] -= (img_tensor.shape[2] - h0 * r0) / 2kps_output[..., 0] /= r0kps_output[..., 1] /= r0kps_output[..., 0].clamp_(0, w0)kps_output[..., 1].clamp_(0, h0)# ---------- 画框 ----------for box in box_output:x1, y1, x2, y2, score, idx = box.cpu().numpy()cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)# ---------- 画关键点和骨架 ----------for kpt in kps_output:for i, k in enumerate(kpt):if len(k) == 3 and k[2] < 0.5:continuecv2.circle(frame, (int(k[0]), int(k[1])), 5, kpt_color[i].tolist(), -1, lineType=cv2.LINE_AA)for i, sk in enumerate(skeleton):pos1 = (int(kpt[sk[0]-1, 0]), int(kpt[sk[0]-1, 1]))pos2 = (int(kpt[sk[1]-1, 0]), int(kpt[sk[1]-1, 1]))if kpt.shape[-1] == 3 and (kpt[sk[0]-1,2]<0.5 or kpt[sk[1]-1,2]<0.5):continuecv2.line(frame, pos1, pos2, limb_color[i].tolist(), 2, lineType=cv2.LINE_AA)# ---------- 写入视频 ----------writer.write(frame)cv2.imshow('Pose Estimation', frame)if cv2.waitKey(10) & 0xFF == ord('q'):breakcap.release()writer.release()cv2.destroyAllWindows()def profile(args, params):model = nn.yolo_v8_n(len(params['names']))shape = (1, 3, args.input_size, args.input_size)model.eval()model(torch.zeros(shape))params = sum(p.numel() for p in model.parameters())if args.local_rank == 0:print(f'Number of parameters: {int(params)}')def main():parser = argparse.ArgumentParser()parser.add_argument('--input-size', default=640, type=int)parser.add_argument('--batch-size', default=32, type=int)parser.add_argument('--local_rank', default=0, type=int)parser.add_argument('--epochs', default=1000, type=int)parser.add_argument('--train', action='store_true')parser.add_argument('--test', action='store_true')parser.add_argument('--demo',default=True, action='store_true')args = parser.parse_args()args.local_rank = int(os.getenv('LOCAL_RANK', 0))args.world_size = int(os.getenv('WORLD_SIZE', 1))if args.world_size > 1:torch.cuda.set_device(device=args.local_rank)torch.distributed.init_process_group(backend='nccl', init_method='env://')if args.local_rank == 0:if not os.path.exists('weights'):os.makedirs('weights')util.setup_seed()util.setup_multi_processes()with open(r'./utils\args.yaml', errors='ignore') as f:params = yaml.safe_load(f)profile(args, params)demo(args)if __name__ == "__main__":main()
部署 关于新版本torch的报错
File "E:\suanfa_diaoyan\YOLOv8-pose-master\YOLOv8-pose-master\nets\nn.py", line 197, in detect_boxself.anchors, self.strides = (x.transpose(0, 1) for x in make_anchors(x_det, self.stride, 0.5))File "E:\suanfa_diaoyan\YOLOv8-pose-master\YOLOv8-pose-master\utils\util.py", line 78, in make_anchorssy, sx = torch.meshgrid(sy, sx)File "E:\miniconda\envs\cuda11.8\lib\site-packages\torch\functional.py", line 489, in meshgridreturn _meshgrid(*tensors, indexing=indexing)File "E:\miniconda\envs\cuda11.8\lib\site-packages\torch\functional.py", line 504, in _meshgridreturn _VF.meshgrid(tensors, **kwargs) # type: ignore[attr-defined]
修改方案
注释报错参数
def forward(self, input: Tensor) -> Tensor:# return F.interpolate(input, self.size, self.scale_factor, self.mode, self.align_corners,# recompute_scale_factor=self.recompute_scale_factor)return F.interpolate(input, self.size, self.scale_factor, self.mode, self.align_corners)