pytorch transforms image padding into a square

At present, I am experimenting with a batch of data. The length and width of each image are different. I tried transforms.Resize(640), which can scale the short side to 640 and scale the long side proportionally. This result will still lead to inconsistent input image size. If transforms.Resize((640,640)) is used, although the image becomes a 640 X 640 square, the internal object ratio is stretched or scaled. My goal is to turn the image into a square without changing the aspect ratio of the object in the original image.

The implementation code is as follows:

class SquarePad:
	def __call__(self, image):
		w, h = image.size
		max_wh = np.max([w, h])
		hp = int((max_wh - w) / 2)
		vp = int((max_wh - h) / 2)
		padding = (hp, vp, hp, vp)
		return F.pad(image, padding, 0, 'constant')

# Data augmentation and normalization for training
# Just normalization for validation
data_transforms = {
    'train': transforms.Compose([
        SquarePad(),
        transforms.Resize(640),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'val': transforms.Compose([
        SquarePad(),
        transforms.Resize(640),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}

SquarePad can padding the short side of the image to be consistent with the long side, and then resize the square to ensure that the proportion of internal objects will not change. How much size to resize can also be determined by yourself~