Spaces:
Runtime error
Runtime error
| #!/usr/bin/env python3 | |
| """ | |
| simple routines used throughout the project, placed here to avoid circular imports | |
| """ | |
| from PIL import Image | |
| import torch | |
| def flip_bottom_half_and_attach(sub_img): | |
| "takes one 256x256 and returns on 512x128 image with the bottom half reversed and attached on the right" | |
| h, w = sub_img.size | |
| new_img = Image.new(sub_img.mode, (w*2, h//2)) | |
| new_img.paste(sub_img.crop((0, 0, w, h//2)), (0, 0)) | |
| new_img.paste(sub_img.crop((0, h//2, w, h)).transpose(Image.FLIP_LEFT_RIGHT), (w, 0)) | |
| return new_img | |
| def square_to_rect_tensor(img_tensor:torch.Tensor): | |
| if len(img_tensor.shape) <= 4: | |
| img_tensor = img_tensor.unsqueeze(0) | |
| channels_last = img_tensor.shape[-1] == 3 | |
| if channels_last: | |
| img_tensor = img_tensor.permute(0, 3, 1, 2) | |
| channels_first = img_tensor.shape[1] == 3 | |
| b,c,w,h = img_tensor.shape | |
| new_img = torch.zeros((b,c,w*2,h//2), dtype=img_tensor.dtype).to(img_tensor.device) | |
| new_img[:,:,:w, :] = img_tensor[:,:,:w,:] | |
| new_img[:,:,w:, :] = torch.flip(img_tensor[:,:,w:, :], dims=[-2]) | |
| return new_img | |
| def square_to_rect(img): | |
| #"""just an alias for flip_bottom_half_and_attach""" | |
| if isinstance(img, torch.Tensor): | |
| return square_to_rect_tensor(img) | |
| return flip_bottom_half_and_attach(img) | |
| def rect_to_square(img): | |
| "takes a 512x128 image and returns a 256x256 image with the bottom half reversed" | |
| w, h = img.size | |
| new_img = Image.new(img.mode, (w//2, h*2)) | |
| new_img.paste(img.crop((0, 0, w//2, h)), (0, 0)) | |
| new_img.paste(img.crop((w//2, 0, w, h)).transpose(Image.FLIP_LEFT_RIGHT), (0, h)) | |
| return new_img | |