When using yolov7 pose,the size of the generated image is smaller than the original,how to fix?

Question

I referenced the code from this 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 )

The size of my original image zidane.jpg is 1280X720,the size of the image result.jpg is 970X576.

How can I not change the size of the generated image result.jpg?

Answer 1

Basically a deep learning model takes an input of a fixed size. Fortunately YOLO is by nature a full convolutional model and so should be able to take a flexible size in input (with predictions considering the same size). In your code the function letterbox resizes your image (see here the code). You can then try to change the size 960 to (1280, 720).