pytorch transforms将图像padding成正方形

目前在实验一批数据，各个图像长宽不一，尝试了transforms.Resize(640)，可以将短边缩放到640，长边按比例缩放。这个结果仍然会导致输入图像size不一致。如果用transforms.Resize((640,640))的话，虽然图像都变成640 X 640的正方形，但是内部物体比例却做了拉伸或缩放。我的目标是在不要改变原图中物体的长宽比例的前提下，将图像变成一个正方形。

实现代码如下：

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可以将图像短的一边padding到与长边一致，再将该正方形做一个resize，就可以保证内部物体比例不会发生改变了。resize到多少大小，也完全可以由自己确定咯~