Fix transpose with non-square images

This commit is contained in:
missionfloyd 2023-04-09 19:19:14 -06:00 committed by GitHub
parent 97c7a402d7
commit d414157701
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -232,6 +232,8 @@ class Transpose:
def transpose(self, image: torch.Tensor, method: str):
batch_size, height, width, _ = image.shape
result = torch.zeros_like(image)
if height != width and method in ("Transpose", "Transverse", "Rotate 90°", "Rotate 270°"):
result = torch.permute(result, (0, 2, 1, 3))
methods = {
"Flip horizontal": Image.Transpose.FLIP_LEFT_RIGHT,