"""Utils for loading images and annotations."""
import os
from functools import lru_cache
from pathlib import Path
from typing import List, Tuple, Union
import numpy as np
from ..data_structure.constants import PathLike
from ..data_structure.models import RenderOutputEnumBlender
os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1'
try:
import cv2
import flow_vis
except ImportError:
raise ImportError('Failed to import necessary packages. Please install by: \npip install "xrfeitoria[vis]"')
__all__ = ['Viewer']
class ExrReader:
"""Utils for exr format data.
Load `.exr` format file.
"""
def __init__(self, exr_path: Union[str, Path]):
"""Initialize with a `.exr` format file.
Args:
exr_path (PathLike): path to `.exr` format file
"""
if isinstance(exr_path, Path):
exr_path = exr_path.as_posix()
exr_mat = cv2.imread(exr_path, cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH)
exr_mat = cv2.cvtColor(exr_mat, cv2.COLOR_BGR2RGB)
self.exr_mat = exr_mat
@staticmethod
def float2int(array: np.ndarray) -> np.ndarray:
"""Convert ndarray to uint8 that can be display as image. Values greater than
255 will be clipped.
Args:
array (np.ndarray): input ndarray
Returns:
np.ndarray: array of dtype=unit8
"""
array = np.round(array * 255)
array = np.clip(array, 0, 255)
return array.astype(np.uint8)
def get_rgb(self) -> np.ndarray:
"""Get RGB channel in `.exr` format.
Returns:
np.ndarray: masks of shape (H, W, 3)
"""
rgb = self.exr_mat
img = self.float2int(rgb)
return img
def get_flow(self) -> np.ndarray:
"""Get optical flow in `.exr` format.
Returns:
np.ndarray: optical flow data of (H, W, 3) converted to colors
"""
flow = self.exr_mat[..., :2]
img = flow_vis.flow_to_color(flow, convert_to_bgr=False)
return img
def get_depth(self, inverse: bool = False, depth_rescale: float = 1.0) -> np.ndarray:
"""Get depth in `.exr` format.
Args:
inverse (bool, optional): whether to inverse the depth.
If True, white (255) represents the farthest, and black (0) represents the nearest.
if False, white (255) represents the nearest, and black (0) represents the farthest.
Defaults to False.
depth_rescale (float, optional): scaling the depth to map it into (0, 255).
``depth = depth / depth_rescale``.
Depth values greater than `depth_rescale` will be clipped. Defaults to 1.0.
Returns:
np.ndarray: depth data of shape (H, W, 3)
"""
depth = self.exr_mat
img = self.float2int(depth / depth_rescale)
if inverse:
img = 255 - img
else:
img[img == 0] = 255
return img
[docs]
class Viewer:
"""Utils for loading images and annotations.
Examples:
>>> viewer = Viewer('/path/to/sequence')
>>> camera_name = 'cam'
>>> frame = 0
>>> img = viewer.get_img(camera_name=camera_name, frame=frame)
>>> mask = viewer.get_mask(camera_name=camera_name, frame=frame)
>>> depth = viewer.get_depth(camera_name=camera_name, frame=frame)
>>> flow = viewer.get_flow(camera_name=camera_name, frame=frame)
>>> normal = viewer.get_normal(camera_name=camera_name, frame=frame)
>>> diffuse = viewer.get_diffuse(camera_name=camera_name, frame=frame)
"""
# folder names of each data modal
IMG = RenderOutputEnumBlender.img.name
MASK = RenderOutputEnumBlender.mask.name
DEPTH = RenderOutputEnumBlender.depth.name
FLOW = RenderOutputEnumBlender.flow.name
NORMAL = RenderOutputEnumBlender.normal.name
DIFFUSE = RenderOutputEnumBlender.diffuse.name
[docs]
def __init__(self, sequence_dir: PathLike) -> None:
"""Initialize with the sequence directory.
Args:
sequence_dir (PathLike): path to the sequence directory
"""
self.sequence_dir = Path(sequence_dir)
@property
@lru_cache
def camera_names(self) -> List[str]:
camera_folders = list(self.sequence_dir.glob(f'{self.IMG}/*'))
return [camera_folder.name for camera_folder in camera_folders]
@property
@lru_cache
def frame_num(self) -> int:
return len(list(self.sequence_dir.glob(f'{self.IMG}/{self.camera_names[0]}/*')))
[docs]
def get_img(self, camera_name: str, frame: int) -> np.ndarray:
"""Get rgb image of the given frame ('img/{frame:04d}.*')
Args:
camera_name (str): the camera name
frame (int): the frame number
Returns:
np.ndarray: image of shape (H, W, 3)
"""
folder = self.sequence_dir / self.IMG / camera_name
if not folder.exists():
raise ValueError(f'Folder of rgb images not found: {folder}')
file_path = next(folder.glob(f'{frame:04d}.*')).resolve()
if not file_path.exists():
raise ValueError(f'Image of {frame}-frame not found: {file_path}')
if file_path.suffix == '.exr':
return ExrReader(file_path).get_rgb()
else:
img = cv2.imread(str(file_path))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
return img
[docs]
def get_diffuse(self, camera_name: str, frame: int) -> np.ndarray:
"""Get diffuse image of the given frame ('diffuse/{frame:04d}.*')
Args:
camera_name (str): the camera name
frame (int): the frame number
Returns:
np.ndarray: image of shape (H, W, 3)
"""
folder = self.sequence_dir / self.DIFFUSE / camera_name
if not folder.exists():
raise ValueError(f'Folder of diffuse images not found: {folder}')
file_path = next(folder.glob(f'{frame:04d}.*')).resolve()
if not file_path.exists():
raise ValueError(f'Diffuse image of {frame}-frame not found: {file_path}')
if file_path.suffix == '.exr':
return ExrReader(file_path).get_rgb()
else:
img = cv2.imread(str(file_path))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
return img
[docs]
def get_mask(self, camera_name: str, frame: int) -> np.ndarray:
"""Get mask of the given frame ('mask/{frame:04d}.*')
Args:
camera_name (str): the camera name
frame (int): the frame number
Returns:
np.ndarray: image of shape (H, W, 3)
"""
folder = self.sequence_dir / self.MASK / camera_name
if not folder.exists():
raise ValueError(f'Folder of masks not found: {folder}')
file_path = next(folder.glob(f'{frame:04d}.*')).resolve()
if not file_path.exists():
raise ValueError(f'Mask of {frame}-frame not found: {file_path}')
if file_path.suffix == '.exr':
return ExrReader(file_path).get_rgb()
else:
img = cv2.imread(str(file_path))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
return img
[docs]
def get_depth(self, camera_name: str, frame: int, inverse: bool = False, depth_rescale: float = 1.0) -> np.ndarray:
"""Get depth of the given frame ('depth/{frame:04d}.*')
Args:
camera_name (str): the camera name
frame (int): the frame number
inverse (bool, optional): whether to inverse the depth.
If True, white (255) represents the farthest, and black (0) represents the nearest.
if False, white (255) represents the nearest, and black (0) represents the farthest.
Defaults to False.
depth_rescale (float, optional): scaling the depth to map it into (0, 255).
``depth = depth / depth_rescale``.
Depth values greater than `depth_rescale` will be clipped. Defaults to 1.0.
Returns:
np.ndarray: depth of shape (H, W, 3)
"""
folder = self.sequence_dir / self.DEPTH / camera_name
if not folder.exists():
raise ValueError(f'Folder of depth not found: {folder}')
file_path = next(folder.glob(f'{frame:04d}.*')).resolve()
if not file_path.exists():
raise ValueError(f'Depth of {frame}-frame not found: {file_path}')
if file_path.suffix == '.exr':
return ExrReader(file_path).get_depth(inverse=inverse, depth_rescale=depth_rescale)
else:
img = cv2.imread(str(file_path))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
return img
[docs]
def get_flow(self, camera_name: str, frame: int) -> np.ndarray:
"""Get optical flow of the given frame ('flow/{frame:04d}.*')
Args:
camera_name (str): the camera name
frame (int): the frame number
Returns:
np.ndarray: optical flow of shape (H, W, 3)
"""
folder = self.sequence_dir / self.FLOW / camera_name
if not folder.exists():
raise ValueError(f'Folder of depth not found: {folder}')
file_path = next(folder.glob(f'{frame:04d}.*')).resolve()
if not file_path.exists():
raise ValueError(f'Depth of {frame}-frame not found: {file_path}')
if file_path.suffix == '.exr':
return ExrReader(file_path).get_flow()
else:
img = cv2.imread(str(file_path))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
return img
[docs]
def get_normal(self, camera_name: str, frame: int) -> np.ndarray:
"""Get normal map of the given frame ('normal/{frame:04d}.*')
Args:
camera_name (str): the camera name
frame (int): the frame number
Returns:
np.ndarray: normal map of shape (H, W, 3)
"""
folder = self.sequence_dir / self.NORMAL / camera_name
if not folder.exists():
raise ValueError(f'Folder of normal mpa not found: {folder}')
file_path = next(folder.glob(f'{frame:04d}.*')).resolve()
if not file_path.exists():
raise ValueError(f'Normal map of {frame}-frame not found: {file_path}')
if file_path.suffix == '.exr':
return ExrReader(file_path).get_rgb()
else:
img = cv2.imread(str(file_path))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
return img