diff --git a/README.md b/README.md index 00d06928..92f14c41 100644 --- a/README.md +++ b/README.md @@ -155,6 +155,8 @@ __Right-handed rule XYZ global coordinate system__. `x-axis` faces forward and ` - `pitch`: counter-clockwise rotation about the `y-axis` - `yaw`: counter-clockwise rotation about the `z-axis` +You can chnage the right-handed coordinate system so that the `z-axis` faces down by adding `z_down=True` as a parameter. + See demo scripts under `scripts`. @@ -187,7 +189,7 @@ Check [CONTRIBUTING.md](./CONTRIBUTING.md) for more information - [ ] Documentations for each transform - [x] Add table and statistics for speed improvements -- [ ] Batch processing for `numpy` +- [ ] Batch processing for `numpy` (uses too much memory; alternative will be `multiprocessing`) - [ ] Mixed precision for `torch` - [ ] `c++` version of grid sampling diff --git a/demo/equi2equi_numpy.py b/demo/equi2equi_numpy.py index 2e5349ac..5a0af855 100644 --- a/demo/equi2equi_numpy.py +++ b/demo/equi2equi_numpy.py @@ -23,7 +23,7 @@ def preprocess( img: Union[np.ndarray, Image.Image], is_cv2: bool = False, ) -> np.ndarray: - r"""Preprocesses image""" + """Preprocesses image""" if isinstance(img, np.ndarray) and is_cv2: img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) if isinstance(img, Image.Image): @@ -37,7 +37,7 @@ def preprocess( def test_video(path: str) -> None: - r"""Test video""" + """Test video""" # Rotation: pi = np.pi inc = pi / 180 @@ -105,7 +105,7 @@ def test_video(path: str) -> None: def test_image(path: str) -> None: - r"""Test single image""" + """Test single image""" # Rotation: rot = { "roll": 0, # diff --git a/demo/equi2equi_torch.py b/demo/equi2equi_torch.py index fcb1dcb6..f9037986 100644 --- a/demo/equi2equi_torch.py +++ b/demo/equi2equi_torch.py @@ -27,7 +27,7 @@ def preprocess( img: Union[np.ndarray, Image.Image], is_cv2: bool = False, ) -> torch.Tensor: - r"""Preprocesses image""" + """Preprocesses image""" if isinstance(img, np.ndarray) and is_cv2: img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) if isinstance(img, Image.Image): @@ -66,7 +66,7 @@ def postprocess( def test_video(path: str) -> None: - r"""Test video""" + """Test video""" # Rotation: pi = np.pi inc = pi / 180 @@ -134,7 +134,7 @@ def test_video(path: str) -> None: def test_image(path: str) -> None: - r"""Test single image""" + """Test single image""" # Rotation: rot = { "roll": 0, # diff --git a/demo/equi2pers_numpy.py b/demo/equi2pers_numpy.py index 8cd04f8d..3945aa68 100644 --- a/demo/equi2pers_numpy.py +++ b/demo/equi2pers_numpy.py @@ -23,7 +23,7 @@ def preprocess( img: Union[np.ndarray, Image.Image], is_cv2: bool = False, ) -> np.ndarray: - r"""Preprocesses image""" + """Preprocesses image""" if isinstance(img, np.ndarray) and is_cv2: img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) if isinstance(img, Image.Image): @@ -42,7 +42,7 @@ def test_video( w_pers: int = 640, fov_x: float = 90.0, ) -> None: - r"""Test video""" + """Test video""" # Rotation: pi = np.pi inc = pi / 180 @@ -116,7 +116,7 @@ def test_image( w_pers: int = 640, fov_x: float = 90.0, ) -> None: - r"""Test single image""" + """Test single image""" # Rotation: rot = { "roll": 0, # diff --git a/demo/equi2pers_torch.py b/demo/equi2pers_torch.py index 51ee5329..fd2603d9 100644 --- a/demo/equi2pers_torch.py +++ b/demo/equi2pers_torch.py @@ -27,7 +27,7 @@ def preprocess( img: Union[np.ndarray, Image.Image], is_cv2: bool = False, ) -> torch.Tensor: - r"""Preprocesses image""" + """Preprocesses image""" if isinstance(img, np.ndarray) and is_cv2: img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) if isinstance(img, Image.Image): @@ -71,7 +71,7 @@ def test_video( w_pers: int = 640, fov_x: float = 90.0, ) -> None: - r"""Test video""" + """Test video""" # Rotation: pi = np.pi inc = pi / 180 @@ -145,7 +145,7 @@ def test_image( w_pers: int = 640, fov_x: float = 90.0, ) -> None: - r"""Test single image""" + """Test single image""" # Rotation: rot = { "roll": 0, # diff --git a/equilib/__init__.py b/equilib/__init__.py index 8c57d52a..4c8c0f75 100644 --- a/equilib/__init__.py +++ b/equilib/__init__.py @@ -4,6 +4,7 @@ from equilib.equi2cube.base import Equi2Cube, equi2cube from equilib.equi2equi.base import Equi2Equi, equi2equi from equilib.equi2pers.base import Equi2Pers, equi2pers +from equilib.info import __version__ __all__ = [ "Cube2Equi", @@ -15,9 +16,3 @@ "equi2equi", "equi2pers", ] - -__version__ = "0.2.1" -__author__ = "Haruya Ishikawa" -__homepage__ = "www.hoge.com" -__description__ = "equirectangular processing" -__url__ = "www.hoge.com/equilib" diff --git a/equilib/common/numpy_utils.py b/equilib/common/numpy_utils.py deleted file mode 100644 index 36a7a6b7..00000000 --- a/equilib/common/numpy_utils.py +++ /dev/null @@ -1,69 +0,0 @@ -#!/usr/bin/env python3 - -import numpy as np - - -def create_coord( - height: int, - width: int, -) -> np.ndarray: - r"""Create mesh coordinate grid with height and width - - `z-axis` scale is `1` - - params: - - height (int) - - width (int) - - return: - - coordinate (np.ndarray) - """ - _xs = np.linspace(0, width - 1, width) - _ys = np.linspace(0, height - 1, height) - xs, ys = np.meshgrid(_xs, _ys) - zs = np.ones_like(xs) - coord = np.stack((xs, ys, zs), axis=2) - return coord - - -def create_rotation_matrix( - x: float, - y: float, - z: float, -) -> np.ndarray: - r"""Create Rotation Matrix - - params: - - x: x-axis rotation float - - y: y-axis rotation float - - z: z-axis rotation float - - return: - - rotation matrix (np.ndarray) - """ - # calculate rotation about the x-axis - R_x = np.array( - [ - [1.0, 0.0, 0.0], - [0.0, np.cos(x), -np.sin(x)], - [0.0, np.sin(x), np.cos(x)], - ] - ) - # calculate rotation about the y-axis - R_y = np.array( - [ - [np.cos(y), 0.0, np.sin(y)], - [0.0, 1.0, 0.0], - [-np.sin(y), 0.0, np.cos(y)], - ] - ) - # calculate rotation about the z-axis - R_z = np.array( - [ - [np.cos(z), -np.sin(z), 0.0], - [np.sin(z), np.cos(z), 0.0], - [0.0, 0.0, 1.0], - ] - ) - - return R_z @ R_y @ R_x diff --git a/equilib/common/torch_utils.py b/equilib/common/torch_utils.py deleted file mode 100644 index 9b882277..00000000 --- a/equilib/common/torch_utils.py +++ /dev/null @@ -1,86 +0,0 @@ -#!/usr/bin/env python3 - -import numpy as np - -import torch - -pi = torch.Tensor([3.14159265358979323846]) - - -def sizeof(tensor: torch.Tensor) -> float: - r"""Get the size of a tensor""" - assert torch.is_tensor(tensor), "ERR: is not tensor" - return tensor.element_size() * tensor.nelement() - - -def get_device(a: torch.Tensor) -> torch.device: - r"""Get device of a Tensor""" - return torch.device(a.get_device() if a.get_device() >= 0 else "cpu") - - -def deg2rad(tensor: torch.Tensor) -> torch.Tensor: - r"""Function that converts angles from degrees to radians.""" - if not torch.is_tensor(tensor): - return tensor * float(pi) / 180.0 - return tensor * pi.to(tensor.device).type(tensor.dtype) / 180.0 - - -def create_coord( - height: int, - width: int, -) -> torch.Tensor: - r"""Create mesh coordinate grid""" - _xs = torch.linspace(0, width - 1, width) - _ys = torch.linspace(0, height - 1, height) - # NOTE: https://github.com/pytorch/pytorch/issues/15301 - # Torch meshgrid behaves differently than numpy - ys, xs = torch.meshgrid([_ys, _xs]) - zs = torch.ones_like(xs) - coord = torch.stack((xs, ys, zs), dim=2) - return coord - - -def create_rotation_matrix( - x: float, - y: float, - z: float, -) -> torch.Tensor: - r"""Create Rotation Matrix - - params: - - x: x-axis rotation float - - y: y-axis rotation float - - z: z-axis rotation float - - return: - - rotation matrix (torch.Tensor) - """ - # calculate rotation about the x-axis - R_x = torch.tensor( - [ - [1.0, 0.0, 0.0], - [0.0, np.cos(x), -np.sin(x)], - [0.0, np.sin(x), np.cos(x)], - ], - dtype=torch.float, - ) - # calculate rotation about the y-axis - R_y = torch.tensor( - [ - [np.cos(y), 0.0, np.sin(y)], - [0.0, 1.0, 0.0], - [-np.sin(y), 0.0, np.cos(y)], - ], - dtype=torch.float, - ) - # calculate rotation about the z-axis - R_z = torch.tensor( - [ - [np.cos(z), -np.sin(z), 0.0], - [np.sin(z), np.cos(z), 0.0], - [0.0, 0.0, 1.0], - ], - dtype=torch.float, - ) - - return R_z @ R_y @ R_x diff --git a/equilib/cube2equi/base.py b/equilib/cube2equi/base.py index 36e6281c..3561f91f 100644 --- a/equilib/cube2equi/base.py +++ b/equilib/cube2equi/base.py @@ -13,7 +13,7 @@ class Cube2Equi(object): - r""" + """ params: - w_out, h_out (int): equirectangular image size - cube_format (str): input cube format("dice", "horizon", "dict", "list") @@ -74,18 +74,19 @@ def cube2equi( sampling_method: str = "default", mode: str = "bilinear", ) -> Union[np.ndarray, torch.Tensor]: - r""" + """ params: - - cubemap: Union[ + - cubemap (Union[ np.ndarray, torch.Tensor, Dict[str, Union[np.ndarray, torch.Tensor]], - List[Union[np.ndarray, torch.Tensor]]] + List[Union[np.ndarray, torch.Tensor]]]) - cube_format (str): ("dice", "horizon", "dict", "list") - w_out (int): - h_out (int): - sampling_method (str): "default" - mode (str): "bilinear" + - z_down (bool): False return: - equi (np.ndarray, torch.Tensor) diff --git a/equilib/cube2equi/cube2equi_numpy.py b/equilib/cube2equi/cube2equi_numpy.py index 9b5b2fdb..f5df1768 100644 --- a/equilib/cube2equi/cube2equi_numpy.py +++ b/equilib/cube2equi/cube2equi_numpy.py @@ -22,7 +22,7 @@ def cube_dict2h( def cube_dice2h(cube_dice: np.ndarray) -> np.ndarray: - r"""dice to horizion + """dice to horizion params: - cube_dice: (C, H, W) @@ -57,7 +57,7 @@ def _to_horizon(cubemap: np.ndarray, cube_format: str) -> np.ndarray: def _equirect_facetype(h: int, w: int) -> np.ndarray: - r"""0F 1R 2B 3L 4U 5D""" + """0F 1R 2B 3L 4U 5D""" tp = np.roll( np.arange(4).repeat(w // 4)[None, :].repeat(h, 0), 3 * w // 8, 1 @@ -92,7 +92,7 @@ def _run_single( sampling_method: str, mode: str, ) -> np.ndarray: - r"""Run a single batch of transformation""" + """Run a single batch of transformation""" w_face = cubemap.shape[-2] theta, phi = create_equi_grid(h_out, w_out) @@ -145,7 +145,7 @@ def run( sampling_method: str, mode: str, ) -> np.ndarray: - r"""Run cube to equirectangular image transformation + """Run cube to equirectangular image transformation params: - cubemap: np.ndarray diff --git a/equilib/cube2equi/cube2equi_torch.py b/equilib/cube2equi/cube2equi_torch.py index 0f4a3d9e..29c1d48b 100644 --- a/equilib/cube2equi/cube2equi_torch.py +++ b/equilib/cube2equi/cube2equi_torch.py @@ -6,7 +6,7 @@ import torch from equilib.grid_sample import torch_func -from equilib.common.torch_utils import get_device +from equilib.utils import torch_utils def cube_list2h(cube_list: List[torch.Tensor]) -> torch.Tensor: @@ -24,7 +24,7 @@ def cube_dict2h( def cube_dice2h(cube_dice: torch.Tensor) -> torch.Tensor: - r"""dice to horizion + """dice to horizion params: - cube_dice: (C, H, W) @@ -62,7 +62,7 @@ def _to_horizon( def _equirect_facetype(h: int, w: int) -> torch.Tensor: - r"""0F 1R 2B 3L 4U 5D""" + """0F 1R 2B 3L 4U 5D""" tp = torch.roll( torch.arange(4) # 1 .repeat_interleave(w // 4) # 2 same as np.repeat @@ -107,13 +107,13 @@ def run( sampling_method: str, mode: str, ) -> torch.Tensor: - r"""Run cube to equirectangular image transformation + """Run cube to equirectangular image transformation params: - - cubemap: np.ndarray - - cube_format: ('dice', 'horizon', 'list', 'dict') - - sampling_method: str - - mode: str + - cubemap (np.ndarray) + - cube_format (str): ('dice', 'horizon', 'list', 'dict') + - sampling_method (str) + - mode (str) """ # Convert all cubemap format to `horizon` and batched @@ -151,7 +151,7 @@ def run( cubemap = torch.stack(cubemap, 0) # get device - device = get_device(cubemap) + device = torch_utils.get_device(cubemap) w_face = cubemap.shape[-2] theta, phi = create_equi_grid(h_out, w_out) diff --git a/equilib/equi2cube/base.py b/equilib/equi2cube/base.py index e084e145..343e58a5 100644 --- a/equilib/equi2cube/base.py +++ b/equilib/equi2cube/base.py @@ -14,16 +14,17 @@ class Equi2Cube(object): - r""" + """ params: - w_face (int): cube face width - cube_format (str): ("dice", "horizon", "dict", "list") - sampling_method (str) - mode (str) + - z_down (bool) inputs: - equi (np.ndarray, torch.Tensor, list, dict) - - rot: (dict, list[dict]): {"roll", "pitch", "yaw"} + - rot (dict, list[dict]): {"roll", "pitch", "yaw"} returns: - cube (np.ndarray, torch.Tensor, list, dict) @@ -35,11 +36,13 @@ def __init__( cube_format: str, sampling_method: str = "default", mode: str = "bilinear", + z_down: bool = False, ) -> None: self.w_face = w_face self.cube_format = cube_format self.sampling_method = sampling_method self.mode = mode + self.z_down = z_down def __call__( self, @@ -65,6 +68,7 @@ def __call__( cube_format=self.cube_format, sampling_method=self.sampling_method, mode=self.mode, + z_down=self.z_down, ) @@ -82,20 +86,22 @@ def equi2cube( cube_format: str, sampling_method: str = "default", mode: str = "bilinear", + z_down: bool = False, ) -> Union[ np.ndarray, torch.Tensor, Dict[str, Union[np.ndarray, torch.Tensor]], List[Union[np.ndarray, torch.Tensor]], ]: - r""" + """ params: - equi (np.ndarray, torch.Tensor, list, dict) - - rot: (dict, list[dict]): {"roll", "pitch", "yaw"} + - rot (dict, list[dict]): {"roll", "pitch", "yaw"} - w_face (int): cube face width - cube_format (str): ("dice", "horizon", "dict", "list") - sampling_method (str) - mode (str) + - z_down (bool) returns: - cube (np.ndarray, torch.Tensor, dict, list) @@ -126,6 +132,7 @@ def equi2cube( cube_format=cube_format, sampling_method=sampling_method, mode=mode, + z_down=z_down, ) elif _type == "torch": return run_torch( @@ -135,6 +142,7 @@ def equi2cube( cube_format=cube_format, sampling_method=sampling_method, mode=mode, + z_down=z_down, ) else: raise ValueError diff --git a/equilib/equi2cube/equi2cube_numpy.py b/equilib/equi2cube/equi2cube_numpy.py index 6d14c1a7..58aa42c0 100644 --- a/equilib/equi2cube/equi2cube_numpy.py +++ b/equilib/equi2cube/equi2cube_numpy.py @@ -1,11 +1,14 @@ #!/usr/bin/env python3 +from functools import partial from typing import Dict, List, Tuple, Union import numpy as np from equilib.grid_sample import numpy_func -from equilib.common.numpy_utils import create_rotation_matrix +from equilib.utils import create_rotation_matrix + +_create_rotation_matrix = partial(create_rotation_matrix, is_torch=False) def cube_h2list(cube_h: np.ndarray) -> List[np.ndarray]: @@ -34,29 +37,8 @@ def cube_h2dice(cube_h: np.ndarray) -> np.ndarray: return cube_dice -def rotation_matrix( - roll: float, - pitch: float, - yaw: float, -) -> np.ndarray: - r"""Create Rotation Matrix - - params: - - roll: x-axis rotation float - - pitch: y-axis rotation float - - yaw: z-axis rotation float - - return: - - rotation matrix (np.ndarray) - - Global coordinates -> x-axis points forward, z-axis points upward - """ - R = create_rotation_matrix(x=roll, y=pitch, z=yaw) - return R - - def create_xyz(w_face: int) -> np.ndarray: - r"""xyz coordinates of the faces of the cube""" + """xyz coordinates of the faces of the cube""" out = np.zeros((w_face, w_face * 6, 3), np.float32) rng = np.linspace(-0.5, 0.5, num=w_face, dtype=np.float32) @@ -100,7 +82,7 @@ def create_xyz(w_face: int) -> np.ndarray: def xyz2rot(xyz) -> Tuple[np.ndarray]: - r"""Return rotation (theta, phi) from xyz""" + """Return rotation (theta, phi) from xyz""" norm = np.linalg.norm(xyz, axis=-1) phi = np.arcsin(xyz[:, :, 2] / norm) theta = np.arctan2(xyz[:, :, 1], xyz[:, :, 0]) @@ -114,8 +96,9 @@ def _run_single( cube_format: str, sampling_method: str, mode: str, + z_down: bool, ) -> np.ndarray: - r"""Call a single run + """Call a single run params: - equi: np.ndarray @@ -123,6 +106,7 @@ def _run_single( - w_face: int - cube_format: str - mode: str + - z """ assert len(equi.shape) == 3, "ERR: {} is not a valid array".format( @@ -135,7 +119,11 @@ def _run_single( xyz = create_xyz(w_face) xyz = xyz[:, :, :, np.newaxis] - xyz_ = rotation_matrix(**rot) @ xyz + # FIXME: not sure why, but z-axis is facing the opposite + # probably I need to change the way I choose the xyz coordinates + # this is a temporary fix for now + z_down = not z_down + xyz_ = _create_rotation_matrix(**rot, z_down=z_down) @ xyz xyz_ = xyz_.squeeze(-1) theta, phi = xyz2rot(xyz_) @@ -174,6 +162,7 @@ def run( cube_format: str, sampling_method: str, mode: str, + z_down: bool, ) -> Union[np.ndarray, List[np.ndarray], List[Dict[str, np.ndarray]]]: r"""Call Equi2Cube @@ -213,6 +202,7 @@ def run( cube_format=cube_format, sampling_method=sampling_method, mode=mode, + z_down=z_down, ) cubemaps.append(cubemap) diff --git a/equilib/equi2cube/equi2cube_torch.py b/equilib/equi2cube/equi2cube_torch.py index 9987e68d..d50cfb61 100644 --- a/equilib/equi2cube/equi2cube_torch.py +++ b/equilib/equi2cube/equi2cube_torch.py @@ -1,17 +1,19 @@ #!/usr/bin/env python3 +from functools import partial import math from typing import Dict, List, Tuple, Union import torch from equilib.grid_sample import torch_func -from equilib.common.torch_utils import ( +from equilib.utils import ( create_rotation_matrix, - get_device, - sizeof, + torch_utils, ) +_create_rotation_matrix = partial(create_rotation_matrix, is_torch=True) + def cube_h2list(cube_h: torch.Tensor) -> List[torch.Tensor]: assert cube_h.shape[-2] * 6 == cube_h.shape[-1] @@ -68,27 +70,8 @@ def cube_h2dice(cube_h: torch.Tensor) -> torch.Tensor: return cube_dice -def rotation_matrix( - roll: float, - pitch: float, - yaw: float, -) -> torch.Tensor: - r"""Create Rotation Matrix - - params: - - roll: x-axis rotation float - - pitch: y-axis rotation float - - yaw: z-axis rotation float - - return: - - rotation matrix (torch.Tensor) - """ - R = create_rotation_matrix(x=roll, y=pitch, z=yaw) - return R - - def create_xyz(w_face: int) -> torch.Tensor: - r"""xyz coordinates of the faces of the cube""" + """xyz coordinates of the faces of the cube""" _dtype = torch.float32 out = torch.zeros((w_face, w_face * 6, 3), dtype=_dtype) rng = torch.linspace(-0.5, 0.5, w_face, dtype=_dtype) @@ -136,7 +119,7 @@ def create_xyz(w_face: int) -> torch.Tensor: def xyz2rot(xyz) -> Tuple[torch.Tensor]: - r"""Return rotation (theta, phi) from xyz""" + """Return rotation (theta, phi) from xyz""" norm = torch.norm(xyz, dim=-1) phi = torch.asin(xyz[:, :, :, 2] / norm) theta = torch.atan2(xyz[:, :, :, 1], xyz[:, :, :, 0]) @@ -150,6 +133,7 @@ def run( cube_format: str, sampling_method: str, mode: str, + z_down: bool, debug: bool = False, ) -> Union[ torch.Tensor, @@ -157,7 +141,7 @@ def run( List[Dict[str, torch.Tensor]], Dict[str, torch.Tensor], ]: - r"""Call Equi2Cube + """Call Equi2Cube params: - equi (Union[torch.Tensor, List[torch.Tensor]]) @@ -166,6 +150,7 @@ def run( - cube_format (str): ('list', 'dict', 'dice') - sampling_method (str) - mode (str) + - z_down (bool) """ assert ( @@ -183,17 +168,22 @@ def run( h_equi, w_equi = equi.shape[-2:] if debug: - print("equi: ", sizeof(equi) / 10e6, "mb") + print("equi: ", torch_utils.sizeof(equi) / 10e6, "mb") # get device - device = get_device(equi) + device = torch_utils.get_device(equi) + + # FIXME: not sure why, but z-axis is facing the opposite + # probably I need to change the way I choose the xyz coordinates + # this is a temporary fix for now + z_down = not z_down # define variables xyz = [] for r in rot: # for each rotations calculate M _xyz = create_xyz(w_face) - xyz_ = rotation_matrix(**r) @ _xyz.unsqueeze(3) + xyz_ = _create_rotation_matrix(**r, z_down=z_down) @ _xyz.unsqueeze(3) xyz_ = xyz_.squeeze(3) xyz.append(xyz_) xyz = torch.stack(xyz, dim=0).to(device) diff --git a/equilib/equi2equi/base.py b/equilib/equi2equi/base.py index ea4d749f..38744e9c 100644 --- a/equilib/equi2equi/base.py +++ b/equilib/equi2equi/base.py @@ -13,11 +13,12 @@ class Equi2Equi(object): - r""" - init params: + """ + params: - w_out, h_out (optional int): equi image size - sampling_method (str): defaults to "default" - mode (str): interpolation mode, defaults to "bilinear" + - z_down (bool) input params: - src (np.ndarray, torch.Tensor, list) @@ -33,11 +34,13 @@ def __init__( h_out: Optional[int] = None, sampling_method: str = "default", mode: str = "bilinear", + z_down: bool = False ) -> None: self.w_out = w_out self.h_out = h_out self.sampling_method = sampling_method self.mode = mode + self.z_down = z_down def __call__( self, @@ -56,6 +59,7 @@ def __call__( rot=rot, sampling_method=self.sampling_method, mode=self.mode, + z_down=self.z_down, ) @@ -71,16 +75,18 @@ def equi2equi( ], sampling_method: str = "default", mode: str = "bilinear", + z_down: bool = False, w_out: Optional[int] = None, h_out: Optional[int] = None, ) -> Union[np.ndarray, torch.Tensor]: - r""" - init params: + """ + params: - src (np.ndarray, torch.Tensor, list) - rot (dict, list[dict]) - w_out, h_out (optional int): equi image size - sampling_method (str): defaults to "default" - mode (str): interpolation mode, defaults to "bilinear" + - z_down (bool) """ # Try and detect which type it is ("numpy" or "torch") @@ -106,6 +112,7 @@ def equi2equi( rot=rot, sampling_method=sampling_method, mode=mode, + z_down=z_down, w_out=w_out, h_out=h_out, ) @@ -115,6 +122,7 @@ def equi2equi( rot=rot, sampling_method=sampling_method, mode=mode, + z_down=z_down, w_out=w_out, h_out=h_out, ) diff --git a/equilib/equi2equi/equi2equi_numpy.py b/equilib/equi2equi/equi2equi_numpy.py index 4454e074..91f60198 100644 --- a/equilib/equi2equi/equi2equi_numpy.py +++ b/equilib/equi2equi/equi2equi_numpy.py @@ -1,15 +1,18 @@ #!/usr/bin/env python3 +from functools import partial from typing import Dict, List, Optional, Tuple, Union import numpy as np from equilib.grid_sample import numpy_func -from equilib.common.numpy_utils import create_rotation_matrix +from equilib.utils import create_rotation_matrix + +_create_rotation_matrix = partial(create_rotation_matrix, is_torch=False) def create_coordinate(h_out: int, w_out: int) -> np.ndarray: - r"""Create mesh coordinate grid with height and width + """Create mesh coordinate grid with height and width return: - coordinate (np.ndarray) @@ -23,29 +26,8 @@ def create_coordinate(h_out: int, w_out: int) -> np.ndarray: return coord -def rotation_matrix( - roll: float, - pitch: float, - yaw: float, -) -> np.ndarray: - r"""Create Rotation Matrix - - params: - - roll: x-axis rotation float - - pitch: y-axis rotation float - - yaw: z-axis rotation float - - return: - - rotation matrix (np.ndarray) - - Global coordinates -> x-axis points forward, z-axis points upward - """ - R = create_rotation_matrix(x=roll, y=pitch, z=yaw) - return R - - def _get_img_size(img: np.ndarray) -> Tuple[int]: - r"""Return height and width""" + """Return height and width""" return img.shape[-2:] @@ -54,6 +36,7 @@ def _run_single( rot: Dict[str, float], sampling_method: str, mode: str, + z_down: bool, w_out: Optional[int] = None, h_out: Optional[int] = None, ) -> np.ndarray: @@ -70,7 +53,7 @@ def _run_single( z = norm_A * np.sin(a[:, :, 1]) A = np.stack((x, y, z), axis=-1) - R = rotation_matrix(**rot) + R = _create_rotation_matrix(**rot, z_down=z_down) # conversion: # B = R @ A @@ -103,16 +86,18 @@ def run( rot: Union[Dict[str, float], List[Dict[str, float]]], sampling_method: str, mode: str, + z_down: bool, w_out: Optional[int] = None, h_out: Optional[int] = None, ) -> np.ndarray: - r"""Run Equi2Pers + """Run Equi2Pers params: - src: equirectangular image np.ndarray[C, H, W] - rot (Dict[str, float]) - sampling_method (str): (default="faster") - mode (str): (default="bilinear") + - z_down (bool) returns: - pers: perspective image np.ndarray[C, H, W] @@ -134,10 +119,11 @@ def run( rot ), "ERR: length of src and rot differs {} vs {}".format(len(src), len(rot)) + # TODO: batch implementation + # TODO: multiprocessing samples = [] for s, r in zip(src, rot): # iterate through batches - # TODO: batch implementation sample = _run_single( src=s, rot=r, @@ -145,6 +131,7 @@ def run( h_out=h_out, sampling_method=sampling_method, mode=mode, + z_down=z_down, ) samples.append(sample) diff --git a/equilib/equi2equi/equi2equi_torch.py b/equilib/equi2equi/equi2equi_torch.py index d5e11294..40fd997b 100644 --- a/equilib/equi2equi/equi2equi_torch.py +++ b/equilib/equi2equi/equi2equi_torch.py @@ -1,22 +1,22 @@ #!/usr/bin/env python3 +from functools import partial import math from typing import Dict, List, Tuple, Union -import numpy as np - import torch from equilib.grid_sample import torch_func -from equilib.common.torch_utils import ( +from equilib.utils import ( create_rotation_matrix, - get_device, - sizeof, + torch_utils, ) +_create_rotation_matrix = partial(create_rotation_matrix, is_torch=True) + def pixel_wise_rot(M: torch.Tensor) -> Tuple[torch.Tensor]: - r"""Rotation coordinates to phi/theta of the equirectangular image + """Rotation coordinates to phi/theta of the equirectangular image params: - M: torch.Tensor @@ -38,8 +38,8 @@ def pixel_wise_rot(M: torch.Tensor) -> Tuple[torch.Tensor]: return phis, thetas -def create_coordinate(h_out: int, w_out: int) -> np.ndarray: - r"""Create mesh coordinate grid with height and width +def create_coordinate(h_out: int, w_out: int) -> torch.Tensor: + """Create mesh coordinate grid with height and width return: - coordinate (np.ndarray) @@ -55,29 +55,8 @@ def create_coordinate(h_out: int, w_out: int) -> np.ndarray: return coord -def rotation_matrix( - roll: float, - pitch: float, - yaw: float, -) -> np.ndarray: - r"""Create Rotation Matrix - - params: - - roll: x-axis rotation float - - pitch: y-axis rotation float - - yaw: z-axis rotation float - - return: - - rotation matrix (np.ndarray) - - Global coordinates -> x-axis points forward, z-axis points upward - """ - R = create_rotation_matrix(x=roll, y=pitch, z=yaw) - return R - - def _get_img_size(img: torch.Tensor) -> Tuple[int]: - r"""Return height and width""" + """Return height and width""" # batch, channel, height, width return img.shape[-2:] @@ -89,9 +68,10 @@ def run( h_out: int, sampling_method: str = "torch", mode: str = "bilinear", + z_down: bool = False, debug: bool = False, ) -> torch.Tensor: - r"""Run Equi2Pers + """Run Equi2Pers params: - src: equirectangular image torch.Tensor[(B), C, H, W] @@ -124,10 +104,10 @@ def run( w_out = w_equi if debug: - print("size of src: ", sizeof(src) / 10e6, "mb") + print("size of src: ", torch_utils.sizeof(src) / 10e6, "mb") # get device - device = get_device(src) + device = torch_utils.get_device(src) # define variables B = [] @@ -138,7 +118,7 @@ def run( y = norm_A * torch.cos(a[:, :, 1]) * torch.sin(a[:, :, 0]) z = norm_A * torch.sin(a[:, :, 1]) A = torch.stack((x, y, z), dim=-1) - R = rotation_matrix(**r) + R = _create_rotation_matrix(**r, z_down=z_down) _B = R @ A.unsqueeze(3) _B = _B.squeeze(3) B.append(_B) diff --git a/equilib/equi2pers/base.py b/equilib/equi2pers/base.py index d8c2b1f7..6f34c4dd 100644 --- a/equilib/equi2pers/base.py +++ b/equilib/equi2pers/base.py @@ -13,11 +13,14 @@ class Equi2Pers(object): - r""" + """ params: - w_pers, h_pers (int): perspective size - fov_x (float): perspective image fov of x-axis - skew (float): skew intrinsic parameter + - sampling_method (str) + - mode (str) + - z_down (bool) inputs: - equi (np.ndarray, torch.Tensor, list) @@ -35,14 +38,15 @@ def __init__( skew: float = 0.0, sampling_method: str = "default", mode: str = "bilinear", + z_down: bool = False, ) -> None: - r""" """ self.w_pers = w_pers self.h_pers = h_pers self.fov_x = fov_x self.skew = skew self.sampling_method = sampling_method self.mode = mode + self.z_down = z_down def __call__( self, @@ -65,6 +69,7 @@ def __call__( skew=self.skew, sampling_method=self.sampling_method, mode=self.mode, + z_down=self.z_down, ) @@ -84,14 +89,16 @@ def equi2pers( skew: float = 0.0, sampling_method: str = "default", mode: str = "bilinear", + z_down: bool = False, ) -> Union[np.ndarray, torch.Tensor]: - r""" + """ params: - equi (np.ndarray, torch.Tensor, list) - rot (dict, list): Dict[str, float] - w_pers, h_pers (int): perspective size - fov_x (float): perspective image fov of x-axis - skew (float): skew intrinsic parameter + - z_down (bool) returns: - pers (np.ndarray, torch.Tensor) @@ -124,6 +131,7 @@ def equi2pers( skew=skew, sampling_method=sampling_method, mode=mode, + z_down=z_down, ) elif _type == "torch": return run_torch( @@ -135,6 +143,7 @@ def equi2pers( skew=skew, sampling_method=sampling_method, mode=mode, + z_down=z_down, ) else: raise ValueError diff --git a/equilib/equi2pers/equi2pers_numpy.py b/equilib/equi2pers/equi2pers_numpy.py index 1fbc23ca..7293ed58 100644 --- a/equilib/equi2pers/equi2pers_numpy.py +++ b/equilib/equi2pers/equi2pers_numpy.py @@ -1,45 +1,27 @@ #!/usr/bin/env python3 +from functools import partial from typing import Dict, List, Tuple, Union import numpy as np from equilib.grid_sample import numpy_func -from equilib.common.numpy_utils import create_rotation_matrix +from equilib.utils import ( + create_global2camera_rotation_matrix, + create_intrinsic_matrix, + create_rotation_matrix +) - -def intrinsic_matrix( - w_pers: int, - h_pers: int, - fov_x: float, - skew: float, -) -> np.ndarray: - r"""Create Intrinsic Matrix - - return: - - K (np.ndarray): 3x3 matrix - - NOTE: - - ref: http://ksimek.github.io/2013/08/13/intrinsic/ - """ - # perspective projection (focal length) - f = w_pers / (2.0 * np.tan(np.radians(fov_x) / 2.0)) - # transform between camera frame and pixel coordinates - K = np.array( - [ - [f, skew, w_pers / 2], - [0.0, f, h_pers / 2], - [0.0, 0.0, 1.0], - ] - ) - return K +_create_global2camera_rotation_matrix = partial(create_global2camera_rotation_matrix, is_torch=False) +_create_intrinsic_matrix = partial(create_intrinsic_matrix, is_torch=False) +_create_rotation_matrix = partial(create_rotation_matrix, is_torch=False) def perspective_coordinate( w_pers: int, h_pers: int, ) -> np.ndarray: - r"""Create mesh coordinate grid with perspective height and width + """Create mesh coordinate grid with perspective height and width return: - coordinate (np.ndarray) @@ -52,48 +34,8 @@ def perspective_coordinate( return coord -def rotation_matrix( - roll: float, - pitch: float, - yaw: float, -) -> np.ndarray: - r"""Create Rotation Matrix - - params: - - roll: x-axis rotation float - - pitch: y-axis rotation float - - yaw: z-axis rotation float - - return: - - rotation matrix (np.ndarray) - - Global coordinates -> x-axis points forward, z-axis points upward - """ - R = create_rotation_matrix(x=roll, y=pitch, z=yaw) - return R - - -def global2camera_rotation_matrix() -> np.ndarray: - r"""Default rotation that changes global to camera coordinates""" - R_XY = np.array( - [ # X <-> Y - [0.0, 1.0, 0.0], - [1.0, 0.0, 0.0], - [0.0, 0.0, 1.0], - ] - ) - R_YZ = np.array( - [ # Y <-> Z - [1.0, 0.0, 0.0], - [0.0, 0.0, 1.0], - [0.0, 1.0, 0.0], - ] - ) - return R_XY @ R_YZ - - def _get_img_size(img: np.ndarray) -> Tuple[int]: - r"""Return height and width""" + """Return height and width""" return img.shape[-2:] @@ -106,20 +48,21 @@ def _run_single( skew: float, sampling_method: str, mode: str, + z_down: bool, ) -> np.ndarray: # NOTE: Precomputable variables m = perspective_coordinate(w_pers=w_pers, h_pers=h_pers) - K = intrinsic_matrix( + K = _create_intrinsic_matrix( w_pers=w_pers, h_pers=h_pers, fov_x=fov_x, skew=skew, ) - g2c_rot = global2camera_rotation_matrix() + g2c_rot = _create_global2camera_rotation_matrix() # Compute variables - R = rotation_matrix(**rot) + R = _create_rotation_matrix(**rot, z_down=z_down) h_equi, w_equi = _get_img_size(equi) # conversion @@ -159,14 +102,16 @@ def run( skew: float, sampling_method: str, mode: str, + z_down: bool, ) -> np.ndarray: - r"""Run Equi2Pers + """Run Equi2Pers params: - equi: equirectangular image np.ndarray[C, H, W] - rot (dict, list): Dict[str, float] - sampling_method (str) - mode (str) + - z_down (bool) returns: - pers: perspective image np.ndarray[C, H, W] @@ -203,6 +148,7 @@ def run( skew=skew, sampling_method=sampling_method, mode=mode, + z_down=z_down, ) samples.append(sample) diff --git a/equilib/equi2pers/equi2pers_torch.py b/equilib/equi2pers/equi2pers_torch.py index ee06086c..3d2ce361 100644 --- a/equilib/equi2pers/equi2pers_torch.py +++ b/equilib/equi2pers/equi2pers_torch.py @@ -1,52 +1,29 @@ #!/usr/bin/env python3 +from functools import partial import math from typing import Dict, List, Tuple, Union -import numpy as np - import torch from equilib.grid_sample import torch_func -from equilib.common.torch_utils import ( +from equilib.utils import ( + create_global2camera_rotation_matrix, + create_intrinsic_matrix, create_rotation_matrix, - deg2rad, - get_device, - sizeof, + torch_utils, ) - -def intrinsic_matrix( - w_pers: int, - h_pers: int, - fov_x: float, - skew: float, -) -> torch.Tensor: - r"""Create Intrinsic Matrix - - return: - - K (torch.Tensor): 3x3 matrix - - NOTE: - - ref: http://ksimek.github.io/2013/08/13/intrinsic/ - """ - fov_x = torch.tensor(fov_x) - f = w_pers / (2 * torch.tan(deg2rad(fov_x) / 2)) - K = torch.tensor( - [ - [f, skew, w_pers / 2], - [0.0, f, h_pers / 2], - [0.0, 0.0, 1.0], - ] - ) - return K +_create_global2camera_rotation_matrix = partial(create_global2camera_rotation_matrix, is_torch=True) +_create_intrinsic_matrix = partial(create_intrinsic_matrix, is_torch=True) +_create_rotation_matrix = partial(create_rotation_matrix, is_torch=True) def perspective_coordinate( w_pers: int, h_pers: int, ) -> torch.Tensor: - r"""Create mesh coordinate grid with perspective height and width + """Create mesh coordinate grid with perspective height and width return: - coordinate (torch.Tensor) @@ -61,48 +38,8 @@ def perspective_coordinate( return coord -def rotation_matrix( - roll: float, - pitch: float, - yaw: float, -) -> np.ndarray: - r"""Create Rotation Matrix - - params: - - roll: x-axis rotation float - - pitch: y-axis rotation float - - yaw: z-axis rotation float - - return: - - rotation matrix (torch.Tensor) - - Global coordinates -> x-axis points forward, z-axis points upward - """ - R = create_rotation_matrix(x=roll, y=pitch, z=yaw) - return R - - -def global2camera_rotation_matrix() -> torch.Tensor: - r"""Default rotation that changes global to camera coordinates""" - R_XY = torch.tensor( - [ # X <-> Y - [0.0, 1.0, 0.0], - [1.0, 0.0, 0.0], - [0.0, 0.0, 1.0], - ] - ) - R_YZ = torch.tensor( - [ # Y <-> Z - [1.0, 0.0, 0.0], - [0.0, 0.0, 1.0], - [0.0, 1.0, 0.0], - ] - ) - return R_XY @ R_YZ - - def _get_img_size(img: torch.Tensor) -> Tuple[int]: - r"""Return height and width""" + """Return height and width""" # batch, channel, height, width return img.shape[-2:] @@ -116,15 +53,17 @@ def run( skew: float, sampling_method: str, mode: str, + z_down: bool, debug: bool = False, ) -> torch.Tensor: - r"""Run Equi2Pers + """Run Equi2Pers params: - equi: equirectangular image torch.Tensor[(B), C, H, W] - rot: Dict[str, float] or List[Dict[str, float]] - sampling_method (str) - mode (str) + - z_down (bool) returns: - pers: perspective image torch.Tensor[C, H, W] @@ -147,10 +86,10 @@ def run( h_equi, w_equi = _get_img_size(equi) if debug: - print("equi: ", sizeof(equi) / 10e6, "mb") + print("equi: ", torch_utils.sizeof(equi) / 10e6, "mb") # get device - device = get_device(equi) + device = torch_utils.get_device(equi) # define variables # FIXME: methods without using loop @@ -161,14 +100,14 @@ def run( w_pers=w_pers, h_pers=h_pers, ) - K = intrinsic_matrix( + K = _create_intrinsic_matrix( w_pers=w_pers, h_pers=h_pers, fov_x=fov_x, skew=skew, ) - g2c_rot = global2camera_rotation_matrix() - R = rotation_matrix(**r) + g2c_rot = _create_global2camera_rotation_matrix() + R = _create_rotation_matrix(**r, z_down=z_down) _M = R @ g2c_rot @ K.inverse() @ m.unsqueeze(3) _M = _M.squeeze(3) M.append(_M) diff --git a/equilib/grid_sample/numpy_grid_sample/faster.py b/equilib/grid_sample/numpy_grid_sample/faster.py index 2f831471..d13fb05b 100644 --- a/equilib/grid_sample/numpy_grid_sample/faster.py +++ b/equilib/grid_sample/numpy_grid_sample/faster.py @@ -13,7 +13,7 @@ def interp2d( dx: np.ndarray, mode: str = "bilinear", ) -> np.ndarray: - r"""Naive Interpolation + """Naive Interpolation (y,x): target pixel mode: interpolation mode """ @@ -31,7 +31,7 @@ def grid_sample( grid: np.ndarray, mode: str = "bilinear", ) -> np.ndarray: - r"""Numpy Grid Sample""" + """Numpy Grid Sample""" channels, h_in, w_in = img.shape _, h_out, w_out = grid.shape diff --git a/equilib/grid_sample/numpy_grid_sample/interp.py b/equilib/grid_sample/numpy_grid_sample/interp.py index 817fee03..3daecaea 100644 --- a/equilib/grid_sample/numpy_grid_sample/interp.py +++ b/equilib/grid_sample/numpy_grid_sample/interp.py @@ -2,5 +2,5 @@ def linear_interp(v0, v1, d, L): - r"""Basic Linear Interpolation""" + """Basic Linear Interpolation""" return v0 * (1 - d) / L + v1 * d / L diff --git a/equilib/grid_sample/numpy_grid_sample/naive.py b/equilib/grid_sample/numpy_grid_sample/naive.py index 977ed0b1..fba6dd6f 100644 --- a/equilib/grid_sample/numpy_grid_sample/naive.py +++ b/equilib/grid_sample/numpy_grid_sample/naive.py @@ -13,7 +13,7 @@ def interp2d( dx: float, mode: str = "bilinear", ) -> np.ndarray: - r"""Naive Interpolation + """Naive Interpolation (y,x): target pixel mode: interpolation mode """ @@ -31,7 +31,7 @@ def grid_sample( grid: np.ndarray, mode: str = "bilinear", ) -> np.ndarray: - r"""Naive grid sample algorithm""" + """Naive grid sample algorithm""" channels, h_in, w_in = img.shape _, h_out, w_out = grid.shape diff --git a/equilib/grid_sample/torch_grid_sample/basic.py b/equilib/grid_sample/torch_grid_sample/basic.py index dffd0499..f84b45b3 100644 --- a/equilib/grid_sample/torch_grid_sample/basic.py +++ b/equilib/grid_sample/torch_grid_sample/basic.py @@ -4,7 +4,7 @@ import torch -from equilib.common.torch_utils import get_device +from equilib.utils import torch_utils from .interp import linear_interp @@ -15,12 +15,12 @@ def interp2d( dx: torch.Tensor, mode: str = "bilinear", ) -> torch.Tensor: - r"""Naive Interpolation + """Naive Interpolation (y,x): target pixel mode: interpolation mode """ q00, q10, q01, q11 = Q - L = torch.tensor(1.0).to(get_device(Q)) + L = torch.tensor(1.0).to(torch_utils.get_device(Q)) if mode == "bilinear": f0 = linear_interp(q00, q01, dx, L) f1 = linear_interp(q10, q11, dx, L) @@ -34,7 +34,7 @@ def grid_sample( grid: torch.Tensor, mode: str = "bilinear", ) -> torch.Tensor: - r"""Torch Grid Sample (Custom) + """Torch Grid Sample (Custom) params: - img: Tensor[B, C, H, W] or Tensor[C, H, W] @@ -53,12 +53,12 @@ def grid_sample( batches, channels, h_in, w_in = img.shape _, _, h_out, w_out = grid.shape - assert get_device(img) == get_device( + assert torch_utils.get_device(img) == torch_utils.get_device( grid ), "ERR: img and grid does not match device " "{} vs {}".format( - get_device(img), get_device(grid) + torch_utils.get_device(img), torch_utils.get_device(grid) ) - device = get_device(img) + device = torch_utils.get_device(img) _dtype = img.dtype _max = torch.tensor(1.0, device=device) diff --git a/equilib/grid_sample/torch_grid_sample/interp.py b/equilib/grid_sample/torch_grid_sample/interp.py index 817fee03..3daecaea 100644 --- a/equilib/grid_sample/torch_grid_sample/interp.py +++ b/equilib/grid_sample/torch_grid_sample/interp.py @@ -2,5 +2,5 @@ def linear_interp(v0, v1, d, L): - r"""Basic Linear Interpolation""" + """Basic Linear Interpolation""" return v0 * (1 - d) / L + v1 * d / L diff --git a/equilib/grid_sample/torch_grid_sample/torch_func.py b/equilib/grid_sample/torch_grid_sample/torch_func.py index fe4c110a..f0814658 100644 --- a/equilib/grid_sample/torch_grid_sample/torch_func.py +++ b/equilib/grid_sample/torch_grid_sample/torch_func.py @@ -7,7 +7,7 @@ def grid_sample( img: torch.Tensor, grid: torch.Tensor, mode: str = "bilinear", **kwargs ) -> torch.Tensor: - r"""Torch Grid Sample (default) + """Torch Grid Sample (default) - Uses `torch.nn.functional.grid_sample` - By far the best way to sample diff --git a/equilib/info.py b/equilib/info.py new file mode 100644 index 00000000..6dfea0f5 --- /dev/null +++ b/equilib/info.py @@ -0,0 +1,8 @@ +#!/usr/bin/env python3 + +__version__ = "0.2.1" +__author__ = "Haruya Ishikawa" +__email__ = "haru.ishi43@gmail.com" +__homepage__ = "www.hoge.com" +__description__ = "equirectangular image processing with python" +__url__ = "https://github.com/haruishi43/equilib" diff --git a/equilib/utils/__init__.py b/equilib/utils/__init__.py new file mode 100644 index 00000000..422bdd98 --- /dev/null +++ b/equilib/utils/__init__.py @@ -0,0 +1,13 @@ +#!/usr/bin/env python3 + +from .rotation import ( + create_global2camera_rotation_matrix, + create_intrinsic_matrix, + create_rotation_matrix, +) + +__all__ = [ + "create_global2camera_rotation_matrix", + "create_intrinsic_matrix", + "create_rotation_matrix", +] diff --git a/equilib/utils/numpy_utils/__init__.py b/equilib/utils/numpy_utils/__init__.py new file mode 100644 index 00000000..420da1db --- /dev/null +++ b/equilib/utils/numpy_utils/__init__.py @@ -0,0 +1,7 @@ +#!/usr/bin/env python3 + +from .rotation import create_rotation_matrix + +__all__ = [ + "create_rotation_matrix", +] diff --git a/equilib/utils/numpy_utils/coord.py b/equilib/utils/numpy_utils/coord.py new file mode 100644 index 00000000..ab4ab7c8 --- /dev/null +++ b/equilib/utils/numpy_utils/coord.py @@ -0,0 +1,26 @@ +#!/usr/bin/env python3 + +import numpy as np + + +def create_coord( + height: int, + width: int, +) -> np.ndarray: + """Create mesh coordinate grid with height and width + + `z-axis` scale is `1` + + params: + - height (int) + - width (int) + + return: + - coordinate (np.ndarray) + """ + _xs = np.linspace(0, width - 1, width) + _ys = np.linspace(0, height - 1, height) + xs, ys = np.meshgrid(_xs, _ys) + zs = np.ones_like(xs) + coord = np.stack((xs, ys, zs), axis=2) + return coord diff --git a/equilib/utils/numpy_utils/rotation.py b/equilib/utils/numpy_utils/rotation.py new file mode 100644 index 00000000..aa47d0ea --- /dev/null +++ b/equilib/utils/numpy_utils/rotation.py @@ -0,0 +1,123 @@ +#!/usr/bin/env python3 + +import numpy as np + + +def create_global2camera_rotation_matrix() -> np.ndarray: + """Rotation from global (world) to camera coordinates""" + R_XY = np.array( + [ # X <-> Y + [0.0, 1.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 0.0, 1.0], + ] + ) + R_YZ = np.array( + [ # Y <-> Z + [1.0, 0.0, 0.0], + [0.0, 0.0, 1.0], + [0.0, 1.0, 0.0], + ] + ) + return R_XY @ R_YZ + + +def create_intrinsic_matrix( + w_pers: int, + h_pers: int, + fov_x: float, + skew: float, +) -> np.ndarray: + """Create Intrinsic Matrix""" + # perspective projection (focal length) + f = w_pers / (2.0 * np.tan(np.radians(fov_x) / 2.0)) + # transform between camera frame and pixel coordinates + K = np.array( + [ + [f, skew, w_pers / 2], + [0.0, f, h_pers / 2], + [0.0, 0.0, 1.0], + ] + ) + return K + + +def create_rotation_matrix( + roll: float, + pitch: float, + yaw: float, + z_down: bool = False, +) -> np.ndarray: + """Create Rotation Matrix""" + # calculate rotation about the x-axis + R_x = np.array( + [ + [1.0, 0.0, 0.0], + [0.0, np.cos(roll), -np.sin(roll)], + [0.0, np.sin(roll), np.cos(roll)], + ] + ) + # calculate rotation about the y-axis + if z_down: + pitch = -pitch + R_y = np.array( + [ + [np.cos(pitch), 0.0, -np.sin(pitch)], + [0.0, 1.0, 0.0], + [np.sin(pitch), 0.0, np.cos(pitch)], + ] + ) + # calculate rotation about the z-axis + if z_down: + yaw = -yaw + R_z = np.array( + [ + [np.cos(yaw), np.sin(yaw), 0.0], + [-np.sin(yaw), np.cos(yaw), 0.0], + [0.0, 0.0, 1.0], + ] + ) + + return R_z @ R_y @ R_x + + +def _create_rotation_matrix( + x: float, + y: float, + z: float, + z_down: bool = False, +) -> np.ndarray: + """Create Rotation Matrix + + NOTE: DEPRECATED + """ + # calculate rotation about the x-axis + R_x = np.array( + [ + [1.0, 0.0, 0.0], + [0.0, np.cos(x), -np.sin(x)], + [0.0, np.sin(x), np.cos(x)], + ] + ) + # calculate rotation about the y-axis + if z_down: + y = -y + R_y = np.array( + [ + [np.cos(y), 0.0, np.sin(y)], + [0.0, 1.0, 0.0], + [-np.sin(y), 0.0, np.cos(y)], + ] + ) + # calculate rotation about the z-axis + if z_down: + z = -z + R_z = np.array( + [ + [np.cos(z), -np.sin(z), 0.0], + [np.sin(z), np.cos(z), 0.0], + [0.0, 0.0, 1.0], + ] + ) + + return R_z @ R_y @ R_x diff --git a/equilib/utils/rotation.py b/equilib/utils/rotation.py new file mode 100644 index 00000000..f82f8b18 --- /dev/null +++ b/equilib/utils/rotation.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 + +from typing import Union + +import numpy as np + +import torch + +from .numpy_utils.rotation import ( + create_global2camera_rotation_matrix as np_grm, + create_intrinsic_matrix as np_im, + create_rotation_matrix as np_rm, +) +from .torch_utils.rotation import ( + create_global2camera_rotation_matrix as th_grm, + create_intrinsic_matrix as th_im, + create_rotation_matrix as th_rm, +) + + +def create_global2camera_rotation_matrix( + is_torch: bool, +) -> Union[np.ndarray, torch.Tensor]: + """Create Rotation Matrix that transform from global to camera coordinate + + params: + - is_torch (bool): return torch.Tensor + + return: + - rotation matrix (np.ndarray, torch.Tensor): 3x3 rotation matrix + """ + if is_torch: + return th_grm() + else: + return np_grm() + + +def create_intrinsic_matrix( + w_pers: int, + h_pers: int, + fov_x: float, + skew: float, + is_torch: bool, +) -> Union[np.ndarray, torch.Tensor]: + """Create Intrinsic Matrix + + params: + - w_pers (int) + - h_pers (int) + - fov_x (float) + - skew (float) + - is_torch (bool): return torch.Tensor + + return: + - K (np.ndarray, torch.Tensor): 3x3 matrix + + NOTE: + - ref: http://ksimek.github.io/2013/08/13/intrinsic/ + """ + if is_torch: + return th_im( + w_pers=w_pers, + h_pers=h_pers, + fov_x=fov_x, + skew=skew, + ) + else: + return np_im( + w_pers=w_pers, + h_pers=h_pers, + fov_x=fov_x, + skew=skew, + ) + + +def create_rotation_matrix( + roll: float, + pitch: float, + yaw: float, + z_down: bool, + is_torch: bool, +) -> Union[np.ndarray, torch.Tensor]: + """Create Rotation Matrix + + params: + - roll (float): x-axis rotation + - pitch (float): y-axis rotation + - yaw (float): z-axis rotation + - z_down (bool): make z-axis face down + - is_torch (bool): return torch.Tensor + + return: + - rotation matrix (np.ndarray, torch.Tensor): 3x3 matrix + """ + if is_torch: + return th_rm( + roll=roll, + pitch=pitch, + yaw=yaw, + z_down=z_down, + ) + else: + return np_rm( + roll=roll, + pitch=pitch, + yaw=yaw, + z_down=z_down, + ) diff --git a/equilib/utils/torch_utils/__init__.py b/equilib/utils/torch_utils/__init__.py new file mode 100644 index 00000000..6e06b243 --- /dev/null +++ b/equilib/utils/torch_utils/__init__.py @@ -0,0 +1,11 @@ +#!/usr/bin/env python3 + +from .rotation import create_rotation_matrix, deg2rad +from .torch_func import get_device, sizeof + +__all__ = [ + "create_rotation_matrix", + "deg2rad", + "get_device", + "sizeof", +] diff --git a/equilib/utils/torch_utils/coord.py b/equilib/utils/torch_utils/coord.py new file mode 100644 index 00000000..1992a67d --- /dev/null +++ b/equilib/utils/torch_utils/coord.py @@ -0,0 +1,18 @@ +#!/usr/bin/env python3 + +import torch + + +def create_coord( + height: int, + width: int, +) -> torch.Tensor: + """Create mesh coordinate grid""" + _xs = torch.linspace(0, width - 1, width) + _ys = torch.linspace(0, height - 1, height) + # NOTE: https://github.com/pytorch/pytorch/issues/15301 + # Torch meshgrid behaves differently than numpy + ys, xs = torch.meshgrid([_ys, _xs]) + zs = torch.ones_like(xs) + coord = torch.stack((xs, ys, zs), dim=2) + return coord diff --git a/equilib/utils/torch_utils/rotation.py b/equilib/utils/torch_utils/rotation.py new file mode 100644 index 00000000..6636011f --- /dev/null +++ b/equilib/utils/torch_utils/rotation.py @@ -0,0 +1,137 @@ +#!/usr/bin/env python3 + +import numpy as np + +import torch + +pi = torch.Tensor([3.14159265358979323846]) + + +def deg2rad(tensor: torch.Tensor) -> torch.Tensor: + """Function that converts angles from degrees to radians""" + return tensor * pi.to(tensor.device).type(tensor.dtype) / 180.0 + + +def create_global2camera_rotation_matrix() -> torch.Tensor: + """Rotation from global (world) to camera coordinates""" + R_XY = torch.tensor( + [ # X <-> Y + [0.0, 1.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 0.0, 1.0], + ] + ) + R_YZ = torch.tensor( + [ # Y <-> Z + [1.0, 0.0, 0.0], + [0.0, 0.0, 1.0], + [0.0, 1.0, 0.0], + ] + ) + return R_XY @ R_YZ + + +def create_intrinsic_matrix( + w_pers: int, + h_pers: int, + fov_x: float, + skew: float, +) -> torch.Tensor: + """Create Intrinsic Matrix""" + fov_x = torch.tensor(fov_x) + f = w_pers / (2 * torch.tan(deg2rad(fov_x) / 2)) + K = torch.tensor( + [ + [f, skew, w_pers / 2], + [0.0, f, h_pers / 2], + [0.0, 0.0, 1.0], + ] + ) + return K + + +def create_rotation_matrix( + roll: float, + pitch: float, + yaw: float, + z_down: bool = False, +) -> torch.Tensor: + """Create Rotation Matrix""" + # calculate rotation about the x-axis + R_x = torch.tensor( + [ + [1.0, 0.0, 0.0], + [0.0, np.cos(roll), -np.sin(roll)], + [0.0, np.sin(roll), np.cos(roll)], + ], + dtype=torch.float, + ) + # calculate rotation about the y-axis + if z_down: + pitch = -pitch + R_y = torch.tensor( + [ + [np.cos(pitch), 0.0, -np.sin(pitch)], + [0.0, 1.0, 0.0], + [np.sin(pitch), 0.0, np.cos(pitch)], + ], + dtype=torch.float, + ) + # calculate rotation about the z-axis + if z_down: + yaw = -yaw + R_z = torch.tensor( + [ + [np.cos(yaw), np.sin(yaw), 0.0], + [-np.sin(yaw), np.cos(yaw), 0.0], + [0.0, 0.0, 1.0], + ], + dtype=torch.float, + ) + + return R_z @ R_y @ R_x + + +def _create_rotation_matrix( + x: float, + y: float, + z: float, + z_down: bool = False, +) -> torch.Tensor: + """Create Rotation Matrix + + NOTE: DEPRECATED + """ + # calculate rotation about the x-axis + R_x = torch.tensor( + [ + [1.0, 0.0, 0.0], + [0.0, np.cos(x), -np.sin(x)], + [0.0, np.sin(x), np.cos(x)], + ], + dtype=torch.float, + ) + # calculate rotation about the y-axis + if z_down: + y = -y + R_y = torch.tensor( + [ + [np.cos(y), 0.0, np.sin(y)], + [0.0, 1.0, 0.0], + [-np.sin(y), 0.0, np.cos(y)], + ], + dtype=torch.float, + ) + # calculate rotation about the z-axis + if z_down: + z = -z + R_z = torch.tensor( + [ + [np.cos(z), -np.sin(z), 0.0], + [np.sin(z), np.cos(z), 0.0], + [0.0, 0.0, 1.0], + ], + dtype=torch.float, + ) + + return R_z @ R_y @ R_x diff --git a/equilib/utils/torch_utils/torch_func.py b/equilib/utils/torch_utils/torch_func.py new file mode 100644 index 00000000..2ae426fd --- /dev/null +++ b/equilib/utils/torch_utils/torch_func.py @@ -0,0 +1,14 @@ +#!/usr/bin/env python3 + +import torch + + +def sizeof(tensor: torch.Tensor) -> float: + """Get the size of a tensor""" + assert torch.is_tensor(tensor), "ERR: is not tensor" + return tensor.element_size() * tensor.nelement() + + +def get_device(a: torch.Tensor) -> torch.device: + """Get device of a Tensor""" + return torch.device(a.get_device() if a.get_device() >= 0 else "cpu") diff --git a/setup.py b/setup.py index f537f1a8..f87757be 100644 --- a/setup.py +++ b/setup.py @@ -16,12 +16,40 @@ def readme(): def find_version(): - version_file = "equilib/__init__.py" + version_file = "equilib/info.py" with open(version_file, "r") as f: exec(compile(f.read(), version_file, "exec")) return locals()["__version__"] +def find_author(): + version_file = "equilib/info.py" + with open(version_file, "r") as f: + exec(compile(f.read(), version_file, "exec")) + return locals()["__author__"] + + +def find_email(): + version_file = "equilib/info.py" + with open(version_file, "r") as f: + exec(compile(f.read(), version_file, "exec")) + return locals()["__email__"] + + +def find_description(): + version_file = "equilib/info.py" + with open(version_file, "r") as f: + exec(compile(f.read(), version_file, "exec")) + return locals()["__description__"] + + +def find_url(): + version_file = "equilib/info.py" + with open(version_file, "r") as f: + exec(compile(f.read(), version_file, "exec")) + return locals()["__url__"] + + def get_requirements(filename="requirements.txt"): here = osp.dirname(osp.realpath(__file__)) with open(osp.join(here, filename), "r") as f: @@ -64,13 +92,13 @@ def get_long_description(): name="pyequilib", version=find_version(), packages=find_packages(exclude=["github2pypi"]), - description="equirectangular image processing with python", + description=find_description(), long_description=get_long_description(), long_description_content_type="text/markdown", - author="Haruya Ishikawa", - author_email="www.haru.ishi43@gmail.com", + author=find_author(), + author_email=find_email(), license="AGPL-3.0", - url="https://github.com/haruishi43/equilib", + url=find_url(), install_requires=get_requirements(), keywords=["Equirectangular", "Computer Vision"], classifiers=[ diff --git a/tests/common/timer.py b/tests/common/timer.py index bb530eff..ba2ec3b0 100644 --- a/tests/common/timer.py +++ b/tests/common/timer.py @@ -1,10 +1,11 @@ #!/usr/bin/env python3 import time +from typing import Callable -def timer(func): - r"""Decorator that reports the execution time.""" +def timer(func: Callable): + """Decorator that reports the execution time""" def wrap(*args, **kwargs): tic = time.perf_counter() diff --git a/tests/test_equi2cube.py b/tests/test_equi2cube.py index dcb9c2f7..69313d73 100644 --- a/tests/test_equi2cube.py +++ b/tests/test_equi2cube.py @@ -20,6 +20,7 @@ CUBE_FORMAT = "dict" # Output cube format SAMPLING_METHOD = "default" # Sampling method MODE = "bilinear" # Sampling mode +Z_DOWN = False # z-axis control USE_CLASS = True # Class or function # Paths @@ -41,6 +42,7 @@ def run_equi2cube( cube_format: str, sampling_method: str, mode: str, + z_down: bool, use_class: bool, ) -> Union[np.ndarray, torch.Tensor]: # h_equi, w_equi = equi.shape[-2:] @@ -51,6 +53,7 @@ def run_equi2cube( cube_format=cube_format, sampling_method=sampling_method, mode=mode, + z_down=z_down, ) sample = equi2cube_instance( equi=equi, @@ -64,6 +67,7 @@ def run_equi2cube( cube_format=cube_format, sampling_method=sampling_method, mode=mode, + z_down=z_down, ) return sample @@ -320,6 +324,7 @@ def test_numpy_single(): cube_format=CUBE_FORMAT, sampling_method=SAMPLING_METHOD, mode=MODE, + z_down=Z_DOWN, use_class=USE_CLASS, ) process_single_numpy_output( @@ -345,6 +350,7 @@ def test_numpy_batch(): cube_format=CUBE_FORMAT, sampling_method=SAMPLING_METHOD, mode=MODE, + z_down=Z_DOWN, use_class=USE_CLASS, ) process_batch_numpy_output( @@ -369,6 +375,7 @@ def test_torch_single(): cube_format=CUBE_FORMAT, sampling_method=SAMPLING_METHOD, mode=MODE, + z_down=Z_DOWN, use_class=USE_CLASS, ) process_single_torch_output( @@ -395,6 +402,7 @@ def test_torch_batch(): cube_format=CUBE_FORMAT, sampling_method=SAMPLING_METHOD, mode=MODE, + z_down=Z_DOWN, use_class=USE_CLASS, ) process_batch_torch_output( diff --git a/tests/test_equi2equi.py b/tests/test_equi2equi.py index 64684c8f..751960ac 100644 --- a/tests/test_equi2equi.py +++ b/tests/test_equi2equi.py @@ -20,6 +20,7 @@ WIDTH, HEIGHT = (640, 320) # Output panorama shape SAMPLING_METHOD = "default" # Sampling method MODE = "bilinear" # Sampling mode +Z_DOWN = True # z-axis control USE_CLASS = True # Class or function # Paths @@ -41,6 +42,7 @@ def run_equi2equi( h_out: int, sampling_method: str, mode: str, + z_down: bool, use_class: bool, ) -> Union[np.ndarray, torch.Tensor]: # h_equi, w_equi = equi.shape[-2:] @@ -51,6 +53,7 @@ def run_equi2equi( h_out=h_out, sampling_method=sampling_method, mode=mode, + z_down=z_down, ) sample = equi2equi_instance( src=equi, @@ -62,6 +65,7 @@ def run_equi2equi( rot=rot, sampling_method=sampling_method, mode=mode, + z_down=z_down, w_out=w_out, h_out=h_out, ) @@ -243,6 +247,7 @@ def test_numpy_single(): h_out=HEIGHT, sampling_method=SAMPLING_METHOD, mode=MODE, + z_down=Z_DOWN, use_class=USE_CLASS, ) process_single_numpy_output( @@ -265,6 +270,7 @@ def test_numpy_batch(): h_out=HEIGHT, sampling_method=SAMPLING_METHOD, mode=MODE, + z_down=Z_DOWN, use_class=USE_CLASS, ) process_batch_numpy_output( @@ -288,6 +294,7 @@ def test_torch_single(): h_out=HEIGHT, sampling_method=SAMPLING_METHOD, mode=MODE, + z_down=Z_DOWN, use_class=USE_CLASS, ) process_single_torch_output( @@ -309,6 +316,7 @@ def test_torch_batch(): h_out=HEIGHT, sampling_method=SAMPLING_METHOD, mode=MODE, + z_down=Z_DOWN, use_class=USE_CLASS, ) process_batch_torch_output( diff --git a/tests/test_equi2pers.py b/tests/test_equi2pers.py index 14cdea92..b5c145b8 100644 --- a/tests/test_equi2pers.py +++ b/tests/test_equi2pers.py @@ -21,6 +21,7 @@ FOV = 90 # Output pers fov SAMPLING_METHOD = "default" # Sampling method MODE = "bilinear" # Sampling mode +Z_DOWN = False # z-axis control USE_CLASS = True # Class or function # Paths @@ -43,6 +44,7 @@ def run_equi2pers( fov: int, sampling_method: str, mode: str, + z_down: bool, use_class: bool, ) -> Union[np.ndarray, torch.Tensor]: # h_equi, w_equi = equi.shape[-2:] @@ -54,6 +56,7 @@ def run_equi2pers( fov_x=fov, sampling_method=sampling_method, mode=mode, + z_down=z_down, ) sample = equi2pers_instance( equi=equi, @@ -68,6 +71,7 @@ def run_equi2pers( fov_x=fov, sampling_method=sampling_method, mode=mode, + z_down=z_down, ) return sample @@ -249,6 +253,7 @@ def test_numpy_single(): fov=FOV, sampling_method=SAMPLING_METHOD, mode=MODE, + z_down=Z_DOWN, use_class=USE_CLASS, ) process_single_numpy_output( @@ -274,6 +279,7 @@ def test_numpy_batch(): fov=FOV, sampling_method=SAMPLING_METHOD, mode=MODE, + z_down=Z_DOWN, use_class=USE_CLASS, ) process_batch_numpy_output( @@ -298,6 +304,7 @@ def test_torch_single(): fov=FOV, sampling_method=SAMPLING_METHOD, mode=MODE, + z_down=Z_DOWN, use_class=USE_CLASS, ) process_single_torch_output( @@ -324,6 +331,7 @@ def test_torch_batch(): fov=FOV, sampling_method=SAMPLING_METHOD, mode=MODE, + z_down=Z_DOWN, use_class=USE_CLASS, ) process_batch_torch_output( diff --git a/equilib/common/__init__.py b/tests/utils/__init__.py similarity index 100% rename from equilib/common/__init__.py rename to tests/utils/__init__.py diff --git a/tests/utils/test_rotation.py b/tests/utils/test_rotation.py new file mode 100644 index 00000000..e5a0d9b4 --- /dev/null +++ b/tests/utils/test_rotation.py @@ -0,0 +1 @@ +#!/usr/bin/env python3