pytorch 当使用yolov7 pose时,生成的图像的大小比原来的小,如何修复?

yqkkidmi  于 2022-11-09  发布在  其他
关注(0)|答案(1)|浏览(629)

我引用了这个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的大小?

z31licg0

z31licg01#

深度学习模型基本上采用固定大小的输入。幸运的是,YOLO本质上是一个全卷积模型,因此应该能够采用灵活的输入大小(预测时考虑相同的大小)。在代码中,函数letterbox会调整图像大小(见此处代码)。然后,您可以尝试将960的大小更改为(1280, 720)

相关问题