Spaces:
Runtime error
Runtime error
| import cv2 | |
| import numpy as np | |
| import torch | |
| import torchvision | |
| class Equirectangular(): | |
| """ | |
| Random sample a panorama image into a perspective view | |
| take https://github.com/fuenwang/Equirec2Perspec/blob/master/Equirec2Perspec.py as a reference | |
| """ | |
| def __init__(self, width = 256, height = 256, FovX = 100, theta = [0, 0]): | |
| """ | |
| width: output image's width | |
| height: output image's height | |
| FovX: perspective camera FOV on x-axis (degree) | |
| theta: theta field where img's theta degree from | |
| """ | |
| self.theta = theta | |
| self.width = width | |
| self.height = height | |
| self.type = type | |
| #create x-axis coordinates and corresponding y-axis coordinates | |
| x = np.arange(width) | |
| y = np.arange(height) | |
| x, y = np.meshgrid(x, y) | |
| #create homogenerous coordinates | |
| z = np.ones_like(x) | |
| xyz = np.concatenate([x[..., None], y[..., None], z[..., None]], axis=-1) | |
| #translation matrix | |
| f = 0.5 * width * 1 / np.tan(np.radians(FovX/2)) | |
| # cx = (width - 1) / 2.0 | |
| # cy = (height - 1) / 2.0 | |
| cx = (width) / 2.0 | |
| cy = (height) / 2.0 | |
| K = np.array([ | |
| [f, 0, cx], | |
| [0, f, cy], | |
| [0, 0, 1], | |
| ], np.float32) | |
| K_inv = np.linalg.inv(K) | |
| xyz = xyz @ K_inv.T | |
| self.xyz = xyz ### self.xyz is the direction of the each ray in the camera space when camera is fixed | |
| def __call__(self, img1): | |
| batch = img1.shape[0] | |
| PHI, THETA = self.getRandomRotation(batch) | |
| y_axis = np.array([0.0, 1.0, 0.0], np.float32) | |
| x_axis = np.array([1.0, 0.0, 0.0], np.float32) | |
| #rotation matrix | |
| xy_grid = [] | |
| for i in range(batch): | |
| R1, _ = cv2.Rodrigues(y_axis * np.radians(PHI[i])) | |
| R2, _ = cv2.Rodrigues(np.dot(R1, x_axis) * np.radians(THETA[i])) | |
| R = R2 @ R1 | |
| #rotate | |
| xyz = self.xyz @ R.T ### ### xyz is the direction of the each ray in the camera space when camera is rotate | |
| norm = np.linalg.norm(xyz, axis=-1, keepdims=True) | |
| xyz_norm = xyz / norm | |
| #transfer to image coordinates | |
| xy = self.xyz2xy(xyz_norm) | |
| device = img1.device | |
| xy = torch.from_numpy(xy).to(device).unsqueeze(0) | |
| xy_grid.append(xy) | |
| xy = torch.cat(xy_grid,dim=0) | |
| #resample | |
| return xy | |
| def xyz2xy(self, xyz_norm): | |
| #normlize | |
| x = xyz_norm[..., 0] | |
| y = xyz_norm[..., 1] | |
| z = xyz_norm[..., 2] | |
| lon = np.arctan2(x, z) | |
| lat = np.arcsin(y) | |
| ### transfer to the lon and lat | |
| X = lon / (np.pi) | |
| Y = lat / (np.pi) * 2 | |
| xy = np.stack([X, Y], axis=-1) | |
| xy = xy.astype(np.float32) | |
| return xy | |
| def getRandomRotation(self,batch_size): | |
| # phi = np.random.rand(batch_size) * 360 -180 | |
| phi = np.random.randint(-180,180,batch_size) | |
| assert(self.theta[0]<self.theta[1]) | |
| theta = np.random.randint(self.theta[0],self.theta[1],batch_size) | |
| # theta = np.random.rand(batch_size)*(self.theta[1]-self.theta[0])-self.theta[0] | |
| return phi, theta | |
| if __name__=='__main__': | |
| # test demo | |
| e = Equirectangular(theta=[0., 40.],width = 64, height = 64,FovX = 100) | |
| img = cv2.imread('dataset/CVACT/streetview/__-DFIFxvZBCn1873qkqXA_grdView.png')[:,:,::-1]/255.0 | |
| img = img.transpose(2, 0, 1).astype(np.float32) | |
| img = torch.from_numpy(img).unsqueeze(0).repeat(10, 1, 1, 1) | |
| equ= e(img) | |
| # print(PHI, THETA) | |
| torchvision.utils.save_image(torch.nn.functional.grid_sample(img, equ.float(), align_corners = True)*0.99, 'test_30.png') |