third_party/AnyBimanual/helpers/utils.py

import numpy as np
import torch
import trimesh
try:
    import pyrender
    from pyrender.trackball import Trackball
except ImportError:
    pyrender = None
    Trackball = None
from rlbench.backend.const import DEPTH_SCALE
from scipy.spatial.transform import Rotation
from rlbench.backend.observation import Observation
from rlbench import CameraConfig, ObservationConfig
from pyrep.const import RenderMode
from typing import List
from PIL import Image
import blosc
import pickle

import einops
from pickle import UnpicklingError
import numpy as np
from scipy.interpolate import CubicSpline, interp1d
import torch
import torchvision.transforms as transforms
import torchvision.transforms.functional as transforms_f

# from diffuser_actor.utils.utils import normalise_quat
SCALE_FACTOR = DEPTH_SCALE
DEFAULT_SCENE_SCALE = 2.0


def loss_weights(replay_sample, beta=1.0):
    loss_weights = 1.0
    if "sampling_probabilities" in replay_sample:
        probs = replay_sample["sampling_probabilities"]
        loss_weights = 1.0 / torch.sqrt(probs + 1e-10)
        loss_weights = (loss_weights / torch.max(loss_weights)) ** beta
    return loss_weights


def soft_updates(net, target_net, tau):
    for param, target_param in zip(net.parameters(), target_net.parameters()):
        target_param.data.copy_(tau * param.data + (1 - tau) * target_param.data)


def stack_on_channel(x):
    # expect (B, T, C, ...)
    return torch.cat(torch.split(x, 1, dim=1), dim=2).squeeze(1)


def normalize_quaternion(quat):
    return np.array(quat) / np.linalg.norm(quat, axis=-1, keepdims=True)


def correct_rotation_instability(disc, resolution):
    # q1 = discrete_euler_to_quaternion(disc, resolution)
    # q2 = discrete_euler_to_quaternion(quaternion_to_discrete_euler(q1, resolution), resolution)
    #
    # d2 = quaternion_to_discrete_euler(q2, resolution)
    #
    # # choose the smallest change
    # if np.any(disc != d2):
    #     if np.sum(disc) < np.sum(d2):
    #         return disc
    #     else:
    #         return d2
    return disc


def check_gimbal_lock(pred_rot_and_grip, gt_rot_and_grip, resolution):
    pred_rot_and_grip_np = pred_rot_and_grip.detach().cpu().numpy()
    gt_rot_and_grip_np = gt_rot_and_grip.detach().cpu().numpy()

    pred_rot = discrete_euler_to_quaternion(pred_rot_and_grip_np[:, :3], resolution)
    gt_rot = discrete_euler_to_quaternion(gt_rot_and_grip_np[:, :3], resolution)
    gimbal_lock_matches = [
        np.all(np.abs(pred_rot[i] - gt_rot[i]) < 1e-10)
        and np.any(pred_rot_and_grip_np[i, :3] != gt_rot_and_grip_np[i, :3])
        for i in range(pred_rot.shape[0])
    ]
    return 0


def quaternion_to_discrete_euler(quaternion, resolution):
    euler = Rotation.from_quat(quaternion).as_euler("xyz", degrees=True) + 180
    assert np.min(euler) >= 0 and np.max(euler) <= 360
    disc = np.around((euler / resolution)).astype(int)
    disc[disc == int(360 / resolution)] = 0
    return disc


def discrete_euler_to_quaternion(discrete_euler, resolution):
    euluer = (discrete_euler * resolution) - 180
    return Rotation.from_euler("xyz", euluer, degrees=True).as_quat()


def point_to_voxel_index(
    point: np.ndarray, voxel_size: np.ndarray, coord_bounds: np.ndarray
):
    bb_mins = np.array(coord_bounds[0:3])
    bb_maxs = np.array(coord_bounds[3:])
    dims_m_one = np.array([voxel_size] * 3) - 1
    bb_ranges = bb_maxs - bb_mins
    res = bb_ranges / (np.array([voxel_size] * 3) + 1e-12)
    voxel_indicy = np.minimum(
        np.floor((point - bb_mins) / (res + 1e-12)).astype(np.int32), dims_m_one
    )
    return voxel_indicy


def voxel_index_to_point(
    voxel_index: torch.Tensor, voxel_size: int, coord_bounds: np.ndarray
):
    res = (coord_bounds[:, 3:] - coord_bounds[:, :3]) / voxel_size
    points = (voxel_index * res) + coord_bounds[:, :3]
    return points


def point_to_pixel_index(
    point: np.ndarray, extrinsics: np.ndarray, intrinsics: np.ndarray
):
    point = np.array([point[0], point[1], point[2], 1])
    world_to_cam = np.linalg.inv(extrinsics)
    point_in_cam_frame = world_to_cam.dot(point)
    px, py, pz = point_in_cam_frame[:3]
    px = 2 * intrinsics[0, 2] - int(-intrinsics[0, 0] * (px / pz) + intrinsics[0, 2])
    py = 2 * intrinsics[1, 2] - int(-intrinsics[1, 1] * (py / pz) + intrinsics[1, 2])
    return px, py


def _compute_initial_camera_pose(scene):
    # Adapted from:
    # https://github.com/mmatl/pyrender/blob/master/pyrender/viewer.py#L1032
    centroid = scene.centroid
    scale = scene.scale
    if scale == 0.0:
        scale = DEFAULT_SCENE_SCALE
    s2 = 1.0 / np.sqrt(2.0)
    cp = np.eye(4)
    cp[:3, :3] = np.array([[0.0, -s2, s2], [1.0, 0.0, 0.0], [0.0, s2, s2]])
    hfov = np.pi / 6.0
    dist = scale / (2.0 * np.tan(hfov))
    cp[:3, 3] = dist * np.array([1.0, 0.0, 1.0]) + centroid
    return cp


def _from_trimesh_scene(trimesh_scene, bg_color=None, ambient_light=None):
    # convert trimesh geometries to pyrender geometries
    geometries = {
        name: pyrender.Mesh.from_trimesh(geom, smooth=False)
        for name, geom in trimesh_scene.geometry.items()
    }
    # create the pyrender scene object
    scene_pr = pyrender.Scene(bg_color=bg_color, ambient_light=ambient_light)
    # add every node with geometry to the pyrender scene
    for node in trimesh_scene.graph.nodes_geometry:
        pose, geom_name = trimesh_scene.graph[node]
        scene_pr.add(geometries[geom_name], pose=pose)
    return scene_pr


def _create_bounding_box(scene, voxel_size, res):
    l = voxel_size * res
    T = np.eye(4)
    w = 0.01
    for trans in [[0, 0, l / 2], [0, l, l / 2], [l, l, l / 2], [l, 0, l / 2]]:
        T[:3, 3] = np.array(trans) - voxel_size / 2
        scene.add_geometry(
            trimesh.creation.box([w, w, l], T, face_colors=[0, 0, 0, 255])
        )
    for trans in [[l / 2, 0, 0], [l / 2, 0, l], [l / 2, l, 0], [l / 2, l, l]]:
        T[:3, 3] = np.array(trans) - voxel_size / 2
        scene.add_geometry(
            trimesh.creation.box([l, w, w], T, face_colors=[0, 0, 0, 255])
        )
    for trans in [[0, l / 2, 0], [0, l / 2, l], [l, l / 2, 0], [l, l / 2, l]]:
        T[:3, 3] = np.array(trans) - voxel_size / 2
        scene.add_geometry(
            trimesh.creation.box([w, l, w], T, face_colors=[0, 0, 0, 255])
        )


def create_voxel_scene(
    voxel_grid: np.ndarray,
    q_attention: np.ndarray = None,
    highlight_coordinate: np.ndarray = None,
    highlight_gt_coordinate: np.ndarray = None,
    highlight_alpha: float = 1.0,
    voxel_size: float = 0.1,
    show_bb: bool = False,
    alpha: float = 0.5,
):
    _, d, h, w = voxel_grid.shape
    v = voxel_grid.transpose((1, 2, 3, 0))
    occupancy = v[:, :, :, -1] != 0
    alpha = np.expand_dims(np.full_like(occupancy, alpha, dtype=np.float32), -1)
    rgb = np.concatenate([(v[:, :, :, 3:6] + 1) / 2.0, alpha], axis=-1)


    if q_attention is not None:
        q = np.max(q_attention, 0)
        q = q / np.max(q)
        show_q = q > 0.75
        occupancy = (show_q + occupancy).astype(bool)
        q = np.expand_dims(q - 0.5, -1)  # Max q can be is 0.9
        q_rgb = np.concatenate(
            [q, np.zeros_like(q), np.zeros_like(q), np.clip(q, 0, 1)], axis=-1
        )
        rgb = np.where(np.expand_dims(show_q, -1), q_rgb, rgb)

    if highlight_coordinate is not None:
        x, y, z = highlight_coordinate
        size = 2
        try:
            occupancy[x-size:x+size+2, y-size:y+size+2, z-size:z+size+2] = True
            rgb[x-size:x+size+2, y-size:y+size+2, z-size:z+size+2] = [0.0, 0.0, 1.0, highlight_alpha]   # blue
        except:
            print("highlight_coordinate is out of bounds")

    if highlight_gt_coordinate is not None:
        x, y, z = highlight_gt_coordinate
        size = 2
        try:
            occupancy[x-size:x+size+1, y-size:y+size+1, z-size:z+size+1] = True
            rgb[x-size:x+size+1, y-size:y+size+1, z-size:z+size+1] = [1.0, 0.0, 0.0, highlight_alpha]   # red
        except:
            print("highlight_gt_coordinate is out of bounds")
    
    transform = trimesh.transformations.scale_and_translate(
        scale=voxel_size, translate=(0.0, 0.0, 0.0)
    )
    trimesh_voxel_grid = trimesh.voxel.VoxelGrid(
        encoding=occupancy, transform=transform
    )
    geometry = trimesh_voxel_grid.as_boxes(colors=rgb)
    scene = trimesh.Scene()
    scene.add_geometry(geometry)
    if show_bb:
        assert d == h == w
        _create_bounding_box(scene, voxel_size, d)
    return scene


def visualise_voxel(
    voxel_grid: np.ndarray,
    q_attention: np.ndarray = None,
    highlight_coordinate: np.ndarray = None,
    highlight_gt_coordinate: np.ndarray = None,
    highlight_alpha: float = 1.0,
    rotation_amount: float = 0.0,
    show: bool = False,
    voxel_size: float = 0.1,
    offscreen_renderer=None,
    show_bb: bool = False,
    alpha: float = 0.5,
    filename: str = ""
):
    if pyrender is None or Trackball is None:
        raise ImportError("pyrender is required for voxel visualisation.")
    scene = create_voxel_scene(
        voxel_grid,
        q_attention,
        highlight_coordinate,
        highlight_gt_coordinate,
        highlight_alpha,
        voxel_size,
        show_bb,
        alpha,
    )
    if show:
        scene.show()
    else:
        r = offscreen_renderer or pyrender.OffscreenRenderer(
            viewport_width=640, viewport_height=480, point_size=1.0
        )
        s = _from_trimesh_scene(
            scene, ambient_light=[0.8, 0.8, 0.8], bg_color=[1.0, 1.0, 1.0]
        )
        cam = pyrender.PerspectiveCamera(
            yfov=np.pi / 4.0, aspectRatio=r.viewport_width / r.viewport_height
        )
        p = _compute_initial_camera_pose(s)
        t = Trackball(p, (r.viewport_width, r.viewport_height), s.scale, s.centroid)
        t.rotate(rotation_amount, np.array([0.0, 0.0, 1.0]))
        s.add(cam, pose=t.pose)
        color, depth = r.render(s)
        return color.copy()


def preprocess(img, dist="transporter"):
    """Pre-process input (subtract mean, divide by std)."""

    transporter_color_mean = [0.18877631, 0.18877631, 0.18877631]
    transporter_color_std = [0.07276466, 0.07276466, 0.07276466]
    transporter_depth_mean = 0.00509261
    transporter_depth_std = 0.00903967

    franka_color_mean = [0.622291933, 0.628313992, 0.623031488]
    franka_color_std = [0.168154213, 0.17626014, 0.184527364]
    franka_depth_mean = 0.872146842
    franka_depth_std = 0.195743116

    clip_color_mean = [0.48145466, 0.4578275, 0.40821073]
    clip_color_std = [0.26862954, 0.26130258, 0.27577711]

    # choose distribution
    if dist == "clip":
        color_mean = clip_color_mean
        color_std = clip_color_std
    elif dist == "franka":
        color_mean = franka_color_mean
        color_std = franka_color_std
    else:
        color_mean = transporter_color_mean
        color_std = transporter_color_std

    if dist == "franka":
        depth_mean = franka_depth_mean
        depth_std = franka_depth_std
    else:
        depth_mean = transporter_depth_mean
        depth_std = transporter_depth_std

    # convert to pytorch tensor (if required)
    if type(img) == torch.Tensor:

        def cast_shape(stat, img):
            tensor = torch.from_numpy(np.array(stat)).to(
                device=img.device, dtype=img.dtype
            )
            tensor = tensor.unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
            tensor = tensor.repeat(img.shape[0], 1, img.shape[-2], img.shape[-1])
            return tensor

        color_mean = cast_shape(color_mean, img)
        color_std = cast_shape(color_std, img)
        depth_mean = cast_shape(depth_mean, img)
        depth_std = cast_shape(depth_std, img)

        # normalize
        img = img.clone()
        img[:, :3, :, :] = (img[:, :3, :, :] / 255 - color_mean) / color_std
        img[:, 3:, :, :] = (img[:, 3:, :, :] - depth_mean) / depth_std
    else:
        # normalize
        img[:, :, :3] = (img[:, :, :3] / 255 - color_mean) / color_std
        img[:, :, 3:] = (img[:, :, 3:] - depth_mean) / depth_std
    return img


def rand_dist(size, min=-1.0, max=1.0):
    return (max - min) * torch.rand(size) + min


def rand_discrete(size, min=0, max=1):
    if min == max:
        return torch.zeros(size)
    return torch.randint(min, max + 1, size)


def split_list(lst, n):
    for i in range(0, len(lst), n):
        yield lst[i : i + n]


def get_device(gpu):
    if gpu is not None and gpu >= 0 and torch.cuda.is_available():
        device = torch.device("cuda:%d" % gpu)
        torch.backends.cudnn.enabled = torch.backends.cudnn.benchmark = True
    else:
        device = torch.device("cpu")
    return device





# def loader(file):
#     if str(file).endswith(".npy"):
#         try:
#             content = np.load(file, allow_pickle=True)
#             return content
#         except UnpicklingError as e:
#             print(f"Can't load {file}: {e}")
#     elif str(file).endswith(".dat"):
#         try:
#             with open(file, "rb") as f:
#                 content = pickle.loads(blosc.decompress(f.read()))
#             return content
#         except UnpicklingError as e:
#             print(f"Can't load {file}: {e}")
#     elif str(file).endswith(".pkl"):
#         try:
#             with open(file, 'rb') as f:
#                 content = pickle.load(f)
#             return content
#         except UnpicklingError as e:
#             print(f"Can't load {file}: {e}")
#     return None


# class Resize:
#     """Resize and pad/crop the image and aligned point cloud."""

#     def __init__(self, scales):
#         self.scales = scales

#     def __call__(self, **kwargs):
#         """Accept tensors as T, N, C, H, W."""
#         keys = list(kwargs.keys())

#         if len(keys) == 0:
#             raise RuntimeError("No args")

#         # Sample resize scale from continuous range
#         sc = np.random.uniform(*self.scales)

#         t, n, c, raw_h, raw_w = kwargs[keys[0]].shape
#         kwargs = {n: arg.flatten(0, 1) for n, arg in kwargs.items()}
#         resized_size = [int(raw_h * sc), int(raw_w * sc)]

#         # Resize
#         kwargs = {
#             n: transforms_f.resize(
#                 arg,
#                 resized_size,
#                 transforms.InterpolationMode.NEAREST
#             )
#             for n, arg in kwargs.items()
#         }

#         # If resized image is smaller than original, pad it with a reflection
#         if raw_h > resized_size[0] or raw_w > resized_size[1]:
#             right_pad, bottom_pad = max(raw_w - resized_size[1], 0), max(
#                 raw_h - resized_size[0], 0
#             )
#             kwargs = {
#                 n: transforms_f.pad(
#                     arg,
#                     padding=[0, 0, right_pad, bottom_pad],
#                     padding_mode="reflect",
#                 )
#                 for n, arg in kwargs.items()
#             }

#         # If resized image is larger than original, crop it
#         i, j, h, w = transforms.RandomCrop.get_params(
#             kwargs[keys[0]], output_size=(raw_h, raw_w)
#         )
#         kwargs = {
#             n: transforms_f.crop(arg, i, j, h, w) for n, arg in kwargs.items()
#         }

#         kwargs = {
#             n: einops.rearrange(arg, "(t n) c h w -> t n c h w", t=t)
#             for n, arg in kwargs.items()
#         }

#         return kwargs


# class TrajectoryInterpolator:
#     """Interpolate a trajectory to have fixed length."""

#     def __init__(self, use=False, interpolation_length=50):
#         self._use = use
#         self._interpolation_length = interpolation_length

#     def __call__(self, trajectory):
#         if not self._use:
#             return trajectory
#         trajectory = trajectory.numpy()
#         # Calculate the current number of steps
#         old_num_steps = len(trajectory)

#         # Create a 1D array for the old and new steps
#         old_steps = np.linspace(0, 1, old_num_steps)
#         new_steps = np.linspace(0, 1, self._interpolation_length)

#         # Interpolate each dimension separately
#         resampled = np.empty((self._interpolation_length, trajectory.shape[1]))
#         for i in range(trajectory.shape[1]):
#             if i == (trajectory.shape[1] - 1):  # gripper opening
#                 interpolator = interp1d(old_steps, trajectory[:, i])
#             else:
#                 interpolator = CubicSpline(old_steps, trajectory[:, i])

#             resampled[:, i] = interpolator(new_steps)

#         resampled = torch.tensor(resampled)
#         if trajectory.shape[1] == 8:
#             resampled[:, 3:7] = normalise_quat(resampled[:, 3:7])
#         return resampled