core/utils.py

from enum import Enum
import torch
from torch import Tensor
import visibility
import numpy as np
from pathlib import Path
from easydict import EasyDict as edict
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.dataloader import default_collate
import random
from torch import initial_seed
import torch.nn as nn
from numpy import random as np_random
from scipy.spatial.transform import Rotation as R
from PIL import Image
import h5py
from numpy import ndarray
import pandas as pd
import torchvision.transforms.functional as TTF
from torchvision import transforms
import abc
from tqdm import tqdm
from torch.backends import cudnn
from torch.nn.parallel import DistributedDataParallel


class CameraIntrinsicParameters(Tensor):
    def __new__(cls, focal_length_x, focal_length_y, principal_point_x, principal_point_y):
        data = torch.tensor([focal_length_x, focal_length_y, principal_point_x, principal_point_y], dtype=torch.float32)
        return Tensor._make_subclass(cls, data, data.requires_grad)

    def __init__(self, focal_length_x, focal_length_y, principal_point_x, principal_point_y):
        pass

    @property
    def focal_length_x(self):
        return self[0]

    @focal_length_x.setter
    def focal_length_x(self, value):
        self[0] = value

    @property
    def focal_length_y(self):
        return self[1]

    @focal_length_y.setter
    def focal_length_y(self, value):
        self[1] = value

    @property
    def principal_point_x(self):
        return self[2]

    @principal_point_x.setter
    def principal_point_x(self, value):
        self[2] = value

    @property
    def principal_point_y(self):
        return self[3]

    @principal_point_y.setter
    def principal_point_y(self, value):
        self[3] = value

    def to_matrix(self) -> Tensor:
        return torch.tensor([[self.focal_length_x, 0, self.principal_point_x],
                             [0, self.focal_length_y, self.principal_point_y],
                             [0, 0, 1.]], dtype=torch.float32)

    def scale(self, w_scale: float, h_scale: float) -> 'CameraIntrinsicParameters':
        self.focal_length_x *= w_scale
        self.focal_length_y *= h_scale
        self.principal_point_x *= w_scale
        self.principal_point_y *= h_scale
        return self

class EngineMode(Enum):
    """
    Enum class representing the different types of data splits.

    Attributes:
        TRAIN (int): Represents the training data split.
        VALID (int): Represents the validation data split.
        TEST (int): Represents the test data split.
    """
    TRAIN = 0
    VALID = 1
    TEST = 2
    
class DepthFlowGenerator:
    def __init__(self, cfg: edict):
        self._real_shape = None
        self._cfg = cfg
        self._occlusion_threshold = cfg.dataset.occlusion_threshold
        self._occlusion_kernel = cfg.dataset.occlusion_kernel

    def gen_depth_img(self, uv, depth, index, cam_params: CameraIntrinsicParameters):
        device = uv.device

        depth_image = torch.zeros(
            self._real_shape[:2], device=device, dtype=torch.float
        )
        depth_image += 1000.0
        mask_image = (-1) * torch.ones(
            self._real_shape[:2], device=device, dtype=torch.float
        )
        index = index.float()
        depth_image, mask_image = visibility.depth_image(
            uv,
            depth,
            index,
            depth_image,
            mask_image,
            uv.shape[0],
            self._real_shape[1],
            self._real_shape[0],
        )
        depth_image[depth_image == 1000.0] = 0.0
        mask_image_deocclusion = (-1) * torch.ones(
            self._real_shape[:2], device=device, dtype=torch.float
        )
        depth_image_no_occlusion = torch.zeros_like(depth_image, device=device)
        depth_image_no_occlusion, mask_image_deocclusion = visibility.visibility2(
            depth_image,
            cam_params,
            mask_image,
            depth_image_no_occlusion,
            mask_image_deocclusion,
            depth_image.shape[1],
            depth_image.shape[0],
            self._occlusion_threshold,
            self._occlusion_kernel,
        )
        return (
            depth_image_no_occlusion,
            mask_image_deocclusion.int(),
            mask_image,
            mask_image,
        )

    def flatten_mask(self, mask_deocclusion, range_mask_uv):
        index_deocclusion = torch.where(mask_deocclusion > 0)
        mask_deocclusion = mask_deocclusion[
            index_deocclusion[0][:], index_deocclusion[1][:]
        ]
        mask_flatten = torch.zeros(
            range_mask_uv.shape[0], device=mask_deocclusion.device, dtype=torch.int32
        )
        mask_flatten[mask_deocclusion.cpu().numpy() - 1] = mask_deocclusion
        return mask_flatten

    def crop_data_from_dict(self, data: dict, patch_shape, engine_mode: EngineMode):
        H, W = patch_shape[:2]
        patch_H, patch_W = patch_shape[2:]
        assert (
            patch_H <= H and patch_W <= W
        ), "Patch size should be smaller than the image size"
        if engine_mode == EngineMode.TRAIN:
            x = np.random.randint(0, H - patch_H) if H > patch_H else 0
            y = np.random.randint(0, W - patch_W) if W > patch_W else 0
        else:
            x = (H - patch_H) // 2
            y = (W - patch_W) // 2
        # Unpack Data from Dict
        return {
            key: value[..., x : x + patch_H, y : y + patch_W]
            for key, value in data.items()
        }

    def push(self, data_dict: dict, engine_mode=EngineMode.TRAIN):
        vision_images = data_dict["vision_image"]
        point_clouds = data_dict["point_cloud"]
        camera_intrinsic_parameters = data_dict["camera_intrinsic_parameters"]
        T_errs = data_dict["tr_error"]
        R_errs = data_dict["rot_error"]
        orders = data_dict["order"]
        device = vision_images[0].device

        vision_images_input = []
        depth_images_input = []
        depth_images_fine = []
        flow_images_gt = []
        valid_masks = []

        for idx in range(len(vision_images)):
            # 1 Unpack Data
            vision_image = vision_images[idx].to(device)
            point_cloud_fine = point_clouds[idx].clone().to(device)
            cam_params = camera_intrinsic_parameters[idx]
            order = orders[idx]
            self._real_shape = [
                int(vision_image.shape[1]),
                int(vision_image.shape[2]),
                vision_image.shape[0],
            ]
            # 2 Transform Point Cloud
            transform_fine2coarse = get_transform_from_rotation_translation(
                R_errs[idx].to(device), T_errs[idx].to(device)
            ).squeeze(0)
            point_cloud_coarse = apply_transform_to_points(
                point_cloud_fine, transform_fine2coarse
            )
            # 3 Project Point Cloud
            uv_fine, depth_fine, mask_fine = project_with_mask(
                point_cloud_fine, self._real_shape, cam_params, order
            )  # (2, N_fine), (N_fine), (N)
            uv_fine = uv_fine.t().int().contiguous()  # (N_fine, 2)
            uv_coarse, depth_coarse, mask_coarse = project_with_mask(
                point_cloud_coarse, self._real_shape, cam_params, order
            )  # (2, N_coarse), (N_coarse), (N)
            uv_coarse = uv_coarse.t().int().contiguous()  # (N_coarse, 2)

            # 4 Get Flow Set
            flow_set, mask_flow = get_flow_set_from_2pixel_sets(
                uv_coarse, uv_fine, mask_coarse, mask_fine
            )  # (N_flow, 2), (N_flow) in (N)

            # 5 Get Depth Image
            ## 5.1 Filter flow points in coarse points
            mask_flow_coarse = mask_coarse[mask_flow]  # (N_flow)
            range_mask_uv_coarse = (
                torch.arange(mask_flow_coarse.shape[0]).to(device) + 1
            )  # (1, ..., N_flow)
            uv_coarse_in_flow = uv_coarse[
                mask_flow[mask_coarse], :
            ]  # N_coarse [(N_flow) in (N_coarse)] -> (N_flow, 2)
            depth_coarse_in_flow = depth_coarse[
                mask_flow[mask_coarse]
            ]  # N_coarse [(N_flow) in (N_coarse)] -> (N_flow)

            ## 5.2 Filter flow points in fine points
            mask_flow_fine = mask_fine[mask_flow]  # (N_flow)
            range_mask_uv_fine = (
                torch.arange(mask_flow_fine.shape[0]).to(device) + 1
            )  # (1, ..., N_flow)
            uv_fine_in_flow = uv_fine[
                mask_flow[mask_fine], :
            ]  # N_fine [(N_flow) in (N_fine)] -> (N_flow, 2)
            depth_fine_in_flow = depth_fine[
                mask_flow[mask_fine]
            ]  # N_fine [(N_flow) in (N_fine)] -> (N_flow)

            ## 5.3 Get Deocclusion Mask
            _, mask_deocclusion_coarse, _, _ = self.gen_depth_img(
                uv_coarse_in_flow,
                depth_coarse_in_flow,
                range_mask_uv_coarse,
                cam_params,
            )
            mask_depth_coarse = self.flatten_mask(
                mask_deocclusion_coarse, range_mask_uv_coarse
            )
            _, mask_deocclusion_fine, _, _ = self.gen_depth_img(
                uv_fine_in_flow, depth_fine_in_flow, range_mask_uv_fine, cam_params
            )
            mask_depth_fine = self.flatten_mask(
                mask_deocclusion_fine, range_mask_uv_fine
            )

            ## 5.4 Get Depth Image for Training
            depth_image, _, _, _ = self.gen_depth_img(
                uv_coarse, depth_coarse, mask_coarse[mask_coarse], cam_params
            )
            depth_image /= 100.0
            depth_image = depth_image.unsqueeze(0)
            mask_depth = (mask_depth_coarse > 0) & (mask_depth_fine > 0)

            depth_image_fine, _, _, _ = self.gen_depth_img(
                uv_fine, depth_fine, mask_fine[mask_fine], cam_params
            )
            depth_image_fine /= 100.0
            depth_image_fine = depth_image_fine.unsqueeze(0)

            # 6 Get Flow Image
            flow_image = get_flow_image_from_flow_set(
                flow_set, uv_coarse_in_flow, mask_depth, self._real_shape[:2]
            )

            # 7 Crop Data
            vision_image, depth_image, depth_image_fine, flow_image = (
                self.crop_data_from_dict(
                    dict(
                        image=vision_image,
                        depth=depth_image,
                        depth_fine=depth_image_fine,
                        flow=flow_image,
                    ),
                    vision_image.shape[-2:] + (320, 960),
                    engine_mode,
                ).values()
            )

            valid_i = (flow_image[0].abs() < 1000) & (flow_image[1].abs() < 1000)

            vision_images_input.append(vision_image)
            depth_images_input.append(depth_image)
            depth_images_fine.append(depth_image_fine)
            flow_images_gt.append(flow_image)
            valid_masks.append(valid_i)

        data_dict.update(
            {
                "vision_images_input": torch.stack(vision_images_input),
                "depth_images_input": torch.stack(depth_images_input),
                "flow_images_gt": torch.stack(flow_images_gt),
                "valid_masks": torch.stack(valid_masks),
                "depth_images_fine": torch.stack(depth_images_fine),
            }
        )

        return data_dict

def adjust_points_shape(points, shape="N3"):
    """Adjust the shape of points to (N, 3) or (N, 4).

    Args:
        points (Tensor): The input points tensor.
        shape (str): The desired shape of the points. It can be "N3" or "N4", "3N" or "4N".

    Returns:
        Tensor: The adjusted points tensor.  
    """
    if shape[0] == "N":
        if points.shape[-2] in [3, 4]:
            points = points.transpose(-1, -2)  # (3, N) -> (N, 3), (4, N) -> (N, 4)
    elif shape[1] == "N":
        if points.shape[-1] in [3, 4]:
            points = points.transpose(-1, -2)  # (N, 3) -> (3, N), (N, 4) -> (4, N)
    if "4" in shape:
        if points.shape[-1] == 3:
            points = torch.cat(
                [points, torch.ones(points.shape[:-1] + (1,), device=points.device)], dim=-1
            )
        elif points.shape[-2] == 3:
            points = torch.cat(
                [points, torch.ones(points.shape[:-2] + (1,), device=points.device)], dim=-2
            )
    elif "3" in shape:
        if points.shape[-1] == 4:
            points = points[..., :3]
        elif points.shape[-2] == 4:
            points = points[..., :3, :]
    return points

def adjust_coordinate(xyz: Tensor, order=[1,2,0]) -> Tensor:
    """
    Adjusts the coordinates of a given tensor based on the specified order.

    Args:
        xyz (torch.Tensor): The input tensor containing coordinates. C * N.
        order (list, optional): A list specifying the new order of the coordinates. Defaults to [1, 2, 0, 3].

    Returns:
        torch.Tensor: The tensor with adjusted coordinates.
    """
    return xyz[order, :]

def project_with_mask(points: Tensor, image_size, camera_params: CameraIntrinsicParameters, adjust_coordinate_order=None, front=True) -> tuple[Tensor, Tensor, Tensor]:
    """
    Projects 3D points onto a 2D image plane using intrinsic camera parameters and returns the projected points,
    their depths, and a mask indicating valid points.
    Args:
        points (torch.Tensor): A 3xN tensor representing the 3D points to be projected.
        image_size (tuple): A tuple (width, height) representing the size of the image.
    Returns:
        tuple: A tuple containing:
            - uv (torch.Tensor): A 2xN_front tensor of the projected 2D points.
            - depth (torch.Tensor): A 1D tensor of the depths of the valid points.
            - mask (torch.Tensor): A 1D tensor indicating which points are valid after projection.
    Raises:
        TypeError: If the input points tensor does not have a shape of 3xN.
    """
    assert points.dim() == 2, f"points must be a 2D matrix. but points is {points.dim()}D with shape {points.shape}"
    points = adjust_points_shape(points, "3N") # (3, N)
    if adjust_coordinate_order is not None:
        points = adjust_coordinate(points, adjust_coordinate_order)  # (3, N)
    mask = torch.ones(points.shape[1], dtype=torch.bool, device=points.device) # (N)
    if front:
        mask_front = mask_pixels_with_front(points[2, :]) # (N)
        points = points[:, mask_front] # (3, N_front)
        mask = mask_front # (N)
    uv = torch.zeros((2, points.shape[1]), device=points.device) # (2, N_front)
    uv[0, :] = camera_params.focal_length_x * points[0, :] / points[2, :] + camera_params.principal_point_x
    uv[1, :] = camera_params.focal_length_y * points[1, :] / points[2, :] + camera_params.principal_point_y
    
    mask_vision = mask_pixels_with_vision(uv, (0.1, image_size[1]), (0.1, image_size[0])) # (N_front)
    # generate complete indexes
    index_front = torch.where(mask == True)[0] # (N_front)
    mask[index_front] = mask[index_front] & mask_vision # (N_front) in (N) & (N_front)

    return uv[:, mask_vision], points[2, mask_vision], mask

def mask_pixels_with_image_size(pixels: Tensor, image_w_range: tuple[float, float], image_h_range: tuple[float, float]) -> Tensor:
    """Compute the masks of the pixels which are within the range of an image.

    Args:
        pixels (Tensor): the pixels in the shape of (..., 2). Note that the pixels are represented as (h, w).
        image_w_range (tuple[float, float]): The range of the image width.
        image_h_range (tuple[float, float]): The range of the image height.

    Returns:
        A BoolTensor of the masks of the pixels in the shape of (..., 2). A pixel is with the image if True.
    """
    masks = torch.logical_and(
        torch.logical_and(torch.ge(pixels[0, ...], image_w_range[0]), torch.lt(pixels[0, ...], image_w_range[1])),
        torch.logical_and(torch.ge(pixels[1, ...], image_h_range[0]), torch.lt(pixels[1, ...], image_h_range[1])),
    )
    return masks

def mask_pixels_with_front(depth: Tensor) -> Tensor:
    """Compute the masks of the pixels which are in the front.

    Args:
        pixels (Tensor): the pixels in the shape of (..., 2). Note that the pixels are represented as (h, w).
        depth (Tensor): the depth tensor.

    Returns:
        A BoolTensor of the masks of the pixels in the shape of (..., 2). A pixel is in the front if True.
    """
    return torch.ge(depth, 0)

def mask_pixels_with_vision(pixels: Tensor, image_w_range: tuple[float, float], image_h_range: tuple[float, float], depth:Tensor=None, front=False):
    """
    Masks the pixels based on vision information.

    Args:
        pixels (Tensor): the pixels in the shape of (..., 2). Note that the pixels are represented as (h, w).
        image_w_range (tuple[float, float]): The range of the image width.
        image_h_range (tuple[float, float]): The range of the image height.
        depth (Tensor, optional): The depth tensor. Defaults to None.
        front (bool, optional): Whether to mask pixels in the front. Defaults to False.

    Returns:
        Tensor: The masked pixels tensor.
    """
    masks = mask_pixels_with_image_size(pixels, image_w_range, image_h_range)
    if front:
        assert depth is not None, "depth should be provided when front is True"
        masks = torch.logical_and(masks, mask_pixels_with_front(depth))
    return masks

def deproject(uv, pc_project_uv, camera_params: CameraIntrinsicParameters):
    index = np.argwhere(uv > 0)
    mask = uv > 0
    z = uv[mask]
    camera_params = camera_params.cpu().numpy()
    x = (index[:, 1] - camera_params[2]) * z / camera_params[0]
    y = (index[:, 0] - camera_params[3]) * z / camera_params[1]
    zxy = np.array([z, x, y])
    zxy = torch.tensor(zxy, dtype=torch.float32)
    zxyw = torch.cat([zxy, torch.ones(1, zxy.shape[1], device=zxy.device)])
    zxy = zxyw[:3, :]
    zxy = zxy.cpu().numpy()
    xyz = zxy[[1, 2, 0], :]

    # apply mask to pc_project_uv
    pc_project_u = pc_project_uv[:, :, 0][mask]
    pc_project_v = pc_project_uv[:, :, 1][mask]
    pc_project = np.array([pc_project_v, pc_project_u])
    match_index = np.array([index[:, 0], index[:, 1]])

    return xyz.transpose(), pc_project.transpose(), match_index.transpose()

def get_flow_image_from_flow_set(projected_points, index, mask, shape):
    """
    Computes the optical flow from projected points.

    Args:
        projected_points (torch.Tensor.int): A tensor containing the projected points.
        index (torch.Tensor): A tensor containing the indices of the projected points.
        mask (torch.Tensor): A boolean tensor used to mask the projected points and indices.
        shape (tuple): The shape of the output flow tensor.

    Returns:
        torch.Tensor: A tensor representing the optical flow with shape (2, *shape).
    """
    projected_points_mask = projected_points[mask, :].float()
    index_mask = index[mask, :].t()
    flow = torch.zeros((2, *shape), device=projected_points.device, dtype=torch.float)
    flow[0].index_put_((index_mask[1], index_mask[0]), projected_points_mask[:, 0])
    flow[1].index_put_((index_mask[1], index_mask[0]), projected_points_mask[:, 1])
    return flow

def get_flow_set_from_2pixel_sets(pixel_source, pixel_target, mask_source, mask_target):
    """
    Computes the flow set from two sets of pixels and their corresponding masks.

    Args:
        pixel_source (numpy.ndarray): Source pixel coordinates of shape (N_source, 2).
        pixel_target (numpy.ndarray): Target pixel coordinates of shape (N_target, 2).
        mask_source (numpy.ndarray): Boolean mask for the source pixels of shape (N_source,) in (N) .
        mask_target (numpy.ndarray): Boolean mask for the target pixels of shape (N_target,) in (N) .

    Returns:
        tuple: A tuple containing:
            - flow_set (numpy.ndarray): The computed flow set of shape (N_flow, 2).
            - mask (numpy.ndarray): The combined mask of shape (N_flow,) in (N) .
    """
    mask = mask_source & mask_target # (N_flow) in (N) 
    index_source = mask[mask_source] # (N_flow) in (N_source)
    index_target = mask[mask_target] # (N_flow) in (N_target)
    flow_set = pixel_target[index_target, :] - pixel_source[index_source, :] # (N_flow, 2)
    return flow_set, mask

def apply_transform_to_points(
    points_source: Tensor, transform_source: Tensor, disentangled: bool = False
) -> Tensor:
    # clone to avoid in-place operation
    points = points_source.clone()
    transform = transform_source.clone()

    # pre shape check
    transform = (
        transform[None, :, :] if transform.dim() == 2 else transform
    )  # (B, 4, 4)
    flag = False
    if points.dim() == 2:
        points = points[None, :, :]
        flag = True
    if points.shape[1] == 3 or points.shape[1] == 4:  # (B, 3, N) -> (B, N, 3), (B, 4, N) -> (B, N, 4)
        points = points.transpose(1, 2)

    if disentangled:
        points_mean = points[..., :3].mean(dim=1)[:, None, :]  # (B, 1, 3)
        points[..., :3] -= points_mean  # (B, N, 3)
    rotation = transform[:, :3, :3]  # (B, 3, 3)
    translation = transform[:, None, :3, 3]  # (B, 1, 3)
    points[..., :3] = torch.matmul(points[..., :3], rotation.transpose(-1, -2))
    if disentangled:
        points[..., :3] += points_mean
    points[..., :3] += translation
    if flag:
        points = points.squeeze(0)
    return points

def get_transform_from_rotation_translation(
    rotation: Tensor, translation: Tensor
) -> Tensor:
    """Compose transformation matrix from rotation matrix and translation vector.
    Args:
        rotation (Tensor): (*, 3, 3)
        translation (Tensor): (*, 3)
    Returns:
        transform (Tensor): (*, 4, 4) float
    """
    input_shape = rotation.shape
    rotation = rotation.view(-1, 3, 3)
    translation = translation.view(-1, 3)
    transform = torch.eye(4).to(rotation).unsqueeze(0).repeat(rotation.shape[0], 1, 1)
    transform[:, :3, :3] = rotation
    transform[:, :3, 3] = translation
    output_shape = input_shape[:-2] + (4, 4)
    transform = transform.view(*output_shape)
    return transform.float()

def get_rotation_translation_from_transform(transform: Tensor) -> tuple[Tensor, Tensor]:
    """Decompose transformation matrix into rotation matrix and translation vector.
    Args:
        transform (Tensor): (*, 4, 4)
    Returns:
        rotation (Tensor): (*, 3, 3)
        translation (Tensor): (*, 3)
    """
    rotation = transform[..., :3, :3]
    translation = transform[..., :3, 3]
    return rotation, translation


evaluations = {}

def register_evaluation(cls):
    evaluations[cls.__name__] = cls
    return cls

def get_evaluation(name: str, cfg: edict):
    assert name in evaluations, f"evaluation {name} is not registered"
    return evaluations[name](cfg)

models = {}

def register_model(cls):
    models[cls.__name__] = cls
    return cls

def create_model(cfg: edict):
    assert cfg.model.name in models, f"model {cfg.model.name} is not registered"
    return models[cfg.model.name](cfg)

datasets = {}

def register_dataset(cls):
    datasets[cls.__name__] = cls
    return cls

def create_dataset(cfg: edict=None, engine_mode: EngineMode=None):
    assert cfg is not None, 'cfg must be provided to create dataset'
    assert 'dataset' in cfg or 'name' in cfg, 'dataset must be provided to create dataset'
    assert engine_mode is not None, 'engine_mode must be provided to create dataset'
    if 'dataset' in cfg:
        assert 'name' in cfg['dataset'], 'dataset name must be provided to create dataset'
        assert cfg['dataset']['name'] in datasets, f"dataset {cfg['dataset']['name']} is not registered"
        return datasets[cfg['dataset']['name']](cfg, engine_mode)
    else:
        return datasets[cfg['name']](cfg, engine_mode)

def merge_inputs(queries):
    # point_clouds = []
    # imgs = []
    calibs = []
    orders = []
    returns = {key: default_collate([d[key] for d in queries]) for key in queries[0]
            if key not in ['camera_intrinsic_parameters', 'order']}
    for input in queries:
        # point_clouds.append(input['point_cloud'])
        # imgs.append(input['vision_image'])
        calibs.append(input['camera_intrinsic_parameters'])
        orders.append(input['order'])
    # returns['point_cloud'] = point_clouds
    # returns['vision_image'] = imgs
    returns['camera_intrinsic_parameters'] = calibs
    returns['order'] = orders
    return returns

def get_test_data_loader(cfg: edict) -> Dataset:
    get_logger().info('Loading test data loader...')
    test_dataset = create_dataset(cfg.dataset, engine_mode=EngineMode.TEST)
    test_data_loader = build_dataloader(test_dataset, 
                                        num_workers=cfg.dataset['num_workers'], 
                                        batch_size=1,
                                        shuffle=False,
                                        collate_fn=merge_inputs
                                        )
    return test_data_loader

def reset_seed_worker_init_fn(worker_id):
    """Reset NumPy and Python seed for data loader worker."""
    seed = initial_seed() % (2 ** 32)
    np_random.seed(seed)
    random.seed(seed)

def build_dataloader(
    dataset,
    batch_size=1,
    num_workers=1,
    shuffle=None,
    collate_fn=None,
    sampler=None,
    pin_memory=True,
    drop_last=False,
):
    data_loader = DataLoader(
        dataset,
        batch_size=batch_size,
        num_workers=num_workers,
        shuffle=shuffle,
        sampler=sampler,
        collate_fn=collate_fn,
        worker_init_fn=reset_seed_worker_init_fn,
        pin_memory=pin_memory,
        drop_last=drop_last,
    )

    return data_loader

from typing import overload

@overload
def ensure_dir(path: str): ...

@overload
def ensure_dir(path: Path): ...

def ensure_dir(path):
    if isinstance(path, str):
        path = Path(path)
    if not path.exists():
        path.mkdir(parents=True, exist_ok=True)
    else:
        assert path.is_dir(), f"'{path}' already exists but is not a directory."

class MetricsManager:
    def __init__(self):
        self._metrics: dict[str, list[float]] = {}
        self._keys: list[str] = []
    
    def register_metric(self, key: str, metric: float):
        self._metrics[key] = []
        self._keys.append(key)
        self._metrics[key].append(metric)
    
    def update_metric(self, key: str, metric: float):
        if key not in self._keys:
            self.register_metric(key, metric)
            return
        self._metrics[key].append(metric)
    
    def update(self, metric_dict: dict[str, float]):
        for key, metric in metric_dict.items():
            self.update_metric(key, metric)
    
    def get_metric_mean_std(self, key: str, threshold_key: str = None, filter_func = None) -> tuple[float, float]:
        assert key in self._keys, f"Key '{key}' not found."
        filter_metrics = self._metrics
        if threshold_key is not None and filter_func is not None:
            filter_metrics = self.filter_metrics(threshold_key, [key], filter_func)
        return np.mean(filter_metrics[key]), np.std(filter_metrics[key])

    def get_metric_mean(self, key: str, threshold_key: str = None, filter_func = None) -> float:
        assert key in self._keys, f"Key '{key}' not found."
        filter_metrics = self._metrics
        if threshold_key is not None and filter_func is not None:
            filter_metrics = self.filter_metrics(threshold_key, [key], filter_func)
        return np.mean(filter_metrics[key])

    def get_metrics_mean_std(self, keys: list[str] = None, threshold_key: str = None, filter_func = None) -> dict[str, tuple[float, float]]:
        if keys is None:
            keys = self._keys
        assert len(keys) > 0, "At least one key must be provided."
        filter_metrics = self._metrics
        mean_std_metrics = {}
        if threshold_key is not None and filter_func is not None:
            filter_metrics = self.filter_metrics(threshold_key, keys, filter_func)
            mean_std_metrics['RR'] = filter_metrics['RR']
            mean_std_metrics['threshold_key'] = threshold_key
            mean_std_metrics['filter_func'] = filter_func.threshold
        for key in keys:
            assert key in self._keys, f"Key '{key}' not found."
            if np.asarray(filter_metrics[key]).size == 0:
                mean_std_metrics[key] = "NoData"
            else:
                mean_std_metrics[key] = (np.mean(filter_metrics[key]), np.std(filter_metrics[key]))
        return mean_std_metrics

    def get_metrics_mean(self, keys: list[str] = None, threshold_key: str = None, filter_func = None) -> dict[str, float]:
        if keys is None:
            keys = self._keys
        assert len(keys) > 0, "At least one key must be provided."
        filter_metrics = self._metrics
        mean_metrics = {}
        if threshold_key is not None and filter_func is not None:
            filter_metrics = self.filter_metrics(threshold_key, keys, filter_func)
            mean_metrics['RR'] = filter_metrics['RR']
            mean_metrics['threshold_key'] = threshold_key
            mean_metrics['filter_func'] = filter_func.threshold
        for key in keys:
            assert key in self._keys, f"Key '{key}' not found."
            if np.asarray(filter_metrics[key]).size == 0:
                mean_metrics[key] = "NoData"
            else:
                mean_metrics[key] = np.mean(filter_metrics[key])
        return mean_metrics

    def get_metrics_median(self, keys: list[str] = None, threshold_key: str = None, filter_func = None) -> dict[str, float]:
        if keys is None:
            keys = self._keys
        assert len(keys) > 0, "At least one key must be provided."
        filter_metrics = self._metrics
        median_metrics = {}
        if threshold_key is not None and filter_func is not None:
            filter_metrics = self.filter_metrics(threshold_key, keys, filter_func)
            median_metrics['RR'] = filter_metrics['RR']
            median_metrics['threshold_key'] = threshold_key
            median_metrics['filter_func'] = filter_func.threshold
        for key in keys:
            assert key in self._keys, f"Key '{key}' not found."
            if np.asarray(filter_metrics[key]).size == 0:
                median_metrics[key] = "NoData"
            else:
                median_metrics[key] = np.median(filter_metrics[key])
        return median_metrics

    def get_metrics_mean_std_median(self, keys: list[str] = None, threshold_key: str = None, filter_func = None) -> dict[str, tuple[float, float, float]]:
        if keys is None:
            keys = self._keys
        assert len(keys) > 0, "At least one key must be provided."
        filter_metrics = self._metrics
        mean_std_median_metrics = {}
        if threshold_key is not None and filter_func is not None:
            filter_metrics = self.filter_metrics(threshold_key, keys, filter_func)
            mean_std_median_metrics['RR'] = filter_metrics['RR']
            mean_std_median_metrics['threshold_key'] = threshold_key
            mean_std_median_metrics['filter_func'] = filter_func.threshold
        for key in keys:
            assert key in self._keys, f"Key '{key}' not found."
            if np.asarray(filter_metrics[key]).size == 0:
                mean_std_median_metrics[key] = "NoData"
            else:
                mean_std_median_metrics[key] = (np.mean(filter_metrics[key]), np.std(filter_metrics[key]), np.median(filter_metrics[key]))
        return mean_std_median_metrics
        
    def filter_metrics(self, threshold_key: str, keys: list[str] = None, filter_func = None) -> dict[str, np.ndarray]:
        if keys is None:
            keys = self._keys
        assert len(keys) > 0, "At least one key must be provided."
        assert threshold_key in self._keys, f"Threshold key '{threshold_key}' not found."
        assert filter_func is not None, "Filter function must be provided."
        filter_indices = np.where(filter_func(np.array(self._metrics[threshold_key])))[0]
        filter_metrics = {}
        filter_metrics['RR'] = (len(filter_indices) / len(self._metrics[threshold_key])) * 100
        for key in keys:
            assert key in self._keys, f"Key '{key}' not found."
            filter_metrics[key] = np.array(self._metrics[key])[filter_indices]
        return filter_metrics
    
    def clear(self):
        self._metrics = {}
        self._keys = []



def setup_engine(seed=None, cudnn_deterministic=True, debug=False):
    if seed is not None:
        random.seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
    if cudnn_deterministic:
        cudnn.benchmark = False
        cudnn.deterministic = True
    else:
        cudnn.benchmark = True
        cudnn.deterministic = False
    torch.autograd.set_detect_anomaly(debug)

def load_state_dict(model, state_dict, strict=False):
    logger = get_logger()

    if isinstance(model, DistributedDataParallel):
        model = model.module

    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)

    if len(missing_keys) > 0:
        logger.warn(f"Missing keys: {missing_keys}")
    if len(unexpected_keys) > 0:
        logger.warn(f"Unexpected keys: {unexpected_keys}")

    if strict and (len(missing_keys) != 0 or len(unexpected_keys) != 0):
        raise RuntimeError("The keys in model and state_dict do not match.")

class BaseTester(abc.ABC):
    def __init__(self):
        # parser
        parser = get_default_parser()
        self._args = parser.parse_args()
        self._cudnn_deterministic = self._args.cudnn_deterministic

        # cuda check
        assert torch.cuda.is_available(), "No CUDA devices available."
        
        cfg = get_config()
        self._cfg = cfg
        # logger
        self._log_file = cfg.experiment.log_dir / "test.log"
        self._logger = get_logger(cfg, self._log_file)

        # find checkpoint
        self._checkpoint = self._args.checkpoint
        assert Path(self._checkpoint).exists(), f"Checkpoint not found: {self._checkpoint}"
        
        # metrics manager
        self._metrics_manager = MetricsManager()
        
        # initialize
        torch.cuda.set_device(*cfg.gpus)
        setup_engine(seed=cfg.experiment.seed, cudnn_deterministic=self._cudnn_deterministic)

        # state
        self.model = None
        self.iteration = None

        # data loader
        self.test_loader = None

    @property
    def args(self):
        return self._args

    @property
    def log_file(self):
        return self._log_file

    def load(self, filename, strict=True):
        self.log('Loading from "{}".'.format(filename))
        state_dict = torch.load(filename, map_location=torch.device("cpu"), weights_only=True)
        assert "model" in state_dict, "No model can be loaded."
        load_state_dict(self.model, state_dict["model"], strict=strict)
        self.log("Model has been loaded.")
        if "metadata" in state_dict:
            epoch = state_dict["metadata"]["epoch"]
            total_steps = state_dict["metadata"]["total_steps"]
            self.log(f"Checkpoint metadata: epoch: {epoch}, total_steps: {total_steps}.")

    def register_model(self, model):
        """Register model."""
        model = model.cuda()
        self.model = model
        message = "Model description:\n" + str(model)
        self.log(message)
        return model

    def register_loader(self, test_loader):
        """Register data loader."""
        self.test_loader = test_loader

    def log(self, message, level="INFO"):
        self._logger.log(message, level=level)

    def write_dict(self, data_dict):
        """Write Wandb event."""
        self._logger.wandb_log(data_dict)
        
    def metrics_clear(self):
        self._metrics_manager.clear()
        
    def metrics_update(self, data_dict):
        self._metrics_manager.update(data_dict)

    def metrics_summary_mean(self):
        return self._metrics_manager.get_metrics_mean()
    
    def metrics_summary_mean_std(self):
        return self._metrics_manager.get_metrics_mean_std()

    def before_test_epoch(self):
        self.metrics_clear()

    def before_test_step(self, iteration, data_dict):
        return data_dict

    @abc.abstractmethod
    def test_step(self, iteration, data_dict) -> dict:
        pass

    @abc.abstractmethod
    def eval_step(self, iteration, data_dict, output_dict) -> dict:
        pass

    def after_test_step(self, iteration, data_dict, output_dict, result_dict):
        pass

    def after_test_epoch(self, summary_dict):
        pass

    def get_log_string(self, iteration, data_dict, output_dict, result_dict) -> str:
        return get_log_string(result_dict)

    @abc.abstractmethod
    def test_epoch(self):
        pass

    def run(self, strict_loading=True):
        assert self.test_loader is not None
        if self._checkpoint is not None:
            self.load(self._checkpoint, strict=strict_loading)
        else:
            self.log("Checkpoint is not specified.", level="WARNING")
        self.model.eval()
        torch.set_grad_enabled(False)
        self.test_epoch()

class SingleTester(BaseTester, abc.ABC):
    def __init__(self):
        super().__init__()

    def test_epoch(self):
        # before epoch
        self.before_test_epoch()
        # setup watcher
        timer = Timer()
        # test loop
        pbar = tqdm(enumerate(self.test_loader), total=len(self.test_loader))
        timer.tic("data")
        for batch_index, data_dict in pbar:
            # on start
            self.iteration = batch_index + 1
            data_dict = move_to_cuda(data_dict)
            data_dict = self.before_test_step(self.iteration, data_dict)
            timer.toc("data")
            # test step
            torch.cuda.synchronize()
            timer.tic("model")
            output_dict = self.test_step(self.iteration, data_dict)
            torch.cuda.synchronize()
            timer.toc("model")
            # eval step
            timer.tic("data")
            timer.tic("eval")
            result_dict = self.eval_step(self.iteration, data_dict, output_dict)
            timer.toc("eval")
            # after step
            self.after_test_step(self.iteration, data_dict, output_dict, result_dict)
            # logging
            result_dict = tensor_to_array(result_dict)
            self.metrics_update(result_dict)
            message = self.get_log_string(self.iteration, data_dict, output_dict, result_dict)
            pbar.set_description(message + ", " + timer.tostring(keys=["data", "model", 'eval'], verbose=False))
            torch.cuda.empty_cache()
        # summary logging
        summary_dict = self.metrics_summary_mean_std()
        self.write_dict(summary_dict)
        message = get_log_string(summary_dict, time_dict=timer.summary(keys=["data", "model", 'eval']))
        self.log(message, level="SUCCESS")
        # after epoch
        self.after_test_epoch(summary_dict)

import time

class Timer:
    def __init__(self):
        self._total_time = {}
        self._count_time = {}
        self._last_time = {}
        self._keys = []

    def register_timer(self, key):
        self._total_time[key] = 0.0
        self._count_time[key] = 0
        self._last_time[key] = None
        self._keys.append(key)

    def tic(self, key):
        if key not in self._keys:
            self.register_timer(key)
        self._last_time[key] = time.time()

    def toc(self, key):
        assert key in self._keys, f"'{key}' is not registered in {self._keys}. Please register it first."
        assert self._last_time[key] is not None, "'tic' must be called before 'toc'."
        duration = time.time() - self._last_time[key]
        self._total_time[key] += duration
        self._count_time[key] += 1
        self._last_time[key] = None

    def get_time(self, key):
        assert key in self._keys, f"'{key}' is not registered in {self._keys}. Please register it first."
        assert self._count_time[key] > 0, f"'toc' must be called at least once for key '{key}'."
        return self._total_time[key] / self._count_time[key]

    def tostring(self, keys=None, verbose=True):
        if keys is None:
            keys = self._keys
        if verbose:
            log_strings = [f"{key}: {self.get_time(key):.3f}s" for key in keys if key in self._keys]
            format_string = ", ".join(log_strings)
        else:
            log_strings = [f"{self.get_time(key):.3f}s" for key in keys if key in self._keys]
            format_string = "time: " + "/".join(log_strings)
        return format_string

    def summary(self, keys=None):
        if keys is None:
            keys = self._keys
        summary_dict = {key: self.get_time(key) for key in keys}
        return summary_dict

import sys
import warnings
import loguru
import wandb

class Logger:
    """Advanced logger with stderr, log file and Wandb support.
    
    When DistributedDataParallel is enabled, only ERROR logs are activated for slave processes.
    """
    def __init__(self, cfg: edict, log_file=None):
        is_master_node = True
        self._logger = loguru.logger
        self._logger.remove()
        fmt_str = "<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level}</level> | <level><n>{message}</n></level>"
        log_level = "DEBUG" if is_master_node else "ERROR"
        self._logger.add(sys.stderr, format=fmt_str, colorize=True, level=log_level)
        self._logger.info("Command executed: " + " ".join(sys.argv))
        self._log_file = log_file if is_master_node else None
        if self._log_file is not None:
            fmt_str = "{time:YYYY-MM-DD HH:mm:ss} | {level} | {message}"
            self._logger.add(self._log_file, format=fmt_str, level="INFO")
            self._logger.info(f"Logs are saved to {self._log_file}.")
        wandb.init(
            project = cfg.title,
            dir = cfg.experiment.project_dir,
            name = cfg.experiment.experiment_name,
            config = cfg
        )
        
    @property
    def log_file(self):
        return self._log_file

    def log(self, message, level="INFO"):
        if level not in ["DEBUG", "INFO", "SUCCESS", "WARNING", "ERROR", "CRITICAL"]:
            self._logger.warning(f"Unsupported logging level: {level}. Fallback to INFO.")
            level = "INFO"
        self._logger.log(level, message)

    def debug(self, message):
        self._logger.debug(message)

    def info(self, message):
        self._logger.info(message)

    def success(self, message):
        self._logger.success(message)

    def warn(self, message):
        self._logger.warning(message)

    def error(self, message):
        self._logger.error(message)

    def critical(self, message):
        self._logger.critical(message)

    def wandb_watch(self, model):
        wandb.watch(model)
    
    def wandb_log(self, data_dict):
        wandb.log(data_dict)

_LOGGER = None

def get_logger(cfg=None, log_file=None):
    """Guarantee only one logger per node is built."""
    global _LOGGER
    if _LOGGER is None:
        _LOGGER = Logger(cfg, log_file=log_file)
    elif log_file is not None:
        log_strings = []
        if log_file is not None:
            log_strings.append(f"log_file={log_file}")
        message = "Logger is already initialized. New parameters (" + ",".join(log_strings) + ") are ignored."
        warnings.warn(message)
    return _LOGGER

def get_print_format(value):
    if isinstance(value, (int, str, tuple)):
        return ""
    if value == 0:
        return ".3f"
    if value < 1e-5:
        return ".3e"
    if value < 1e-2:
        return ".6f"
    return ".3f"

def get_format_strings(result_dict):
    """Get format string for a list of key-value pairs."""
    format_strings = []
    if "metadata" in result_dict:
        # handle special key "metadata"
        format_strings.append(result_dict["metadata"])
    for key, value in result_dict.items():
        if key == "metadata":
            continue
        if isinstance(value, (tuple)):
            format_string = f"{key}: " + "/".join([f"{item:{get_print_format(item)}}" for item in value])
        else:
            format_string = f"{key}: {value:{get_print_format(value)}}"
        format_strings.append(format_string)
    return format_strings

def get_log_string(
    result_dict, epoch=None, max_epoch=None, iteration=None, max_iteration=None, lr=None, time_dict=None
):
    log_strings = []
    if epoch is not None:
        epoch_string = f"epoch: {epoch}"
        if max_epoch is not None:
            epoch_string += f"/{max_epoch}"
        log_strings.append(epoch_string)
    if iteration is not None:
        iter_string = f"iter: {iteration}"
        if max_iteration is not None:
            iter_string += f"/{max_iteration}"
        log_strings.append(iter_string)
    log_strings += get_format_strings(result_dict)
    if lr is not None:
        log_strings.append("lr: {:.3e}".format(lr))
    if time_dict is not None:
        time_string = "time: " + "/".join([f"{time_dict[key]:.3f}s" for key in time_dict])
        log_strings.append(time_string)
    message = ", ".join(log_strings)
    return message

def move_to_cuda(x):
    """Move all tensors to cuda."""
    if isinstance(x, list):
        x = [move_to_cuda(item) for item in x]
    elif isinstance(x, tuple):
        x = tuple([move_to_cuda(item) for item in x])
    elif isinstance(x, dict):
        x = {key: move_to_cuda(value) for key, value in x.items()}
    elif isinstance(x, Tensor):
        x = x.cuda()
    return x

def tensor_to_array(x):
    """Release all pytorch tensors to item or numpy arrays."""
    if isinstance(x, list):
        x = [tensor_to_array(item) for item in x]
    elif isinstance(x, tuple):
        x = tuple([tensor_to_array(item) for item in x])
    elif isinstance(x, dict):
        x = {key: tensor_to_array(value) for key, value in x.items()}
    elif isinstance(x, Tensor):
        if x.numel() == 1:
            x = x.item()
        else:
            x = x.detach().cpu().numpy()
    return x

import argparse

_PARSER = None

def get_default_parser():
    global _PARSER
    if _PARSER is None:
        _PARSER = argparse.ArgumentParser()
    return _PARSER


def parse_args():
    parser = get_default_parser()
    args = parser.parse_args()
    return args

def add_base_args():
    parser = get_default_parser()
    parser.add_argument("--cfg", type=str, required=True, default='./test.toml', help="config file path")
    parser.add_argument("--checkpoint", type=str, default=None, help="load from checkpoint")
    parser.add_argument("--cudnn_deterministic", type=bool, default=True, help="use deterministic method")

def add_trainer_args():
    parser = get_default_parser()
    parser.add_argument_group("trainer", "trainer arguments")
    parser.add_argument("--resume", action="store_true", help="resume training from the latest checkpoint")
    parser.add_argument("--log_steps", type=int, default=100, help="logging steps")
    parser.add_argument("--debug", action="store_true", help="debug mode with grad check")
    parser.add_argument("--detect_anomaly", action="store_true", help="detect anomaly with autograd")
    parser.add_argument("--save_latest_n_models", type=int, default=-1, help="save latest n models")
    parser.add_argument("--watch_model", action="store_true", help="watch model with wandb")

def add_tester_args():
    parser = get_default_parser()
    parser.add_argument_group("tester", "tester arguments")

add_base_args()

from easydict import EasyDict
import tomllib as tml
def read_toml_file(file_name: str) -> EasyDict:
    """
    Read a TOML file and return its contents as a dictionary.

    Args:
        file_path (str): The path to the TOML file.

    Returns:
        dict[str, any]: The contents of the TOML file as a dictionary.

    """
    with open(file_name, 'rb') as toml_file:
        return EasyDict(tml.load(toml_file))

_CONFIG = None

def get_deafult_config():
    global _CONFIG
    if _CONFIG is None:
        _CONFIG = Config()
    return _CONFIG

def get_config():
    return get_deafult_config().cfg

import datetime
import threading

class SingletonType(type):
    _instance_lock = threading.Lock()
    def __call__(cls, *args, **kwargs):
        if not hasattr(cls, "_instance"):
            with SingletonType._instance_lock:
                if not hasattr(cls, "_instance"):
                    cls._instance = super(SingletonType,cls).__call__(*args, **kwargs)
        return cls._instance
    
class Config(metaclass=SingletonType):
    def __init__(self):
        self.cfg = read_toml_file(parse_args().cfg)
        self.add_experiment_cfg()
    
    def add_experiment_cfg(self):
        """
        Adds experiment configuration details to the given configuration dictionary.

        Returns:
            edict: The updated configuration dictionary with added experiment details.

        The function performs the following actions:
            - Sets the experiment name to the title from the configuration.
            - Sets the experiment time to the current datetime in the format YYYYMMDD_HHMMSS.
            - Sets the working directory to the parent directory of the given filename.
            - Sets the project directory to a subdirectory named after the title within the working directory.
            - Sets the output directory to a subdirectory named after the experiment name within the working directory.
            - Sets the checkpoint directory to a "checkpoints" subdirectory within the output directory.
            - Sets the log directory to a "logs" subdirectory within the output directory.
            - Ensures that all directories ending with "_dir" exist by creating them if necessary.
        """
        if "experiment" not in self.cfg:
            self.cfg.experiment = edict()
        self.cfg.experiment.time = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        self.cfg.experiment.working_dir = Path(parse_args().cfg).resolve().parent.parent # ./working_dir/cfg/cfg.toml
        self.cfg.experiment.project_dir = self.cfg.experiment.working_dir / self.cfg.title
        self.cfg.experiment.name_dir = self.cfg.experiment.project_dir / ("train" if self.cfg.mode == "train" else "test")
        self.cfg.experiment.experiment_name = self.cfg.mode + "_" + self.cfg.experiment.time
        self.cfg.experiment.output_dir = self.cfg.experiment.name_dir / self.cfg.experiment.experiment_name
        self.cfg.experiment.checkpoint_dir = self.cfg.experiment.output_dir / "checkpoints"
        self.cfg.experiment.log_dir = self.cfg.experiment.output_dir / "logs"
        self.cfg.experiment.result_dir = self.cfg.experiment.output_dir / "result"
        for dir in self.cfg.experiment:
            if dir.endswith("_dir"):
                ensure_dir(self.cfg.experiment[dir])

    def __str__(self) -> str:
        return f"Configuration details:\n{self.cfg}"

class Evaluation(nn.Module, abc.ABC):
    def __init__(self, cfg: edict):
        self._cfg = cfg
        super(Evaluation, self).__init__()
        
    @abc.abstractmethod
    def evaluation_fn(self, data_dict: dict, output_dict: dict) -> dict:
        raise NotImplementedError
    
    def forward(self, data_dict: dict, output_dict: dict):
        result_dict = self.evaluation_fn(data_dict, output_dict)
        return result_dict

def inverse_rotation_translation(
    rotation: Tensor, translation: Tensor
) -> tuple[Tensor, Tensor]:
    """Inverse rotation and translation.
    Args:
        rotation (Tensor): (*, 3, 3)
        translation (Tensor): (*, 3)
    Returns:
        inv_rotation (Tensor): (*, 3, 3). float
        inv_translation (Tensor): (*, 3). float
    """
    inv_rotation = rotation.transpose(-1, -2).float()  # (*, 3, 3)
    inv_translation = -torch.matmul(
        inv_rotation, translation.view(-1, 3, 1).float()
    ).squeeze(
        -1
    )  # (*, 3)
    return inv_rotation, inv_translation

class TransformDistanceType(Enum):
    """Distance type for rigid transformations."""

    COMMON = 0
    I2D_LOC = 1

def rotation_matrix_distance(
    rotation_matrix1: Tensor, rotation_matrix2: Tensor
) -> Tensor:
    """Compute the distance between two rotation matrices. The error unit of the calculation is the rotation angle error.
    Args:
        rotation_matrix1 (Tensor): (*, 3, 3)
        rotation_matrix2 (Tensor): (*, 3, 3)
    Returns:
        distance (Tensor): (*)
    """
    rotation_matrix1 = rotation_matrix1.view(-1, 3, 3)
    rotation_matrix2 = rotation_matrix2.view(-1, 3, 3)
    return torch.tensor(
        [
            abs(
                torch.acos(
                    (
                        torch.trace(
                            torch.mm(
                                torch.inverse(rotation_matrix1).view(3, 3),
                                rotation_matrix2.view(3, 3),
                            )
                        )
                        - 1
                    )
                    / 2
                )
            )
            * 180.0
            / np.pi
            for i in range(rotation_matrix1.shape[0])
        ]
    )

def rotation_matrix_to_quaternion(rotation_matrix: Tensor) -> Tensor:
    """Convert rotation matrix to quaternion.
    Args:
        rotation_matrix (Tensor): (*, 3, 3)
    Returns:
        quaternion (Tensor): (*, 4) 【xyzw】
    """
    rotation_matrix = rotation_matrix.view(-1, 3, 3)
    return torch.tensor(
        np.array(
            [
                R.from_matrix(rotation_matrix[i].detach().cpu().numpy()).as_quat()
                for i in range(rotation_matrix.shape[0])
            ]
        )
    )

def transform_distance(
    transform1: Tensor, transform2: Tensor, flag=TransformDistanceType.I2D_LOC
) -> tuple[Tensor, Tensor]:
    """Compute distance between two rigid transformations.
    Args:
        transform1 (Tensor): (*, 4, 4)
        transform2 (Tensor): (*, 4, 4)
    Returns:
        distance(Tensor, Tensor): rotation distance, translation distance
    """
    match flag:
        case TransformDistanceType.COMMON:
            rotation1, translation1 = get_rotation_translation_from_transform(
                transform1
            )  # (*, 3, 3), (*, 3)
            rotation2, translation2 = get_rotation_translation_from_transform(
                transform2
            )  # (*, 3, 3), (*, 3)
            rotation_distance = rotation_matrix_distance(
                rotation1, rotation2
            )  # (*,)
            translation_distance = torch.norm(
                translation1 - translation2, dim=-1
            )  # (*,)
        case TransformDistanceType.I2D_LOC:
            rotation, translation = get_rotation_translation_from_transform(
                transform=torch.matmul(inverse_transform(transform2), transform1)
            )
            rotation_distance = (
                quaternion_distance(
                    rotation_matrix_to_quaternion(rotation),
                    torch.tensor([[0.0, 0.0, 0.0, 1.0]]),
                )
                * 180.0
                / torch.pi
            )
            translation_distance = torch.norm(translation) * 100  # (*,)
    return rotation_distance, translation_distance

def inverse_transform(transform: Tensor) -> Tensor:
    """Inverse rigid transform.
    Args:
        transform (Tensor): (*, 4, 4)
    Return:
        inv_transform (Tensor): (*, 4, 4)
    """
    rotation, translation = get_rotation_translation_from_transform(
        transform
    )  # (*, 3, 3), (*, 3)
    inv_rotation = rotation.transpose(-1, -2)  # (*, 3, 3)
    inv_translation = -torch.matmul(inv_rotation, translation.unsqueeze(-1)).squeeze(
        -1
    )  # (*, 3)
    inv_transform = get_transform_from_rotation_translation(
        inv_rotation, inv_translation
    )  # (*, 4, 4)
    return inv_transform

def quaternion_inverse(quaternion: Tensor) -> Tensor:
    """Compute the inverse of a quaternion.
    Args:
        quaternion (Tensor): (*, 4)
    Returns:
        quaternion_inv (Tensor): (*, 4)
    """
    quaternion = quaternion.view(-1, 4)
    return torch.tensor(
        np.array(
            [
                R.from_quat(quaternion[i].detach().cpu().numpy()).inv().as_quat()
                for i in range(quaternion.shape[0])
            ]
        )
    )

def quaternion_multiply(quaternion1: Tensor, quaternion2: Tensor) -> Tensor:
    """Compute the multiplication of two quaternions.
    Args:
        quaternion1 (Tensor): (*, 4) 【xyzw】
        quaternion2 (Tensor): (*, 4) 【xyzw】
    Returns:
        quaternion (Tensor): (*, 4) 【xyzw】
    """
    q = quaternion1.view(-1, 4)[:, [3, 0, 1, 2]]
    r = quaternion2.view(-1, 4)[:, [3, 0, 1, 2]]
    t = torch.zeros(q.shape[0], 4, device=q.device)
    t[:, 0] = (
        r[:, 0] * q[:, 0] - r[:, 1] * q[:, 1] - r[:, 2] * q[:, 2] - r[:, 3] * q[:, 3]
    )
    t[:, 1] = (
        r[:, 0] * q[:, 1] + r[:, 1] * q[:, 0] - r[:, 2] * q[:, 3] + r[:, 3] * q[:, 2]
    )
    t[:, 2] = (
        r[:, 0] * q[:, 2] + r[:, 1] * q[:, 3] + r[:, 2] * q[:, 0] - r[:, 3] * q[:, 1]
    )
    t[:, 3] = (
        r[:, 0] * q[:, 3] - r[:, 1] * q[:, 2] + r[:, 2] * q[:, 1] + r[:, 3] * q[:, 0]
    )
    return t

def quaternion_distance(quaternion1: Tensor, quaternion2: Tensor) -> Tensor:
    """Compute the distance between two quaternions. The error unit of the calculation is the rotation angle error.
    Args:
        quaternion1 (Tensor): (*, 4)
        quaternion2 (Tensor): (*, 4)
    Returns:
        distance (Tensor): (*) 弧度制
    """
    t = quaternion_multiply(quaternion1, quaternion_inverse(quaternion2))
    return 2 * torch.atan2(torch.norm(t[:, 1:], dim=1), torch.abs(t[:, 0]))

def quaternion_to_rotation_matrix(quaternion: Tensor) -> Tensor:
    """Convert quaternion to rotation matrix.
    Args:
        quaternion (Tensor): (*, 4)
    Returns:
        rotation_matrix (Tensor): (*, 3, 3) float
    """
    quaternion = quaternion.view(-1, 4)
    return torch.tensor(
        np.array(
            [
                R.from_quat(quaternion[i].detach().cpu().numpy()).as_matrix()
                for i in range(quaternion.shape[0])
            ]
        )
    ).float()
    
def rotation_matrix_to_quaternion(rotation_matrix: Tensor) -> Tensor:
    """Convert rotation matrix to quaternion.
    Args:
        rotation_matrix (Tensor): (*, 3, 3)
    Returns:
        quaternion (Tensor): (*, 4) 【xyzw】
    """
    rotation_matrix = rotation_matrix.view(-1, 3, 3)
    return torch.tensor(
        np.array(
            [
                R.from_matrix(rotation_matrix[i].detach().cpu().numpy()).as_quat()
                for i in range(rotation_matrix.shape[0])
            ]
        )
    )

def rotation_vector_to_rotation_matrix(rotation_vector: Tensor) -> Tensor:
    """Convert rotation vector to rotation matrix.
    Args:
        rotation_vector (Tensor): (*, 3)
    Returns:
        rotation_matrix (Tensor): (*, 3, 3)
    """
    rotation_vector = torch.tensor(rotation_vector).view(-1, 3)
    return torch.tensor(
        np.array(
            [
                R.from_rotvec(rotation_vector[i].detach().cpu().numpy()).as_matrix()
                for i in range(rotation_vector.shape[0])
            ]
        )
    )

def angle_to_rotation_matrix(angle: Tensor, degrees: bool = True) -> Tensor:
    """Convert angle to rotation matrix.
    Args:
        angle (Tensor): (*, 3)
        degrees (bool): True means degrees, False means radians.
    Returns:
        rotation_matrix (Tensor): (*, 3, 3) float
    """
    angle = angle.view(-1, 3)
    return torch.tensor(
        np.array(
            [
                R.from_euler(
                    "xyz", angle[i].detach().cpu().numpy(), degrees=degrees
                ).as_matrix()
                for i in range(angle.shape[0])
            ]
        ), device=angle.device
    ).float()

def is_path_exist(*file_parts: str) -> bool:
    """
    Check if a file or directory exists by joining multiple path components using pathlib.

    Args:
        *file_parts (str): Multiple components of the file or directory path.

    Returns:
        bool: True if the file or directory exists, False otherwise.
    """
    path = Path(*file_parts)
    return path.exists()

def read_image_file(file_name: str) -> Image.Image:
    """
    Load an image from a given file path.

    Args:
        img_path (str): The path to the image file.

    Returns:
        PIL.Image.Image: The loaded image.
    """
    try:
        img = Image.open(file_name)
        return img
    except IOError as e:
        raise IOError(f"Error opening image file {file_name}: {e}")

def read_h5_file(file_name: str, keys: list[str] | str) -> dict[str, ndarray]:
    """
    Read data from an HDF5 file.

    Args:
        file_name (str): The path to the HDF5 file.
        keys (list[str] | str): The key(s) of the dataset(s) to read from the file. If a single key is provided as a string,
                                it will be converted to a list with a single element.

    Returns:
        dict[str, ndarray]: A dictionary where the keys are the provided dataset keys and the values are the corresponding
                            NumPy arrays containing the data.

    Raises:
        Exception: If there is an error reading the file.

    """
    try:
        keys = [keys] if isinstance(keys, str) else keys
        with h5py.File(file_name, 'r') as hf:
            h5_data = {key: hf[key][:] for key in keys}
        return h5_data
    except Exception as e:
        print(f'File Broken: {file_name}')
        raise e
    
def create_df(data: dict[str, list[any]]) -> pd.DataFrame:
    """
    Create a DataFrame from a dictionary.

    Args:
        data (dict[str, list[any]]): A dictionary where the keys are the column names and the values are lists of data.

    Returns:
        pd.DataFrame: A DataFrame containing the data from the dictionary.

    """
    return pd.DataFrame(data)

def read_csv_file(file_parts: list[str] | str, **kwargs) -> pd.DataFrame:
    """
    Read data from a CSV file.

    Args:
        file_name (str): The path to the CSV file.
        **kwargs: Additional keyword arguments to pass to `pd.read_csv`.

    Returns:
        pd.DataFrame: A DataFrame containing the data from the CSV file.

    Raises:
        Exception: If there is an error reading the file.

    """
    try:
        file_name = Path(*([file_parts] if isinstance(file_parts, str) else file_parts))
        return pd.read_csv(file_name, **kwargs)
    except Exception as e:
        print(f'File Broken: {file_name}')
        raise e

def write_csv_file(data: pd.DataFrame, file_parts: list[str] | str, **kwargs) -> None:
    """
    Write data to a CSV file.

    Args:
        data (pd.DataFrame): The data to write to the CSV file.
        file_parts (str): The path to the CSV file.
        **kwargs: Additional keyword arguments to pass to `pd.to_csv`.

    """
    file_name = Path(*([file_parts] if isinstance(file_parts, str) else file_parts))
    data.to_csv(file_name, **kwargs)
    
class Dataset_I2P(Dataset, abc.ABC):
    def __init__(self, cfg: edict, engine_mode:EngineMode=EngineMode.TRAIN) -> None:
        super(Dataset_I2P, self).__init__()
        self._cfg = cfg
        self._engine_mode = engine_mode
        self._w_scale = cfg.w_scale
        self._h_scale = cfg.h_scale
        self._adjust_coordinate_order = cfg.adjust_coordinate_order if cfg.adjust_coordinate_order != "" else None
        self.GTs_R = {}
        self.GTs_T = {}
        self.all_files = []
        if engine_mode == EngineMode.TRAIN:
            assert 'train_sequences' in cfg, 'train_sequences must be provided in the configuration'
            for sequence in cfg['train_sequences']:
                self.process_sequence(sequence)
        else:
            self.process_sequence(cfg['test_sequence'])
        self.test_RT = self.get_test_RT()
    
    @abc.abstractmethod
    def process_sequence(self, sequence: str) -> None:
        pass

    def get_test_RT(self) -> list:
        test_RT = []
        if self._engine_mode == EngineMode.TRAIN:
            return test_RT
        test_RT_file = '/'.join([self._cfg['root_folder'], f'test_RT_seq{self._cfg.test_sequence}_{self._cfg.max_r:.2f}_{self._cfg.max_t:.2f}.csv'])
        if not is_path_exist(test_RT_file):
            get_logger().success(f'TEST SET - Not found: {test_RT_file}, Generating a new one')
            rad_factor = np.pi / 180.0
            len_files = len(self.all_files)
            data = {
                'tx': np.random.uniform(-self._cfg['max_t'], self._cfg['max_t'], len_files),
                'ty': np.random.uniform(-self._cfg['max_t'], self._cfg['max_t'], len_files),
                'tz': np.random.uniform(-self._cfg['max_t'], min(self._cfg['max_t'], 1.0), len_files),
                'rx': np.random.uniform(-self._cfg['max_r'], self._cfg['max_r'], len_files) * rad_factor,
                'ry': np.random.uniform(-self._cfg['max_r'], self._cfg['max_r'], len_files) * rad_factor,
                'rz': np.random.uniform(-self._cfg['max_r'], self._cfg['max_r'], len_files) * rad_factor
            }
            write_csv_file(create_df(data), test_RT_file)
        get_logger().success(f'TEST SET: Using this file: {test_RT_file}')
        test_RT.extend(read_csv_file(test_RT_file, sep=',').values.tolist())

        assert len(test_RT) == len(self.all_files), f"Something wrong {len(test_RT)} != {len(self.all_files)}"
        return test_RT

    def custom_transform(self, rgb, img_rotation=0., flip=False): # TODO: Update this function
        to_tensor = transforms.ToTensor()
        normalization = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                             std=[0.229, 0.224, 0.225])
        if self._engine_mode == EngineMode.TRAIN:
            color_transform = transforms.ColorJitter(0.1, 0.1, 0.1)
            rgb = color_transform(rgb)
            if flip:
                rgb = TTF.hflip(rgb)
            rgb = TTF.rotate(rgb, img_rotation)
        rgb = to_tensor(rgb)
        rgb = normalization(rgb)
        return rgb

    def adjust_point_cloud(self, pc: np.ndarray) -> torch.Tensor:
        """
        Preprocess a point cloud to ensure it is in a homogeneous coordinate system.

        Args:
            pc (numpy.ndarray): The input point cloud, expected to be of shape (N, 3) or (N, 4) or (3, N) or (4, N).

        Returns:
            torch.Tensor: The preprocessed point cloud in homogeneous coordinates. (N, 4)
        """
        if(isinstance(pc, torch.Tensor)):
            pc_in = pc
        else:
            pc_in = torch.from_numpy(pc.astype(np.float32))
        assert pc_in.dim() == 2, f"PointCloud must be a 2D matrix, but got {pc_in.dim()}D with shape {pc_in.shape}"
        if pc_in.shape[0] == 3 or pc_in.shape[0] == 4: # (3, N) or (4, N) -> (N, 3) or (N, 4)
            pc_in = pc_in.t()
        if pc_in.shape[1] == 3:
            pc_in = torch.cat((pc_in, torch.ones(pc_in.shape[0]).unsqueeze(0)), 1) # (N, 4)
        elif pc_in.shape[1] == 4:
            if torch.all(pc_in[:,3] == 1.):
                pc_in[:,3] = 1.
        else:
            raise TypeError("Wrong PointCloud shape", pc_in.shape)
        return pc_in

    def augment_data(self, image, point_cloud, 
            camera_intrinsic_parameters: CameraIntrinsicParameters, 
            camera_extrinsic_parameters: torch.Tensor) -> tuple:
        """
        Augment data by applying random horizontal mirroring and random rotation.

        Args:
            img_path (str): Path to the image file.
            pc_in (numpy.ndarray): The input point cloud data.
            train_mode (bool, optional): Whether the augmentation is for training. Defaults to True.

        Returns:
            tuple: The augmented image and point cloud.
        """
        image_rotation = 0
        h_mirror = False
        if self._engine_mode == EngineMode.TRAIN:
            # Random horizontal mirroring && Random rotation
            image_rotation = np.random.uniform(-5, 5)
            h_mirror = np.random.rand() > 0.5
        image = self.custom_transform(image, image_rotation, h_mirror)
        # Downsample Point Cloud 204800. N*4 -> 204800 * 4
        if self._engine_mode == EngineMode.TRAIN:
            if h_mirror:
                point_cloud[:, 1] *= -1  
                camera_intrinsic_parameters.principal_point_x = image.shape[2] - camera_intrinsic_parameters.principal_point_x
            R = angle_to_rotation_matrix(torch.tensor([image_rotation, 0, 0]))
            T = torch.tensor([0., 0., 0.]).float()
            transform = inverse_transform(get_transform_from_rotation_translation(R, T))
            point_cloud = apply_transform_to_points(point_cloud, transform) 
        image, camera_intrinsic_parameters = self.scale_image(image, camera_intrinsic_parameters)
        if camera_extrinsic_parameters is not None:
            point_cloud = self.adjust_point_cloud(point_cloud)
            point_cloud = apply_transform_to_points(point_cloud[:, :3], camera_extrinsic_parameters)
        point_cloud = self.adjust_point_cloud(point_cloud)
        return image, point_cloud, camera_intrinsic_parameters

    def generate_random_transforms(self, idx: int=None) -> tuple[torch.Tensor, torch.Tensor]:
        if self._engine_mode == EngineMode.TRAIN:
            R, T = get_rotation_translation_from_transform(generate_random_transforms(self._cfg['max_r'], self._cfg['max_t']))
        else:
            R = angle_to_rotation_matrix(torch.tensor(self.test_RT[idx][4:]), False)
            T = torch.tensor(self.test_RT[idx][1:4])
        return inverse_rotation_translation(R, T)

    def scale_image(self, image, cam_params: CameraIntrinsicParameters):
        self._real_shape = [int(image.shape[1] * self._h_scale), int(image.shape[2] * self._w_scale), image.shape[0]] # H, W, C
        if self._w_scale == 1 and self._h_scale == 1:
            return image, cam_params
        downsample = transforms.Resize(self._real_shape[:2], interpolation=Image.NEAREST)
        image = downsample(image)
        cam_params.scale(self._w_scale, self._h_scale)
        return image, cam_params

    @abc.abstractmethod
    def get_camera_parameters(self, path: str) -> tuple[CameraIntrinsicParameters, torch.Tensor]:
        pass
    
    @abc.abstractmethod
    def get_point_cloud_path(self, idx) -> str:
        pass
    
    @abc.abstractmethod
    def get_image_path(self, idx) -> str:
        pass
    
    @abc.abstractmethod
    def get_camera_parameters_path(self, idx) -> str:
        pass
    
    def __len__(self) -> int:
        return len(self.all_files)
    
    def __getitem__(self, idx):
        image_path = self.get_image_path(idx)
        image = read_image_file(image_path)
        point_cloud_path = self.get_point_cloud_path(idx)
        point_cloud = self.adjust_point_cloud(read_h5_file(point_cloud_path, 'PC')['PC'])
        camera_parameters_path = self.get_camera_parameters_path(idx)
        camera_intrinsic_parameters, camera_extrinsic_parameters = self.get_camera_parameters(camera_parameters_path)
        
        image, point_cloud, camera_intrinsic_parameters = self.augment_data(
            image, point_cloud, camera_intrinsic_parameters, camera_extrinsic_parameters
            )
        R, T = self.generate_random_transforms(idx)
        return {'vision_image': image, 'point_cloud': point_cloud, 'camera_intrinsic_parameters': camera_intrinsic_parameters,
                    'tr_error': T, 'rot_error': R, 'order': self._adjust_coordinate_order}

def generate_random_translation(max_offset: float) -> list[float]:
    """
    Generate a random translation vector within the specified maximum offset.

    Args:
        max_offset (float): The maximum offset for each translation component.

    Returns:
        list[float]: A list containing the randomly generated translation vector [transl_x, transl_y, transl_z].
    """
    transl_x = np.random.uniform(-max_offset, max_offset)
    transl_y = np.random.uniform(-max_offset, max_offset)
    transl_z = np.random.uniform(-max_offset, min(max_offset, 1.0))
    return [transl_x, transl_y, transl_z]

def generate_random_rotation_euler(max_angle: float) -> list[float]:
    """
    Generates a random rotation in Euler angles representation.

    Args:
        max_angle (float): The maximum angle in radians for each Euler angle.

    Returns:
        Tensor: A tensor representing the rotation in Euler angles.

    """
    rotation_euler = [np.random.uniform(-max_angle, max_angle) for _ in range(3)]
    return rotation_euler

def generate_random_transforms(max_angle: float, max_offset: float) -> Tensor:
    """Generate random rotation and translation.
    Args:
        batch_size (int): number of samples
        device (torch.device): device
    Returns:
        rotation (Tensor): (*, 3, 3)
        translation (Tensor): (*, 3)
    """
    rotation_euler = generate_random_rotation_euler(max_angle)
    rotation_matrix_tensor = angle_to_rotation_matrix(torch.tensor(rotation_euler))
    translation_tensor = torch.tensor(generate_random_translation(max_offset))
    return get_transform_from_rotation_translation(
        rotation_matrix_tensor, translation_tensor
    )