''' Author: Manu Gond (manu.gond@miun.se) Date: Nov-15-2022 Objective: Accumulation of some general functions which I use daily in my code realted to image relasted task. The function names and parameters are self explanetory. Requirements: Installed python libraries which have been imported. ''' import torch from torchvision.utils import save_image from torchvision.transforms import transforms import torchmetrics import cv2 import numpy as np from PIL import Image import utils #======================= Read and Write =====================# def readImage(location): image = Image.open(location).convert("RGB") return image def writeImage(image, location): image.save(location) def writeTensorImage(image, filename): save_image(image, filename) def removeChannel(sourceLocation, targetLocation): img = readImage(sourceLocation) writeImage(img, targetLocation) def getImageTransform(width, height): transform = transforms.Compose([transforms.Resize((height,width)), transforms.ToTensor()]) return transform def convertTensor(image): transform = getImageTransform(image.size[0], image.size[1]) image = transform(image) return image #=================== 360 Images =======================# def rotateERP180(image): ''' :param image: PIL Image :return: BxHxW Torch Tensor Image ''' W = image.size[0] H = image.size[1] transform = getImageTransform(W, H) image = transform(image) image1 = image[:, :, 0:(W//2)] image2 = image[:, :, (W//2):W] image3 = torch.zeros(image.size()) image3[:, :, 0:(W//2)] = image2 image3[:, :, (W//2):W] = image1 return image3 def convertERP2Cube(e_img, face_w=256, mode='bilinear', cube_format='dice'): ''' e_img: ndarray in shape of [H, W, *] face_w: int, the length of each face of the cubemap ''' assert len(e_img.shape) == 3 h, w = e_img.shape[:2] if mode == 'bilinear': order = 1 elif mode == 'nearest': order = 0 else: raise NotImplementedError('unknown mode') xyz = utils.xyzcube(face_w) uv = utils.xyz2uv(xyz) coor_xy = utils.uv2coor(uv, h, w) cubemap = np.stack([ utils.sample_equirec(e_img[..., i], coor_xy, order=order) for i in range(e_img.shape[2]) ], axis=-1) if cube_format == 'horizon': pass elif cube_format == 'list': cubemap = utils.cube_h2list(cubemap) elif cube_format == 'dict': cubemap = utils.cube_h2dict(cubemap) elif cube_format == 'dice': cubemap = utils.cube_h2dice(cubemap) else: raise NotImplementedError() return cubemap def convertCube2ERP(cubemap, h, w, mode='bilinear', cube_format='dice'): if mode == 'bilinear': order = 1 elif mode == 'nearest': order = 0 else: raise NotImplementedError('unknown mode') if cube_format == 'horizon': pass elif cube_format == 'list': cubemap = utils.cube_list2h(cubemap) elif cube_format == 'dict': cubemap = utils.cube_dict2h(cubemap) elif cube_format == 'dice': cubemap = utils.cube_dice2h(cubemap) else: raise NotImplementedError('unknown cube_format') assert len(cubemap.shape) == 3 assert cubemap.shape[0] * 6 == cubemap.shape[1] assert w % 8 == 0 face_w = cubemap.shape[0] uv = utils.equirect_uvgrid(h, w) u, v = np.split(uv, 2, axis=-1) u = u[..., 0] v = v[..., 0] cube_faces = np.stack(np.split(cubemap, 6, 1), 0) # Get face id to each pixel: 0F 1R 2B 3L 4U 5D tp = utils.equirect_facetype(h, w) coor_x = np.zeros((h, w)) coor_y = np.zeros((h, w)) for i in range(4): mask = (tp == i) coor_x[mask] = 0.5 * np.tan(u[mask] - np.pi * i / 2) coor_y[mask] = -0.5 * np.tan(v[mask]) / np.cos(u[mask] - np.pi * i / 2) mask = (tp == 4) c = 0.5 * np.tan(np.pi / 2 - v[mask]) coor_x[mask] = c * np.sin(u[mask]) coor_y[mask] = c * np.cos(u[mask]) mask = (tp == 5) c = 0.5 * np.tan(np.pi / 2 - np.abs(v[mask])) coor_x[mask] = c * np.sin(u[mask]) coor_y[mask] = -c * np.cos(u[mask]) # Final renormalize coor_x = (np.clip(coor_x, -0.5, 0.5) + 0.5) * face_w coor_y = (np.clip(coor_y, -0.5, 0.5) + 0.5) * face_w equirec = np.stack([ utils.sample_cubefaces(cube_faces[..., i], tp, coor_y, coor_x, order=order) for i in range(cube_faces.shape[3]) ], axis=-1) return equirec def convertCube2Slices(image): ''' :param image: Image numpy array :return: List of Torch Tensors, CxHxW ''' image = convertTensor(image) C, H, W = image.size() #print(C,H,W) top = torch.zeros((C,W//4,W//4)) left = torch.zeros(top.size()) front = torch.zeros(top.size()) right = torch.zeros(top.size()) back = torch.zeros(top.size()) bottom = torch.zeros(top.size()) top = image[:, 0:H//3, (W//4):(W//4)*2] left = image[:, (H//3):(H//3)*2, 0:W//4] front = image[:, (H//3):(H//3)*2, (W//4):(W//4)*2] right = image[:, (H//3):(H//3)*2, (W//4)*2:(W//4)*3] back = image[:, (H // 3):(H // 3) * 2, (W // 4) * 3:] bottom = image[:, (H//3)*2:, (W//4):(W//4)*2] ''' save_image(top, 'top.png') save_image(left, 'left.png') save_image(front, 'front.png') save_image(right, 'right.png') save_image(back, 'back.png') save_image(bottom, 'bottom.png') ''' return [top, left, front, right, back, bottom] def convertSlicesToCube(imageList): ''' top = convertTensor(readImage(imageList[0])) left = convertTensor(readImage(imageList[1])) front = convertTensor(readImage(imageList[2])) right = convertTensor(readImage(imageList[3])) back = convertTensor(readImage(imageList[4])) bottom = convertTensor(readImage(imageList[5])) ''' top = imageList[0] left = imageList[1] front = imageList[2] right = imageList[3] back = imageList[4] bottom = imageList[5] C, H, W = 3, top.size()[1]*3, top.size()[2]*4 cube = torch.zeros((C, H, W)) cube[:, 0:H//3, (W//4):(W//4)*2] = top cube[:, (H // 3):(H // 3) * 2, 0:W // 4] = left cube[:, (H // 3):(H // 3) * 2, (W // 4):(W // 4) * 2] = front cube[:, (H // 3):(H // 3) * 2, (W // 4) * 2:(W // 4) * 3] = right cube[:, (H // 3):(H // 3) * 2, (W // 4) * 3:] = back cube[:, (H // 3) * 2:, (W // 4):(W // 4) * 2] = bottom return cube #=================== Quality Measures =======================# ''' Predicted Shape : BxCxHxW Original Shape : BxCxHxW Data Type: Torch Tensor ''' def getSSIM(predicted, original): SSIM = torchmetrics.StructuralSimilarityIndexMeasure() return SSIM(predicted, original).item() def getPSNR(predicted, original): PSNR = torchmetrics.PeakSignalNoiseRatio() return PSNR(predicted, original).item() def getMSE(predicted, original): MSE = torchmetrics.MeanSquaredError() return MSE(predicted, original).item() def getMAE(predicted, original): MAE = torchmetrics.MeanAbsoluteError() return MAE(predicted, original).item() if __name__ == "__main__": ''' img = readImage("31_image_0_0.png") img = convertERP2Cube(e_img=np.asarray(img), face_w=256) img = Image.fromarray(img.astype('uint8'),'RGB') convertCube2Slices(img) ''' #image = convertSlicesToCube(["top.png", "left.png", "front.png", "right.png", "back.png", "bottom.png"]) #writeTensorImage(image,'this.png') ''' writeImage(img, 'cube.png') img = readImage('cube.png') img = convertCube2ERP(np.asarray(img),512,1024) img = Image.fromarray(img.astype('uint8'),'RGB') writeImage(img, 'cubeERP.png') img1 = readImage("31_image_0_0.png") img2 = readImage("cubeERP.png") img1 = convertTensor(img1) img2 = convertTensor(img2) print(getSSIM(img1.unsqueeze(0), img2.unsqueeze(0))) ''' #img = rotateERP180(img) #writeTensorImage(img, 'rotated_image.png') #img = convertTensor(img) #print(getMAE(img.unsqueeze(0),img.unsqueeze(0)))