pytorch transforma el relleno de la imagen en un cuadrado

Actualmente, estoy experimentando con un lote de datos. La longitud y el ancho de cada imagen son diferentes. Intenté transforms.Resize(640), que puede escalar el lado corto a 640 y escalar el lado largo proporcionalmente. Este resultado aún conducirá a un tamaño de imagen de entrada inconsistente. Si se usa transforms.Resize((640,640)), aunque la imagen se convierte en un cuadrado de 640 X 640, la proporción del objeto interno se estira o escala. Mi objetivo es convertir la imagen en un cuadrado sin cambiar la relación de aspecto del objeto en la imagen original.

El código de implementación es el siguiente:

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 puede rellenar el lado corto de la imagen para que sea consistente con el lado largo y luego cambiar el tamaño del cuadrado para garantizar que la proporción de objetos internos no cambie. Usted mismo también puede determinar cuánto tamaño cambiar el tamaño ~