def __init__(self, model_path):
self.model = bodypose_model()
if torch.cuda.is_available():
self.model = self.model.cuda()
model_dict = util.transfer(self.model, torch.load(model_path))
self.model.load_state_dict(model_dict)
self.model.eval()
def __init__(self, model_path):
self.model = bodypose_model() # 创建模型
if torch.cuda.is_available():
self.model = self.model.cuda()
model_dict = util.transfer(self.model,
torch.load(model_path)) # 参数格式转换
self.model.load_state_dict(model_dict) # 导入参数
self.model.eval() # 使网络中bn层和dropout层失效。
type=float)
parser.add_argument(
"-show",
nargs='+',
help="types to show: -1 shows the skeletons, or idx for specific part",
default=(-1, 2),
type=int)
parser.add_argument("-thre",
help="threshold for heatmap part",
default=0.2,
type=str)
args = parser.parse_args()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if args.backbone == 'CMU':
model = bodypose_model().to(device)
model.load_state_dict(
torch.load('weights/bodypose_model',
map_location=lambda storage, loc: storage))
elif args.backbone == 'Mobilenet':
model = PoseEstimationWithMobileNet().to(device)
model.load_state_dict(
torch.load('weights/MobileNet_bodypose_model',
map_location=lambda storage, loc: storage))
model.eval()
print('openpose {} model is successfully loaded...'.format(args.backbone))
cap = cv2.VideoCapture(args.video)
_, frame = cap.read()