Skip to content

GitLab

Commit 5d1d329d authored by Nianchen Deng's avatar Nianchen Deng
Browse files

sync

parent f6604bd2
Hide whitespace changes
Inline Side-by-side
my/color_mode.py deleted 100644 → 0
View file @ f6604bd2
RGB = 0
GRAY = 1
YCbCr = 2
def to_str(color_mode):
return "gray" if color_mode == GRAY \
else ("ybr" if color_mode == YCbCr
else "rgb")
def from_str(color_str):
return GRAY if color_str == 'gray' \
else (YCbCr if color_str == 'ybr'
else RGB)
my/conf.py deleted 100644 → 0
View file @ f6604bd2
import torch
import numpy as np
from . import util
class Conf(object):
def __init__(self):
self.pupil_size = 0.02
self.retinal_res = torch.tensor([ 320, 320 ])
self.layer_res = torch.tensor([ 320, 320 ])
self.layer_hfov = 90 # layers' horizontal FOV
self.eye_hfov = 80 # eye's horizontal FOV (ignored in foveated rendering)
self.eye_enable_fovea = False # enable foveated rendering
self.eye_fovea_angles = [ 40, 80 ] # eye's foveation layers' angles
self.eye_fovea_downsamples = [ 1, 2 ] # eye's foveation layers' downsamples
self.d_layer = [ 1, 3 ] # layers' distance
self.eye_fovea_blend = [ self._GenFoveaLayerBlend(0) ]
# blend maps of fovea layers
self.light_field_dim = 5
def GetNLayers(self):
return len(self.d_layer)
def GetLayerSize(self, i):
w = util.Fov2Length(self.layer_hfov)
h = w * self.layer_res[0] / self.layer_res[1]
return torch.tensor([ h, w ]) * self.d_layer[i]
def GetPixelSizeOfLayer(self, i):
'''
Get pixel size of layer i
'''
return util.Fov2Length(self.layer_hfov) * self.d_layer[i] / self.layer_res[0]
def GetEyeViewportSize(self):
fov = self.eye_fovea_angles[-1] if self.eye_enable_fovea else self.eye_hfov
w = util.Fov2Length(fov)
h = w * self.retinal_res[0] / self.retinal_res[1]
return torch.tensor([ h, w ])
def GetRegionOfFoveaLayer(self, i):
'''
Get region of fovea layer i in retinal image
Returns
--------
slice object stores the start and end of region
'''
roi_size = int(np.ceil(self.retinal_res[0] * self.eye_fovea_angles[i] / self.eye_fovea_angles[-1]))
roi_offset = int((self.retinal_res[0] - roi_size) / 2)
return slice(roi_offset, roi_offset + roi_size)
def _GenFoveaLayerBlend(self, i):
'''
Generate blend map for fovea layer i
Parameters
--------
i - index of fovea layer
Returns
--------
H[i] x W[i], blend map
'''
region = self.GetRegionOfFoveaLayer(i)
width = region.stop - region.start
R = width / 2
p = util.MeshGrid([ width, width ])
r = torch.linalg.norm(p - R, 2, dim=2, keepdim=False)
return util.SmoothStep(R, R * 0.6, r)
my/flow.py deleted 100644 → 0
View file @ f6604bd2
import matplotlib.pyplot as plt
import torch
import util
import numpy as np
def FlowMap(b_last_frame, b_map):
'''
Map images using the flow data.
Parameters
--------
b_last_frame - B x 3 x H x W tensor, batch of images
b_map - B x H x W x 2, batch of map data records pixel coords in last frames
Returns
--------
B x 3 x H x W tensor, batch of images mapped by flow data
'''
return torch.nn.functional.grid_sample(b_last_frame, b_map, align_corners=False)
class Flow(object):
'''
Class representating optical flow
Properties
--------
b_data - B x H x W x 2, batch of flow data
b_map - B x H x W x 2, batch of map data records pixel coords in last frames
b_invalid_mask - B x H x W, batch of masks, indicate invalid elements in corresponding flow data
'''
def Load(paths):
'''
Create a Flow instance using a batch of encoded data images loaded from paths
Parameters
--------
paths - list of encoded data image paths
Returns
--------
Flow instance
'''
b_encoded_image = util.ReadImageTensor(paths, rgb_only=False, permute=False, batch_dim=True)
return Flow(b_encoded_image)
def __init__(self, b_encoded_image):
'''
Initialize a Flow instance from a batch of encoded data images
Parameters
--------
b_encoded_image - batch of encoded data images
'''
b_encoded_image = b_encoded_image.mul(255)
# print("b_encoded_image:",b_encoded_image.shape)
self.b_invalid_mask = (b_encoded_image[:, :, :, 0] == 255)
self.b_data = (b_encoded_image[:, :, :, 0:2] / 254 + b_encoded_image[:, :, :, 2:4] - 127) / 127
self.b_data[:, :, :, 1] = -self.b_data[:, :, :, 1]
D = self.b_data.size()
grid = util.MeshGrid((D[1], D[2]), True)
self.b_map = (grid - self.b_data - 0.5) * 2
self.b_map[self.b_invalid_mask] = torch.tensor([ -2.0, -2.0 ])
def getMap(self):
return self.b_map
def Visualize(self, scale_factor = 1):
'''
Visualize the flow data by "color wheel".
Parameters
--------
scale_factor - scale factor of flow data to visualize, default is 1
Returns
--------
B x 3 x H x W tensor, visualization of flow data
'''
try:
Flow.b_color_wheel
except AttributeError:
Flow.b_color_wheel = util.ReadImageTensor('color_wheel.png')
return torch.nn.functional.grid_sample(Flow.b_color_wheel.expand(self.b_data.size()[0], -1, -1, -1),
(self.b_data * scale_factor), align_corners=False)
\ No newline at end of file
my/imgio.py deleted 100644 → 0
View file @ f6604bd2
from typing import List, NoReturn
import torch
from torchvision.transforms.functional import convert_image_dtype
from torchvision.io.image import read_image
from torchvision.utils import save_image
def ReadImages(*args, paths: List[str] = None, dtype=torch.float) -> torch.Tensor:
raise NotImplementedError('The method has bug. Use util.ReadImageTensor instead')
if not paths:
paths = args
images = torch.stack([read_image(path) for path in paths], dim=0)
return convert_image_dtype(images, dtype)
def SaveImages(input, *args, paths: List[str] = None) -> NoReturn:
raise NotImplementedError('The method has bug. Use util.WriteImageTensor instead')
if not paths:
paths = args
if input.size(0) != len(paths):
raise ValueError('batch size of input images is not same as length of paths')
for i, path in enumerate(range(paths)):
save_image(input[i], path)
\ No newline at end of file
my/netio.py deleted 100644 → 0
View file @ f6604bd2
from typing import Mapping
import torch
import numpy as np
def PrintNet(net):
model_parameters = filter(lambda p: p.requires_grad, net.parameters())
params = sum([np.prod(p.size()) for p in model_parameters])
print("%d" % params)
def LoadNet(path, model, solver=None, discriminator=None):
print('Load net from %s ...' % path)
whole_dict: Mapping = torch.load(path)
model.load_state_dict(whole_dict['model'])
if solver:
solver.load_state_dict(whole_dict['solver'])
if discriminator:
discriminator.load_state_dict(whole_dict['discriminator'])
return whole_dict['iters'] if 'iters' in whole_dict else 0
def SaveNet(path, model, solver=None, discriminator=None, iters=None):
print('Saving net to %s ...' % path)
whole_dict = {
'model': model.state_dict()
}
if solver:
whole_dict.update({'solver': solver.state_dict()})
if discriminator:
whole_dict.update({'discriminator': discriminator.state_dict()})
if iters:
whole_dict.update({'iters': iters})
torch.save(whole_dict, path)
\ No newline at end of file
my/sample_in_pupil.py deleted 100644 → 0
View file @ f6604bd2
import torch
import numpy as np
def RandomGen(pupil_size: float, n_samples: int) -> torch.Tensor:
"""
Random sample n_samples positions in pupil region
:param pupil_size: multi-layers' parameters configuration
:param n_samples: number of samples to generate
:return: n_samples x 3, with 3D sample position in each row
"""
samples = torch.empty(n_samples, 3)
i = 0
while i < n_samples:
s = (torch.rand(2) - 0.5) * pupil_size
if np.linalg.norm(s) > pupil_size / 2.:
continue
samples[i, :] = [s[0], s[1], 0]
i += 1
return samples
def CircleGen(pupil_size: float, circles: int) -> torch.Tensor:
"""
Sample positions on circles in pupil region
:param pupil_size: diameter of pupil
:param circles: number of circles to sample
:return: M x 3, with 3D sample position in each row
"""
samples = torch.zeros(1, 3)
for i in range(1, circles):
r = pupil_size / 2. / (circles - 1) * i
n = 4 * i
for j in range(0, n):
angle = 2 * np.pi / n * j
samples = torch.cat([samples, torch.tensor([[r * np.cos(angle), r * np.sin(angle), 0]])], 0)
return samples
my/util.py deleted 100644 → 0
View file @ f6604bd2
from typing import List, Tuple, Union
import os
import math
import torch
import torchvision
import torchvision.transforms.functional as trans_func
import glm
import csv
import numpy as np
import matplotlib.pyplot as plt
from torch.types import Number
from torchvision.utils import save_image
gvec_type = [glm.dvec1, glm.dvec2, glm.dvec3, glm.dvec4]
gmat_type = [[glm.dmat2, glm.dmat2x3, glm.dmat2x4],
[glm.dmat3x2, glm.dmat3, glm.dmat3x4],
[glm.dmat4x2, glm.dmat4x3, glm.dmat4]]
def Fov2Length(angle):
return math.tan(math.radians(angle) / 2) * 2
def SmoothStep(x0, x1, x):
y = torch.clamp((x - x0) / (x1 - x0), 0, 1)
return y * y * (3 - 2 * y)
def MatImg2Tensor(img, permute=True, batch_dim=True):
batch_input = len(img.shape) == 4
if permute:
t = torch.from_numpy(np.transpose(img,
[0, 3, 1, 2] if batch_input else [2, 0, 1]))
else:
t = torch.from_numpy(img)
if not batch_input and batch_dim:
t = t.unsqueeze(0)
return t
def MatImg2Numpy(img, permute=True, batch_dim=True):
batch_input = len(img.shape) == 4
if permute:
t = np.transpose(img, [0, 3, 1, 2] if batch_input else [2, 0, 1])
else:
t = img
if not batch_input and batch_dim:
t = t.unsqueeze(0)
return t
def Tensor2MatImg(t: torch.Tensor) -> np.ndarray:
"""
Convert image tensor to numpy ndarray suitable for matplotlib
:param t: 2D (HW), 3D (CHW/HWC) or 4D (BCHW/BHWC) tensor
:return: numpy ndarray (...C), with channel transposed to the last dim
"""
img = t.squeeze().cpu().detach().numpy()
if len(img.shape) == 2: # Single channel image
return img
batch_input = len(img.shape) == 4
if t.size()[batch_input] <= 4:
return np.transpose(img, [0, 2, 3, 1] if batch_input else [1, 2, 0])
return img
def ReadImageTensor(path, permute=True, rgb_only=True, batch_dim=True):
channels = 3 if rgb_only else 4
if isinstance(path, list):
first_image = plt.imread(path[0])[:, :, 0:channels]
b_image = np.empty(
(len(path), first_image.shape[0], first_image.shape[1], channels), dtype=np.float32)
b_image[0] = first_image
for i in range(1, len(path)):
b_image[i] = plt.imread(path[i])[:, :, 0:channels]
return MatImg2Tensor(b_image, permute)
return MatImg2Tensor(plt.imread(path)[:, :, 0:channels], permute, batch_dim)
def ReadImageNumpyArray(path, permute=True, rgb_only=True, batch_dim=True):
channels = 3 if rgb_only else 4
if isinstance(path, list):
first_image = plt.imread(path[0])[:, :, 0:channels]
b_image = np.empty(
(len(path), first_image.shape[0], first_image.shape[1], channels), dtype=np.float32)
b_image[0] = first_image
for i in range(1, len(path)):
b_image[i] = plt.imread(path[i])[:, :, 0:channels]
return MatImg2Numpy(b_image, permute)
return MatImg2Numpy(plt.imread(path)[:, :, 0:channels], permute, batch_dim)
def WriteImageTensor(t, path):
#image = Tensor2MatImg(t)
if isinstance(path, list):
if (len(t.size()) != 4 and len(path) != 1) or t.size()[0] != len(path):
raise ValueError
for i in range(len(path)):
save_image(t[i], path[i])
#plt.imsave(path[i], image[i])
else:
if len(t.squeeze().size()) >= 4:
raise ValueError
#plt.imsave(path, image)
save_image(t, path)
def PlotImageTensor(t: torch.Tensor, *, ax: plt.Axes = None):
"""
Plot a image tensor using matplotlib
:param t: 2D (single channel image), 3D (multiple channels image) or 4D (3D image with batch dim) tensor
:param ax: (Optional) Specify the axes to plot image
"""
return plt.imshow(Tensor2MatImg(t)) if ax is None else ax.imshow(Tensor2MatImg(t))
def Tensor2Glm(t):
t = t.squeeze()
size = t.size()
if len(size) == 1:
if size[0] <= 0 or size[0] > 4:
raise ValueError
return gvec_type[size[0] - 1](t.cpu().numpy())
if len(size) == 2:
if size[0] <= 1 or size[0] > 4 or size[1] <= 1 or size[1] > 4:
raise ValueError
return gmat_type[size[1] - 2][size[0] - 2](t.cpu().numpy())
raise ValueError
def Glm2Tensor(val):
return torch.from_numpy(np.array(val))
def MeshGrid(size: Tuple[int, int], normalize: bool = False, swap_dim: bool = False):
"""
Generate a mesh grid
:param size: grid size (rows, columns)
:param normalize: return coords in normalized space? defaults to False
:param swap_dim: if True, return coords in (y, x) order, defaults to False
:return: rows x columns x 2 tensor
"""
y, x = torch.meshgrid(torch.tensor(range(size[0])),
torch.tensor(range(size[1])))
if swap_dim:
if normalize:
return torch.stack([y / (size[0] - 1.), x / (size[1] - 1.)], 2)
else:
return torch.stack([y, x], 2)
if normalize:
return torch.stack([x / (size[1] - 1.), y / (size[0] - 1.)], 2)
else:
return torch.stack([x, y], 2)
def CreateDirIfNeed(path):
if not os.path.exists(path):
os.makedirs(path)
def get_angle(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
angle = -torch.atan(x / y) + (y < 0) * math.pi + 0.5 * math.pi
return angle
def CartesianToSpherical(cart: torch.Tensor, inverse_r: bool = False) -> torch.Tensor:
"""
Convert coordinates from Cartesian to Spherical
:param cart ```Tensor(..., 3)```: coordinates in Cartesian
:param inverse_r: whether to inverse r
:return ```Tensor(..., 3)```: coordinates in Spherical (r, theta, phi)
"""
rho = torch.sqrt(torch.sum(cart * cart, dim=-1))
theta = get_angle(cart[..., 0], cart[..., 2])
if inverse_r:
rho = rho.reciprocal()
phi = torch.acos(cart[..., 1] * rho)
else:
phi = torch.acos(cart[..., 1] / rho)
return torch.stack([rho, theta, phi], dim=-1)
def SphericalToCartesian(spher: torch.Tensor) -> torch.Tensor:
"""
Convert coordinates from Spherical to Cartesian
:param spher: ... x 3, coordinates in Spherical
:return: ... x 3, coordinates in Cartesian (r, theta, phi)
"""
rho = spher[..., 0]
sin_theta_phi = torch.sin(spher[..., 1:3])
cos_theta_phi = torch.cos(spher[..., 1:3])
x = rho * cos_theta_phi[..., 0] * sin_theta_phi[..., 1]
y = rho * cos_theta_phi[..., 1]
z = rho * sin_theta_phi[..., 0] * sin_theta_phi[..., 1]
return torch.stack([x, y, z], dim=-1)
def RaySphereIntersect(p: torch.Tensor, v: torch.Tensor, r: torch.Tensor) -> torch.Tensor:
"""
Calculate intersections of each rays and each spheres
:param p ```Tensor(B, 3)```: positions of rays
:param v ```Tensor(B, 3)```: directions of rays
:param r ```Tensor(N)```: , radius of spheres
:return ```Tensor(B, N, 3)```: points of intersection
:return ```Tensor(B, N)```: depths of intersection along ray
"""
# p, v: Expand to (B, 1, 3)
p = p.unsqueeze(1)
v = v.unsqueeze(1)
# pp, vv, pv: (B, 1)
pp = (p * p).sum(dim=2)
vv = (v * v).sum(dim=2)
pv = (p * v).sum(dim=2)
depths = (((pv * pv - vv * (pp - r * r)).sqrt() - pv) / vv)
return p + depths[..., None] * v, depths
def GetDepthLayers(depth_range: Tuple[float, float], n_layers: int) -> List[float]:
"""
Get [n_layers] foreground layers whose diopters are distributed uniformly
in [depth_range] plus a background layer
:param depth_range: depth range of foreground layers
:param n_layers: number of foreground layers
:return: list of [n_layers+1] depths
"""
diopter_range = (1 / depth_range[1], 1 / depth_range[0])
depths = [1e5] # Background layer
depths += list(1.0 /
np.linspace(diopter_range[0], diopter_range[1], n_layers))
return depths
def GetRotMatrix(theta: Union[float, torch.Tensor], phi: Union[float, torch.Tensor]) -> torch.Tensor:
"""
Get rotation matrix from angles in spherical space
:param theta ```Tensor(..., 1) | float```: rotation angles around y axis
:param phi ```Tensor(..., 1) | float```: rotation angles around x axis
:return: ```Tensor(..., 3, 3)``` rotation matrices
"""
if not isinstance(theta, torch.Tensor):
theta = torch.tensor([theta])
if not isinstance(phi, torch.Tensor):
phi = torch.tensor([phi])
spher = torch.cat([torch.ones_like(theta), theta, phi], dim=-1)
print(spher)
forward = SphericalToCartesian(spher) # (..., 3)
up = torch.tensor([0.0, 1.0, 0.0])
forward, up = torch.broadcast_tensors(forward, up)
print(forward, up)
right = torch.cross(forward, up, dim=-1) # (..., 3)
up = torch.cross(right, forward, dim=-1) # (..., 3)
print(right, up, forward)
return torch.stack([right, up, forward], dim=-2) # (..., 3, 3)
def broadcast_cat(input: torch.Tensor,
s: Union[Number, List[Number], torch.Tensor],
dim=-1,
append: bool = True) -> torch.Tensor:
"""
Concatenate a tensor with a scalar along last dimension
:param input ```Tensor(..., N)```: input tensor
:param s: scalar
:param append: append or prepend the scalar to input tensor
:return: ```Tensor(..., N+1)```
"""
if dim != -1:
raise NotImplementedError('currently only support the last dimension')
if isinstance(s, torch.Tensor):
x = s
elif isinstance(s, list):
x = torch.tensor(s, dtype=input.dtype, device=input.device)
else:
x = torch.tensor([s], dtype=input.dtype, device=input.device)
expand_shape = list(input.size())
expand_shape[dim] = -1
x = x.expand(expand_shape)
return torch.cat([input, x] if append else [x, input], dim)
def generate_video(frames: torch.Tensor, path: str, fps: float,
repeat: int = 1, pingpong: bool = False,
video_codec: str = 'libx264'):
"""
Generate video from a sequence of frames after converting type and
permuting channels to meet the requirement of ```torchvision.io.write_video()```
:param frames ```Tensor(B, C, H, W)```: a sequence of frames
:param path: video path
:param fps: frames per second
:param repeat: repeat times
:param pingpong: whether repeat sequence in pinpong form
:param video_codec: video codec
"""
frames = trans_func.convert_image_dtype(frames, torch.uint8)
frames = frames.detach().cpu().permute(0, 2, 3, 1)
if pingpong:
frames = torch.cat([frames, frames.flip(0)], 0)
frames = frames.expand(repeat, -1, -1, -1, 3).flatten(0, 1)
torchvision.io.write_video(path, frames, fps, video_codec)
def is_image_file(filename):
return any(filename.endswith(extension) for extension in [".png", ".jpg", ".jpeg"])
def save_2d_tensor(path, x):
with open(path, 'w', encoding='utf-8', newline='') as f:
csv_writer = csv.writer(f)
for i in range(x.shape[0]):
csv_writer.writerow(x[i])
def view_like(input: torch.Tensor, ref: torch.Tensor) -> torch.Tensor:
"""
Reshape input to be the same size as ref except the last dimension
:param input ```Tensor(..., C)```: input tensor
:param ref ```Tensor(B.., *): reference tensor
:return ```Tensor(B.., C)```: reshaped tensor
"""
out_shape = list(ref.size())
out_shape[-1] = -1
return input.view(out_shape)
def rgb2ycbcr(input: torch.Tensor) -> torch.Tensor:
"""
Convert input tensor from RGB to YCbCr
:param input ```Tensor(..., 3) | Tensor(..., 3, H, W)```:
:return ```Tensor(..., 3) | Tensor(..., 3, H, W)```:
"""
if input.size(-1) == 3:
r = input[..., 0:1]
g = input[..., 1:2]
b = input[..., 2:3]
dim_c = -1
else:
r = input[..., 0:1, :, :]
g = input[..., 1:2, :, :]
b = input[..., 2:3, :, :]
dim_c = -3
y = r * 0.25678824 + g * 0.50412941 + b * 0.09790588 + 0.0625
cb = r * -0.14822353 + g * -0.29099216 + b * 0.43921569 + 0.5
cr = r * 0.43921569 + g * -0.36778824 + b * -0.07142745 + 0.5
return torch.cat([y, cb, cr], dim_c)
def rgb2ycbcr(input: torch.Tensor) -> torch.Tensor:
"""
Convert input tensor from RGB to YCbCr
:param input ```Tensor(..., 3) | Tensor(..., 3, H, W)```:
:return ```Tensor(..., 3) | Tensor(..., 3, H, W)```:
"""
if input.size(-1) == 3:
r = input[..., 0:1]
g = input[..., 1:2]
b = input[..., 2:3]
dim_c = -1
else:
r = input[..., 0:1, :, :]
g = input[..., 1:2, :, :]
b = input[..., 2:3, :, :]
dim_c = -3
y = r * 0.257 + g * 0.504 + b * 0.098 + 0.0625
cb = r * -0.148 + g * -0.291 + b * 0.439 + 0.5
cr = r * 0.439 + g * -0.368 + b * -0.071 + 0.5
return torch.cat([cb, cr, y], dim_c)
def ycbcr2rgb(input: torch.Tensor) -> torch.Tensor:
"""
Convert input tensor from YCbCr to RGB
:param input ```Tensor(..., 3) | Tensor(..., 3, H, W)```:
:return ```Tensor(..., 3) | Tensor(..., 3, H, W)```:
"""
if input.size(-1) == 3:
cb = input[..., 0:1]
cr = input[..., 1:2]
y = input[..., 2:3]
dim_c = -1
else:
cb = input[..., 0:1, :, :]
cr = input[..., 1:2, :, :]
y = input[..., 2:3, :, :]
dim_c = -3
y = y - 0.0625
cb = cb - 0.5
cr = cr - 0.5
r = y * 1.164 + cr * 1.596
g = y * 1.164 + cb * -0.392 + cr * -0.813
b = y * 1.164 + cb * 2.017
return torch.cat([r, g, b], dim_c)
def horizontal_shift_image(input: torch.Tensor, shift: int, dim=-1) -> torch.Tensor:
if shift == 0:
return input
shifted = torch.zeros_like(input)
if dim == -1:
if shift > 0:
shifted[..., shift:] = input[..., :-shift]
else:
shifted[..., :shift] = input[..., -shift:]
elif dim == -2:
if shift > 0:
shifted[..., shift:, :] = input[..., :-shift, :]
else:
shifted[..., :shift, :] = input[..., -shift:, :]
else:
raise NotImplementedError
return shifted
def depth_sample(depth_range: Tuple[float, float], n: int, lindisp: bool) -> torch.Tensor:
if lindisp:
depth_range = (1 / depth_range[0], 1 / depth_range[1])
samples = torch.linspace(depth_range[0], depth_range[1], n)
return samples
nerf++ @ a30f1a5a
Subproject commit a30f1a5ad116e43aad90c426a966b2a3fcedaf7e
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment