pom/utils.py

#!/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