我引用了这个link中的代码。
import torch
import cv2
from torchvision import transforms
import numpy as np
from utils.datasets import letterbox
from utils.general import non_max_suppression_kpt
from utils.plots import output_to_keypoint, plot_skeleton_kpts
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
weigths = torch.load('weights/yolov7-w6-pose.pt', map_location=device)
model = weigths['model']
_ = model.float().eval()
if torch.cuda.is_available():
model.half().to(device)
image = cv2.imread('image/zidane.jpg')
image = letterbox(image, 960, stride=64, auto=True)[0]
image_ = image.copy()
image = transforms.ToTensor()(image)
image = torch.tensor(np.array([image.numpy()]))
if torch.cuda.is_available():
image = image.half().to(device)
output, _ = model(image)
output = non_max_suppression_kpt(output, 0.25, 0.65, nc=model.yaml['nc'], nkpt=model.yaml['nkpt'], kpt_label=True)
with torch.no_grad():
output = output_to_keypoint(output)
nimg = image[0].permute(1, 2, 0) * 255
nimg = nimg.cpu().numpy().astype(np.uint8)
nimg = cv2.cvtColor(nimg, cv2.COLOR_RGB2BGR)
for idx in range(output.shape[0]):
plot_skeleton_kpts(nimg, output[idx, 7:].T, 3)
nimg2 = cv2.cvtColor(nimg, cv2.COLOR_BGR2RGB)
cv2.imwrite('E:/result.jpg',nimg2 )
我的原始图像zidane.jpg
的大小是1280X720,图像result.jpg
的大小是970X576。
如何不更改生成的图像result.jpg
的大小?
1条答案
按热度按时间z31licg01#
深度学习模型基本上采用固定大小的输入。幸运的是,YOLO本质上是一个全卷积模型,因此应该能够采用灵活的输入大小(预测时考虑相同的大小)。在代码中,函数
letterbox
会调整图像大小(见此处代码)。然后,您可以尝试将960
的大小更改为(1280, 720)
。