Official code release for ConceptGraphs
"""
PyTorch dataset classes for datasets in the NICE-SLAM format.
Large chunks of code stolen and adapted from:
https://github.com/cvg/nice-slam/blob/645b53af3dc95b4b348de70e759943f7228a61ca/src/utils/datasets.py

Support for Replica (sequences from the iMAP paper), TUM RGB-D, NICE-SLAM Apartment.
TODO: Add Azure Kinect dataset support
"""

import abc
import glob
import json
import os
from pathlib import Path
from typing import Dict, List, Optional, Union

import cv2
import imageio
import numpy as np
import torch
import torch.nn.functional as F
import yaml
from natsort import natsorted
from scipy.spatial.transform import Rotation as R

from gradslam.datasets import datautils
from gradslam.geometry.geometryutils import relative_transformation
from gradslam.slam.pointfusion import PointFusion
from gradslam.structures.rgbdimages import RGBDImages

from conceptgraph.utils.general_utils import to_scalar


def as_intrinsics_matrix(intrinsics):
    """
    Get matrix representation of intrinsics.

    """
    K = np.eye(3)
    K[0, 0] = intrinsics[0]
    K[1, 1] = intrinsics[1]
    K[0, 2] = intrinsics[2]
    K[1, 2] = intrinsics[3]
    return K

def from_intrinsics_matrix(K: torch.Tensor) -> tuple[float, float, float, float]:
    '''
    Get fx, fy, cx, cy from the intrinsics matrix
    
    return 4 scalars
    '''
    fx = to_scalar(K[0, 0])
    fy = to_scalar(K[1, 1])
    cx = to_scalar(K[0, 2])
    cy = to_scalar(K[1, 2])
    return fx, fy, cx, cy


def readEXR_onlydepth(filename):
    """
    Read depth data from EXR image file.

    Args:
        filename (str): File path.

    Returns:
        Y (numpy.array): Depth buffer in float32 format.
    """
    # move the import here since only CoFusion needs these package
    # sometimes installation of openexr is hard, you can run all other datasets
    # even without openexr
    import Imath
    import OpenEXR as exr

    exrfile = exr.InputFile(filename)
    header = exrfile.header()
    dw = header["dataWindow"]
    isize = (dw.max.y - dw.min.y + 1, dw.max.x - dw.min.x + 1)

    channelData = dict()

    for c in header["channels"]:
        C = exrfile.channel(c, Imath.PixelType(Imath.PixelType.FLOAT))
        C = np.fromstring(C, dtype=np.float32)
        C = np.reshape(C, isize)

        channelData[c] = C

    Y = None if "Y" not in header["channels"] else channelData["Y"]

    return Y


class GradSLAMDataset(torch.utils.data.Dataset):
    def __init__(
        self,
        config_dict,
        stride: Optional[int] = 1,
        start: Optional[int] = 0,
        end: Optional[int] = -1,
        desired_height: int = 480,
        desired_width: int = 640,
        channels_first: bool = False,
        normalize_color: bool = False,
        device="cuda:0",
        dtype=torch.float,
        load_embeddings: bool = False,
        embedding_dir: str = "feat_lseg_240_320",
        embedding_dim: int = 512,
        relative_pose: bool = True, # If True, the pose is relative to the first frame
        **kwargs,
    ):
        super().__init__()
        self.name = config_dict["dataset_name"]
        self.device = device
        self.png_depth_scale = config_dict["camera_params"]["png_depth_scale"]

        self.orig_height = config_dict["camera_params"]["image_height"]
        self.orig_width = config_dict["camera_params"]["image_width"]
        self.fx = config_dict["camera_params"]["fx"]
        self.fy = config_dict["camera_params"]["fy"]
        self.cx = config_dict["camera_params"]["cx"]
        self.cy = config_dict["camera_params"]["cy"]

        self.dtype = dtype

        self.desired_height = desired_height
        self.desired_width = desired_width
        self.height_downsample_ratio = float(self.desired_height) / self.orig_height
        self.width_downsample_ratio = float(self.desired_width) / self.orig_width
        self.channels_first = channels_first
        self.normalize_color = normalize_color

        self.load_embeddings = load_embeddings
        self.embedding_dir = embedding_dir
        self.embedding_dim = embedding_dim
        self.relative_pose = relative_pose

        self.start = start
        self.end = end
        if start < 0:
            raise ValueError("start must be positive. Got {0}.".format(stride))
        if not (end == -1 or end > start):
            raise ValueError(
                "end ({0}) must be -1 (use all images) or greater than start ({1})".format(end, start)
            )

        self.distortion = (
            np.array(config_dict["camera_params"]["distortion"])
            if "distortion" in config_dict["camera_params"]
            else None
        )
        self.crop_size = (
            config_dict["camera_params"]["crop_size"]
            if "crop_size" in config_dict["camera_params"]
            else None
        )

        self.crop_edge = None
        if "crop_edge" in config_dict["camera_params"].keys():
            self.crop_edge = config_dict["camera_params"]["crop_edge"]

        self.color_paths, self.depth_paths, self.embedding_paths = self.get_filepaths()
        if len(self.color_paths) != len(self.depth_paths):
            raise ValueError("Number of color and depth images must be the same.")
        if self.load_embeddings:
            if len(self.color_paths) != len(self.embedding_paths):
                raise ValueError(
                    "Mismatch between number of color images and number of embedding files."
                )
        self.num_imgs = len(self.color_paths)
        self.poses = self.load_poses()
        
        if self.end == -1:
            self.end = self.num_imgs

        self.color_paths = self.color_paths[self.start : self.end : stride]
        self.depth_paths = self.depth_paths[self.start : self.end : stride]
        if self.load_embeddings:
            self.embedding_paths = self.embedding_paths[self.start : self.end : stride]
        self.poses = self.poses[self.start : self.end : stride]
        # Tensor of retained indices (indices of frames and poses that were retained)
        self.retained_inds = torch.arange(self.num_imgs)[self.start : self.end : stride]
        # Update self.num_images after subsampling the dataset
        self.num_imgs = len(self.color_paths)

        # self.transformed_poses = datautils.poses_to_transforms(self.poses)
        self.poses = torch.stack(self.poses)
        if self.relative_pose:
            self.transformed_poses = self._preprocess_poses(self.poses)
        else:
            self.transformed_poses = self.poses

    def __len__(self):
        return self.num_imgs

    def get_filepaths(self):
        """Return paths to color images, depth images. Implement in subclass."""
        raise NotImplementedError

    def load_poses(self):
        """Load camera poses. Implement in subclass."""
        raise NotImplementedError

    def _preprocess_color(self, color: np.ndarray):
        r"""Preprocesses the color image by resizing to :math:`(H, W, C)`, (optionally) normalizing values to
        :math:`[0, 1]`, and (optionally) using channels first :math:`(C, H, W)` representation.

        Args:
            color (np.ndarray): Raw input rgb image

        Retruns:
            np.ndarray: Preprocessed rgb image

        Shape:
            - Input: :math:`(H_\text{old}, W_\text{old}, C)`
            - Output: :math:`(H, W, C)` if `self.channels_first == False`, else :math:`(C, H, W)`.
        """
        color = cv2.resize(
            color,
            (self.desired_width, self.desired_height),
            interpolation=cv2.INTER_LINEAR,
        )
        if self.normalize_color:
            color = datautils.normalize_image(color)
        if self.channels_first:
            color = datautils.channels_first(color)
        return color

    def _preprocess_depth(self, depth: np.ndarray):
        r"""Preprocesses the depth image by resizing, adding channel dimension, and scaling values to meters. Optionally
        converts depth from channels last :math:`(H, W, 1)` to channels first :math:`(1, H, W)` representation.

        Args:
            depth (np.ndarray): Raw depth image

        Returns:
            np.ndarray: Preprocessed depth

        Shape:
            - depth: :math:`(H_\text{old}, W_\text{old})`
            - Output: :math:`(H, W, 1)` if `self.channels_first == False`, else :math:`(1, H, W)`.
        """
        depth = cv2.resize(
            depth.astype(float),
            (self.desired_width, self.desired_height),
            interpolation=cv2.INTER_NEAREST,
        )
        depth = np.expand_dims(depth, -1)
        if self.channels_first:
            depth = datautils.channels_first(depth)
        return depth / self.png_depth_scale
    
    def _preprocess_poses(self, poses: torch.Tensor):
        r"""Preprocesses the poses by setting first pose in a sequence to identity and computing the relative
        homogenous transformation for all other poses.

        Args:
            poses (torch.Tensor): Pose matrices to be preprocessed

        Returns:
            Output (torch.Tensor): Preprocessed poses

        Shape:
            - poses: :math:`(L, 4, 4)` where :math:`L` denotes sequence length.
            - Output: :math:`(L, 4, 4)` where :math:`L` denotes sequence length.
        """
        return relative_transformation(
            poses[0].unsqueeze(0).repeat(poses.shape[0], 1, 1),
            poses,
            orthogonal_rotations=False,
        )
        
    def get_cam_K(self):
        '''
        Return camera intrinsics matrix K
        
        Returns:
            K (torch.Tensor): Camera intrinsics matrix, of shape (3, 3)
        '''
        K = as_intrinsics_matrix([self.fx, self.fy, self.cx, self.cy])
        K = torch.from_numpy(K)
        return K
    
    def read_embedding_from_file(self, embedding_path: str):
        '''
        Read embedding from file and process it. To be implemented in subclass for each dataset separately.
        '''
        raise NotImplementedError

    def __getitem__(self, index):
        color_path = self.color_paths[index]
        depth_path = self.depth_paths[index]
        color = np.asarray(imageio.imread(color_path), dtype=float)
        color = self._preprocess_color(color)
        color = torch.from_numpy(color)
        if ".png" in depth_path:
            # depth_data = cv2.imread(depth_path, cv2.IMREAD_UNCHANGED)
            depth = np.asarray(imageio.imread(depth_path), dtype=np.int64)
        elif ".exr" in depth_path:
            depth = readEXR_onlydepth(depth_path)
        elif ".npy" in depth_path:
            depth = np.load(depth_path)
        else:
            raise NotImplementedError

        K = as_intrinsics_matrix([self.fx, self.fy, self.cx, self.cy])
        K = torch.from_numpy(K)
        if self.distortion is not None:
            # undistortion is only applied on color image, not depth!
            color = cv2.undistort(color, K, self.distortion)

        depth = self._preprocess_depth(depth)
        depth = torch.from_numpy(depth)

        K = datautils.scale_intrinsics(
            K, self.height_downsample_ratio, self.width_downsample_ratio
        )
        intrinsics = torch.eye(4).to(K)
        intrinsics[:3, :3] = K

        pose = self.transformed_poses[index]

        if self.load_embeddings:
            embedding = self.read_embedding_from_file(self.embedding_paths[index])
            return (
                color.to(self.device).type(self.dtype),
                depth.to(self.device).type(self.dtype),
                intrinsics.to(self.device).type(self.dtype),
                pose.to(self.device).type(self.dtype),
                embedding.to(self.device),  # Allow embedding to be another dtype
                # self.retained_inds[index].item(),
            )

        return (
            color.to(self.device).type(self.dtype),
            depth.to(self.device).type(self.dtype),
            intrinsics.to(self.device).type(self.dtype),
            pose.to(self.device).type(self.dtype),
            # self.retained_inds[index].item(),
        )


class ICLDataset(GradSLAMDataset):
    def __init__(
        self,
        config_dict: Dict,
        basedir: Union[Path, str],
        sequence: Union[Path, str],
        stride: Optional[int] = 1,
        start: Optional[int] = 0,
        end: Optional[int] = -1,
        desired_height: Optional[int] = 480,
        desired_width: Optional[int] = 640,
        load_embeddings: Optional[bool] = False,
        embedding_dir: Optional[Union[Path, str]] = "embeddings",
        embedding_dim: Optional[int] = 512,
        embedding_file_extension: Optional[str] = "pt",
        **kwargs,
    ):
        self.input_folder = os.path.join(basedir, sequence)
        # Attempt to find pose file (*.gt.sim)
        self.pose_path = glob.glob(os.path.join(self.input_folder, "*.gt.sim"))
        if self.pose_path == 0:
            raise ValueError("Need pose file ending in extension `*.gt.sim`")
        self.pose_path = self.pose_path[0]
        self.embedding_file_extension = embedding_file_extension
        super().__init__(
            config_dict,
            stride=stride,
            start=start,
            end=end,
            desired_height=desired_height,
            desired_width=desired_width,
            load_embeddings=load_embeddings,
            embedding_dir=embedding_dir,
            embedding_dim=embedding_dim,
            **kwargs,
        )

    def get_filepaths(self):
        color_paths = natsorted(glob.glob(f"{self.input_folder}/rgb/*.png"))
        depth_paths = natsorted(glob.glob(f"{self.input_folder}/depth/*.png"))
        embedding_paths = None
        if self.load_embeddings:
            embedding_paths = natsorted(
                glob.glob(
                    f"{self.input_folder}/{self.embedding_dir}/*.{self.embedding_file_extension}"
                )
            )
        return color_paths, depth_paths, embedding_paths

    def load_poses(self):
        poses = []

        lines = []
        with open(self.pose_path, "r") as f:
            lines = f.readlines()

        _posearr = []
        for line in lines:
            line = line.strip().split()
            if len(line) == 0:
                continue
            _npvec = np.asarray(
                [float(line[0]), float(line[1]), float(line[2]), float(line[3])]
            )
            _posearr.append(_npvec)
        _posearr = np.stack(_posearr)

        for pose_line_idx in range(0, _posearr.shape[0], 3):
            _curpose = np.zeros((4, 4))
            _curpose[3, 3] = 3
            _curpose[0] = _posearr[pose_line_idx]
            _curpose[1] = _posearr[pose_line_idx + 1]
            _curpose[2] = _posearr[pose_line_idx + 2]
            poses.append(torch.from_numpy(_curpose).float())

        return poses

    def read_embedding_from_file(self, embedding_file_path):
        embedding = torch.load(embedding_file_path)
        return embedding.permute(0, 2, 3, 1)  # (1, H, W, embedding_dim)


class ReplicaDataset(GradSLAMDataset):
    def __init__(
        self,
        config_dict,
        basedir,
        sequence,
        stride: Optional[int] = None,
        start: Optional[int] = 0,
        end: Optional[int] = -1,
        desired_height: Optional[int] = 480,
        desired_width: Optional[int] = 640,
        load_embeddings: Optional[bool] = False,
        embedding_dir: Optional[str] = "embeddings",
        embedding_dim: Optional[int] = 512,
        **kwargs,
    ):
        self.input_folder = os.path.join(basedir, sequence)
        self.pose_path = os.path.join(self.input_folder, "traj.txt")
        super().__init__(
            config_dict,
            stride=stride,
            start=start,
            end=end,
            desired_height=desired_height,
            desired_width=desired_width,
            load_embeddings=load_embeddings,
            embedding_dir=embedding_dir,
            embedding_dim=embedding_dim,
            **kwargs,
        )

    def get_filepaths(self):
        color_paths = natsorted(glob.glob(f"{self.input_folder}/results/frame*.jpg"))
        depth_paths = natsorted(glob.glob(f"{self.input_folder}/results/depth*.png"))
        embedding_paths = None
        if self.load_embeddings:
            embedding_paths = natsorted(
                glob.glob(f"{self.input_folder}/{self.embedding_dir}/*.pt")
            )
        return color_paths, depth_paths, embedding_paths

    def load_poses(self):
        poses = []
        with open(self.pose_path, "r") as f:
            lines = f.readlines()
        for i in range(self.num_imgs):
            line = lines[i]
            c2w = np.array(list(map(float, line.split()))).reshape(4, 4)
            # c2w[:3, 1] *= -1
            # c2w[:3, 2] *= -1
            c2w = torch.from_numpy(c2w).float()
            poses.append(c2w)
        return poses

    def read_embedding_from_file(self, embedding_file_path):
        embedding = torch.load(embedding_file_path)
        return embedding.permute(0, 2, 3, 1)  # (1, H, W, embedding_dim)


class ScannetDataset(GradSLAMDataset):
    def __init__(
        self,
        config_dict,
        basedir,
        sequence,
        stride: Optional[int] = None,
        start: Optional[int] = 0,
        end: Optional[int] = -1,
        desired_height: Optional[int] = 968,
        desired_width: Optional[int] = 1296,
        load_embeddings: Optional[bool] = False,
        embedding_dir: Optional[str] = "embeddings",
        embedding_dim: Optional[int] = 512,
        **kwargs,
    ):
        self.input_folder = os.path.join(basedir, sequence)
        self.pose_path = None
        super().__init__(
            config_dict,
            stride=stride,
            start=start,
            end=end,
            desired_height=desired_height,
            desired_width=desired_width,
            load_embeddings=load_embeddings,
            embedding_dir=embedding_dir,
            embedding_dim=embedding_dim,
            **kwargs,
        )

    def get_filepaths(self):
        color_paths = natsorted(glob.glob(f"{self.input_folder}/color/*.jpg"))
        depth_paths = natsorted(glob.glob(f"{self.input_folder}/depth/*.png"))
        embedding_paths = None
        if self.load_embeddings:
            embedding_paths = natsorted(
                glob.glob(f"{self.input_folder}/{self.embedding_dir}/*.pt")
            )
        return color_paths, depth_paths, embedding_paths

    def load_poses(self):
        poses = []
        posefiles = natsorted(glob.glob(f"{self.input_folder}/pose/*.txt"))
        for posefile in posefiles:
            _pose = torch.from_numpy(np.loadtxt(posefile))
            poses.append(_pose)
        return poses

    def read_embedding_from_file(self, embedding_file_path):
        print(embedding_file_path)
        embedding = torch.load(embedding_file_path, map_location="cpu")
        return embedding.permute(0, 2, 3, 1)  # (1, H, W, embedding_dim)


class AzureKinectDataset(GradSLAMDataset):
    def __init__(
        self,
        config_dict,
        basedir,
        sequence,
        stride: Optional[int] = None,
        start: Optional[int] = 0,
        end: Optional[int] = -1,
        desired_height: Optional[int] = 480,
        desired_width: Optional[int] = 640,
        load_embeddings: Optional[bool] = False,
        embedding_dir: Optional[str] = "embeddings",
        embedding_dim: Optional[int] = 512,
        **kwargs,
    ):
        self.input_folder = os.path.join(basedir, sequence)
        self.pose_path = None

        # check if a file named 'poses_global_dvo.txt' exists in the basedir / sequence folder 
        if os.path.isfile(os.path.join(basedir, sequence, 'poses_global_dvo.txt')):
            self.pose_path = os.path.join(basedir, sequence, 'poses_global_dvo.txt')
            
        # if "odomfile" in kwargs.keys():
        #     self.pose_path = kwargs["odomfile"]
        super().__init__(
            config_dict,
            stride=stride,
            start=start,
            end=end,
            desired_height=desired_height,
            desired_width=desired_width,
            load_embeddings=load_embeddings,
            embedding_dir=embedding_dir,
            embedding_dim=embedding_dim,
            **kwargs,
        )

    def get_filepaths(self):
        color_paths = natsorted(glob.glob(f"{self.input_folder}/color/*.jpg"))
        depth_paths = natsorted(glob.glob(f"{self.input_folder}/depth/*.png"))
        embedding_paths = None
        if self.load_embeddings:
            embedding_paths = natsorted(
                glob.glob(f"{self.input_folder}/{self.embedding_dir}/*.pt")
            )
        return color_paths, depth_paths, embedding_paths

    def load_poses(self):
        if self.pose_path is None:
            print(
                "WARNING: Dataset does not contain poses. Returning identity transform."
            )
            return [torch.eye(4).float() for _ in range(self.num_imgs)]
        else:
            # Determine whether the posefile ends in ".log"
            # a .log file has the following format for each frame
            # frame_idx frame_idx+1
            # row 1 of 4x4 transform
            # row 2 of 4x4 transform
            # row 3 of 4x4 transform
            # row 4 of 4x4 transform
            # [repeat for all frames]
            #
            # on the other hand, the "poses_o3d.txt" or "poses_dvo.txt" files have the format
            # 16 entries of 4x4 transform
            # [repeat for all frames]
            if self.pose_path.endswith(".log"):
                # print("Loading poses from .log format")
                poses = []
                lines = None
                with open(self.pose_path, "r") as f:
                    lines = f.readlines()
                if len(lines) % 5 != 0:
                    raise ValueError(
                        "Incorrect file format for .log odom file "
                        "Number of non-empty lines must be a multiple of 5"
                    )
                num_lines = len(lines) // 5
                for i in range(0, num_lines):
                    _curpose = []
                    _curpose.append(list(map(float, lines[5 * i + 1].split())))
                    _curpose.append(list(map(float, lines[5 * i + 2].split())))
                    _curpose.append(list(map(float, lines[5 * i + 3].split())))
                    _curpose.append(list(map(float, lines[5 * i + 4].split())))
                    _curpose = np.array(_curpose).reshape(4, 4)
                    poses.append(torch.from_numpy(_curpose))
            else:
                poses = []
                lines = None
                with open(self.pose_path, "r") as f:
                    lines = f.readlines()
                for line in lines:
                    if len(line.split()) == 0:
                        continue
                    c2w = np.array(list(map(float, line.split()))).reshape(4, 4)
                    poses.append(torch.from_numpy(c2w))
            return poses

    def read_embedding_from_file(self, embedding_file_path):
        embedding = torch.load(embedding_file_path)
        return embedding  # .permute(0, 2, 3, 1)  # (1, H, W, embedding_dim)


class RealsenseDataset(GradSLAMDataset):
    """
    Dataset class to process depth images captured by realsense camera on the tabletop manipulator 
    """
    def __init__(
        self,
        config_dict,
        basedir,
        sequence,
        stride: Optional[int] = None,
        start: Optional[int] = 0,
        end: Optional[int] = -1,
        desired_height: Optional[int] = 480,
        desired_width: Optional[int] = 640,
        load_embeddings: Optional[bool] = False,
        embedding_dir: Optional[str] = "embeddings",
        embedding_dim: Optional[int] = 512,
        **kwargs,
    ):
        self.input_folder = os.path.join(basedir, sequence)
        # only poses/images/depth corresponding to the realsense_camera_order are read/used
        self.pose_path = os.path.join(self.input_folder, "poses")
        super().__init__(
            config_dict,
            stride=stride,
            start=start,
            end=end,
            desired_height=desired_height,
            desired_width=desired_width,
            load_embeddings=load_embeddings,
            embedding_dir=embedding_dir,
            embedding_dim=embedding_dim,
            **kwargs,
        )

    def get_filepaths(self):
        color_paths = natsorted(
            glob.glob(os.path.join(self.input_folder, "rgb", "*.jpg"))
        )
        depth_paths = natsorted(
            glob.glob(os.path.join(self.input_folder, "depth", "*.png"))
        )
        embedding_paths = None
        if self.load_embeddings:
            embedding_paths = natsorted(
                glob.glob(f"{self.input_folder}/{self.embedding_dir}/*.pt")
            )
        return color_paths, depth_paths, embedding_paths

    def load_poses(self):
        posefiles = natsorted(glob.glob(os.path.join(self.pose_path, "*.npy")))
        poses = []
        P = torch.tensor(
            [
                [1, 0, 0, 0],
                [0, -1, 0, 0],
                [0, 0, -1, 0],
                [0, 0, 0, 1]
            ]
        ).float()
        for posefile in posefiles:
            c2w = torch.from_numpy(np.load(posefile)).float()
            _R = c2w[:3, :3]
            _t = c2w[:3, 3]
            _pose = P @ c2w @ P.T
            poses.append(_pose)
        return poses

    def read_embedding_from_file(self, embedding_file_path):
        embedding = torch.load(embedding_file_path)
        return embedding.permute(0, 2, 3, 1)  # (1, H, W, embedding_dim)


class Record3DDataset(GradSLAMDataset):
    """
    Dataset class to read in saved files from the structure created by our
    `save_record3d_stream.py` script
    """
    def __init__(
        self,
        config_dict,
        basedir,
        sequence,
        stride: Optional[int] = None,
        start: Optional[int] = 0,
        end: Optional[int] = -1,
        desired_height: Optional[int] = 480,
        desired_width: Optional[int] = 640,
        load_embeddings: Optional[bool] = False,
        embedding_dir: Optional[str] = "embeddings",
        embedding_dim: Optional[int] = 512,
        **kwargs,
    ):
        self.input_folder = os.path.join(basedir, sequence)
        self.pose_path = os.path.join(self.input_folder, "poses")
        super().__init__(
            config_dict,
            stride=stride,
            start=start,
            end=end,
            desired_height=desired_height,
            desired_width=desired_width,
            load_embeddings=load_embeddings,
            embedding_dir=embedding_dir,
            embedding_dim=embedding_dim,
            **kwargs,
        )

    def get_filepaths(self):
        color_paths = natsorted(
            glob.glob(os.path.join(self.input_folder, "rgb", "*.png"))
        )
        depth_paths = natsorted(
            glob.glob(os.path.join(self.input_folder, "depth", "*.png"))
        )
        embedding_paths = None
        if self.load_embeddings:
            embedding_paths = natsorted(
                glob.glob(f"{self.input_folder}/{self.embedding_dir}/*.pt")
            )
        return color_paths, depth_paths, embedding_paths

    def load_poses(self):
        posefiles = natsorted(glob.glob(os.path.join(self.pose_path, "*.npy")))
        poses = []
        P = torch.tensor(
            [
                [1, 0, 0, 0],
                [0, -1, 0, 0],
                [0, 0, -1, 0],
                [0, 0, 0, 1]
            ]
        ).float()
        for posefile in posefiles:
            c2w = torch.from_numpy(np.load(posefile)).float()
            _R = c2w[:3, :3]
            _t = c2w[:3, 3]
            _pose = P @ c2w @ P.T
            poses.append(_pose)
        return poses

    def read_embedding_from_file(self, embedding_file_path):
        embedding = torch.load(embedding_file_path)
        return embedding.permute(0, 2, 3, 1)  # (1, H, W, embedding_dim)


class MultiscanDataset(GradSLAMDataset):
    def __init__(
        self,
        config_dict,
        basedir,
        sequence,
        stride: Optional[int] = None,
        start: Optional[int] = 0,
        end: Optional[int] = -1,
        desired_height: Optional[int] = 480,
        desired_width: Optional[int] = 640,
        load_embeddings: Optional[bool] = False,
        embedding_dir: Optional[str] = "embeddings",
        embedding_dim: Optional[int] = 512,
        **kwargs,
    ):
        self.input_folder = os.path.join(basedir, sequence)
        self.pose_path = os.path.join(self.input_folder, f"{sequence}.jsonl")
        
        scene_meta = json.load(
            open(os.path.join(self.input_folder, f"{sequence}.json"), "r")
        )
        cam_K = scene_meta['streams'][0]['intrinsics']
        cam_K = np.array(cam_K).reshape(3, 3).T
        
        config_dict['camera_params']['fx'] = cam_K[0, 0]
        config_dict['camera_params']['fy'] = cam_K[1, 1]
        config_dict['camera_params']['cx'] = cam_K[0, 2]
        config_dict['camera_params']['cy'] = cam_K[1, 2]
        config_dict["camera_params"]["image_height"] = scene_meta['streams'][0]['resolution'][0]
        config_dict["camera_params"]["image_width"] = scene_meta['streams'][0]['resolution'][1]
        
        super().__init__(
            config_dict,
            stride=stride,
            start=start,
            end=end,
            desired_height=desired_height,
            desired_width=desired_width,
            load_embeddings=load_embeddings,
            embedding_dir=embedding_dir,
            embedding_dim=embedding_dim,
            **kwargs,
        )
        
    def get_filepaths(self):
        color_paths = natsorted(glob.glob(f"{self.input_folder}/outputs/color/*.png"))
        depth_paths = natsorted(glob.glob(f"{self.input_folder}/outputs/depth/*.png"))
        embedding_paths = None
        if self.load_embeddings:
            embedding_paths = natsorted(
                glob.glob(f"{self.input_folder}/{self.embedding_dir}/*.pt")
            )
            
        return color_paths, depth_paths, embedding_paths
        
    def load_poses(self):
        poses = []
        with open(self.pose_path, "r") as f:
            lines = f.readlines()
        n_sampled = len(glob.glob(f"{self.input_folder}/outputs/color/*.png"))
        step = round(len(lines) / float(n_sampled))

        poses = []
        for i in range(0, len(lines), step):
            line = lines[i]
            info = json.loads(line)
            transform = np.asarray(info.get('transform'))
            transform = np.reshape(transform, (4, 4), order='F')
            transform = np.dot(transform, np.diag([1, -1, -1, 1]))
            transform = transform / transform[3][3]
            poses.append(torch.from_numpy(transform).float())
            
        return poses
        
    def read_embedding_from_file(self, embedding_file_path):
        embedding = torch.load(embedding_file_path)
        return embedding.permute(0, 2, 3, 1)  # (1, H, W, embedding_dim)


class Hm3dDataset(GradSLAMDataset):
    def __init__(
        self,
        config_dict,
        basedir,
        sequence,
        stride: Optional[int] = None,
        start: Optional[int] = 0,
        end: Optional[int] = -1,
        desired_height: Optional[int] = 480,
        desired_width: Optional[int] = 640,
        load_embeddings: Optional[bool] = False,
        embedding_dir: Optional[str] = "embeddings",
        embedding_dim: Optional[int] = 512,
        **kwargs,
    ):
        self.input_folder = os.path.join(basedir, sequence)
        self.pose_path = None
        super().__init__(
            config_dict,
            stride=stride,
            start=start,
            end=end,
            desired_height=desired_height,
            desired_width=desired_width,
            load_embeddings=load_embeddings,
            embedding_dir=embedding_dir,
            embedding_dim=embedding_dim,
            **kwargs,
        )
        
    def get_filepaths(self):
        color_paths = natsorted(glob.glob(f"{self.input_folder}/*.png"))
        depth_paths = natsorted(glob.glob(f"{self.input_folder}/*_depth.npy"))
        embedding_paths = None
        if self.load_embeddings:
            embedding_paths = natsorted(
                glob.glob(f"{self.input_folder}/{self.embedding_dir}/*.pt")
            )
        return color_paths, depth_paths, embedding_paths
    
    def load_poses(self):
        poses = []
        posefiles = natsorted(glob.glob(f"{self.input_folder}/*.json"))
        
        P = torch.tensor(
            [
                [1, 0, 0, 0],
                [0, -1, 0, 0],
                [0, 0, -1, 0],
                [0, 0, 0, 1]
            ]
        ).float()
        
        for posefile in posefiles:
            with open(posefile, 'r') as f:
                pose_raw = json.load(f)
            pose = np.asarray(pose_raw['pose'])
            
            pose = torch.from_numpy(pose).float()
            pose = P @ pose @ P.T
            
            poses.append(pose)
            
        return poses

def load_dataset_config(path, default_path=None):
    """
    Loads config file.

    Args:
        path (str): path to config file.
        default_path (str, optional): whether to use default path. Defaults to None.

    Returns:
        cfg (dict): config dict.

    """
    # load configuration from file itself
    with open(path, "r") as f:
        cfg_special = yaml.full_load(f)

    # check if we should inherit from a config
    inherit_from = cfg_special.get("inherit_from")

    # if yes, load this config first as default
    # if no, use the default_path
    if inherit_from is not None:
        cfg = load_dataset_config(inherit_from, default_path)
    elif default_path is not None:
        with open(default_path, "r") as f:
            cfg = yaml.full_load(f)
    else:
        cfg = dict()

    # include main configuration
    update_recursive(cfg, cfg_special)

    return cfg


def update_recursive(dict1, dict2):
    """
    Update two config dictionaries recursively.

    Args:
        dict1 (dict): first dictionary to be updated.
        dict2 (dict): second dictionary which entries should be used.
    """
    for k, v in dict2.items():
        if k not in dict1:
            dict1[k] = dict()
        if isinstance(v, dict):
            update_recursive(dict1[k], v)
        else:
            dict1[k] = v


def common_dataset_to_batch(dataset):
    colors, depths, poses = [], [], []
    intrinsics, embeddings = None, None
    for idx in range(len(dataset)):
        _color, _depth, intrinsics, _pose, _embedding = dataset[idx]
        colors.append(_color)
        depths.append(_depth)
        poses.append(_pose)
        if _embedding is not None:
            if embeddings is None:
                embeddings = [_embedding]
            else:
                embeddings.append(_embedding)
    colors = torch.stack(colors)
    depths = torch.stack(depths)
    poses = torch.stack(poses)
    if embeddings is not None:
        embeddings = torch.stack(embeddings, dim=1)
        # # (1, NUM_IMG, DIM_EMBED, H, W) -> (1, NUM_IMG, H, W, DIM_EMBED)
        # embeddings = embeddings.permute(0, 1, 3, 4, 2)
    colors = colors.unsqueeze(0)
    depths = depths.unsqueeze(0)
    intrinsics = intrinsics.unsqueeze(0).unsqueeze(0)
    poses = poses.unsqueeze(0)
    colors = colors.float()
    depths = depths.float()
    intrinsics = intrinsics.float()
    poses = poses.float()
    if embeddings is not None:
        embeddings = embeddings.float()
    return colors, depths, intrinsics, poses, embeddings


def get_dataset(dataconfig, basedir, sequence, **kwargs):
    config_dict = load_dataset_config(dataconfig)
    if config_dict["dataset_name"].lower() in ["icl"]:
        return ICLDataset(config_dict, basedir, sequence, **kwargs)
    elif config_dict["dataset_name"].lower() in ["replica"]:
        return ReplicaDataset(config_dict, basedir, sequence, **kwargs)
    elif config_dict["dataset_name"].lower() in ["azure", "azurekinect"]:
        return AzureKinectDataset(config_dict, basedir, sequence, **kwargs)
    elif config_dict["dataset_name"].lower() in ["scannet"]:
        return ScannetDataset(config_dict, basedir, sequence, **kwargs)
    elif config_dict["dataset_name"].lower() in ["record3d"]:
        return Record3DDataset(config_dict, basedir, sequence, **kwargs)
    elif config_dict["dataset_name"].lower() in ["realsense"]:
        return RealsenseDataset(config_dict, basedir, sequence, **kwargs)
    elif config_dict["dataset_name"].lower() in ["multiscan"]:
        return MultiscanDataset(config_dict, basedir, sequence, **kwargs)
    elif config_dict['dataset_name'].lower() in ['hm3d']:
        return Hm3dDataset(config_dict, basedir, sequence, **kwargs)
    else:
        raise ValueError(f"Unknown dataset name {config_dict['dataset_name']}")


if __name__ == "__main__":
    cfg = load_dataset_config(
        "/home/qiao/src/gradslam-foundation/examples/dataconfigs/replica/replica.yaml"
    )
    dataset = ReplicaDataset(
        config_dict=cfg,
        basedir="/home/qiao/src/nice-slam/Datasets/Replica",
        sequence="office0",
        start=0,
        end=1900,
        stride=100,
        # desired_height=680,
        # desired_width=1200,
        desired_height=240,
        desired_width=320,
    )

    colors, depths, poses = [], [], []
    intrinsics = None
    for idx in range(len(dataset)):
        _color, _depth, intrinsics, _pose = dataset[idx]
        colors.append(_color)
        depths.append(_depth)
        poses.append(_pose)
    colors = torch.stack(colors)
    depths = torch.stack(depths)
    poses = torch.stack(poses)
    colors = colors.unsqueeze(0)
    depths = depths.unsqueeze(0)
    intrinsics = intrinsics.unsqueeze(0).unsqueeze(0)
    poses = poses.unsqueeze(0)
    colors = colors.float()
    depths = depths.float()
    intrinsics = intrinsics.float()
    poses = poses.float()

    # create rgbdimages object
    rgbdimages = RGBDImages(
        colors,
        depths,
        intrinsics,
        poses,
        channels_first=False,
        has_embeddings=False,  # KM
    )

    # SLAM
    slam = PointFusion(odom="gt", dsratio=1, device="cuda:0", use_embeddings=False)
    pointclouds, recovered_poses = slam(rgbdimages)

    import open3d as o3d

    print(pointclouds.colors_padded.shape)
    pcd = pointclouds.open3d(0)
    o3d.visualization.draw_geometries([pcd])

    # from icl_dataset import ICLWithCLIPEmbeddings

    # dataset = ICLWithCLIPEmbeddings(
    #     os.path.join("/home/krishna/data/icl/"),
    #     trajectories="living_room_traj1_frei_png",
    #     stride=10,
    #     height=480,
    #     width=640,
    #     embedding_dir="feat_lseg_240_320",
    # )
    # colors, depths, intrinsics, poses, _, _, embeddings = dataset[
    #     0
    # ]  # next(iter(loader))
    # print(colors.shape, depths.shape, intrinsics.shape, poses.shape)

conceptgraph/dataset/datasets_common.py
concept-graphs/concept-graphs
"""
Preprocess an unzipped .r3d file to the Record3DDataset format.
"""

import glob
import json
import os
import yaml
from dataclasses import dataclass
from pathlib import Path
from typing import List, Tuple, Union

import cv2
import liblzfse  # https://pypi.org/project/pyliblzfse/
import numpy as np
import png  # pip install pypng
import torch
import tyro
from natsort import natsorted
from PIL import Image
from scipy.spatial.transform import Rotation
from tqdm import tqdm, trange

@dataclass
class ProgramArgs:
    datapath: str = "/home/krishna/data/record3d/krishna-bcs-room"


def load_depth(filepath):
    with open(filepath, 'rb') as depth_fh:
        raw_bytes = depth_fh.read()
        decompressed_bytes = liblzfse.decompress(raw_bytes)
        depth_img = np.frombuffer(decompressed_bytes, dtype=np.float32)

    # depth_img = depth_img.reshape((640, 480))  # For a FaceID camera 3D Video
    depth_img = depth_img.reshape((256, 192))  # For a LiDAR 3D Video

    return depth_img


def load_conf(filepath):
    with open(filepath, 'rb') as conf_fh:
        raw_bytes = conf_fh.read()
        decompressed_bytes = liblzfse.decompress(raw_bytes)
        conf_img = np.frombuffer(decompressed_bytes, dtype=np.uint8)

    # depth_img = depth_img.reshape((640, 480))  # For a FaceID camera 3D Video
    conf_img = conf_img.reshape((256, 192))  # For a LiDAR 3D Video

    return conf_img


def load_color(filepath):
    img = cv2.imread(filepath)
    return cv2.resize(img, (192, 256))


def write_color(outpath, img):
    cv2.imwrite(outpath, img)


def write_depth(outpath, depth):
    depth = depth * 1000
    depth = depth.astype(np.uint16)
    depth = Image.fromarray(depth)
    depth.save(outpath)

def write_conf(outpath, conf):
    np.save(outpath, conf)

def write_pose(outpath, pose):
    np.save(outpath, pose.astype(np.float32))


def get_poses(metadata_dict: dict) -> int:
    """Converts Record3D's metadata dict into pose matrices needed by nerfstudio
    Args:
        metadata_dict: Dict containing Record3D metadata
    Returns:
        np.array of pose matrices for each image of shape: (num_images, 4, 4)
    """

    poses_data = np.array(metadata_dict["poses"])  # (N, 3, 4)
    # NB: Record3D / scipy use "scalar-last" format quaternions (x y z w)
    # https://fzheng.me/2017/11/12/quaternion_conventions_en/
    camera_to_worlds = np.concatenate(
        [Rotation.from_quat(poses_data[:, :4]).as_matrix(), poses_data[:, 4:, None]],
        axis=-1,
    ).astype(np.float32)

    homogeneous_coord = np.zeros_like(camera_to_worlds[..., :1, :])
    homogeneous_coord[..., :, 3] = 1
    camera_to_worlds = np.concatenate([camera_to_worlds, homogeneous_coord], -2)
    
    return camera_to_worlds


def get_intrinsics(metadata_dict: dict, downscale_factor: float = 7.5) -> int:
    """Converts Record3D metadata dict into intrinsic info needed by nerfstudio
    Args:
        metadata_dict: Dict containing Record3D metadata
        downscale_factor: factor to scale RGB image by (usually scale factor is
            set to 7.5 for record3d 1.8 or higher -- this is the factor that downscales
            RGB images to lidar)
    Returns:
        dict with camera intrinsics keys needed by nerfstudio
    """

    # Camera intrinsics
    K = np.array(metadata_dict["K"]).reshape((3, 3)).T
    K = K / downscale_factor
    K[2, 2] = 1.0
    focal_length = K[0, 0]

    H = metadata_dict["h"]
    W = metadata_dict["w"]

    H = int(H / downscale_factor)
    W = int(W / downscale_factor)

    # # TODO(akristoffersen): The metadata dict comes with principle points,
    # # but caused errors in image coord indexing. Should update once that is fixed.
    # cx, cy = W / 2, H / 2
    cx, cy = K[0, 2], K[1, 2]

    intrinsics_dict = {
        "fx": focal_length,
        "fy": focal_length,
        "cx": cx,
        "cy": cy,
        "w": W,
        "h": H,
    }

    return intrinsics_dict


def main():
    args = tyro.cli(ProgramArgs)
    
    metadata = None
    with open(os.path.join(args.datapath, "metadata"), "r") as f:
        metadata = json.load(f)
    
    # Keys in metadata dict
    # h, w, K, fps, dw, dh, initPose, poses, cameraType, frameTimestamps
    # print(metadata.keys())

    poses = get_poses(metadata)
    intrinsics_dict = get_intrinsics(metadata)
    
    color_paths = natsorted(glob.glob(os.path.join(args.datapath, "rgbd", "*.jpg")))
    depth_paths = natsorted(glob.glob(os.path.join(args.datapath, "rgbd", "*.depth")))
    conf_paths = natsorted(glob.glob(os.path.join(args.datapath, "rgbd", "*.conf")))

    os.makedirs(os.path.join(args.datapath, "rgb"), exist_ok=True)
    os.makedirs(os.path.join(args.datapath, "conf"), exist_ok=True)
    os.makedirs(os.path.join(args.datapath, "depth"), exist_ok=True)
    os.makedirs(os.path.join(args.datapath, "poses"), exist_ok=True)

    cfg = {}
    cfg["dataset_name"] = "record3d"
    cfg["camera_params"] = {}
    cfg["camera_params"]["image_height"] = intrinsics_dict["h"]
    cfg["camera_params"]["image_width"] = intrinsics_dict["w"]
    cfg["camera_params"]["fx"] = intrinsics_dict["fx"].item()
    cfg["camera_params"]["fy"] = intrinsics_dict["fy"].item()
    cfg["camera_params"]["cx"] = intrinsics_dict["cx"].item()
    cfg["camera_params"]["cy"] = intrinsics_dict["cy"].item()
    cfg["camera_params"]["png_depth_scale"] = 1000.0
    print(cfg)
    with open(os.path.join(args.datapath, "dataconfig.yaml"), "w") as f:
        yaml.dump(cfg, f)

    for i in trange(len(color_paths)):
        color = load_color(color_paths[i])
        depth = load_depth(depth_paths[i])
        conf = load_conf(conf_paths[i])
        basename = os.path.splitext(os.path.basename(color_paths[i]))[0]
        # color_path = os.path.splitext(os.path.basename(color_paths[i]))[0] + ".png"
        write_color(os.path.join(args.datapath, "rgb", basename + ".png"), color)
        # depth_path = os.path.splitext(os.path.basename(depth_paths[i]))[0] + ".png"
        write_depth(os.path.join(args.datapath, "depth", basename + ".png"), depth)
        # conf_path = os.path.splitext(os.path.basename(conf_paths[i]))[0] + ".npy"
        write_conf(os.path.join(args.datapath, "conf", basename + ".npy"), conf)
        write_pose(os.path.join(args.datapath, "poses", basename + ".npy"), poses[i])
        # c2w = poses[i]
        # frame = {
        #     "file_path": os.path.join("rgb", color_path),
        #     "depth_path": os.path.join("depth", depth_path),
        #     "conf_path": os.path.join("conf", conf_path),
        #     "transform_matrix": c2w.tolist(),
        # }
        # frames.append(frame)


if __name__ == "__main__":
    main()

conceptgraph/dataset/preprocess_r3d_file.py
concept-graphs/concept-graphs
REPLICA_CLASSES = [
    "other",
    "backpack",
    "base-cabinet",
    "basket",
    "bathtub",
    "beam",
    "beanbag",
    "bed",
    "bench",
    "bike",
    "bin",
    "blanket",
    "blinds",
    "book",
    "bottle",
    "box",
    "bowl",
    "camera",
    "cabinet",
    "candle",
    "chair",
    "chopping-board",
    "clock",
    "cloth",
    "clothing",
    "coaster",
    "comforter",
    "computer-keyboard",
    "cup",
    "cushion",
    "curtain",
    "ceiling",
    "cooktop",
    "countertop",
    "desk",
    "desk-organizer",
    "desktop-computer",
    "door",
    "exercise-ball",
    "faucet",
    "floor",
    "handbag",
    "hair-dryer",
    "handrail",
    "indoor-plant",
    "knife-block",
    "kitchen-utensil",
    "lamp",
    "laptop",
    "major-appliance",
    "mat",
    "microwave",
    "monitor",
    "mouse",
    "nightstand",
    "pan",
    "panel",
    "paper-towel",
    "phone",
    "picture",
    "pillar",
    "pillow",
    "pipe",
    "plant-stand",
    "plate",
    "pot",
    "rack",
    "refrigerator",
    "remote-control",
    "scarf",
    "sculpture",
    "shelf",
    "shoe",
    "shower-stall",
    "sink",
    "small-appliance",
    "sofa",
    "stair",
    "stool",
    "switch",
    "table",
    "table-runner",
    "tablet",
    "tissue-paper",
    "toilet",
    "toothbrush",
    "towel",
    "tv-screen",
    "tv-stand",
    "umbrella",
    "utensil-holder",
    "vase",
    "vent",
    "wall",
    "wall-cabinet",
    "wall-plug",
    "wardrobe",
    "window",
    "rug",
    "logo",
    "bag",
    "set-of-clothing",
]

REPLICA_SCENE_IDS = [
    "room0",
    "room1",
    "room2",
    "office0",
    "office1",
    "office2",
    "office3",
    "office4",
]

# with _ in REPLICA_SCENE_IDS_
REPLICA_SCENE_IDS_ = [
    "room_0",
    "room_1",
    "room_2",
    "office_0",
    "office_1",
    "office_2",
    "office_3",
    "office_4",
]

REPLICA_EXISTING_CLASSES = [
    0,
    3,
    7,
    8,
    10,
    11,
    12,
    13,
    14,
    15,
    16,
    17,
    18,
    19,
    20,
    22,
    23,
    26,
    29,
    31,
    34,
    35,
    37,
    40,
    44,
    47,
    52,
    54,
    56,
    59,
    60,
    61,
    62,
    63,
    64,
    65,
    70,
    71,
    76,
    78,
    79,
    80,
    82,
    83,
    87,
    88,
    91,
    92,
    93,
    95,
    97,
    98,
]

conceptgraph/dataset/replica_constants.py
concept-graphs/concept-graphs
import argparse
import os
from threading import Event
import time

import cv2
import numpy as np
import torch
import yaml
from PIL import Image
from record3d import Record3DStream
from scipy.spatial.transform import Rotation
from tqdm import trange


def write_color(outpath, img):
    cv2.imwrite(outpath, img)


def write_depth(outpath, depth):
    depth = depth * 1000
    depth = depth.astype(np.uint16)
    depth = Image.fromarray(depth)
    depth.save(outpath)


def write_pose(outpath, pose):
    # quat_trans: Quaternion + world position (accessible via camera_pose.[qx|qy|qz|qw|tx|ty|tz])
    # NB: Record3D / scipy use "scalar-last" format quaternions (x y z w)
    # https://fzheng.me/2017/11/12/quaternion_conventions_en/
    # pose = np.asarray([pose.qx, pose.qy, pose.qz, pose.qw, pose.tx, pose.ty, pose.tz])

    c2w = np.zeros((4, 4))
    c2w[3, 3] = 1.0
    c2w[:3, :3] = Rotation.from_quat(pose[:4]).as_matrix()
    c2w[:3, 3] = pose[4:]
    np.save(outpath, c2w.astype(np.float32))


class DemoApp:
    def __init__(self, savedir_name="saved-record3d", seq_name="debug"):
        self.event = Event()
        self.session = None
        self.DEVICE_TYPE__TRUEDEPTH = 0
        self.DEVICE_TYPE__LIDAR = 1
        self.imgs = []
        self.depths = []
        self.intrinsics = None
        self.poses = []
        self.savedir_name = savedir_name
        self.seq_name = seq_name

    def on_new_frame(self):
        """
        This method is called from non-main thread, therefore cannot be used for presenting UI.
        """
        self.event.set()  # Notify the main thread to stop waiting and process new frame.

    def on_stream_stopped(self):
        print("Stream stopped")
        # seq_ts = f"{self.seq_name}_{str(int(time.time()))}"
        # savedir = os.path.join(self.savedir_name, seq_ts)
        savedir = os.path.join(self.savedir_name, self.seq_name)
        # add timestamp to savedir
        # savedir = os.path.join(savedir, str(int(time.time())))

        savedir_rgb = os.path.join(savedir, "rgb")
        savedir_depth = os.path.join(savedir, "depth")
        savedir_poses = os.path.join(savedir, "poses")

        os.makedirs(savedir, exist_ok=True)
        os.makedirs(savedir_rgb, exist_ok=True)
        os.makedirs(savedir_depth, exist_ok=True)
        os.makedirs(savedir_poses, exist_ok=True)
        # print(f"Line 69, savedir_poses: {savedir_poses}")

        self.cfg = {}
        self.cfg["dataset_name"] = "record3d"
        self.cfg["camera_params"] = {}
        self.cfg["camera_params"]["image_height"] = self.imgs[0].shape[0]
        self.cfg["camera_params"]["image_width"] = self.imgs[0].shape[1]
        self.cfg["camera_params"]["fx"] = self.intrinsics[0, 0].item()
        self.cfg["camera_params"]["fy"] = self.intrinsics[1, 1].item()
        self.cfg["camera_params"]["cx"] = self.intrinsics[0, 2].item()
        self.cfg["camera_params"]["cy"] = self.intrinsics[1, 2].item()
        self.cfg["camera_params"]["png_depth_scale"] = 1000.0
        print(self.cfg)
        with open(os.path.join(savedir, "dataconfig.yaml"), "w") as f:
            yaml.dump(self.cfg, f)

        for _i in trange(len(self.imgs)):
            fname = str(_i).zfill(5)
            savefile = os.path.join(savedir_rgb, f"{fname}.png")
            write_color(savefile, self.imgs[_i])
            savefile = os.path.join(savedir_depth, f"{fname}.png")
            write_depth(savefile, self.depths[_i])
            savefile = os.path.join(savedir_poses, f"{fname}.npy")
            write_pose(savefile, self.poses[_i])
        
        cv2.destroyallwindows()
        exit()

    def connect_to_device(self, dev_idx):
        print("Searching for devices")
        devs = Record3DStream.get_connected_devices()
        print("{} device(s) found".format(len(devs)))
        for dev in devs:
            print("\tID: {}\n\tUDID: {}\n".format(dev.product_id, dev.udid))

        if len(devs) <= dev_idx:
            raise RuntimeError(
                "Cannot connect to device #{}, try different index.".format(dev_idx)
            )

        dev = devs[dev_idx]
        self.session = Record3DStream()
        self.session.on_new_frame = self.on_new_frame
        self.session.on_stream_stopped = self.on_stream_stopped
        self.session.connect(dev)  # Initiate connection and start capturing

    def get_intrinsic_mat_from_coeffs(self, coeffs):
        return np.array(
            [[coeffs.fx, 0, coeffs.tx], [0, coeffs.fy, coeffs.ty], [0, 0, 1]]
        )

    def start_processing_stream(self):
        while True:
            # print(f"Line 118, True: {True}")
            self.event.wait()  # Wait for new frame to arrive

            # Copy the newly arrived RGBD frame
            depth = self.session.get_depth_frame()
            # print(f"Line 123, depth: {depth}")
            rgb = self.session.get_rgb_frame()
            intrinsic_mat = self.get_intrinsic_mat_from_coeffs(
                self.session.get_intrinsic_mat()
            )
            pose = (
                self.session.get_camera_pose()
            )  # Quaternion + world position (accessible via camera_pose.[qx|qy|qz|qw|tx|ty|tz])

            # print(intrinsic_mat)

            # You can now e.g. create point cloud by projecting the depth map using the intrinsic matrix.

            # Postprocess it
            if self.session.get_device_type() == self.DEVICE_TYPE__TRUEDEPTH:
                depth = cv2.flip(depth, 1)
                rgb = cv2.flip(rgb, 1)

            rgb = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)

            if self.intrinsics is None:
                self.intrinsics = intrinsic_mat
            self.imgs.append(rgb)
            self.depths.append(depth)
            self.poses.append(
                [pose.qx, pose.qy, pose.qz, pose.qw, pose.tx, pose.ty, pose.tz]
            )

            # Show the RGBD Stream
            cv2.imshow("RGB", rgb)
            cv2.imshow("Depth", depth)
            cv2.waitKey(1)

            self.event.clear()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--savedir_name", type=str, default="saved-record3d", help="Dir to save data to")
    parser.add_argument("--seq_name", type=str, default="debug", help="Unique name for captured trajectory")
    args = parser.parse_args()

    app = DemoApp(savedir_name=args.savedir_name, seq_name=args.seq_name)
    
    app.connect_to_device(dev_idx=0)
    app.start_processing_stream()

conceptgraph/dataset/save_record3d.py
concept-graphs/concept-graphs
'''
Acknowledgements: 
This code is adapted from the LLaVA repository by Krishna Murthy Jatavallabhula 
'''

import argparse
import os
import pickle as pkl
import sys
from dataclasses import dataclass
from io import BytesIO
from pathlib import Path
from typing import List, Optional, Tuple, Union

import cv2
import matplotlib
import matplotlib.patches as patches
import matplotlib.pyplot as plt
import numpy as np
import requests
import torch
import torch.nn as nn
import torch.nn.functional as F
from mpl_toolkits.axes_grid1 import make_axes_locatable
from PIL import Image
from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    AutoTokenizer,
    CLIPImageProcessor,
    CLIPVisionModel,
    LlamaConfig,
    LlamaForCausalLM,
    LlamaModel,
)
from transformers import logging as hf_logging
from transformers.modeling_outputs import (
    BaseModelOutputWithPast,
    CausalLMOutputWithPast,
)

try:
    LLAVA_PYTHON_PATH = os.environ["LLAVA_PYTHON_PATH"]
except KeyError:
    print("Please set the environment variable LLAVA_PYTHON_PATH to the path of the LLaVA repository")
    sys.exit(1)
    
sys.path.append(LLAVA_PYTHON_PATH)

from llava.conversation import SeparatorStyle, conv_templates
# from llava.model.utils import KeywordsStoppingCriteria
from llava.mm_utils import KeywordsStoppingCriteria
from llava.utils import disable_torch_init

torch.autograd.set_grad_enabled(False)

# Set logging verbosity for the transformers package to only log errors
hf_logging.set_verbosity_error()


DEFAULT_IMAGE_TOKEN = "<image>"
DEFAULT_IMAGE_PATCH_TOKEN = "<im_patch>"
DEFAULT_IM_START_TOKEN = "<im_start>"
DEFAULT_IM_END_TOKEN = "<im_end>"


class LlavaConfig(LlamaConfig):
    model_type = "llava"


class LlavaLlamaModelTweaked(LlamaModel):
    config_class = LlavaConfig

    def __init__(self, config: LlamaConfig):
        super().__init__(config)

        # Instantiate the multimodal vision tower
        if hasattr(config, "mm_vision_tower"):
            # HACK: for FSDP (fully-sharded data parallel); cast to a list
            self.vision_tower = [
                CLIPVisionModel.from_pretrained(config.mm_vision_tower)
            ]

        # Instantiate the multimodal linear projection (vision-language)
        if hasattr(config, "use_mm_proj"):
            self.mm_projector = nn.Linear(config.mm_hidden_size, config.hidden_size)

    def get_vision_tower(self):
        # Return the vision tower
        vision_tower = getattr(self, "vision_tower", None)
        if type(vision_tower) is list:
            vision_tower = vision_tower[0]
        return vision_tower

    def initialize_vision_modules(
        self,
        vision_tower,
        mm_vision_select_layer,
        pretrain_mm_mlp_adapter=None,
        fsdp=None,
    ):
        self.config.mm_vision_tower = vision_tower

        image_processor = CLIPImageProcessor.from_pretrained(vision_tower)

        if not hasattr(self, "vision_tower"):
            vision_tower = CLIPVisionModel.from_pretrained(vision_tower)
        else:
            vision_tower = self.vision_tower[0]
        # Vision tower is frozen
        vision_tower.requires_grad_(False)

        if fsdp is not None and len(fsdp) > 0:
            self.vision_tower = [vision_tower]
        else:
            self.vision_tower = vision_tower

        vision_config = vision_tower.config
        num_patches = (vision_config.image_size // vision_config.patch_size) ** 2

        # If using a vision tower, we need a multimodal projection (from vision to text)
        self.config.use_mm_proj = True
        self.config.mm_hidden_size = vision_config.hidden_size
        self.config.mm_vision_select_layer = mm_vision_select_layer

        if not hasattr(self, "mm_projector"):
            self.mm_projector = nn.Linear(
                vision_config.hidden_size, self.config.hidden_size
            )

        if pretrain_mm_mlp_adapter is not None:
            mm_projector_weights = torch.load(
                pretrain_mm_mlp_adapter, map_location="cpu"
            )
            self.mm_projector.load_state_dict(
                {k.split(".")[-1]: v for k, v in mm_projector_weights.items()}
            )

        return dict(
            image_processor=image_processor,
            image_token_len=num_patches,
            vision_config=vision_config,
        )

    def encode_image(
        self,
        images: Union[List, torch.FloatTensor],
    ) -> Union[list, torch.FloatTensor]:
        with torch.no_grad():
            vision_tower = self.get_vision_tower()
            if vision_tower is not None and images is not None:
                if type(images) is list:
                    image_features = []
                    for image in images:
                        image_forward_out = vision_tower(
                            image.unsqueeze(0), output_hidden_states=True
                        )
                        select_hidden_state_layer = getattr(
                            self.config, "mm_vision_select_layer", -1
                        )
                        # TODO (Krishna): Figure out why the zero-th dim is dropped (perhaps that is the
                        # global token?)
                        image_feature = select_hidden_state[:, 1:]
                        image_features.append(image_feature)
                    # image_features = [self.mm_projector(image_feature)[0] for image_feature in image_features]
                else:
                    image_forward_outs = vision_tower(
                        images.to(vision_tower.dtype), output_hidden_states=True
                    )
                    select_hidden_state_layer = getattr(
                        self.config, "mm_vision_select_layer", -1
                    )
                    select_hidden_state = image_forward_outs.hidden_states[
                        select_hidden_state_layer
                    ]
                    image_features = select_hidden_state[:, 1:].to(images.dtype)
                    # image_features = self.mm_projector(image_features)
            return image_features

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        images: Optional[Union[list, torch.FloatTensor]] = None,
        image_features: Optional[Union[list, torch.FloatTensor]] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutputWithPast]:
        # HACK: replace back original embeddings for LLaVA pretraining
        orig_embeds_params = getattr(self, "orig_embeds_params", None)
        # if orig_embeds_params is not None:
        #     orig_embeds_params = orig_embeds_params[0]
        #     with torch.no_grad():
        #         self.get_input_embeddings().weight.data[:-2] = orig_embeds_params[:-2].data

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        vision_tower = self.get_vision_tower()
        if vision_tower is not None and (input_ids.shape[1] != 1 or self.training):
            # If images are provided, but image features is not provided, compute the image_features
            # from the images passed in. (else, skip this step)
            if image_features is None and images is not None:
                # TODO: this is a modified multimodal LLM -- Haotian Liu
                with torch.no_grad():
                    if type(images) is list:
                        # variable length images
                        image_features = []
                        for image in images:
                            image_forward_out = vision_tower(
                                image.unsqueeze(0), output_hidden_states=True
                            )
                            select_hidden_state_layer = getattr(
                                self.config, "mm_vision_select_layer", -1
                            )
                            select_hidden_state = image_forward_out.hidden_states[
                                select_hidden_state_layer
                            ]
                            image_feature = select_hidden_state[:, 1:]
                            image_features.append(image_feature)
                    else:
                        image_forward_outs = vision_tower(
                            images.to(vision_tower.dtype), output_hidden_states=True
                        )
                        select_hidden_state_layer = getattr(
                            self.config, "mm_vision_select_layer", -1
                        )
                        select_hidden_state = image_forward_outs.hidden_states[
                            select_hidden_state_layer
                        ]
                        image_features = select_hidden_state[:, 1:].to(images.dtype)
                if type(images) is list:
                    image_features = [
                        self.mm_projector(image_feature)[0]
                        for image_feature in image_features
                    ]
                else:
                    image_features = self.mm_projector(image_features)

            if type(image_features) is list:
                image_features = [
                    self.mm_projector(image_feature)[0]
                    for image_feature in image_features
                ]
            elif image_features is not None:  # image_features is torch.tensor
                image_features = self.mm_projector(image_features)

            # TODO (Krishna): The following line from the original LLaVA impl seems like a hardcoded param
            # size. Replace this by reading in these dims from the config file in future.
            dummy_image_features = torch.zeros(
                256, 1024, device=inputs_embeds.device, dtype=inputs_embeds.dtype
            )
            dummy_image_features = self.mm_projector(dummy_image_features)
            # (Krishna): If no image_features are passed as input, feed the dummy image features through
            if image_features is None:
                image_features = dummy_image_features

            new_input_embeds = []
            cur_image_idx = 0
            for cur_input_ids, cur_input_embeds in zip(input_ids, inputs_embeds):
                if (cur_input_ids == vision_tower.config.im_patch_token).sum() == 0:
                    # multimodal LLM, but the current sample is not multimodal
                    cur_input_embeds = (
                        cur_input_embeds + (0.0 * dummy_image_features).sum()
                    )
                    new_input_embeds.append(cur_input_embeds)
                    cur_image_idx += 1
                    continue
                if vision_tower.config.use_im_start_end:
                    cur_image_features = image_features[cur_image_idx]
                    num_patches = cur_image_features.shape[0]
                    if (cur_input_ids == vision_tower.config.im_start_token).sum() != (
                        cur_input_ids == vision_tower.config.im_end_token
                    ).sum():
                        raise ValueError(
                            "The number of image start tokens and image end tokens should be the same."
                        )
                    image_start_tokens = torch.where(
                        cur_input_ids == vision_tower.config.im_start_token
                    )[0]
                    for image_start_token_pos in image_start_tokens:
                        cur_image_features = image_features[cur_image_idx].to(
                            device=cur_input_embeds.device
                        )
                        num_patches = cur_image_features.shape[0]
                        if (
                            cur_input_ids[image_start_token_pos + num_patches + 1]
                            != vision_tower.config.im_end_token
                        ):
                            raise ValueError(
                                "The image end token should follow the image start token."
                            )
                        if orig_embeds_params is not None:
                            cur_new_input_embeds = torch.cat(
                                (
                                    cur_input_embeds[:image_start_token_pos].detach(),
                                    cur_input_embeds[
                                        image_start_token_pos : image_start_token_pos
                                        + 1
                                    ],
                                    cur_image_features,
                                    cur_input_embeds[
                                        image_start_token_pos
                                        + num_patches
                                        + 1 : image_start_token_pos
                                        + num_patches
                                        + 2
                                    ],
                                    cur_input_embeds[
                                        image_start_token_pos + num_patches + 2 :
                                    ].detach(),
                                ),
                                dim=0,
                            )
                        else:
                            cur_new_input_embeds = torch.cat(
                                (
                                    cur_input_embeds[: image_start_token_pos + 1],
                                    cur_image_features,
                                    cur_input_embeds[
                                        image_start_token_pos + num_patches + 1 :
                                    ],
                                ),
                                dim=0,
                            )
                        cur_image_idx += 1
                    new_input_embeds.append(cur_new_input_embeds)
                else:
                    cur_image_features = image_features[cur_image_idx]
                    num_patches = cur_image_features.shape[0]
                    if (
                        cur_input_ids == vision_tower.config.im_patch_token
                    ).sum() != num_patches:
                        raise ValueError(
                            "The number of image patch tokens should be the same as the number of image patches."
                        )
                    masked_indices = torch.where(
                        cur_input_ids == vision_tower.config.im_patch_token
                    )[0]
                    mask_index_start = masked_indices[0]
                    if (
                        masked_indices
                        != torch.arange(
                            mask_index_start,
                            mask_index_start + num_patches,
                            device=masked_indices.device,
                            dtype=masked_indices.dtype,
                        )
                    ).any():
                        raise ValueError(
                            "The image patch tokens should be consecutive."
                        )
                    if orig_embeds_params is not None:
                        cur_new_input_embeds = torch.cat(
                            (
                                cur_input_embeds[:mask_index_start].detach(),
                                cur_image_features,
                                cur_input_embeds[
                                    mask_index_start + num_patches :
                                ].detach(),
                            ),
                            dim=0,
                        )
                    else:
                        cur_new_input_embeds = torch.cat(
                            (
                                cur_input_embeds[:mask_index_start],
                                cur_image_features,
                                cur_input_embeds[mask_index_start + num_patches :],
                            ),
                            dim=0,
                        )
                    new_input_embeds.append(cur_new_input_embeds)
                    cur_image_idx += 1
            inputs_embeds = torch.stack(new_input_embeds, dim=0)

        return super().forward(
            input_ids=None,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )


class LlavaLlamaForCausalLMTweaked(LlamaForCausalLM):
    config_class = LlavaConfig

    def __init__(self, config):
        super().__init__(config)
        self.model = LlavaLlamaModelTweaked(config)

        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        # Initialize weights and apply final processing
        self.post_init()

    def get_model(self):
        return self.model

    def get_vision_tower(self):
        return self.get_model().get_vision_tower()

    def get_vision_tower(self):
        model = self.get_model()
        vision_tower = model.vision_tower
        if type(vision_tower) is list:
            vision_tower = vision_tower[0]
        return vision_tower

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        images: Optional[Union[list, torch.FloatTensor]] = None,
        image_features: Optional[Union[list, torch.FloatTensor]] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        output_attentions = (
            output_attentions
            if output_attentions is not None
            else self.config.output_attentions
        )
        output_hidden_states = (
            output_hidden_states
            if output_hidden_states is not None
            else self.config.output_hidden_states
        )
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            images=images,
            image_features=image_features,
        )

        hidden_states = outputs[0]
        logits = self.lm_head(hidden_states)

        loss = None
        if labels is not None:
            # Shift so that tokens < n predict n
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            # Flatten the tokens
            loss_fct = torch.nn.CrossEntropyLoss()
            shift_logits = shift_logits.view(-1, self.config.vocab_size)
            shift_labels = shift_labels.view(-1)
            # Enable model/pipeline parallelism
            shift_labels = shift_labels.to(shift_logits.device)
            loss = loss_fct(shift_logits, shift_labels)

        if not return_dict:
            output = (logits,) + outputs[1:]
            return (loss,) + output if loss is not None else output

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

    def prepare_inputs_for_generation(
        self,
        input_ids,
        past_key_values=None,
        attention_mask=None,
        inputs_embeds=None,
        **kwargs,
    ):
        if past_key_values:
            input_ids = input_ids[:, -1:]

        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
        if inputs_embeds is not None and past_key_values is None:
            model_inputs = {"inputs_embeds": inputs_embeds}
        else:
            model_inputs = {"input_ids": input_ids}

        model_inputs.update(
            {
                "past_key_values": past_key_values,
                "use_cache": kwargs.get("use_cache"),
                "attention_mask": attention_mask,
                "images": kwargs.get("images", None),
                "image_features": kwargs.get("image_features", None),
            }
        )
        return model_inputs

    def initialize_vision_tokenizer(
        self,
        mm_use_im_start_end,
        tokenizer,
        device,
        tune_mm_mlp_adapter=False,
        pretrain_mm_mlp_adapter=None,
    ):
        vision_config = self.get_vision_tower().config
        vision_config.use_im_start_end = mm_use_im_start_end
        tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
        self.resize_token_embeddings(len(tokenizer))

        if mm_use_im_start_end:
            num_new_tokens = tokenizer.add_tokens(
                [DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True
            )
            self.resize_token_embeddings(len(tokenizer))
            (
                vision_config.im_start_token,
                vision_config.im_end_token,
            ) = tokenizer.convert_tokens_to_ids(
                [DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN]
            )

            if num_new_tokens > 0:
                input_embeddings = self.get_input_embeddings().weight.data
                output_embeddings = self.get_output_embeddings().weight.data

                input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(
                    dim=0, keepdim=True
                )
                output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(
                    dim=0, keepdim=True
                )

                input_embeddings[-num_new_tokens:] = input_embeddings_avg
                output_embeddings[-num_new_tokens:] = output_embeddings_avg

            if tune_mm_mlp_adapter:
                self.get_model().orig_embeds_params = [
                    self.get_input_embeddings().weight.data.clone().to(device=device)
                ]
                for p in self.get_input_embeddings().parameters():
                    p.requires_grad = True
                for p in self.get_output_embeddings().parameters():
                    p.requires_grad = False

            if pretrain_mm_mlp_adapter:
                mm_projector_weights = torch.load(
                    pretrain_mm_mlp_adapter, map_location="cpu"
                )
                embed_tokens_weight = mm_projector_weights["model.embed_tokens.weight"]
                assert num_new_tokens == 2
                if input_embeddings.shape == embed_tokens_weight.shape:
                    input_embeddings[-num_new_tokens:] = embed_tokens_weight[
                        -num_new_tokens:
                    ]
                elif embed_tokens_weight.shape[0] == num_new_tokens:
                    input_embeddings[-num_new_tokens:] = embed_tokens_weight
                else:
                    raise ValueError(
                        f"Unexpected embed_tokens_weight shape. Pretrained: {embed_tokens_weight.shape}. "
                        f"Current: {input_embeddings.shape}. Numer of new tokens: {num_new_tokens}."
                    )

        vision_config.im_patch_token = tokenizer.convert_tokens_to_ids(
            [DEFAULT_IMAGE_PATCH_TOKEN]
        )[0]


class LLaVaChat(object):
    def __init__(self, model_path, conv_mode="multimodal", num_gpus=1):
        self.model_path = model_path
        self.conv_mode = conv_mode
        self.num_gpus = num_gpus

        # Handle multi-gpu config
        if self.num_gpus == 1:
            kwargs = {}
        else:
            kwargs = {
                "device_map": "auto",
                "max_memory": {i: "13GiB" for i in range(self.num_gpus)},
            }

        # pytorch spends a substantial amount of time initializing default weights for
        # each linear layer and layernorm layer in the model. Since we will load weights
        # from disk anyways, disable this redundant default init.
        # This accelerates model creation.
        disable_torch_init()

        # Initialize the tokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
        # Initialize the model
        self.model = LlavaLlamaForCausalLMTweaked.from_pretrained(
            self.model_path, torch_dtype=torch.float16, low_cpu_mem_usage=True, **kwargs
        )
        self.model.cuda()

        # Image preprocessor
        self.image_processor = CLIPImageProcessor.from_pretrained(
            self.model.config.mm_vision_tower, torch_dtype=torch.float16
        )

        self.mm_use_im_start_end = getattr(
            self.model.config, "mm_use_im_start_end", False
        )
        self.tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
        if self.mm_use_im_start_end:
            self.tokenizer.add_tokens(
                [DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True
            )

        self.vision_tower = self.model.get_model().vision_tower[0]
        if self.vision_tower.device.type == "meta":
            self.vision_tower = CLIPVisionModel.from_pretrained(
                self.vision_tower.config._name_or_path,
                torch_dtype=torch.float16,
                low_cpu_mem_usage=True,
            ).cuda()
            self.model.get_model().vision_tower[0] = self.vision_tower
        else:
            self.vision_tower.to(device="cuda", dtype=torch.float16)
        self.vision_config = self.vision_tower.config
        self.vision_config.im_patch_token = self.tokenizer.convert_tokens_to_ids(
            [DEFAULT_IMAGE_PATCH_TOKEN]
        )[0]
        self.vision_config.use_im_start_end = self.mm_use_im_start_end
        if self.mm_use_im_start_end:
            (
                self.vision_config.im_start_token,
                self.vision_config.im_end_token,
            ) = self.tokenizer.convert_tokens_to_ids(
                [DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN]
            )
        self.image_token_len = (
            self.vision_config.image_size // self.vision_config.patch_size
        ) ** 2

        # # Initialize a conversation from template (default conv_mode is "multimodal")
        # # (conv_mode determines the conversation template to use from llava.conversation module)
        # self.conv = conv_templates[self.conv_mode].copy()

        # # Cache for image features
        # self.image_features = None
        
        self.reset()
        
    def reset(self):
        # Initialize a conversation from template (default conv_mode is "multimodal")
        # (conv_mode determines the conversation template to use from llava.conversation module)
        self.conv = conv_templates[self.conv_mode].copy()
        
        # Cache for image features
        self.image_features = None

    def __call__(self, query, image_features=None):
        # Given this query, and the image_featurese, prompt LLaVA with the query,
        # using the image_features as context.

        qs = query
        if image_features is not None:
            if self.mm_use_im_start_end:
                qs = (
                    qs
                    + "\n"
                    + DEFAULT_IM_START_TOKEN
                    + DEFAULT_IMAGE_PATCH_TOKEN * self.image_token_len
                    + DEFAULT_IM_END_TOKEN
                )
            else:
                qs = qs + "\n" + DEFAULT_IMAGE_PATCH_TOKEN * self.image_token_len
            if self.image_features is None:
                self.image_features = image_features
        else:
            qs = qs + "\n"

        self.conv.append_message(self.conv.roles[0], qs)
        self.conv.append_message(self.conv.roles[1], None)

        input_ids = None

        # Get the prompt
        prompt = self.conv.get_prompt()
        # Tokenize this prompt
        inputs = self.tokenizer([prompt])
        # Cast to torch tensor and to GPU
        input_ids = torch.as_tensor(inputs.input_ids).cuda()

        stop_str = (
            self.conv.sep
            if self.conv.sep_style != SeparatorStyle.TWO
            else self.conv.sep2
        )
        keywords = [stop_str]
        stopping_criteria = KeywordsStoppingCriteria(
            keywords, self.tokenizer, input_ids
        )

        with torch.inference_mode():
            output_ids = self.model.generate(
                input_ids,
                image_features=self.image_features,
                do_sample=True,
                temperature=0.2,
                max_new_tokens=1024,
                stopping_criteria=[stopping_criteria],
            )

        input_token_len = input_ids.shape[1]
        n_diff_input_output = (
            (input_ids != output_ids[:, :input_token_len]).sum().item()
        )
        if n_diff_input_output > 0:
            print(
                f"[Warning] {n_diff_input_output} output_ids are not the same as the input_ids"
            )
        outputs = self.tokenizer.batch_decode(
            output_ids[:, input_token_len:], skip_special_tokens=True
        )[0]
        self.conv.append_message(self.conv.roles[1], outputs)
        outputs = outputs.strip()
        if outputs.endswith(stop_str):
            outputs = outputs[: -len(stop_str)]
        outputs = outputs.strip()
        return outputs

    def load_image(self, image_file):
        if image_file.startswith("http") or image_file.startswith("https"):
            response = requests.get(image_file)
            image = Image.open(BytesIO(response.content)).convert("RGB")
        else:
            image = Image.open(image_file).convert("RGB")
        return image

    def encode_image(self, image_tensor_half_cuda):
        return self.model.model.encode_image(image_tensor_half_cuda)
    



if __name__ == "__main__":
    # from types import SimpleNamespace

    # args = SimpleNamespace()
    # args.model_path = "/home/krishna/code/LLaVA/LLaVA-7b-v0"
    # args.image_file = "/home/krishna/Downloads/figurines/images/frame_00150.jpg"
    # args.conv_mode = "multimodal"
    # args.num_gpus = 1
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_path", type=str, default=None)
    parser.add_argument("--image_file", type=str, required=True)
    parser.add_argument("--conv_mode", type=str, default="multimodal")
    parser.add_argument("--num_gpus", type=int, default=1)
    args = parser.parse_args()
    
    if args.model_path is None:
        try: 
            args.model_path = os.environ["LLAVA_CKPT_PATH"]
        except KeyError:
            print("Please provide a model path or set the environment variable LLAVA_CKPT_PATH")
            exit(1)

    chat = LLaVaChat(args.model_path, args.conv_mode, args.num_gpus)
    print("LLaVA chat initialized...")

    image = chat.load_image(args.image_file)
    image_tensor = chat.image_processor.preprocess(image, return_tensors="pt")[
        "pixel_values"
    ][0]
    image_features = chat.encode_image(image_tensor[None, ...].half().cuda())
    # query = input("Enter a query ('q' to quit): ")
    query = "List the set of objects in this image."
    outputs = chat(query=query, image_features=image_features)
    print(outputs)

    query = "List potential uses for each."
    outputs = chat(query=query, image_features=None)
    print(outputs)

    # while True:
    #     query = input("Enter a query ('q' to quit): ")
    #     # query = "Describe this image"
    #     # print(query)
    #     outputs = chat(query=query, image_features=None)
    #     print(outputs)

conceptgraph/llava/llava_model.py
concept-graphs/concept-graphs
class GPTPrompt:
    def __init__(self):
        
        self.old_system_prompt = """
        The input is a list of JSONs describing multiple predictions of a single object. Each JSON has four fields: 
        1. id: a unique identifier for the object
        2. bbox_extent: the 3D bounding box extents of the object 
        3. bbox_center: the 3D bounding box center of the object 
        4. caption: a caption predicted by an image captioning model referencing that object. This caption should be brief and in sparse prose. For example, the caption "the object described appears to be described as a electric bicycle, it is sitting alongside a red suitcase which is nearby" should be shortened to "electric bike, near red suitcase".
        There may be upto 10 such bounding boxes and captions in each input.
        The captions may not always be accurate or consistent (often, predictions may just be wrong). 
        The only valid objects are ones that we commonly find in indoor scenes. Any predictions that reference 
        people, animals, or objects that are impossible to find indoors, must be tagged as "invalid".

        Output a brief, informative language tag for each object being referenced to. If the captions are grossly
        inconsistent, output "invalid" for that object.

        The output must be a single JSON containing just the following fields. "summary" indicating a brief 
        summary of your understanding of the object being referenced to. "possible_tags" indicating a list of 
        possible tags that you think the object could be. "object_tag" indicating the final tag that you think 
        this object should be, considering everything else in the scene (particularly, nearby other objects). 
        Before suggesting the final tag, consider the actual size of the object (you have this in the "bbox") 
        and identify the best possible tag this object could be. Verify that the output is in valid JSON format with
        the fields "summary", "possible_tags", and "object_tag". The object_tag must be supported by the captions.
        It is very important that the output must be valid JSON, and nothing else.
        """
        
        self.new_but_long_system_prompt = """You are a helpful assistant that helps identify and describe objects in a scene. Your input is a in JSON format, and you should always reply in JSON format. Your input will contain the field "caption" which is a list of captions of an image attempting to identify the objects in the image. Your response should contain the fields "summary" containing a concise summary of the object(s) in the image. If an object is mentioned more than once, that prediction is likely accurate. If many objects are mentioned more than once, and a container or surface is mentioned, it is likely that the image was of those objects (mentioned more than once) on that container or surface. If no object is mentioned more than once, and each caption is unrelated to the rest, it could be a blank or too small image or blurry image, the captions are likely incorrect, so just say "conflicting captions about [objects] in the summary field and put "invalid" in the "object_tag" field. The field "possible_tags" should contain a list of possible tags that you think the object(s) could be. The field "object_tag" should contain the final tag that you think this object should be, considering all the information given. These are based on scans of indoor scenes so most objects will be those found in indoor spaces. """
        
        self.system_prompt = """Identify and describe objects in scenes. Input and output must be in JSON format. The input field 'captions' contains a list of image captions aiming to identify objects. Output 'summary' as a concise description of the identified object(s). An object mentioned multiple times is likely accurate. If various objects are repeated and a container/surface is noted such as a shelf or table, assume the (repeated) objects are on that container/surface. For unrelated, non-repeating (or empty) captions, summarize as 'conflicting (or empty) captions about [objects]' and set 'object_tag' to 'invalid'. Output 'possible_tags' listing potential object categories. Set 'object_tag' as the conclusive identification. Focus on indoor object types, as the input captions are from indoor scans."""
        
        self.example_1 = """
        "id": 1,
        "captions": [
        "a jacket hanging on a wall, either on a hook or a rack.",
        "a jacket, which is hanging on a wall or a rack.",
        "a jacket, which is either being worn or draped over a person's shoulders.",
        "a sweater, which is hanging on a clothes hanger.",
        "a hooded jacket, which is either hanging on a hook or draped over a shower rail.",
        "a mannequin, which is wearing a yellow shirt and a red jacket.",
        "a jacket, which is hanging on a hook or a rack.",
        "a hooded sweatshirt, which is either being held by someone or hanging on a hook or a rack.",
        "a hanger with a yellow and black jacket hanging on it.",
        "a yellow and black striped umbrella.",
        "a hanging coat, which is either yellow or red depending on the specific description.",
        "a cat.",
        "a person wearing a yellow and red jacket."
        ]
        }"
        """
        
        self.response_1 = """{
        "summary": "a jacket hanging on a rack",
        "possible_tags": ["jacket", "sweater", "hooded jacket", "hooded sweatshirt", "coat", "hanging clothing", "hanger", "wall", "rack"],
        "object_tag": "hanging jacket"
        }"""
        
        self.example_2 = """{
        "id": 12,
        "captions": [
        "a bookshelf filled with books",
        "a bicycle helmet"
        ]
        }"""
        
        self.response_2 = """{
        "summary": "conflicting captions of a bookshelf and a bicycle helmet",
        "possible_tags": ["bookshelf", "bicycle helmet", "helmet", "bookcase", "book", "books", "shelf", "shelves"],
        "object_tag": "invalid"
        }"""
        
        self.example_3 = """{
        "id": 304,
        "captions": [
        "a pair of scissors.",
        "a sewing machine.",
        "a white shelf or rack, which is filled with various boxes and files.",
        "a white shelf or bookshelf that is filled with various items."
        ]
        }"""
        
        self.response_3 = """{
        "summary": "a white shelf or rack filled with various items",
        "possible_tags": ["shelf", "rack", "bookshelf", "bookcase", "box", "file", "scissors", "sewing machine"],
        "object_tag": "shelf"
        }"""
        
        self.example_4 = """{
        "id": 433,
        "captions": [
        "a white toilet.",
        "a pile of various exercise equipment, including a set of tennis balls.",
        "a white table.",
        "a barbell.",
        "a microwave.",
        "a tall, thin black bottle.",
        "a barbell, which is situated on a table.",
        "a pile of various exercise equipment, including a bench, situated in a room with desks and a whiteboard.",
        "a tall vase or pitcher.",
        "a tennis ball.",
        "a shirt with a space design on it.",
        "a tennis ball.",
        "a gray shirt with the NASA logo on it."
        ]
        }"""
        
        self.response_4 = """{
        "summary": "a white table containing a barbell and possibly with some other exercise equipment",
        "possible_tags": ["exercise equipment", "tennis ball", "barbell", "bench", "table", "microwave", "bottle", "vase", "pitcher", "shirt"],
        "object_tag": "white table"
        }"""
        
        self.example_5 = """{
        "id": 231,
        "captions": [
            "a teddy bear.",
            "a doorknob.",
            "a television set.",
            "a laptop computer."
        ]
        }"""
        
        self.response_5 = """{
        "summary": "conflicting captions of a teddy bear, a doorknob, a television set, and a laptop computer",
        "possible_tags": ["teddy bear", "doorknob", "television set", "laptop computer", "computer"],
        "object_tag": "invalid"
        }"""


    def get_json(self):
        prompt_json = [
            {
                "role": "system",
                "content": self.system_prompt
            },
            {
                "role": "user",
                "content": self.example_1
            },
            {
                "role": "assistant",
                "content": self.response_1
            },
            {
                "role": "user",
                "content": self.example_2
            },
            {
                "role": "assistant",
                "content": self.response_2
            },
            {
                "role": "user",
                "content": self.example_3
            },
            {
                "role": "assistant",
                "content": self.response_3
            },
            {
                "role": "user",
                "content": self.example_4
            },
            {
                "role": "assistant",
                "content": self.response_4
            },
            {
                "role": "user",
                "content": self.example_5
            },
            {
                "role": "assistant",
                "content": self.response_5
            }
        ]
        return prompt_json

# Usage example
if __name__ == "__main__":
    prompt_obj = GPTPrompt()
    json_data = prompt_obj.get_json()
    print(json_data)
conceptgraph/scenegraph/GPTPrompt.py
concept-graphs/concept-graphs
"""
Build a scene graph from the segment-based map and captions from LLaVA.
"""

import gc
import gzip
import json
import os
import pickle as pkl
import time
from dataclasses import dataclass
from pathlib import Path
from types import SimpleNamespace
from typing import List, Literal, Union
from textwrap import wrap

import cv2
import matplotlib.pyplot as plt

import numpy as np
import rich
import torch
import tyro
from PIL import Image
from scipy.sparse import csr_matrix
from scipy.sparse.csgraph import connected_components, minimum_spanning_tree
from tqdm import tqdm, trange
from transformers import logging as hf_logging

# from mappingutils import (
#     MapObjectList,
#     compute_3d_giou_accuracte_batch,
#     compute_3d_iou_accuracte_batch,
#     compute_iou_batch,
#     compute_overlap_matrix_faiss,
#     num_points_closer_than_threshold_batch,
# )

torch.autograd.set_grad_enabled(False)
hf_logging.set_verbosity_error()

# Import OpenAI API
import openai

openai.api_key = os.getenv("OPENAI_API_KEY")


@dataclass
class ProgramArgs:
    mode: Literal[
        "extract-node-captions",
        "refine-node-captions",
        "build-scenegraph",
        "generate-scenegraph-json",
        "annotate-scenegraph",
    ]

    # Path to cache directory
    cachedir: str = "saved/room0"
    
    prompts_path: str = "prompts/gpt_prompts.json"

    # Path to map file
    mapfile: str = "saved/room0/map/scene_map_cfslam.pkl.gz"

    # Path to file storing segment class names
    class_names_file: str = "saved/room0/gsa_classes_ram.json"

    # Device to use
    device: str = "cuda:0"

    # Voxel size for downsampling
    downsample_voxel_size: float = 0.025

    # Maximum number of detections to consider, per object
    max_detections_per_object: int = 10

    # Suppress objects with less than this number of observations
    min_views_per_object: int = 2

    # List of objects to annotate (default: all objects)
    annot_inds: Union[List[int], None] = None

    # Masking option
    masking_option: Literal["blackout", "red_outline", "none"] = "none"

def load_scene_map(args, scene_map):
    """
    Loads a scene map from a gzip-compressed pickle file. This is a function because depending whether the mapfile was made using cfslam_pipeline_batch.py or merge_duplicate_objects.py, the file format is different (see below). So this function handles that case.
    
    The function checks the structure of the deserialized object to determine
    the correct way to load it into the `scene_map` object. There are two
    expected formats:
    1. A dictionary containing an "objects" key.
    2. A list or a dictionary (replace with your expected type).
    """
    
    with gzip.open(Path(args.mapfile), "rb") as f:
        loaded_data = pkl.load(f)
        
        # Check the type of the loaded data to decide how to proceed
        if isinstance(loaded_data, dict) and "objects" in loaded_data:
            scene_map.load_serializable(loaded_data["objects"])
        elif isinstance(loaded_data, list) or isinstance(loaded_data, dict):  # Replace with your expected type
            scene_map.load_serializable(loaded_data)
        else:
            raise ValueError("Unexpected data format in map file.")
        print(f"Loaded {len(scene_map)} objects")


def prjson(input_json, indent=0):
    """ Pretty print a json object """
    if not isinstance(input_json, list):
        input_json = [input_json]
        
    print("[")
    for i, entry in enumerate(input_json):
        print("  {")
        for j, (key, value) in enumerate(entry.items()):
            terminator = "," if j < len(entry) - 1 else ""
            if isinstance(value, str):
                formatted_value = value.replace("\\n", "\n").replace("\\t", "\t")
                print('    "{}": "{}"{}'.format(key, formatted_value, terminator))
            else:
                print(f'    "{key}": {value}{terminator}')
        print("  }" + ("," if i < len(input_json) - 1 else ""))
    print("]")

def crop_image_pil(image: Image, x1: int, y1: int, x2: int, y2: int, padding: int = 0) -> Image:
    """
    Crop the image with some padding

    Args:
        image: PIL image
        x1, y1, x2, y2: bounding box coordinates
        padding: padding around the bounding box

    Returns:
        image_crop: PIL image

    Implementation from the CFSLAM repo
    """
    image_width, image_height = image.size
    x1 = max(0, x1 - padding)
    y1 = max(0, y1 - padding)
    x2 = min(image_width, x2 + padding)
    y2 = min(image_height, y2 + padding)

    image_crop = image.crop((x1, y1, x2, y2))
    return image_crop


def draw_red_outline(image, mask):
    """ Draw a red outline around the object i nan image"""
    # Convert PIL Image to numpy array
    image_np = np.array(image)

    red_outline = [255, 0, 0]

    # Find contours in the binary mask
    contours, _ = cv2.findContours(mask.astype(np.uint8) * 255, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # Draw red outlines around the object. The last argument "3" indicates the thickness of the outline.
    cv2.drawContours(image_np, contours, -1, red_outline, 3)

    # Optionally, add padding around the object by dilating the drawn contours
    kernel = np.ones((5, 5), np.uint8)
    image_np = cv2.dilate(image_np, kernel, iterations=1)
    
    image_pil = Image.fromarray(image_np)

    return image_pil


def crop_image_and_mask(image: Image, mask: np.ndarray, x1: int, y1: int, x2: int, y2: int, padding: int = 0):
    """ Crop the image and mask with some padding. I made a single function that crops both the image and the mask at the same time because I was getting shape mismatches when I cropped them separately.This way I can check that they are the same shape."""
    
    image = np.array(image)
    # Verify initial dimensions
    if image.shape[:2] != mask.shape:
        print("Initial shape mismatch: Image shape {} != Mask shape {}".format(image.shape, mask.shape))
        return None, None

    # Define the cropping coordinates
    x1 = max(0, x1 - padding)
    y1 = max(0, y1 - padding)
    x2 = min(image.shape[1], x2 + padding)
    y2 = min(image.shape[0], y2 + padding)
    # round the coordinates to integers
    x1, y1, x2, y2 = round(x1), round(y1), round(x2), round(y2)

    # Crop the image and the mask
    image_crop = image[y1:y2, x1:x2]
    mask_crop = mask[y1:y2, x1:x2]

    # Verify cropped dimensions
    if image_crop.shape[:2] != mask_crop.shape:
        print("Cropped shape mismatch: Image crop shape {} != Mask crop shape {}".format(image_crop.shape, mask_crop.shape))
        return None, None
    
    # convert the image back to a pil image
    image_crop = Image.fromarray(image_crop)

    return image_crop, mask_crop

def blackout_nonmasked_area(image_pil, mask):
    """ Blackout the non-masked area of an image"""
    # convert image to numpy array
    image_np = np.array(image_pil)
    # Create an all-black image of the same shape as the input image
    black_image = np.zeros_like(image_np)
    # Wherever the mask is True, replace the black image pixel with the original image pixel
    black_image[mask] = image_np[mask]
    # convert back to pil image
    black_image = Image.fromarray(black_image)
    return black_image

def plot_images_with_captions(images, captions, confidences, low_confidences, masks, savedir, idx_obj):
    """ This is debug helper function that plots the images with the captions and masks overlaid and saves them to a directory. This way you can inspect exactly what the LLaVA model is captioning which image with the mask, and the mask confidence scores overlaid."""
    
    n = min(9, len(images))  # Only plot up to 9 images
    nrows = int(np.ceil(n / 3))
    ncols = 3 if n > 1 else 1
    fig, axarr = plt.subplots(nrows, ncols, figsize=(10, 5 * nrows), squeeze=False)  # Adjusted figsize

    for i in range(n):
        row, col = divmod(i, 3)
        ax = axarr[row][col]
        ax.imshow(images[i])

        # Apply the mask to the image
        img_array = np.array(images[i])
        if img_array.shape[:2] != masks[i].shape:
            ax.text(0.5, 0.5, "Plotting error: Shape mismatch between image and mask", ha='center', va='center')
        else:
            green_mask = np.zeros((*masks[i].shape, 3), dtype=np.uint8)
            green_mask[masks[i]] = [0, 255, 0]  # Green color where mask is True
            ax.imshow(green_mask, alpha=0.15)  # Overlay with transparency

        title_text = f"Caption: {captions[i]}\nConfidence: {confidences[i]:.2f}"
        if low_confidences[i]:
            title_text += "\nLow Confidence"
        
        # Wrap the caption text
        wrapped_title = '\n'.join(wrap(title_text, 30))
        
        ax.set_title(wrapped_title, fontsize=12)  # Reduced font size for better fitting
        ax.axis('off')

    # Remove any unused subplots
    for i in range(n, nrows * ncols):
        row, col = divmod(i, 3)
        axarr[row][col].axis('off')
    
    plt.tight_layout()
    plt.savefig(savedir / f"{idx_obj}.png")
    plt.close()



def extract_node_captions(args):
    from conceptgraph.llava.llava_model import LLaVaChat

    # NOTE: args.mapfile is in cfslam format
    from conceptgraph.slam.slam_classes import MapObjectList

    # Load the scene map
    scene_map = MapObjectList()
    load_scene_map(args, scene_map)
    
    # Scene map is in CFSLAM format
    # keys: 'image_idx', 'mask_idx', 'color_path', 'class_id', 'num_detections',
    # 'mask', 'xyxy', 'conf', 'n_points', 'pixel_area', 'contain_number', 'clip_ft',
    # 'text_ft', 'pcd_np', 'bbox_np', 'pcd_color_np'

    # Imports to help with feature extraction
    # from extract_mask_level_features import (
    #     crop_bbox_from_img,
    #     get_model_and_preprocessor,
    #     preprocess_and_encode_pil_image,
    # )

    # Load class names from the json file
    class_names = None
    with open(Path(args.class_names_file), "r") as f:
        class_names = json.load(f)
    print(class_names)

    # Creating a namespace object to pass args to the LLaVA chat object
    chat_args = SimpleNamespace()
    chat_args.model_path = os.getenv("LLAVA_MODEL_PATH")
    chat_args.conv_mode = "v0_mmtag" # "multimodal"
    chat_args.num_gpus = 1

    # rich console for pretty printing
    console = rich.console.Console()

    # Initialize LLaVA chat
    chat = LLaVaChat(chat_args.model_path, chat_args.conv_mode, chat_args.num_gpus)
    # chat = LLaVaChat(chat_args)
    print("LLaVA chat initialized...")
    query = "Describe the central object in the image."
    # query = "Describe the object in the image that is outlined in red."

    # Directories to save features and captions
    savedir_feat = Path(args.cachedir) / "cfslam_feat_llava"
    savedir_feat.mkdir(exist_ok=True, parents=True)
    savedir_captions = Path(args.cachedir) / "cfslam_captions_llava"
    savedir_captions.mkdir(exist_ok=True, parents=True)
    savedir_debug = Path(args.cachedir) / "cfslam_captions_llava_debug"
    savedir_debug.mkdir(exist_ok=True, parents=True)

    caption_dict_list = []

    for idx_obj, obj in tqdm(enumerate(scene_map), total=len(scene_map)):
        conf = obj["conf"]
        conf = np.array(conf)
        idx_most_conf = np.argsort(conf)[::-1]

        features = []
        captions = []
        low_confidences = []
        
        image_list = []
        caption_list = []
        confidences_list = []
        low_confidences_list = []
        mask_list = []  # New list for masks
        if len(idx_most_conf) < 2:
            continue 
        idx_most_conf = idx_most_conf[:args.max_detections_per_object]

        for idx_det in tqdm(idx_most_conf):
            # image = Image.open(correct_path).convert("RGB")
            image = Image.open(obj["color_path"][idx_det]).convert("RGB")
            xyxy = obj["xyxy"][idx_det]
            class_id = obj["class_id"][idx_det]
            class_name = class_names[class_id]
            # Retrieve and crop mask
            mask = obj["mask"][idx_det]

            padding = 10
            x1, y1, x2, y2 = xyxy
            # image_crop = crop_image_pil(image, x1, y1, x2, y2, padding=padding)
            image_crop, mask_crop = crop_image_and_mask(image, mask, x1, y1, x2, y2, padding=padding)
            if args.masking_option == "blackout":
                image_crop_modified = blackout_nonmasked_area(image_crop, mask_crop)
            elif args.masking_option == "red_outline":
                image_crop_modified = draw_red_outline(image_crop, mask_crop)
            else:
                image_crop_modified = image_crop  # No modification

            _w, _h = image_crop.size
            if _w * _h < 70 * 70:
                # captions.append("small object")
                print("small object. Skipping LLaVA captioning...")
                low_confidences.append(True)
                continue
            else:
                low_confidences.append(False)

            # image_tensor = chat.image_processor.preprocess(image_crop, return_tensors="pt")["pixel_values"][0]
            image_tensor = chat.image_processor.preprocess(image_crop_modified, return_tensors="pt")["pixel_values"][0]

            image_features = chat.encode_image(image_tensor[None, ...].half().cuda())
            features.append(image_features.detach().cpu())

            chat.reset()
            console.print("[bold red]User:[/bold red] " + query)
            outputs = chat(query=query, image_features=image_features)
            console.print("[bold green]LLaVA:[/bold green] " + outputs)
            captions.append(outputs)
        
            # print(f"Line 274, obj['mask'][idx_det].shape: {obj['mask'][idx_det].shape}")
            # print(f"Line 276, image.size: {image.size}")
            
            # For the LLava debug folder
            conf_value = conf[idx_det]
            image_list.append(image_crop)
            caption_list.append(outputs)
            confidences_list.append(conf_value)
            low_confidences_list.append(low_confidences[-1])
            mask_list.append(mask_crop)  # Add the cropped mask

        caption_dict_list.append(
            {
                "id": idx_obj,
                "captions": captions,
                "low_confidences": low_confidences,
            }
        )

        # Concatenate the features
        if len(features) > 0:
            features = torch.cat(features, dim=0)

        # Save the feature descriptors
        torch.save(features, savedir_feat / f"{idx_obj}.pt")
        
        # Again for the LLava debug folder
        if len(image_list) > 0:
            plot_images_with_captions(image_list, caption_list, confidences_list, low_confidences_list, mask_list, savedir_debug, idx_obj)

    # Save the captions
    # Remove the "The central object in the image is " prefix from 
    # the captions as it doesnt convey and actual info
    for item in caption_dict_list:
        item["captions"] = [caption.replace("The central object in the image is ", "") for caption in item["captions"]]
    # Save the captions to a json file
    with open(Path(args.cachedir) / "cfslam_llava_captions.json", "w", encoding="utf-8") as f:
        json.dump(caption_dict_list, f, indent=4, sort_keys=False)


def save_json_to_file(json_str, filename):
    with open(filename, "w", encoding="utf-8") as f:
        json.dump(json_str, f, indent=4, sort_keys=False)


def refine_node_captions(args):
    # NOTE: args.mapfile is in cfslam format
    from conceptgraph.slam.slam_classes import MapObjectList
    from conceptgraph.scenegraph.GPTPrompt import GPTPrompt

    # Load the captions for each segment
    caption_file = Path(args.cachedir) / "cfslam_llava_captions.json"
    captions = None
    with open(caption_file, "r") as f:
        captions = json.load(f)
    # loaddir_captions = Path(args.cachedir) / "cfslam_captions_llava"
    # captions = []
    # for idx_obj in range(len(os.listdir(loaddir_captions))):
    #     with open(loaddir_captions / f"{idx_obj}.pkl", "rb") as f:
    #         captions.append(pkl.load(f))

    # Load the scene map
    scene_map = MapObjectList()
    load_scene_map(args, scene_map)
    
    # load the prompt
    gpt_messages = GPTPrompt().get_json()

    TIMEOUT = 25  # Timeout in seconds

    responses_savedir = Path(args.cachedir) / "cfslam_gpt-4_responses"
    responses_savedir.mkdir(exist_ok=True, parents=True)

    responses = []
    unsucessful_responses = 0

    # loop over every object
    for _i in trange(len(captions)):
        if len(captions[_i]) == 0:
            continue
        
        # Prepare the object prompt 
        _dict = {}
        _caption = captions[_i]
        _bbox = scene_map[_i]["bbox"]
        # _bbox = o3d.geometry.OrientedBoundingBox.create_from_points(o3d.utility.Vector3dVector(scene_map[_i]["bbox"]))
        _dict["id"] = _caption["id"]
        # _dict["bbox_extent"] = np.round(_bbox.extent, 1).tolist()
        # _dict["bbox_center"] = np.round(_bbox.center, 1).tolist()
        _dict["captions"] = _caption["captions"]
        # _dict["low_confidences"] = _caption["low_confidences"]
        # Convert to printable string
        
        # Make and format the full prompt
        preds = json.dumps(_dict, indent=0)

        start_time = time.time()
    
        curr_chat_messages = gpt_messages[:]
        curr_chat_messages.append({"role": "user", "content": preds})
        chat_completion = openai.ChatCompletion.create(
            # model="gpt-3.5-turbo",
            model="gpt-4",
            messages=curr_chat_messages,
            timeout=TIMEOUT,  # Timeout in seconds
        )
        elapsed_time = time.time() - start_time
        if elapsed_time > TIMEOUT:
            print("Timed out exceeded!")
            _dict["response"] = "FAIL"
            # responses.append('{"object_tag": "FAIL"}')
            save_json_to_file(_dict, responses_savedir / f"{_caption['id']}.json")
            responses.append(json.dumps(_dict))
            unsucessful_responses += 1
            continue
        
        # count unsucessful responses
        if "invalid" in chat_completion["choices"][0]["message"]["content"].strip("\n"):
            unsucessful_responses += 1
            
        # print output
        prjson([{"role": "user", "content": preds}])
        print(chat_completion["choices"][0]["message"]["content"])
        print(f"Unsucessful responses so far: {unsucessful_responses}")
        _dict["response"] = chat_completion["choices"][0]["message"]["content"].strip("\n")
        
        # save the response
        responses.append(json.dumps(_dict))
        save_json_to_file(_dict, responses_savedir / f"{_caption['id']}.json")
        # responses.append(chat_completion["choices"][0]["message"]["content"].strip("\n"))

    # tags = []
    # for response in responses:
    #     try:
    #         parsed = json.loads(response)
    #         tags.append(parsed["object_tag"])
    #     except:
    #         tags.append("FAIL")

    # Save the responses to a text file
    # with open(Path(args.cachedir) / "gpt-3.5-turbo_responses.txt", "w") as f:
    #     for response in responses:
    #         f.write(response + "\n")
    with open(Path(args.cachedir) / "cfslam_gpt-4_responses.pkl", "wb") as f:
        pkl.dump(responses, f)


def extract_object_tag_from_json_str(json_str):
    start_str_found = False
    is_object_tag = False
    object_tag_complete = False
    object_tag = ""
    r = json_str.strip().split()
    for _idx, _r in enumerate(r):
        if not start_str_found:
            # Searching for open parenthesis of JSON
            if _r == "{":
                start_str_found = True
                continue
            else:
                continue
        # Start string found. Now skip everything until the object_tag field
        if not is_object_tag:
            if _r == '"object_tag":':
                is_object_tag = True
                continue
            else:
                continue
        # object_tag field found. Read it
        if is_object_tag and not object_tag_complete:
            if _r == '"':
                continue
            else:
                if _r.strip() in [",", "}"]:
                    break
                object_tag += f" {_r}"
                continue
    return object_tag


def build_scenegraph(args):
    from conceptgraph.slam.slam_classes import MapObjectList
    from conceptgraph.slam.utils import compute_overlap_matrix

    # Load the scene map
    scene_map = MapObjectList()
    load_scene_map(args, scene_map)

    response_dir = Path(args.cachedir) / "cfslam_gpt-4_responses"
    responses = []
    object_tags = []
    also_indices_to_remove = [] # indices to remove if the json file does not exist
    for idx in range(len(scene_map)):
        # check if the json file exists first 
        if not (response_dir / f"{idx}.json").exists():
            also_indices_to_remove.append(idx)
            continue
        with open(response_dir / f"{idx}.json", "r") as f:
            _d = json.load(f)
            try:
                _d["response"] = json.loads(_d["response"])
            except json.JSONDecodeError:
                _d["response"] = {
                    'summary': f'GPT4 json reply failed: Here is the invalid response {_d["response"]}',
                    'possible_tags': ['possible_tag_json_failed'],
                    'object_tag': 'invalid'
                }
            responses.append(_d)
            object_tags.append(_d["response"]["object_tag"])

    # # Load the responses from the json file
    # responses = None
    # # Load json file into a list
    # with open(Path(args.cachedir) / "cfslam_gpt-4_responses.pkl", "rb") as f:
    #     responses = pkl.load(f)
    # object_tags = []
    # for response in responses:
    #     # _tag = extract_object_tag_from_json_str(response)
    #     # _tag = _tag.lower().replace('"', "").strip()
    #     _d = json.loads(json.loads(response)["responses"])
    #     object_tags.append(_d["object_tag"])

    # Remove segments that correspond to "invalid" tags
    indices_to_remove = [i for i in range(len(responses)) if object_tags[i].lower() in ["fail", "invalid"]]
    # Also remove segments that do not have a minimum number of observations
    indices_to_remove = set(indices_to_remove)
    for obj_idx in range(len(scene_map)):
        conf = scene_map[obj_idx]["conf"]
        # Remove objects with less than args.min_views_per_object observations
        if len(conf) < args.min_views_per_object:
            indices_to_remove.add(obj_idx)
    indices_to_remove = list(indices_to_remove)
    # combine with also_indices_to_remove and sort the list
    indices_to_remove = list(set(indices_to_remove + also_indices_to_remove))
    
    # List of tags in original scene map that are in the pruned scene map
    segment_ids_to_retain = [i for i in range(len(scene_map)) if i not in indices_to_remove]
    with open(Path(args.cachedir) / "cfslam_scenegraph_invalid_indices.pkl", "wb") as f:
        pkl.dump(indices_to_remove, f)
    print(f"Removed {len(indices_to_remove)} segments")
    
    # Filtering responses based on segment_ids_to_retain
    responses = [resp for resp in responses if resp['id'] in segment_ids_to_retain]

    # Assuming each response dictionary contains an 'object_tag' key for the object tag.
    # Extract filtered object tags based on filtered_responses
    object_tags = [resp['response']['object_tag'] for resp in responses]


    pruned_scene_map = []
    pruned_object_tags = []
    for _idx, segmentidx in enumerate(segment_ids_to_retain):
        pruned_scene_map.append(scene_map[segmentidx])
        pruned_object_tags.append(object_tags[_idx])
    scene_map = MapObjectList(pruned_scene_map)
    object_tags = pruned_object_tags
    del pruned_scene_map
    # del pruned_object_tags
    gc.collect()
    num_segments = len(scene_map)

    for i in range(num_segments):
        scene_map[i]["caption_dict"] = responses[i]
        # scene_map[i]["object_tag"] = object_tags[i]

    # Save the pruned scene map (create the directory if needed)
    if not (Path(args.cachedir) / "map").exists():
        (Path(args.cachedir) / "map").mkdir(parents=True, exist_ok=True)
    with gzip.open(Path(args.cachedir) / "map" / "scene_map_cfslam_pruned.pkl.gz", "wb") as f:
        pkl.dump(scene_map.to_serializable(), f)

    print("Computing bounding box overlaps...")
    bbox_overlaps = compute_overlap_matrix(args, scene_map)

    # Construct a weighted adjacency matrix based on similarity scores
    weights = []
    rows = []
    cols = []
    for i in range(num_segments):
        for j in range(i + 1, num_segments):
            if i == j:
                continue
            if bbox_overlaps[i, j] > 0.01:
                weights.append(bbox_overlaps[i, j])
                rows.append(i)
                cols.append(j)
                weights.append(bbox_overlaps[i, j])
                rows.append(j)
                cols.append(i)

    adjacency_matrix = csr_matrix((weights, (rows, cols)), shape=(num_segments, num_segments))

    # Find the minimum spanning tree of the weighted adjacency matrix
    mst = minimum_spanning_tree(adjacency_matrix)

    # Find connected components in the minimum spanning tree
    _, labels = connected_components(mst)

    components = []
    _total = 0
    if len(labels) != 0:
        for label in range(labels.max() + 1):
            indices = np.where(labels == label)[0]
            _total += len(indices.tolist())
            components.append(indices.tolist())

    with open(Path(args.cachedir) / "cfslam_scenegraph_components.pkl", "wb") as f:
        pkl.dump(components, f)

    # Initialize a list to store the minimum spanning trees of connected components
    minimum_spanning_trees = []
    relations = []
    if len(labels) != 0:
        # Iterate over each connected component
        for label in range(labels.max() + 1):
            component_indices = np.where(labels == label)[0]
            # Extract the subgraph for the connected component
            subgraph = adjacency_matrix[component_indices][:, component_indices]
            # Find the minimum spanning tree of the connected component subgraph
            _mst = minimum_spanning_tree(subgraph)
            # Add the minimum spanning tree to the list
            minimum_spanning_trees.append(_mst)

        TIMEOUT = 25  # timeout in seconds

        if not (Path(args.cachedir) / "cfslam_object_relations.json").exists():
            relation_queries = []
            for componentidx, component in enumerate(components):
                if len(component) <= 1:
                    continue
                for u, v in zip(
                    minimum_spanning_trees[componentidx].nonzero()[0], minimum_spanning_trees[componentidx].nonzero()[1]
                ):
                    segmentidx1 = component[u]
                    segmentidx2 = component[v]
                    _bbox1 = scene_map[segmentidx1]["bbox"]
                    _bbox2 = scene_map[segmentidx2]["bbox"]

                    input_dict = {
                        "object1": {
                            "id": segmentidx1,
                            "bbox_extent": np.round(_bbox1.extent, 1).tolist(),
                            "bbox_center": np.round(_bbox1.center, 1).tolist(),
                            "object_tag": object_tags[segmentidx1],
                        },
                        "object2": {
                            "id": segmentidx2,
                            "bbox_extent": np.round(_bbox2.extent, 1).tolist(),
                            "bbox_center": np.round(_bbox2.center, 1).tolist(),
                            "object_tag": object_tags[segmentidx2],
                        },
                    }
                    print(f"{input_dict['object1']['object_tag']}, {input_dict['object2']['object_tag']}")

                    relation_queries.append(input_dict)

                    input_json_str = json.dumps(input_dict)

                    # Default prompt
                    DEFAULT_PROMPT = """
                    The input is a list of JSONs describing two objects "object1" and "object2". You need to produce a JSON
                    string (and nothing else), with two keys: "object_relation", and "reason".

                    Each of the JSON fields "object1" and "object2" will have the following fields:
                    1. bbox_extent: the 3D bounding box extents of the object
                    2. bbox_center: the 3D bounding box center of the object
                    3. object_tag: an extremely brief description of the object

                    Produce an "object_relation" field that best describes the relationship between the two objects. The
                    "object_relation" field must be one of the following (verbatim):
                    1. "a on b": if object a is an object commonly placed on top of object b
                    2. "b on a": if object b is an object commonly placed on top of object a
                    3. "a in b": if object a is an object commonly placed inside object b
                    4. "b in a": if object b is an object commonly placed inside object a
                    5. "none of these": if none of the above best describe the relationship between the two objects

                    Before producing the "object_relation" field, produce a "reason" field that explains why
                    the chosen "object_relation" field is the best.
                    """

                    start_time = time.time()
                    chat_completion = openai.ChatCompletion.create(
                        # model="gpt-3.5-turbo",
                        model="gpt-4",
                        messages=[{"role": "user", "content": DEFAULT_PROMPT + "\n\n" + input_json_str}],
                        timeout=TIMEOUT,  # Timeout in seconds
                    )
                    elapsed_time = time.time() - start_time
                    output_dict = input_dict
                    if elapsed_time > TIMEOUT:
                        print("Timed out exceeded!")
                        output_dict["object_relation"] = "FAIL"
                        continue
                    else:
                        try:
                            # Attempt to parse the output as a JSON
                            chat_output_json = json.loads(chat_completion["choices"][0]["message"]["content"])
                            # If the output is a valid JSON, then add it to the output dictionary
                            output_dict["object_relation"] = chat_output_json["object_relation"]
                            output_dict["reason"] = chat_output_json["reason"]
                        except:
                            output_dict["object_relation"] = "FAIL"
                            output_dict["reason"] = "FAIL"
                    relations.append(output_dict)

                    # print(chat_completion["choices"][0]["message"]["content"])

            # Save the query JSON to file
            print("Saving query JSON to file...")
            with open(Path(args.cachedir) / "cfslam_object_relation_queries.json", "w") as f:
                json.dump(relation_queries, f, indent=4)

            # Saving the output
            print("Saving object relations to file...")
            with open(Path(args.cachedir) / "cfslam_object_relations.json", "w") as f:
                json.dump(relations, f, indent=4)
        else:
            relations = json.load(open(Path(args.cachedir) / "cfslam_object_relations.json", "r"))

    scenegraph_edges = []

    _idx = 0
    for componentidx, component in enumerate(components):
        if len(component) <= 1:
            continue
        for u, v in zip(
            minimum_spanning_trees[componentidx].nonzero()[0], minimum_spanning_trees[componentidx].nonzero()[1]
        ):
            segmentidx1 = component[u]
            segmentidx2 = component[v]
            # print(f"{segmentidx1}, {segmentidx2}, {relations[_idx]['object_relation']}")
            if relations[_idx]["object_relation"] != "none of these":
                scenegraph_edges.append((segmentidx1, segmentidx2, relations[_idx]["object_relation"]))
            _idx += 1
    print(f"Created 3D scenegraph with {num_segments} nodes and {len(scenegraph_edges)} edges")

    with open(Path(args.cachedir) / "cfslam_scenegraph_edges.pkl", "wb") as f:
        pkl.dump(scenegraph_edges, f)


def generate_scenegraph_json(args):
    from conceptgraph.slam.slam_classes import MapObjectList
    

    # Generate the JSON file summarizing the scene, if it doesn't exist already
    # or if the --recopmute_scenegraph_json flag is set
    scene_desc = []
    print("Generating scene graph JSON file...")

    # Load the pruned scene map
    scene_map = MapObjectList()
    with gzip.open(Path(args.cachedir) / "map" / "scene_map_cfslam_pruned.pkl.gz", "rb") as f:
        scene_map.load_serializable(pkl.load(f))
    print(f"Loaded scene map with {len(scene_map)} objects")

    for i, segment in enumerate(scene_map):
        _d = {
            "id": segment["caption_dict"]["id"],
            "bbox_extent": np.round(segment['bbox'].extent, 1).tolist(),
            "bbox_center": np.round(segment['bbox'].center, 1).tolist(),
            "possible_tags": segment["caption_dict"]["response"]["possible_tags"],
            "object_tag": segment["caption_dict"]["response"]["object_tag"],
            "caption": segment["caption_dict"]["response"]["summary"],
        }
        scene_desc.append(_d)
    with open(Path(args.cachedir) / "scene_graph.json", "w") as f:
        json.dump(scene_desc, f, indent=4)


def display_images(image_list):
    num_images = len(image_list)
    cols = 2  # Number of columns for the subplots (you can change this as needed)
    rows = (num_images + cols - 1) // cols

    _, axes = plt.subplots(rows, cols, figsize=(10, 5))

    for i, ax in enumerate(axes.flat):
        if i < num_images:
            img = image_list[i]
            ax.imshow(img)
            ax.axis("off")
        else:
            ax.axis("off")

    plt.tight_layout()
    plt.show()


def annotate_scenegraph(args):
    from conceptgraph.slam.slam_classes import MapObjectList

    # Load the pruned scene map
    scene_map = MapObjectList()
    with gzip.open(Path(args.cachedir) / "map" / "scene_map_cfslam_pruned.pkl.gz", "rb") as f:
        scene_map.load_serializable(pkl.load(f))

    annot_inds = None
    if args.annot_inds is not None:
        annot_inds = args.annot_inds
    # If annot_inds is not None, we also need to load the annotation json file and only
    # annotate the objects that are specified in the annot_inds list
    annots = []
    if annot_inds is not None:
        annots = json.load(open(Path(args.cachedir) / "annotated_scenegraph.json", "r"))

    if annot_inds is None:
        annot_inds = list(range(len(scene_map)))

    for idx in annot_inds:
        print(f"Object {idx} out of {len(annot_inds)}...")
        obj = scene_map[idx]

        prev_annot = None
        if len(annots) >= idx + 1:
            prev_annot = annots[idx]

        annot = {}
        annot["id"] = idx

        conf = obj["conf"]
        conf = np.array(conf)
        idx_most_conf = np.argsort(conf)[::-1]
        print(obj.keys())

        imgs = []

        for idx_det in idx_most_conf:
            image = Image.open(obj["color_path"][idx_det]).convert("RGB")
            xyxy = obj["xyxy"][idx_det]
            mask = obj["mask"][idx_det]

            padding = 10
            x1, y1, x2, y2 = xyxy
            image_crop = crop_image_pil(image, x1, y1, x2, y2, padding=padding)
            mask_crop = crop_image_pil(Image.fromarray(mask), x1, y1, x2, y2, padding=padding)
            mask_crop = np.array(mask_crop)[..., None]
            mask_crop[mask_crop == 0] = 0.05
            image_crop = np.array(image_crop) * mask_crop
            imgs.append(image_crop)
            if len(imgs) >= 5:
                break
            # if idx_det >= 5:
            #     break

        # Display the images
        display_images(imgs)
        plt.close("all")

        # Ask the user to annotate the object
        if prev_annot is not None:
            print("Previous annotation:")
            print(prev_annot)
        annot["object_tags"] = input("Enter object tags (comma-separated): ")
        annot["colors"] = input("Enter colors (comma-separated): ")
        annot["materials"] = input("Enter materials (comma-separated): ")

        if prev_annot is not None:
            annots[idx] = annot
        else:
            annots.append(annot)

        go_on = input("Continue? (y/n): ")
        if go_on == "n":
            break

    # Save the annotations
    with open(Path(args.cachedir) / "annotated_scenegraph.json", "w") as f:
        json.dump(annots, f, indent=4)


def main():
    # Process command-line args (if any)
    args = tyro.cli(ProgramArgs)
    
    # print using masking option
    print(f"args.masking_option: {args.masking_option}")

    if args.mode == "extract-node-captions":
        extract_node_captions(args)
    elif args.mode == "refine-node-captions":
        refine_node_captions(args)
    elif args.mode == "build-scenegraph":
        build_scenegraph(args)
    elif args.mode == "generate-scenegraph-json":
        generate_scenegraph_json(args)
    elif args.mode == "annotate-scenegraph":
        annotate_scenegraph(args)
    else:
        raise ValueError(f"Unknown mode: {args.mode}")


if __name__ == "__main__":
    main()

conceptgraph/scenegraph/build_scenegraph_cfslam.py
concept-graphs/concept-graphs
"""
Post-process the results from the SLAM to merge duplicate objects and filter out
objects with too few points or detections.
"""

import argparse
import copy
import os
import pickle
import gzip

import matplotlib
import numpy as np

from conceptgraph.slam.slam_classes import MapObjectList

from conceptgraph.slam.utils import filter_objects, merge_objects

def get_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument("--result_path", type=str, required=True)
    parser.add_argument("--savefile", type=str, required=True)
    
    # To inspect the results of merge_overlap_objects
    # This is mainly to quickly try out different thresholds
    parser.add_argument("--merge_overlap_thresh", type=float, default=-1)
    parser.add_argument("--merge_visual_sim_thresh", type=float, default=-1)
    parser.add_argument("--merge_text_sim_thresh", type=float, default=-1)
    parser.add_argument("--obj_min_points", type=int, default=0)
    parser.add_argument("--obj_min_detections", type=int, default=0)
    
    return parser

if __name__ == "__main__":
    parser = get_parser()
    args = parser.parse_args()

    result_path = args.result_path

    with gzip.open(result_path, "rb") as f:
        results = pickle.load(f)
        
    cfg = results['cfg']
    class_names = results['class_names']
    class_colors = results['class_colors']
        
    objects = MapObjectList()
    objects.load_serializable(results['objects'])
    
    # Run the post-processing filtering and merging in instructed to do so
    cfg.obj_min_points = args.obj_min_points
    cfg.obj_min_detections = args.obj_min_detections
    cfg.merge_overlap_thresh = args.merge_overlap_thresh
    cfg.merge_visual_sim_thresh = args.merge_visual_sim_thresh
    cfg.merge_text_sim_thresh = args.merge_text_sim_thresh
    objects = filter_objects(cfg, objects)
    objects = merge_objects(cfg, objects)
    
    # if not args.no_clip:
    #     print("Initializing CLIP model...")
    #     clip_model, _, clip_preprocess = open_clip.create_model_and_transforms("ViT-H-14", "laion2b_s32b_b79k")
    #     clip_model = clip_model.to("cuda")
    #     clip_tokenizer = open_clip.get_tokenizer("ViT-H-14")
    #     print("Done initializing CLIP model.")

    cmap = matplotlib.colormaps.get_cmap("turbo")
    
    # Handle the background objects if they exist
    print(results.keys())
    bg_objects = None
    if 'bg_objects' in results and results['bg_objects'] is not None:
        bg_objects = MapObjectList()
        bg_objects.load_serializable(results['bg_objects'])

        pcds_bg = copy.deepcopy(bg_objects.get_values("pcd"))
        bboxes_bg = copy.deepcopy(bg_objects.get_values("bbox"))
        
        indices_bg = np.arange(len(objects), len(objects) + len(bg_objects))
        
        objects.extend(bg_objects)
    
    # pcds = copy.deepcopy(objects.get_values("pcd"))
    # bboxes = copy.deepcopy(objects.get_values("bbox"))

    # # similarities = objects.compute_similarities(text_query_ft)
    # objects_clip_fts = objects.get_stacked_values_torch("clip_ft")

    objects = objects.to_serializable()
    # print(len(objects))
    
    print(f"\nStarted with {len(results['objects'])} objects then filtered to {len(objects)} objects\n")

    # Save the results
    # If the parent directory doesn't exist, create it
    if not os.path.exists(os.path.dirname(args.savefile)):
        os.makedirs(os.path.dirname(args.savefile), exist_ok=True)
    with gzip.open(args.savefile, "wb") as f:
        pickle.dump(objects, f)
    
conceptgraph/scenegraph/merge_duplicate_objects.py
concept-graphs/concept-graphs
import argparse
import json
import os, glob
import imageio
import natsort
import gzip, pickle
from enum import Enum

import open3d as o3d
import numpy as np
import time

from tqdm import trange

from conceptgraph.slam.slam_classes import MapObjectList
from conceptgraph.utils.general_utils import to_numpy
from conceptgraph.utils.vis import better_camera_frustum, LineMesh

import distinctipy

class COLOR_MODE(Enum):
    RGB = 0
    CLASS = 1
    INSTANCE = 2
    
def create_ball_mesh(center, radius, color=(0, 1, 0)):
    """
    Create a colored mesh sphere.
    
    Args:
    - center (tuple): (x, y, z) coordinates for the center of the sphere.
    - radius (float): Radius of the sphere.
    - color (tuple): RGB values in the range [0, 1] for the color of the sphere.
    
    Returns:
    - o3d.geometry.TriangleMesh: Colored mesh sphere.
    """
    mesh_sphere = o3d.geometry.TriangleMesh.create_sphere(radius=radius)
    mesh_sphere.translate(center)
    mesh_sphere.paint_uniform_color(color)
    return mesh_sphere

def get_parser():
    parser = argparse.ArgumentParser(description="Visualize a series of point clouds as an animation.")
    parser.add_argument("--input_folder", type=str, required=True, help="Folder where the objects of the mapping process are stored.")
    parser.add_argument("--meta_folder", type=str, default=None, help="Folder where the meta information is stored. Default the same as input_folder.")
    parser.add_argument("--edge_file", type=str, default=None, help="Path to the scene graph relationship json file. If not provided, the edges will not be shown.")
    parser.add_argument("--sleep_time", type=float, default=0.1, help="Time to sleep between each frame.")
    parser.add_argument("--follow_cam", action="store_true", help="If set, follow the camera pose.")
    parser.add_argument("--use_original_color", action="store_true", help="If set, will use color scheme from the CFSLAM pipeline.")
    parser.add_argument("--height_cutoff", type=float, default=np.inf, help="Object nodes above this height will not be shown.")
    return parser

cached_frame = None
def load_frame(path):
    global cached_frame
    if cached_frame is None:
        with gzip.open(path, "rb") as f:
            frame = pickle.load(f)
    else:
        frame = cached_frame
        
    if isinstance(frame, dict):
        camera_pose = frame["camera_pose"]
        objects = MapObjectList()
        objects.load_serializable(frame["objects"])
        
        if frame['bg_objects'] is None:
            bg_objects = None
        else:
            bg_objects = MapObjectList()
            bg_objects.load_serializable(frame["bg_objects"])
    elif isinstance(frame, list):
        cached_frame = frame
        objects = MapObjectList()
        objects.load_serializable(frame)

        bg_objects = None
        camera_pose = None
    else:
        raise ValueError("Unknown frame type: ", type(frame))
        
    return camera_pose, objects, bg_objects

def main(args):
    # Load metadata
    if args.meta_folder is None:
        args.meta_folder = args.input_folder
        
    meta_path = os.path.join(args.meta_folder, "meta.pkl.gz")
    with gzip.open(meta_path, "rb") as f:
        meta_info = pickle.load(f)
        
    cfg = meta_info["cfg"]
    class_names = meta_info["class_names"]
    main.class_colors = meta_info["class_colors"]
    
    if not args.use_original_color:
        distinct_colors = distinctipy.get_colors(len(main.class_colors), pastel_factor=0.5)
        for k, c in zip(main.class_colors.keys(), distinct_colors):
            main.class_colors[k] = c
    
    # Scan the mapping results from all frames
    frame_paths = glob.glob(os.path.join(args.input_folder, "*.pkl.gz"))
    frame_paths = [path for path in frame_paths if path != meta_path]
    frame_paths = natsort.natsorted(frame_paths)
    
    # Load edge files
    if args.edge_file is not None:
        with open(args.edge_file, "r") as f:
            edges = json.load(f)

    # Create the visualization
    vis = o3d.visualization.VisualizerWithKeyCallback()
    vis.create_window()
    
    view_ctrl = vis.get_view_control()
    camera_param = view_ctrl.convert_to_pinhole_camera_parameters()
    
    main.color_mode = COLOR_MODE.RGB
    if "class_agnostic" in cfg and cfg.class_agnostic:
        main.color_mode = COLOR_MODE.INSTANCE
        
    main.paused = False
    main.show_bg = True
    main.forward = True
    main.show_top = True
    main.show_graph = True
    main.show_cam = True
    main.show_bbox = True
    main.frame_idx = 0
    
    # Define the key callback functions
    def color_mode_rgb(vis):
        main.color_mode = COLOR_MODE.RGB
        
    def color_mode_class(vis):
        main.color_mode = COLOR_MODE.CLASS
        
    def color_mode_inst(vis):
        main.color_mode = COLOR_MODE.INSTANCE
        
    def save_color_scheme(vis):
        filename = input("Enter the filename to save the color scheme: ")
        path = os.path.join("./tmp/cfslam_vis_color/", f"{filename}.json")
        path = os.path.abspath(path)
        os.makedirs(os.path.dirname(path), exist_ok=True)
        with open(path, "w") as f:
            json.dump(main.class_colors, f)
        print("Color scheme saved to: ", path)
        
    def load_color_scheme(vis):
        filename = input("Enter the filename to load the color scheme: ")
        path = os.path.join("./tmp/cfslam_vis_color/", f"{filename}.json")
        path = os.path.abspath(path)
        
        if not os.path.exists(path):
            print("File does not exist:", path)
            return
        
        with open(path, "r") as f:
            load_colors = json.load(f)
            
        if len(load_colors) != len(main.class_colors):
            print("Color scheme does not match the class names.")
            return
        main.class_colors = load_colors
        print("Color scheme loaded from: ", path)
    
    def pause_resume(vis):
        main.paused = not main.paused
        
    def toggle_bg(vis):
        main.show_bg = not main.show_bg
        
    def toggle_forward(vis):
        main.forward = not main.forward
        
    def toggle_top(vis):
        main.show_top = not main.show_top
        
    def toggle_cam(vis):
        main.show_cam = not main.show_cam

    def toggle_bbox(vis):
        main.show_bbox = not main.show_bbox
        
    def save_vis_capture(vis):
        vis.poll_events()
        vis.update_renderer()
        
        rendered_image = vis.capture_screen_float_buffer(False)
        rendered_image = np.asarray(rendered_image)
        
        rendered_image = (rendered_image * 255).astype(np.uint8)
        
        filename = input("Enter the filename to save the visualization: ")
        path = os.path.join("./tmp/cfslam_vis/", f"{filename}.png")
        os.makedirs(os.path.dirname(path), exist_ok=True)
        imageio.imwrite(path, rendered_image)
        
    def save_camera_params(vis):
        filename = input("Enter the filename to save the camera parameters: ")
        path = os.path.join("./tmp/cfslam_vis_cam/", f"{filename}.json")
        path = os.path.abspath(path)
        os.makedirs(os.path.dirname(path), exist_ok=True)

        vis_ctrl = vis.get_view_control()
        params = vis_ctrl.convert_to_pinhole_camera_parameters()
        o3d.io.write_pinhole_camera_parameters(path, params)
        print("Camera parameters saved to: ", path)

    def load_camera_params(vis):
        filename = input("Enter the filename to load the camera parameters: ")
        path = os.path.join("./tmp/cfslam_vis_cam/", f"{filename}.json")
        path = os.path.abspath(path)
        
        if not os.path.exists(path):
            print("File does not exist:", path)
            return
        
        vis_ctrl = vis.get_view_control()
        params = o3d.io.read_pinhole_camera_parameters(path)
        vis_ctrl.convert_from_pinhole_camera_parameters(params)
    
    def quit(vis):
        vis.destroy_window()
        print(main.frame_idx)
        exit()
        
    def go_to_frame_idx(vis):
        idx = int(input("Enter the frame index to go to: "))
        if not (0 <= main.frame_idx < len(frame_paths)):
            print("Invalid frame index: ", idx)
            return
        main.frame_idx = idx
        print("Frame index set to: ", main.frame_idx)
    
    vis.register_key_callback(ord("R"), color_mode_rgb)
    vis.register_key_callback(ord("C"), color_mode_class)
    vis.register_key_callback(ord("I"), color_mode_inst)
    
    vis.register_key_callback(ord("O"), save_color_scheme)
    vis.register_key_callback(ord("P"), load_color_scheme)
    
    vis.register_key_callback(ord(" "), pause_resume)
    vis.register_key_callback(ord("B"), toggle_bg)
    vis.register_key_callback(ord("F"), toggle_forward)
    vis.register_key_callback(ord("K"), toggle_cam)
    vis.register_key_callback(ord("X"), toggle_bbox)
    vis.register_key_callback(ord("N"), go_to_frame_idx)

    vis.register_key_callback(ord("T"), toggle_top)
    vis.register_key_callback(ord("S"), save_vis_capture)
    vis.register_key_callback(ord("W"), save_camera_params)
    vis.register_key_callback(ord("E"), load_camera_params)
    vis.register_key_callback(ord("Q"), quit)
    
    t = 0 # How many frames have been rendered
    while True:
        tic = time.time()
        camera_pose, objects, bg_objects = load_frame(frame_paths[main.frame_idx])

        # Hide background objects by setting it to None
        if not main.show_bg:
            bg_objects = None
            
        if bg_objects is not None and not main.show_top:
            # Hide the ceiling
            bg_classes = bg_objects.get_most_common_class()
            bg_objects = bg_objects.slice_by_mask(
                np.asarray(bg_classes) != class_names.index("ceiling")
            )
        
        # Give proper colors to the objects
        if main.color_mode == COLOR_MODE.RGB:
            pass # The default color is RGB, do nothing
        elif main.color_mode == COLOR_MODE.CLASS:
            objects.color_by_most_common_classes(main.class_colors, color_bbox=True)
            if bg_objects is not None:
                bg_objects.color_by_most_common_classes(main.class_colors, color_bbox=True)
        elif main.color_mode == COLOR_MODE.INSTANCE:
            objects.color_by_instance()
            if bg_objects is not None:
                bg_objects.color_by_instance()
                
        # Add geometries from the current frame
        vis.clear_geometries()
        for pcd in objects.get_values("pcd"):
            vis.add_geometry(pcd, reset_bounding_box = t<=3)
            
        if main.show_bbox:
            # for bbox in objects.get_values("bbox"):
            #     vis.add_geometry(bbox, reset_bounding_box = t<=3)
            for obj in objects:
                pcd = obj['pcd']
                bbox = pcd.get_axis_aligned_bounding_box()
                # bbox.color = obj['bbox'].color
                bbox.color = (0, 1, 0)
                vis.add_geometry(bbox, reset_bounding_box = t<=3)
        
        if bg_objects is not None:
            for pcd in bg_objects.get_values("pcd"):
                vis.add_geometry(pcd, reset_bounding_box = t<=3)

        # Set the camera pose
        if args.follow_cam:
            # Set the rendering camera pose to the observation camera pose
            camera_param.extrinsic = np.linalg.inv(to_numpy(camera_pose))
            view_ctrl.convert_from_pinhole_camera_parameters(camera_param)
            view_ctrl.camera_local_translate(forward=-0.4, right=0.0, up=0.0)
        else:
            if main.show_cam and camera_pose is not None:
                camera_frustum = better_camera_frustum(
                    camera_pose, 680//2, 1200//2, scale=0.2, color=[0, 1., 0]
                )
                vis.add_geometry(camera_frustum, reset_bounding_box = t<=3)
        
        vis.poll_events()
        vis.update_renderer()
        
        t += 1
        
        if main.paused:
            continue
        else:
            if main.forward:
                main.frame_idx += 1
                main.frame_idx = min(main.frame_idx, len(frame_paths)-1)
            else:
                main.frame_idx -= 1
                main.frame_idx = max(main.frame_idx, 0)
                
            toc = time.time()
            # print(toc - tic)
            if toc - tic < args.sleep_time:
                time.sleep(args.sleep_time - (toc-tic))
                
    

if __name__ == "__main__":
    parser = get_parser()
    args = parser.parse_args()
    main(args)
conceptgraph/scripts/animate_mapping_interactive.py
concept-graphs/concept-graphs
import argparse
import math
import os, glob
from pathlib import Path
import imageio
from matplotlib import pyplot as plt
import natsort
import gzip, pickle
from enum import Enum

from PIL import Image
from tqdm import trange
import open3d as o3d
import numpy as np
import time
import torch
import cv2

from supervision.detection.core import Detections
from supervision.draw.color import Color, ColorPalette

from conceptgraph.dataset.datasets_common import get_dataset

from conceptgraph.utils.general_utils import to_numpy
from conceptgraph.utils.vis import vis_result_fast, vis_result_slow_caption
from conceptgraph.slam.slam_classes import MapObjectList
from conceptgraph.slam.utils import filter_gobs

def get_parser():
    parser = argparse.ArgumentParser(description="Visualize a series of point clouds as an animation.")
    parser.add_argument("--input_folder", type=str, help="Folder where the objects of the mapping process are stored.")
    return parser

def main(args):
    meta_path = os.path.join(args.input_folder, "meta.pkl.gz")
    frame_paths = glob.glob(os.path.join(args.input_folder, "*.pkl.gz"))
    frame_paths = [path for path in frame_paths if path != meta_path]
    frame_paths = natsort.natsorted(frame_paths)

    with gzip.open(meta_path, "rb") as f:
        meta_info = pickle.load(f)
        
    cfg = meta_info["cfg"]
    class_names = meta_info["class_names"]
    class_colors = meta_info["class_colors"]
    
    color_mode_instance = "class_agnostic" in cfg and cfg.class_agnostic
    
    video_save_path = Path(args.input_folder).parent / f"{cfg.gsa_variant}_{cfg.save_suffix}.mp4"

    # Save each video separately in this folder
    sep_video_save_folder = Path(args.input_folder).parent.parent / "anime_all_frames" / f"{cfg.gsa_variant}_{cfg.save_suffix}"
    os.makedirs(sep_video_save_folder, exist_ok=True)
    
    # The sv annotation module use BGR color and expects color to be in [0, 255]
    class_name_to_sv_color = {
        n: Color(round(c[2]*255), round(c[1]*255), round(c[0]*255))
        for n, c in zip(class_names, class_colors.values())
    }
    
    # Initialize the dataset
    dataset = get_dataset(
        dataconfig=cfg.dataset_config,
        start=cfg.start,
        end=cfg.end,
        stride=cfg.stride,
        basedir=cfg.dataset_root,
        sequence=cfg.scene_id,
        device="cpu",
        dtype=torch.float,
    )
    cam_K = dataset.get_cam_K()
    
    color_path = dataset.color_paths[0]
    image_pil = Image.open(color_path)
    
    # Create the visualization
    vis = o3d.visualization.VisualizerWithKeyCallback()
    vis.create_window(
        window_name = "Mapping",
        width = image_pil.width,
        height = image_pil.height,
    )
    
    view_ctrl = vis.get_view_control()
    view_ctrl.change_field_of_view(10)
    camera_param = view_ctrl.convert_to_pinhole_camera_parameters()
    
    main.show_bg = True
    
    result_frames = []
    result_frames_sep = {
        "image_rgb": [],
        "image_class": [],
        "render_rgb": [],
        "render_class": [],
    }
    
    i = 0
    for i in trange(len(frame_paths)):
        color_path = dataset.color_paths[i]
        image_pil = Image.open(color_path)
        image_rgb = np.array(image_pil)
        
        color_path = Path(color_path)
        detections_path = color_path.parent.parent / cfg.detection_folder_name / color_path.name
        detections_path = detections_path.with_suffix(".pkl.gz")
        color_path = str(color_path)
        detections_path = str(detections_path)
        
        # Load the detection results at this frame
        with gzip.open(detections_path, "rb") as f:
            gobs = pickle.load(f)
            
        # Filter out the objects as performed during the mapping process
        gobs = filter_gobs(cfg, gobs, image_rgb)
        
        # Annotate the image with the remaining detection results
        detections = Detections(
            xyxy = gobs['xyxy'],
            confidence = gobs['confidence'],
            class_id = gobs['class_id'],
            mask = gobs['mask'],
        )
        
        frame_class_names = gobs['classes']
        frame_classid_to_svcolor = [
            class_name_to_sv_color[class_name]
            for class_name in frame_class_names
        ]
        frame_color_palette = ColorPalette(frame_classid_to_svcolor)
        if color_mode_instance:
            image_class, labels = vis_result_fast(
                image_rgb, detections, frame_class_names,
                instance_random_color=True,
                draw_bbox=False,
            )
        else:
            image_class, labels = vis_result_fast(
                image_rgb, detections, frame_class_names, frame_color_palette,
                instance_random_color=False,
            )

        # Load the mapping results up to this frame
        with gzip.open(frame_paths[i], "rb") as f:
            frame = pickle.load(f)
            
        if i > 0:
            vis.clear_geometries()
            
        camera_pose = frame["camera_pose"]
        
        objects = MapObjectList()
        objects.load_serializable(frame["objects"])
        
        bg_objects = None
        if main.show_bg:
            if frame['bg_objects'] is None:
                # print("No background objects found.")
                pass
            else:
                bg_objects = MapObjectList()
                bg_objects.load_serializable(frame["bg_objects"])

        # First render the objects in RGB color
        pcds = objects.get_values("pcd")
        bboxes = objects.get_values("bbox")
        for geom in pcds + bboxes:
            vis.add_geometry(geom, reset_bounding_box = i<2)
        
        if bg_objects is not None:
            for geom in bg_objects.get_values("pcd"):
                vis.add_geometry(geom, reset_bounding_box = i<2)
        
        camera_param.extrinsic = np.linalg.inv(to_numpy(camera_pose))
        view_ctrl.convert_from_pinhole_camera_parameters(camera_param)
        view_ctrl.camera_local_translate(forward=-0.4, right=0.0, up=0.0)
        vis.poll_events()
        vis.update_renderer()
        
        render_rgb = vis.capture_screen_float_buffer(False)
        render_rgb = np.asarray(render_rgb)
        
        # Then render the objects in class-coded colors
        if color_mode_instance:
            objects.color_by_instance()
            if bg_objects is not None:
                bg_objects.color_by_instance()
        else:
            objects.color_by_most_common_classes(class_colors, color_bbox=True)
            if bg_objects is not None:
                bg_objects.color_by_most_common_classes(class_colors, color_bbox=True)
            
        pcds = objects.get_values("pcd")
        bboxes = objects.get_values("bbox")
        for geom in pcds + bboxes:
            vis.add_geometry(geom, reset_bounding_box = i<2)
        
        if bg_objects is not None:
            for geom in bg_objects.get_values("pcd"):
                vis.add_geometry(geom, reset_bounding_box = i<2)
        
        camera_param.extrinsic = np.linalg.inv(to_numpy(camera_pose))
        view_ctrl.convert_from_pinhole_camera_parameters(camera_param)
        view_ctrl.camera_local_translate(forward=-0.4, right=0.0, up=0.0)
        vis.poll_events()
        vis.update_renderer()
        
        render_class = vis.capture_screen_float_buffer(False)
        render_class = np.asarray(render_class)
        
        # plt.subplot(2, 2, 1)
        # plt.imshow(image_pil)
        # plt.subplot(2, 2, 2)
        # plt.imshow(image_class)
        # plt.subplot(2, 2, 3)
        # plt.imshow(render_rgb)
        # plt.subplot(2, 2, 4)
        # plt.imshow(render_class)
        # plt.show()
        
        result_frames_sep["image_rgb"].append(image_rgb)
        result_frames_sep["image_class"].append(image_class)
        result_frames_sep["render_rgb"].append((render_rgb * 255).astype(np.uint8))
        result_frames_sep["render_class"].append((render_class * 255).astype(np.uint8))
        
        # Stack the two renders side-by-side, on the bottom of the image and annotated image
        image_stack = np.concatenate([image_rgb, image_class], axis=1)
        render_stack = np.concatenate([render_rgb, render_class], axis=1)
        render_stack = (render_stack * 255).astype(np.uint8)
        vis_stack = np.concatenate([image_stack, render_stack], axis=0)
        
        # plt.figure(figsize=(20, 10))
        # plt.imshow(vis_stack)
        # plt.axis("off")
        # plt.show()
        
        result_frames.append(vis_stack)
        
    vis.destroy_window()
    
    # Save the result as a video
    imageio.mimwrite(video_save_path, result_frames, fps=10)
    
    for k, v in result_frames_sep.items():
        sep_video_save_path = sep_video_save_folder / f"{k}.mp4"
        imageio.mimwrite(sep_video_save_path, v, fps=10)

if __name__ == "__main__":
    parser = get_parser()
    args = parser.parse_args()
    main(args)
conceptgraph/scripts/animate_mapping_save.py
concept-graphs/concept-graphs
'''
The script is used to extract Grounded SAM results on a posed RGB-D dataset. 
The results will be dumped to a folder under the scene folder. 
'''

import os
import argparse
from pathlib import Path
import re
from typing import Any, List
from PIL import Image
import cv2
import json
import imageio
import matplotlib
matplotlib.use("TkAgg")
from matplotlib import pyplot as plt
import numpy as np
import pickle
import gzip
import open_clip

import torch
import torchvision
from torch.utils.data import Dataset
import supervision as sv
from tqdm import trange

from conceptgraph.dataset.datasets_common import get_dataset
from conceptgraph.utils.vis import vis_result_fast, vis_result_slow_caption
import torch.nn.functional as F


try: 
    from groundingdino.util.inference import Model
    from segment_anything import sam_model_registry, SamPredictor, SamAutomaticMaskGenerator
except ImportError as e:
    print("Import Error: Please install Grounded Segment Anything following the instructions in README.")
    raise e

# Set up some path used in this script
# Assuming all checkpoint files are downloaded as instructed by the original GSA repo
if "GSA_PATH" in os.environ:
    GSA_PATH = os.environ["GSA_PATH"]
else:
    raise ValueError("Please set the GSA_PATH environment variable to the path of the GSA repo. ")
    
import sys
TAG2TEXT_PATH = os.path.join(GSA_PATH, "Tag2Text")
EFFICIENTSAM_PATH = os.path.join(GSA_PATH, "EfficientSAM")
sys.path.append(GSA_PATH) # This is needed for the following imports in this file
sys.path.append(TAG2TEXT_PATH) # This is needed for some imports in the Tag2Text files
sys.path.append(EFFICIENTSAM_PATH)
try:
    from Tag2Text.models import tag2text
    from Tag2Text import inference_tag2text, inference_ram
    import torchvision.transforms as TS
except ImportError as e:
    print("Tag2text sub-package not found. Please check your GSA_PATH. ")
    raise e

# Disable torch gradient computation
torch.set_grad_enabled(False)
    
# GroundingDINO config and checkpoint
GROUNDING_DINO_CONFIG_PATH = os.path.join(GSA_PATH, "GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py")
GROUNDING_DINO_CHECKPOINT_PATH = os.path.join(GSA_PATH, "./groundingdino_swint_ogc.pth")

# Segment-Anything checkpoint
SAM_ENCODER_VERSION = "vit_h"
SAM_CHECKPOINT_PATH = os.path.join(GSA_PATH, "./sam_vit_h_4b8939.pth")

# Tag2Text checkpoint
TAG2TEXT_CHECKPOINT_PATH = os.path.join(TAG2TEXT_PATH, "./tag2text_swin_14m.pth")
RAM_CHECKPOINT_PATH = os.path.join(TAG2TEXT_PATH, "./ram_swin_large_14m.pth")

FOREGROUND_GENERIC_CLASSES = [
    "item", "furniture", "object", "electronics", "wall decoration", "door"
]

FOREGROUND_MINIMAL_CLASSES = [
    "item"
]

def get_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--dataset_root", type=Path, required=True,
    )
    parser.add_argument(
        "--dataset_config", type=str, required=True,
        help="This path may need to be changed depending on where you run this script. "
    )
    
    parser.add_argument("--scene_id", type=str, default="train_3")
    
    parser.add_argument("--start", type=int, default=0)
    parser.add_argument("--end", type=int, default=-1)
    parser.add_argument("--stride", type=int, default=1)

    parser.add_argument("--desired-height", type=int, default=480)
    parser.add_argument("--desired-width", type=int, default=640)

    parser.add_argument("--box_threshold", type=float, default=0.25)
    parser.add_argument("--text_threshold", type=float, default=0.25)
    parser.add_argument("--nms_threshold", type=float, default=0.5)

    parser.add_argument("--class_set", type=str, default="scene", 
                        choices=["scene", "generic", "minimal", "tag2text", "ram", "none"], 
                        help="If none, no tagging and detection will be used and the SAM will be run in dense sampling mode. ")
    parser.add_argument("--add_bg_classes", action="store_true", 
                        help="If set, add background classes (wall, floor, ceiling) to the class set. ")
    parser.add_argument("--accumu_classes", action="store_true",
                        help="if set, the class set will be accumulated over frames")

    parser.add_argument("--sam_variant", type=str, default="sam",
                        choices=['fastsam', 'mobilesam', "lighthqsam"])
    
    parser.add_argument("--save_video", action="store_true")
    
    parser.add_argument("--device", type=str, default="cuda")
    
    parser.add_argument("--use_slow_vis", action="store_true", 
                        help="If set, use vis_result_slow_caption. Only effective when using ram/tag2text. ")
    
    parser.add_argument("--exp_suffix", type=str, default=None,
                        help="The suffix of the folder that the results will be saved to. ")
    
    return parser


def compute_clip_features(image, detections, clip_model, clip_preprocess, clip_tokenizer, classes, device):
    backup_image = image.copy()
    
    image = Image.fromarray(image)
    
    # padding = args.clip_padding  # Adjust the padding amount as needed
    padding = 20  # Adjust the padding amount as needed
    
    image_crops = []
    image_feats = []
    text_feats = []

    
    for idx in range(len(detections.xyxy)):
        # Get the crop of the mask with padding
        x_min, y_min, x_max, y_max = detections.xyxy[idx]

        # Check and adjust padding to avoid going beyond the image borders
        image_width, image_height = image.size
        left_padding = min(padding, x_min)
        top_padding = min(padding, y_min)
        right_padding = min(padding, image_width - x_max)
        bottom_padding = min(padding, image_height - y_max)

        # Apply the adjusted padding
        x_min -= left_padding
        y_min -= top_padding
        x_max += right_padding
        y_max += bottom_padding

        cropped_image = image.crop((x_min, y_min, x_max, y_max))
        
        # Get the preprocessed image for clip from the crop 
        preprocessed_image = clip_preprocess(cropped_image).unsqueeze(0).to("cuda")

        crop_feat = clip_model.encode_image(preprocessed_image)
        crop_feat /= crop_feat.norm(dim=-1, keepdim=True)
        
        class_id = detections.class_id[idx]
        tokenized_text = clip_tokenizer([classes[class_id]]).to("cuda")
        text_feat = clip_model.encode_text(tokenized_text)
        text_feat /= text_feat.norm(dim=-1, keepdim=True)
        
        crop_feat = crop_feat.cpu().numpy()
        text_feat = text_feat.cpu().numpy()

        image_crops.append(cropped_image)
        image_feats.append(crop_feat)
        text_feats.append(text_feat)
        
    # turn the list of feats into np matrices
    image_feats = np.concatenate(image_feats, axis=0)
    text_feats = np.concatenate(text_feats, axis=0)

    return image_crops, image_feats, text_feats


# Prompting SAM with detected boxes
def get_sam_segmentation_from_xyxy(sam_predictor: SamPredictor, image: np.ndarray, xyxy: np.ndarray) -> np.ndarray:
    sam_predictor.set_image(image)
    result_masks = []
    for box in xyxy:
        masks, scores, logits = sam_predictor.predict(
            box=box,
            multimask_output=True
        )
        index = np.argmax(scores)
        result_masks.append(masks[index])
    return np.array(result_masks)


def get_sam_predictor(variant: str, device: str | int) -> SamPredictor:
    if variant == "sam":
        sam = sam_model_registry[SAM_ENCODER_VERSION](checkpoint=SAM_CHECKPOINT_PATH)
        sam.to(device)
        sam_predictor = SamPredictor(sam)
        return sam_predictor
    
    if variant == "mobilesam":
        from MobileSAM.setup_mobile_sam import setup_model
        MOBILE_SAM_CHECKPOINT_PATH = os.path.join(GSA_PATH, "./EfficientSAM/mobile_sam.pt")
        checkpoint = torch.load(MOBILE_SAM_CHECKPOINT_PATH)
        mobile_sam = setup_model()
        mobile_sam.load_state_dict(checkpoint, strict=True)
        mobile_sam.to(device=device)
        
        sam_predictor = SamPredictor(mobile_sam)
        return sam_predictor

    elif variant == "lighthqsam":
        from LightHQSAM.setup_light_hqsam import setup_model
        HQSAM_CHECKPOINT_PATH = os.path.join(GSA_PATH, "./EfficientSAM/sam_hq_vit_tiny.pth")
        checkpoint = torch.load(HQSAM_CHECKPOINT_PATH)
        light_hqsam = setup_model()
        light_hqsam.load_state_dict(checkpoint, strict=True)
        light_hqsam.to(device=device)
        
        sam_predictor = SamPredictor(light_hqsam)
        return sam_predictor
        
    elif variant == "fastsam":
        raise NotImplementedError
    else:
        raise NotImplementedError
    


# The SAM based on automatic mask generation, without bbox prompting
def get_sam_segmentation_dense(
    variant:str, model: Any, image: np.ndarray
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    '''
    The SAM based on automatic mask generation, without bbox prompting
    
    Args:
        model: The mask generator or the YOLO model
        image: )H, W, 3), in RGB color space, in range [0, 255]
        
    Returns:
        mask: (N, H, W)
        xyxy: (N, 4)
        conf: (N,)
    '''
    if variant == "sam":
        results = model.generate(image)
        mask = []
        xyxy = []
        conf = []
        for r in results:
            mask.append(r["segmentation"])
            r_xyxy = r["bbox"].copy()
            # Convert from xyhw format to xyxy format
            r_xyxy[2] += r_xyxy[0]
            r_xyxy[3] += r_xyxy[1]
            xyxy.append(r_xyxy)
            conf.append(r["predicted_iou"])
        mask = np.array(mask)
        xyxy = np.array(xyxy)
        conf = np.array(conf)
        return mask, xyxy, conf
    elif variant == "fastsam":
        # The arguments are directly copied from the GSA repo
        results = model(
            image,
            imgsz=1024,
            device="cuda",
            retina_masks=True,
            iou=0.9,
            conf=0.4,
            max_det=100,
        )
        raise NotImplementedError
    else:
        raise NotImplementedError


def get_sam_mask_generator(variant:str, device: str | int) -> SamAutomaticMaskGenerator:
    if variant == "sam":
        sam = sam_model_registry[SAM_ENCODER_VERSION](checkpoint=SAM_CHECKPOINT_PATH)
        sam.to(device)
        mask_generator = SamAutomaticMaskGenerator(
            model=sam,
            points_per_side=12,
            points_per_batch=144,
            pred_iou_thresh=0.88,
            stability_score_thresh=0.95,
            crop_n_layers=0,
            min_mask_region_area=100,
        )
        return mask_generator
    elif variant == "fastsam":
        raise NotImplementedError
        # from ultralytics import YOLO
        # from FastSAM.tools import *
        # FASTSAM_CHECKPOINT_PATH = os.path.join(GSA_PATH, "./EfficientSAM/FastSAM-x.pt")
        # model = YOLO(args.model_path)
        # return model
    else:
        raise NotImplementedError


def process_tag_classes(text_prompt:str, add_classes:List[str]=[], remove_classes:List[str]=[]) -> list[str]:
    '''
    Convert a text prompt from Tag2Text to a list of classes. 
    '''
    classes = text_prompt.split(',')
    classes = [obj_class.strip() for obj_class in classes]
    classes = [obj_class for obj_class in classes if obj_class != '']
    
    for c in add_classes:
        if c not in classes:
            classes.append(c)
    
    for c in remove_classes:
        classes = [obj_class for obj_class in classes if c not in obj_class.lower()]
    
    return classes


def process_ai2thor_classes(classes: List[str], add_classes:List[str]=[], remove_classes:List[str]=[]) -> List[str]:
    '''
    Some pre-processing on AI2Thor objectTypes in a scene
    '''
    classes = list(set(classes))
    
    for c in add_classes:
        classes.append(c)
        
    for c in remove_classes:
        classes = [obj_class for obj_class in classes if c not in obj_class.lower()]

    # Split the element in classes by captical letters
    classes = [obj_class.replace("TV", "Tv") for obj_class in classes]
    classes = [re.findall('[A-Z][^A-Z]*', obj_class) for obj_class in classes]
    # Join the elements in classes by space
    classes = [" ".join(obj_class) for obj_class in classes]
    
    return classes
    
    
def main(args: argparse.Namespace):
    ### Initialize the Grounding DINO model ###
    grounding_dino_model = Model(
        model_config_path=GROUNDING_DINO_CONFIG_PATH, 
        model_checkpoint_path=GROUNDING_DINO_CHECKPOINT_PATH, 
        device=args.device
    )

    ### Initialize the SAM model ###
    if args.class_set == "none":
        mask_generator = get_sam_mask_generator(args.sam_variant, args.device)
    else:
        sam_predictor = get_sam_predictor(args.sam_variant, args.device)
    
    ###
    # Initialize the CLIP model
    clip_model, _, clip_preprocess = open_clip.create_model_and_transforms(
        "ViT-H-14", "laion2b_s32b_b79k"
    )
    clip_model = clip_model.to(args.device)
    clip_tokenizer = open_clip.get_tokenizer("ViT-H-14")
    
    # Initialize the dataset
    dataset = get_dataset(
        dataconfig=args.dataset_config,
        start=args.start,
        end=args.end,
        stride=args.stride,
        basedir=args.dataset_root,
        sequence=args.scene_id,
        desired_height=args.desired_height,
        desired_width=args.desired_width,
        device="cpu",
        dtype=torch.float,
    )

    global_classes = set()
    
    if args.class_set == "scene":
        # Load the object meta information
        obj_meta_path = args.dataset_root / args.scene_id / "obj_meta.json"
        with open(obj_meta_path, "r") as f:
            obj_meta = json.load(f)
        # Get a list of object classes in the scene
        classes = process_ai2thor_classes(
            [obj["objectType"] for obj in obj_meta],
            add_classes=[],
            remove_classes=['wall', 'floor', 'room', 'ceiling']
        )
    elif args.class_set == "generic":
        classes = FOREGROUND_GENERIC_CLASSES
    elif args.class_set == "minimal":
        classes = FOREGROUND_MINIMAL_CLASSES
    elif args.class_set in ["tag2text", "ram"]:
        ### Initialize the Tag2Text or RAM model ###
        
        if args.class_set == "tag2text":
            # The class set will be computed by tag2text on each image
            # filter out attributes and action categories which are difficult to grounding
            delete_tag_index = []
            for i in range(3012, 3429):
                delete_tag_index.append(i)

            specified_tags='None'
            # load model
            tagging_model = tag2text.tag2text_caption(pretrained=TAG2TEXT_CHECKPOINT_PATH,
                                                    image_size=384,
                                                    vit='swin_b',
                                                    delete_tag_index=delete_tag_index)
            # threshold for tagging
            # we reduce the threshold to obtain more tags
            tagging_model.threshold = 0.64 
        elif args.class_set == "ram":
            tagging_model = tag2text.ram(pretrained=RAM_CHECKPOINT_PATH,
                                         image_size=384,
                                         vit='swin_l')
            
        tagging_model = tagging_model.eval().to(args.device)
        
        # initialize Tag2Text
        tagging_transform = TS.Compose([
            TS.Resize((384, 384)),
            TS.ToTensor(), 
            TS.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
        ])
        
        classes = None
    elif args.class_set == "none":
        classes = ['item']
    else:
        raise ValueError("Unknown args.class_set: ", args.class_set)

    if args.class_set not in ["ram", "tag2text"]:
        print("There are total", len(classes), "classes to detect. ")
    elif args.class_set == "none":
        print("Skipping tagging and detection models. ")
    else:
        print(f"{args.class_set} will be used to detect classes. ")
        
    save_name = f"{args.class_set}"
    if args.sam_variant != "sam": # For backward compatibility
        save_name += f"_{args.sam_variant}"
    if args.exp_suffix:
        save_name += f"_{args.exp_suffix}"
    
    if args.save_video:
        video_save_path = args.dataset_root / args.scene_id / f"gsa_vis_{save_name}.mp4"
        frames = []
    
    for idx in trange(len(dataset)):
        ### Relevant paths and load image ###
        color_path = dataset.color_paths[idx]

        color_path = Path(color_path)
        
        vis_save_path = color_path.parent.parent / f"gsa_vis_{save_name}" / color_path.name
        detections_save_path = color_path.parent.parent / f"gsa_detections_{save_name}" / color_path.name
        detections_save_path = detections_save_path.with_suffix(".pkl.gz")
        
        os.makedirs(os.path.dirname(vis_save_path), exist_ok=True)
        os.makedirs(os.path.dirname(detections_save_path), exist_ok=True)
        
        # opencv can't read Path objects... sigh...
        color_path = str(color_path)
        vis_save_path = str(vis_save_path)
        detections_save_path = str(detections_save_path)
        
        image = cv2.imread(color_path) # This will in BGR color space
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # Convert to RGB color space
        image_pil = Image.fromarray(image_rgb)
        
        ### Tag2Text ###
        if args.class_set in ["ram", "tag2text"]:
            raw_image = image_pil.resize((384, 384))
            raw_image = tagging_transform(raw_image).unsqueeze(0).to(args.device)
            
            if args.class_set == "ram":
                res = inference_ram.inference(raw_image , tagging_model)
                caption="NA"
            elif args.class_set == "tag2text":
                res = inference_tag2text.inference(raw_image , tagging_model, specified_tags)
                caption=res[2]

            # Currently ", " is better for detecting single tags
            # while ". " is a little worse in some case
            text_prompt=res[0].replace(' |', ',')
            
            # Add "other item" to capture objects not in the tag2text captions. 
            # Remove "xxx room", otherwise it will simply include the entire image
            # Also hide "wall" and "floor" for now...
            add_classes = ["other item"]
            remove_classes = [
                "room", "kitchen", "office", "house", "home", "building", "corner",
                "shadow", "carpet", "photo", "shade", "stall", "space", "aquarium", 
                "apartment", "image", "city", "blue", "skylight", "hallway", 
                "bureau", "modern", "salon", "doorway", "wall lamp"
            ]
            bg_classes = ["wall", "floor", "ceiling"]

            if args.add_bg_classes:
                add_classes += bg_classes
            else:
                remove_classes += bg_classes

            classes = process_tag_classes(
                text_prompt, 
                add_classes = add_classes,
                remove_classes = remove_classes,
            )
            
        # add classes to global classes
        global_classes.update(classes)
        
        if args.accumu_classes:
            # Use all the classes that have been seen so far
            classes = list(global_classes)
            
        ### Detection and segmentation ###
        if args.class_set == "none":
            # Directly use SAM in dense sampling mode to get segmentation
            mask, xyxy, conf = get_sam_segmentation_dense(
                args.sam_variant, mask_generator, image_rgb)
            detections = sv.Detections(
                xyxy=xyxy,
                confidence=conf,
                class_id=np.zeros_like(conf).astype(int),
                mask=mask,
            )
            image_crops, image_feats, text_feats = compute_clip_features(
                image_rgb, detections, clip_model, clip_preprocess, clip_tokenizer, classes, args.device)

            ### Visualize results ###
            annotated_image, labels = vis_result_fast(
                image, detections, classes, instance_random_color=True)
            
            cv2.imwrite(vis_save_path, annotated_image)
        else:
            # Using GroundingDINO to detect and SAM to segment
            detections = grounding_dino_model.predict_with_classes(
                image=image, # This function expects a BGR image...
                classes=classes,
                box_threshold=args.box_threshold,
                text_threshold=args.text_threshold,
            )
            
            if len(detections.class_id) > 0:
                ### Non-maximum suppression ###
                # print(f"Before NMS: {len(detections.xyxy)} boxes")
                nms_idx = torchvision.ops.nms(
                    torch.from_numpy(detections.xyxy), 
                    torch.from_numpy(detections.confidence), 
                    args.nms_threshold
                ).numpy().tolist()
                # print(f"After NMS: {len(detections.xyxy)} boxes")

                detections.xyxy = detections.xyxy[nms_idx]
                detections.confidence = detections.confidence[nms_idx]
                detections.class_id = detections.class_id[nms_idx]
                
                # Somehow some detections will have class_id=-1, remove them
                valid_idx = detections.class_id != -1
                detections.xyxy = detections.xyxy[valid_idx]
                detections.confidence = detections.confidence[valid_idx]
                detections.class_id = detections.class_id[valid_idx]
                
                ### Segment Anything ###
                detections.mask = get_sam_segmentation_from_xyxy(
                    sam_predictor=sam_predictor,
                    image=image_rgb,
                    xyxy=detections.xyxy
                )

                # Compute and save the clip features of detections  
                image_crops, image_feats, text_feats = compute_clip_features(
                    image_rgb, detections, clip_model, clip_preprocess, clip_tokenizer, classes, args.device)
            else:
                image_crops, image_feats, text_feats = [], [], []
            
            ### Visualize results ###
            annotated_image, labels = vis_result_fast(image, detections, classes)
            
            # save the annotated grounded-sam image
            if args.class_set in ["ram", "tag2text"] and args.use_slow_vis:
                annotated_image_caption = vis_result_slow_caption(
                    image_rgb, detections.mask, detections.xyxy, labels, caption, text_prompt)
                Image.fromarray(annotated_image_caption).save(vis_save_path)
            else:
                cv2.imwrite(vis_save_path, annotated_image)
        
        if args.save_video:
            frames.append(annotated_image)
        
        # Convert the detections to a dict. The elements are in np.array
        results = {
            "xyxy": detections.xyxy,
            "confidence": detections.confidence,
            "class_id": detections.class_id,
            "mask": detections.mask,
            "classes": classes,
            "image_crops": image_crops,
            "image_feats": image_feats,
            "text_feats": text_feats,
        }
        
        if args.class_set in ["ram", "tag2text"]:
            results["tagging_caption"] = caption
            results["tagging_text_prompt"] = text_prompt
        
        # save the detections using pickle
        # Here we use gzip to compress the file, which could reduce the file size by 500x
        with gzip.open(detections_save_path, "wb") as f:
            pickle.dump(results, f)
    
    # save global classes
    with open(args.dataset_root / args.scene_id / f"gsa_classes_{save_name}.json", "w") as f:
        json.dump(list(global_classes), f)
            
    if args.save_video:
        imageio.mimsave(video_save_path, frames, fps=10)
        print(f"Video saved to {video_save_path}")
        

if __name__ == "__main__":
    parser = get_parser()
    args = parser.parse_args()
    main(args)
conceptgraph/scripts/generate_gsa_results.py
concept-graphs/concept-graphs
import copy
import os
import pickle
import gzip
import argparse

import matplotlib
import numpy as np
import open3d as o3d
import torch
import torch.nn.functional as F
import open_clip
from gradslam.structures.pointclouds import Pointclouds

from conceptgraph.slam.slam_classes import MapObjectList

from conceptgraph.slam.utils import filter_objects, merge_objects

def get_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument("--result_path", type=str, required=True)
    parser.add_argument("--rgb_pcd_path", type=str, default=None)
    
    # To inspect the results of merge_overlap_objects
    # This is mainly to quickly try out different thresholds
    parser.add_argument("--merge_overlap_thresh", type=float, default=-1)
    parser.add_argument("--merge_visual_sim_thresh", type=float, default=-1)
    parser.add_argument("--merge_text_sim_thresh", type=float, default=-1)
    parser.add_argument("--obj_min_points", type=int, default=0)
    parser.add_argument("--obj_min_detections", type=int, default=0)
    
    return parser

if __name__ == "__main__":
    parser = get_parser()
    args = parser.parse_args()

    result_path = args.result_path

    with gzip.open(result_path, "rb") as f:
        results = pickle.load(f)
        
    objects = MapObjectList()
    objects.load_serializable(results['objects'])
    
    # Run the post-processing filtering and merging in instructed to do so
    cfg = copy.deepcopy(results['cfg'])
    cfg.obj_min_points = args.obj_min_points
    cfg.obj_min_detections = args.obj_min_detections
    cfg.merge_overlap_thresh = args.merge_overlap_thresh
    cfg.merge_visual_sim_thresh = args.merge_visual_sim_thresh
    cfg.merge_text_sim_thresh = args.merge_text_sim_thresh
    objects = filter_objects(cfg, objects)
    objects = merge_objects(cfg, objects)
    
    updated_results = {
        'objects': objects.to_serializable(),
        'cfg': results['cfg'],
        'class_names': results['class_names'],
        'class_colors': results['class_colors'],
    }
    
    if 'bg_objects' in results and results['bg_objects'] is not None:
        updated_results['bg_objects'] = results['bg_objects']
    
    save_path = result_path[:-7] + "_post.pkl.gz"
    
    with gzip.open(save_path, "wb") as f:
        pickle.dump(updated_results, f)
    print(f"Saved full point cloud to {save_path}")
    
conceptgraph/scripts/run_post_filter_merge.py
concept-graphs/concept-graphs
import argparse
import json
import os
from pathlib import Path
import numpy as np
import torch
import torch.nn.functional as F
from tqdm import trange

import open3d as o3d

from conceptgraph.dataset.datasets_common import get_dataset

from gradslam.slam.pointfusion import PointFusion
from gradslam.structures.pointclouds import Pointclouds
from gradslam.structures.rgbdimages import RGBDImages

def get_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--dataset_root", type=Path, required=True,
    )
    parser.add_argument(
        "--dataset_config", type=str, required=True,
        help="This path may need to be changed depending on where you run this script. "
    )
    parser.add_argument("--scene_id", type=str, default="train_3")
    parser.add_argument("--image_height", type=int, default=480)
    parser.add_argument("--image_width", type=int, default=640)
    
    parser.add_argument("--start", type=int, default=0)
    parser.add_argument("--end", type=int, default=-1)
    parser.add_argument("--stride", type=int, default=1)
    
    parser.add_argument("--visualize", action="store_true")
    parser.add_argument("--save_pcd", action="store_true")
    
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--device", type=str, default="cuda")
    
    parser.add_argument("--load_semseg", action="store_true",
                        help="Load GT semantic segmentation and run fusion on them. ")
    
    return parser

def main(args: argparse.Namespace):
    if args.load_semseg:
        load_embeddings = True
        embedding_dir = "embed_semseg"
        semseg_classes = json.load(open(
            args.dataset_root / args.scene_id / "embed_semseg_classes.json", "r"
        ))
        embedding_dim = len(semseg_classes)
    else:
        load_embeddings = False
        embedding_dir = "embeddings"
        embedding_dim = 512
    
    dataset = get_dataset(
        dataconfig=args.dataset_config,
        basedir=args.dataset_root,
        sequence=args.scene_id,
        desired_height=args.image_height,
        desired_width=args.image_width,
        start=args.start,
        end=args.end,
        stride=args.stride,
        load_embeddings=load_embeddings,
        embedding_dir = embedding_dir,
        embedding_dim = embedding_dim,
        relative_pose=False,
    )
    
    slam = PointFusion(
        odom="gt", 
        dsratio=1, 
        device=args.device
    )
    
    frame_cur, frame_prev = None, None
    pointclouds = Pointclouds(
        device=args.device
    )
    
    for idx in trange(len(dataset)):
        if load_embeddings:
            _color, _depth, intrinsics, _pose, _embedding = dataset[idx]
            _embedding = _embedding.unsqueeze(0).half()
            _confidence = torch.ones_like(_embedding)
        else:
            _color, _depth, intrinsics, _pose = dataset[idx]
            _embedding = None
            _confidence = None

        pose_np = _pose.cpu().numpy()

        _pose = torch.from_numpy(pose_np).to(_pose.device).to(_pose.dtype)
        
        frame_cur = RGBDImages(
            _color.unsqueeze(0).unsqueeze(0),
            _depth.unsqueeze(0).unsqueeze(0),
            intrinsics.unsqueeze(0).unsqueeze(0),
            _pose.unsqueeze(0).unsqueeze(0),
            embeddings = _embedding,
            confidence_image = _confidence,
        )
        
        pointclouds, _ = slam.step(pointclouds, frame_cur, frame_prev)
        # frame_prev = frame_cur # Keep it None when we use the gt odom
        torch.cuda.empty_cache()
        
    dir_to_save_map = os.path.join(args.dataset_root, args.scene_id, "rgb_cloud")
    print(f"Saving the map to {dir_to_save_map}")
    os.makedirs(dir_to_save_map, exist_ok=True)
    pointclouds.save_to_h5(dir_to_save_map)

    # Set the filename for the PCD file
    pcd_file_path = os.path.join(dir_to_save_map, "pointcloud.pcd")
    pcd = pointclouds.open3d(0)

    if args.save_pcd:
        o3d.io.write_point_cloud(pcd_file_path, pcd)  # Saving as PCD

    if args.visualize:
        o3d.visualization.draw_geometries([pcd])


if __name__ == "__main__":
    args = get_parser().parse_args()
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    
    main(args)
conceptgraph/scripts/run_slam_rgb.py
concept-graphs/concept-graphs

import os, struct
import numpy as np
import zlib
import imageio
import cv2
import png

COMPRESSION_TYPE_COLOR = {-1:'unknown', 0:'raw', 1:'png', 2:'jpeg'}
COMPRESSION_TYPE_DEPTH = {-1:'unknown', 0:'raw_ushort', 1:'zlib_ushort', 2:'occi_ushort'}

class RGBDFrame():

  def load(self, file_handle):
    self.camera_to_world = np.asarray(struct.unpack('f'*16, file_handle.read(16*4)), dtype=np.float32).reshape(4, 4)
    self.timestamp_color = struct.unpack('Q', file_handle.read(8))[0]
    self.timestamp_depth = struct.unpack('Q', file_handle.read(8))[0]
    self.color_size_bytes = struct.unpack('Q', file_handle.read(8))[0]
    self.depth_size_bytes = struct.unpack('Q', file_handle.read(8))[0]
    self.color_data = b''.join(struct.unpack('c'*self.color_size_bytes, file_handle.read(self.color_size_bytes)))
    self.depth_data = b''.join(struct.unpack('c'*self.depth_size_bytes, file_handle.read(self.depth_size_bytes)))


  def decompress_depth(self, compression_type):
    if compression_type == 'zlib_ushort':
       return self.decompress_depth_zlib()
    else:
       raise ValueError("Invalid compression type")


  def decompress_depth_zlib(self):
    return zlib.decompress(self.depth_data)


  def decompress_color(self, compression_type):
    if compression_type == 'jpeg':
       return self.decompress_color_jpeg()
    else:
       raise ValueError("Invalid compression type")


  def decompress_color_jpeg(self):
    return imageio.imread(self.color_data)


class SensorData:

  def __init__(self, filename):
    self.version = 4
    self.load(filename)


  def load(self, filename):
    with open(filename, 'rb') as f:
      version = struct.unpack('I', f.read(4))[0]
      assert self.version == version
      strlen = struct.unpack('Q', f.read(8))[0]
      self.sensor_name = b''.join(struct.unpack('c'*strlen, f.read(strlen)))
      self.intrinsic_color = np.asarray(struct.unpack('f'*16, f.read(16*4)), dtype=np.float32).reshape(4, 4)
      self.extrinsic_color = np.asarray(struct.unpack('f'*16, f.read(16*4)), dtype=np.float32).reshape(4, 4)
      self.intrinsic_depth = np.asarray(struct.unpack('f'*16, f.read(16*4)), dtype=np.float32).reshape(4, 4)
      self.extrinsic_depth = np.asarray(struct.unpack('f'*16, f.read(16*4)), dtype=np.float32).reshape(4, 4)
      self.color_compression_type = COMPRESSION_TYPE_COLOR[struct.unpack('i', f.read(4))[0]]
      self.depth_compression_type = COMPRESSION_TYPE_DEPTH[struct.unpack('i', f.read(4))[0]]
      self.color_width = struct.unpack('I', f.read(4))[0]
      self.color_height =  struct.unpack('I', f.read(4))[0]
      self.depth_width = struct.unpack('I', f.read(4))[0]
      self.depth_height =  struct.unpack('I', f.read(4))[0]
      self.depth_shift =  struct.unpack('f', f.read(4))[0]
      num_frames =  struct.unpack('Q', f.read(8))[0]
      self.frames = []
      for i in range(num_frames):
        frame = RGBDFrame()
        frame.load(f)
        self.frames.append(frame)


  def export_depth_images(self, output_path, image_size=None, frame_skip=1):
    if not os.path.exists(output_path):
      os.makedirs(output_path)
    print('exporting', len(self.frames)//frame_skip, ' depth frames to', output_path)
    for f in range(0, len(self.frames), frame_skip):
      depth_data = self.frames[f].decompress_depth(self.depth_compression_type)
      depth = np.fromstring(depth_data, dtype=np.uint16).reshape(self.depth_height, self.depth_width)
      if image_size is not None:
        depth = cv2.resize(depth, (image_size[1], image_size[0]), interpolation=cv2.INTER_NEAREST)
      #imageio.imwrite(os.path.join(output_path, str(f) + '.png'), depth)
      with open(os.path.join(output_path, str(f) + '.png'), 'wb') as f: # write 16-bit
        writer = png.Writer(width=depth.shape[1], height=depth.shape[0], bitdepth=16)
        depth = depth.reshape(-1, depth.shape[1]).tolist()
        writer.write(f, depth)

  def export_color_images(self, output_path, image_size=None, frame_skip=1):
    if not os.path.exists(output_path):
      os.makedirs(output_path)
    print('exporting', len(self.frames)//frame_skip, 'color frames to', output_path)
    for f in range(0, len(self.frames), frame_skip):
      color = self.frames[f].decompress_color(self.color_compression_type)
      if image_size is not None:
        color = cv2.resize(color, (image_size[1], image_size[0]), interpolation=cv2.INTER_NEAREST)
      imageio.imwrite(os.path.join(output_path, str(f) + '.jpg'), color)


  def save_mat_to_file(self, matrix, filename):
    with open(filename, 'w') as f:
      for line in matrix:
        np.savetxt(f, line[np.newaxis], fmt='%f')


  def export_poses(self, output_path, frame_skip=1):
    if not os.path.exists(output_path):
      os.makedirs(output_path)
    print('exporting', len(self.frames)//frame_skip, 'camera poses to', output_path)
    for f in range(0, len(self.frames), frame_skip):
      self.save_mat_to_file(self.frames[f].camera_to_world, os.path.join(output_path, str(f) + '.txt'))


  def export_intrinsics(self, output_path):
    if not os.path.exists(output_path):
      os.makedirs(output_path)
    print('exporting camera intrinsics to', output_path)
    self.save_mat_to_file(self.intrinsic_color, os.path.join(output_path, 'intrinsic_color.txt'))
    self.save_mat_to_file(self.extrinsic_color, os.path.join(output_path, 'extrinsic_color.txt'))
    self.save_mat_to_file(self.intrinsic_depth, os.path.join(output_path, 'intrinsic_depth.txt'))
    self.save_mat_to_file(self.extrinsic_depth, os.path.join(output_path, 'extrinsic_depth.txt'))

conceptgraph/scripts/scannet_process/SensorData.py
concept-graphs/concept-graphs
import copy
import gc
import os
import pickle
import gzip
import argparse

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import rich

import open3d as o3d
from PIL import Image
import torch
import torch.nn.functional as F
import open_clip
from gradslam.structures.pointclouds import Pointclouds

from conceptgraph.slam.slam_classes import MapObjectList

from conceptgraph.llava.llava_model import LLaVaChat
from conceptgraph.utils.image import crop_image_pil
try: LLAVA_CKPT_PATH = os.environ["LLAVA_CKPT_PATH"]
except KeyError: raise ValueError("Please set the environment variable LLAVA_CKPT_PATH to the path of the LLaVA checkpoint")

def get_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument("--result_path", type=str, required=True)
    parser.add_argument("--rgb_pcd_path", type=str, default=None)
    return parser

if __name__ == "__main__":
    parser = get_parser()
    args = parser.parse_args()
    
    # rich console for pretty printing
    # console = rich.console.Console()

    result_path = args.result_path
    rgb_pcd_path = args.rgb_pcd_path

    if rgb_pcd_path is not None:        
        pointclouds = Pointclouds.load_pointcloud_from_h5(rgb_pcd_path)
        global_pcd = pointclouds.open3d(0, include_colors=True)

    with gzip.open(result_path, "rb") as f:
        results = pickle.load(f)
        
    class_names = results['class_names']
    class_colors = results['class_colors']
        
    objects = MapObjectList()
    objects.load_serializable(results['objects'])
    
    cmap = matplotlib.colormaps.get_cmap("turbo")
    
    pcds = copy.deepcopy(objects.get_values("pcd"))
    bboxes = copy.deepcopy(objects.get_values("bbox"))
    
    # Get the color for each object when colored by their class
    object_classes = []
    for i in range(len(objects)):
        obj = objects[i]
        pcd = pcds[i]
        obj_classes = np.asarray(obj['class_id'])
        # Get the most common class for this object as the class
        values, counts = np.unique(obj_classes, return_counts=True)
        obj_class = values[np.argmax(counts)]
        object_classes.append(obj_class)
        
    # Color the object by their classes
    for i in range(len(objects)):
        pcd = pcds[i]
        obj_class = object_classes[i]
        pcd.colors = o3d.utility.Vector3dVector(
            np.tile(
                class_colors[str(obj_class)],
                (len(pcd.points), 1)
            )
        )
        
    # Build a global point cloud and record the object indices
    global_pcd = o3d.geometry.PointCloud()
    global_pcd.points = o3d.utility.Vector3dVector(np.concatenate([pcd.points for pcd in pcds]))
    global_pcd.colors = o3d.utility.Vector3dVector(np.concatenate([pcd.colors for pcd in pcds]))
    object_indices = np.concatenate([np.full(len(pcd.points), i) for i, pcd in enumerate(pcds)])
    
    # chat = LLaVaChat(LLAVA_CKPT_PATH, "multimodal", 1)
    chat = LLaVaChat(LLAVA_CKPT_PATH, "default", 1)
    print("LLaVA chat initialized...")
    
    while True:
        # Set the title of the window
        vis = o3d.visualization.VisualizerWithEditing()

        if result_path is not None:
            vis.create_window(window_name=f'Open3D - {os.path.basename(result_path)}', width=1280, height=720)
        else:
            vis.create_window(window_name=f'Open3D', width=1280, height=720)

        vis.add_geometry(global_pcd)
            
        print("")
        print("==> Please points using [shift + left click]")
        print("   Press [shift + right click] to undo point picking")
        print("   Only the last point picking will be considered")
        print("==> Afther picking points, press q for close the window")
        
        # Render the scene
        vis.run()
        
        vis.destroy_window()
        
        picked_indices = vis.get_picked_points()
        if len(picked_indices) == 0:
            print("No points picked, exiting...")
            exit()
            
        chat.reset()
            
        picked_obj_idx = object_indices[picked_indices[-1]]
        picked_obj = objects[picked_obj_idx]
        print(f"Object index: {picked_obj_idx}")

        picked_obj_classes = np.asarray(picked_obj['class_id'])
        values, counts = np.unique(picked_obj_classes, return_counts=True)
        print("Object class:", end=" ")
        for i in counts.argsort()[::-1]:
            print(f"{class_names[values[i]]}({counts[i]})", end=" ")
        print("")
        picked_obj_class_name = class_names[values[np.argmax(counts)]]
        
        conf = picked_obj['conf']
        conf = picked_obj['n_points']
        conf = np.asarray(conf)
        idx_sort = np.argsort(conf)[::-1]
        
        object_crops = []
        crop_features = []
        for i, idx_det in enumerate(idx_sort[:5]):
            # wrong_path = picked_obj['color_path'][idx_det]
            # correct_prefix = '/home/sacha/data/scan12/'
            # wrong_prefix = '/home/kuwajerw/local_data/azure/liam_lab_w_objects/scan12/'
            # correct_path = wrong_path.replace(wrong_prefix, correct_prefix)

            # image = Image.open(correct_path)

            image = Image.open(picked_obj['color_path'][idx_det])
            mask = picked_obj['mask'][idx_det]
            class_id = picked_obj['class_id'][idx_det]
            xyxy = picked_obj['xyxy'][idx_det]
            class_name = class_names[class_id]
            
            x1, y1, x2, y2 = xyxy
            image_crop = crop_image_pil(image, x1, y1, x2, y2, 10)

            image_tensor = chat.image_processor.preprocess(image_crop, return_tensors="pt")[
                "pixel_values"
            ][0]
            image_feature = chat.encode_image(image_tensor[None, ...].half().cuda())
            
            object_crops.append(image_crop)
            crop_features.append(image_feature)
            
        # Show the crops using plt
        fig, axes = plt.subplots(1, len(object_crops), figsize=(10, 3))
        if len(object_crops) == 1:
            axes = [axes]
        for i, image_crop in enumerate(object_crops):
            axes[i].imshow(image_crop)
            axes[i].set_title(f"{picked_obj_class_name} {i}")
            axes[i].axis("off")
        fig.suptitle("Image crops used for LLaVA")
        plt.show(block=False)
        
        # print("Using the LLaVA feature from the most confident detection...")
        # llava_feat = crop_features[0]
        
        print(f"Using the mean feature from {len(crop_features)} image crops...")
        print(crop_features[0].shape)
        llava_feat = torch.stack(crop_features, dim=0).mean(dim=0)
        print(llava_feat.shape)
        
        query = "What is the central object in this image?"
        outputs = chat(query=query, image_features=llava_feat.half().cuda())
        # console.print("[bold red]User:[/bold red] " + query)
        print("User: " + query)
        # print(outputs)
        print("LLaVA: " + outputs)

        while True:
            query = input("Enter a query ('q' to quit): ")
            if query == "q":
                break
            # console.print("[bold red]User:[/bold red] " + query)
            print("User: " + query)
            outputs = chat(query=query, image_features=None)
            # print(outputs)
            print("LLaVA: " + outputs)

            
        plt.close(fig)
        
        
            
            
        
            
            
conceptgraph/scripts/visualize_cfslam_interact_llava.py
concept-graphs/concept-graphs
import copy
import json
import os
import pickle
import gzip
import argparse

import matplotlib
import numpy as np
import pandas as pd
import open3d as o3d
import torch
import torch.nn.functional as F
import open_clip

import distinctipy

from gradslam.structures.pointclouds import Pointclouds

from conceptgraph.slam.slam_classes import MapObjectList
from conceptgraph.utils.vis import LineMesh
from conceptgraph.slam.utils import filter_objects, merge_objects

def create_ball_mesh(center, radius, color=(0, 1, 0)):
    """
    Create a colored mesh sphere.
    
    Args:
    - center (tuple): (x, y, z) coordinates for the center of the sphere.
    - radius (float): Radius of the sphere.
    - color (tuple): RGB values in the range [0, 1] for the color of the sphere.
    
    Returns:
    - o3d.geometry.TriangleMesh: Colored mesh sphere.
    """
    mesh_sphere = o3d.geometry.TriangleMesh.create_sphere(radius=radius)
    mesh_sphere.translate(center)
    mesh_sphere.paint_uniform_color(color)
    return mesh_sphere

def get_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument("--result_path", type=str, default=None)
    parser.add_argument("--rgb_pcd_path", type=str, default=None)
    parser.add_argument("--edge_file", type=str, default=None)
    
    parser.add_argument("--no_clip", action="store_true", 
                        help="If set, the CLIP model will not init for fast debugging.")
    
    # To inspect the results of merge_overlap_objects
    # This is mainly to quickly try out different thresholds
    parser.add_argument("--merge_overlap_thresh", type=float, default=-1)
    parser.add_argument("--merge_visual_sim_thresh", type=float, default=-1)
    parser.add_argument("--merge_text_sim_thresh", type=float, default=-1)
    parser.add_argument("--obj_min_points", type=int, default=0)
    parser.add_argument("--obj_min_detections", type=int, default=0)
    
    return parser

def load_result(result_path):
    with gzip.open(result_path, "rb") as f:
        results = pickle.load(f)
    
    if isinstance(results, dict):
        objects = MapObjectList()
        objects.load_serializable(results["objects"])
        
        if results['bg_objects'] is None:
            bg_objects = None
        else:
            bg_objects = MapObjectList()
            bg_objects.load_serializable(results["bg_objects"])

        class_colors = results['class_colors']
    elif isinstance(results, list):
        objects = MapObjectList()
        objects.load_serializable(results)

        bg_objects = None
        class_colors = distinctipy.get_colors(len(objects), pastel_factor=0.5)
        class_colors = {str(i): c for i, c in enumerate(class_colors)}
    else:
        raise ValueError("Unknown results type: ", type(results))
        
    return objects, bg_objects, class_colors

def main(args):
    result_path = args.result_path
    rgb_pcd_path = args.rgb_pcd_path
    
    assert not (result_path is None and rgb_pcd_path is None), \
        "Either result_path or rgb_pcd_path must be provided."

    if rgb_pcd_path is not None:        
        pointclouds = Pointclouds.load_pointcloud_from_h5(rgb_pcd_path)
        global_pcd = pointclouds.open3d(0, include_colors=True)
        
        if result_path is None:
            print("Only visualizing the pointcloud...")
            o3d.visualization.draw_geometries([global_pcd])
            exit()
        
    objects, bg_objects, class_colors = load_result(result_path)
    
    if args.edge_file is not None:
        # Load edge files and create meshes for the scene graph
        scene_graph_geometries = []
        with open(args.edge_file, "r") as f:
            edges = json.load(f)
        
        classes = objects.get_most_common_class()
        colors = [class_colors[str(c)] for c in classes]
        obj_centers = []
        for obj, c in zip(objects, colors):
            pcd = obj['pcd']
            bbox = obj['bbox']
            points = np.asarray(pcd.points)
            center = np.mean(points, axis=0)
            extent = bbox.get_max_bound()
            extent = np.linalg.norm(extent)
            # radius = extent ** 0.5 / 25
            radius = 0.10
            obj_centers.append(center)

            # remove the nodes on the ceiling, for better visualization
            ball = create_ball_mesh(center, radius, c)
            scene_graph_geometries.append(ball)
            
        for edge in edges:
            if edge['object_relation'] == "none of these":
                continue
            id1 = edge["object1"]['id']
            id2 = edge["object2"]['id']

            line_mesh = LineMesh(
                points = np.array([obj_centers[id1], obj_centers[id2]]),
                lines = np.array([[0, 1]]),
                colors = [1, 0, 0],
                radius=0.02
            )

            scene_graph_geometries.extend(line_mesh.cylinder_segments)
    
    if not args.no_clip:
        print("Initializing CLIP model...")
        clip_model, _, clip_preprocess = open_clip.create_model_and_transforms("ViT-H-14", "laion2b_s32b_b79k")
        clip_model = clip_model.to("cuda")
        clip_tokenizer = open_clip.get_tokenizer("ViT-H-14")
        print("Done initializing CLIP model.")

    cmap = matplotlib.colormaps.get_cmap("turbo")
    
    if bg_objects is not None:
        indices_bg = np.arange(len(objects), len(objects) + len(bg_objects))
        objects.extend(bg_objects)
        
    # Sub-sample the point cloud for better interactive experience
    for i in range(len(objects)):
        pcd = objects[i]['pcd']
        pcd = pcd.voxel_down_sample(0.05)
        objects[i]['pcd'] = pcd
    
    pcds = copy.deepcopy(objects.get_values("pcd"))
    bboxes = copy.deepcopy(objects.get_values("bbox"))
    
    # Get the color for each object when colored by their class
    object_classes = []
    for i in range(len(objects)):
        obj = objects[i]
        pcd = pcds[i]
        obj_classes = np.asarray(obj['class_id'])
        # Get the most common class for this object as the class
        values, counts = np.unique(obj_classes, return_counts=True)
        obj_class = values[np.argmax(counts)]
        object_classes.append(obj_class)
    
    # Set the title of the window
    vis = o3d.visualization.VisualizerWithKeyCallback()

    if result_path is not None:
        vis.create_window(window_name=f'Open3D - {os.path.basename(result_path)}', width=1280, height=720)
    else:
        vis.create_window(window_name=f'Open3D', width=1280, height=720)

    # Add geometry to the scene
    for geometry in pcds + bboxes:
        vis.add_geometry(geometry)
        
    main.show_bg_pcd = True
    def toggle_bg_pcd(vis):
        if bg_objects is None:
            print("No background objects found.")
            return
        
        for idx in indices_bg:
            if main.show_bg_pcd:
                vis.remove_geometry(pcds[idx], reset_bounding_box=False)
                vis.remove_geometry(bboxes[idx], reset_bounding_box=False)
            else:
                vis.add_geometry(pcds[idx], reset_bounding_box=False)
                vis.add_geometry(bboxes[idx], reset_bounding_box=False)
        
        main.show_bg_pcd = not main.show_bg_pcd
        
    main.show_global_pcd = False
    def toggle_global_pcd(vis):
        if args.rgb_pcd_path is None:
            print("No RGB pcd path provided.")
            return
        
        if main.show_global_pcd:
            vis.remove_geometry(global_pcd, reset_bounding_box=False)
        else:
            vis.add_geometry(global_pcd, reset_bounding_box=False)
        
        main.show_global_pcd = not main.show_global_pcd
        
    main.show_scene_graph = False
    def toggle_scene_graph(vis):
        if args.edge_file is None:
            print("No edge file provided.")
            return
        
        if main.show_scene_graph:
            for geometry in scene_graph_geometries:
                vis.remove_geometry(geometry, reset_bounding_box=False)
        else:
            for geometry in scene_graph_geometries:
                vis.add_geometry(geometry, reset_bounding_box=False)
        
        main.show_scene_graph = not main.show_scene_graph
        
    def color_by_class(vis):
        for i in range(len(objects)):
            pcd = pcds[i]
            obj_class = object_classes[i]
            pcd.colors = o3d.utility.Vector3dVector(
                np.tile(
                    class_colors[str(obj_class)],
                    (len(pcd.points), 1)
                )
            )

        for pcd in pcds:
            vis.update_geometry(pcd)
            
    def color_by_rgb(vis):
        for i in range(len(pcds)):
            pcd = pcds[i]
            pcd.colors = objects[i]['pcd'].colors
        
        for pcd in pcds:
            vis.update_geometry(pcd)
            
    def color_by_instance(vis):
        instance_colors = cmap(np.linspace(0, 1, len(pcds)))
        for i in range(len(pcds)):
            pcd = pcds[i]
            pcd.colors = o3d.utility.Vector3dVector(
                np.tile(
                    instance_colors[i, :3],
                    (len(pcd.points), 1)
                )
            )
            
        for pcd in pcds:
            vis.update_geometry(pcd)
        
    def color_by_clip_sim(vis):
        if args.no_clip:
            print("CLIP model is not initialized.")
            return

        text_query = input("Enter your query: ")
        text_queries = [text_query]
        
        text_queries_tokenized = clip_tokenizer(text_queries).to("cuda")
        text_query_ft = clip_model.encode_text(text_queries_tokenized)
        text_query_ft = text_query_ft / text_query_ft.norm(dim=-1, keepdim=True)
        text_query_ft = text_query_ft.squeeze()
        
        # similarities = objects.compute_similarities(text_query_ft)
        objects_clip_fts = objects.get_stacked_values_torch("clip_ft")
        objects_clip_fts = objects_clip_fts.to("cuda")
        similarities = F.cosine_similarity(
            text_query_ft.unsqueeze(0), objects_clip_fts, dim=-1
        )
        max_value = similarities.max()
        min_value = similarities.min()
        normalized_similarities = (similarities - min_value) / (max_value - min_value)
        probs = F.softmax(similarities, dim=0)
        max_prob_idx = torch.argmax(probs)
        similarity_colors = cmap(normalized_similarities.detach().cpu().numpy())[..., :3]
        
        for i in range(len(objects)):
            pcd = pcds[i]
            map_colors = np.asarray(pcd.colors)
            pcd.colors = o3d.utility.Vector3dVector(
                np.tile(
                    [
                        similarity_colors[i, 0].item(),
                        similarity_colors[i, 1].item(),
                        similarity_colors[i, 2].item()
                    ], 
                    (len(pcd.points), 1)
                )
            )

        for pcd in pcds:
            vis.update_geometry(pcd)
            
    def save_view_params(vis):
        param = vis.get_view_control().convert_to_pinhole_camera_parameters()
        o3d.io.write_pinhole_camera_parameters("temp.json", param)
        
    vis.register_key_callback(ord("B"), toggle_bg_pcd)
    vis.register_key_callback(ord("S"), toggle_global_pcd)
    vis.register_key_callback(ord("C"), color_by_class)
    vis.register_key_callback(ord("R"), color_by_rgb)
    vis.register_key_callback(ord("F"), color_by_clip_sim)
    vis.register_key_callback(ord("I"), color_by_instance)
    vis.register_key_callback(ord("V"), save_view_params)
    vis.register_key_callback(ord("G"), toggle_scene_graph)
    
    # Render the scene
    vis.run()
    
if __name__ == "__main__":
    parser = get_parser()
    args = parser.parse_args()
    main(args)
conceptgraph/scripts/visualize_cfslam_results.py
concept-graphs/concept-graphs
'''
The script is used to model Grounded SAM detections in 3D, it assumes the tag2text classes are avaialable. It also assumes the dataset has Clip features saved for each object/mask.
'''

# Standard library imports
import copy
from datetime import datetime
import os
from pathlib import Path
import gzip
import pickle

# Related third party imports
from PIL import Image
import cv2
import imageio
import numpy as np
import open3d as o3d
import torch
import torch.nn.functional as F
from tqdm import tqdm, trange

import hydra
import omegaconf
from omegaconf import DictConfig

# Local application/library specific imports
from conceptgraph.dataset.datasets_common import get_dataset
from conceptgraph.utils.vis import OnlineObjectRenderer
from conceptgraph.utils.ious import (
    compute_2d_box_contained_batch
)
from conceptgraph.utils.general_utils import to_tensor

from conceptgraph.slam.slam_classes import MapObjectList, DetectionList
from conceptgraph.slam.utils import (
    create_or_load_colors,
    merge_obj2_into_obj1, 
    denoise_objects,
    filter_objects,
    merge_objects, 
    gobs_to_detection_list,
)
from conceptgraph.slam.mapping import (
    compute_spatial_similarities,
    compute_visual_similarities,
    aggregate_similarities,
    merge_detections_to_objects
)

BG_CLASSES = ["wall", "floor", "ceiling"]

# Disable torch gradient computation
torch.set_grad_enabled(False)

def compute_match_batch(cfg, spatial_sim: torch.Tensor, visual_sim: torch.Tensor) -> torch.Tensor:
    '''
    Compute object association based on spatial and visual similarities
    
    Args:
        spatial_sim: a MxN tensor of spatial similarities
        visual_sim: a MxN tensor of visual similarities
    Returns:
        A MxN tensor of binary values, indicating whether a detection is associate with an object. 
        Each row has at most one 1, indicating one detection can be associated with at most one existing object.
        One existing object can receive multiple new detections
    '''
    assign_mat = torch.zeros_like(spatial_sim)
    if cfg.match_method == "sim_sum":
        sims = (1 + cfg.phys_bias) * spatial_sim + (1 - cfg.phys_bias) * visual_sim # (M, N)
        row_max, row_argmax = torch.max(sims, dim=1) # (M,), (M,)
        for i in row_max.argsort(descending=True):
            if row_max[i] > cfg.sim_threshold:
                assign_mat[i, row_argmax[i]] = 1
            else:
                break
    else:
        raise ValueError(f"Unknown matching method: {cfg.match_method}")
    
    return assign_mat

def prepare_objects_save_vis(objects: MapObjectList, downsample_size: float=0.025):
    objects_to_save = copy.deepcopy(objects)
            
    # Downsample the point cloud
    for i in range(len(objects_to_save)):
        objects_to_save[i]['pcd'] = objects_to_save[i]['pcd'].voxel_down_sample(downsample_size)

    # Remove unnecessary keys
    for i in range(len(objects_to_save)):
        for k in list(objects_to_save[i].keys()):
            if k not in [
                'pcd', 'bbox', 'clip_ft', 'text_ft', 'class_id', 'num_detections', 'inst_color'
            ]:
                del objects_to_save[i][k]
                
    return objects_to_save.to_serializable()
    
def process_cfg(cfg: DictConfig):
    cfg.dataset_root = Path(cfg.dataset_root)
    cfg.dataset_config = Path(cfg.dataset_config)
    
    if cfg.dataset_config.name != "multiscan.yaml":
        # For datasets whose depth and RGB have the same resolution
        # Set the desired image heights and width from the dataset config
        dataset_cfg = omegaconf.OmegaConf.load(cfg.dataset_config)
        if cfg.image_height is None:
            cfg.image_height = dataset_cfg.camera_params.image_height
        if cfg.image_width is None:
            cfg.image_width = dataset_cfg.camera_params.image_width
        print(f"Setting image height and width to {cfg.image_height} x {cfg.image_width}")
    else:
        # For dataset whose depth and RGB have different resolutions
        assert cfg.image_height is not None and cfg.image_width is not None, \
            "For multiscan dataset, image height and width must be specified"

    return cfg
    
@hydra.main(version_base=None, config_path="../configs/slam_pipeline", config_name="base")
def main(cfg : DictConfig):
    cfg = process_cfg(cfg)
    
    # Initialize the dataset
    dataset = get_dataset(
        dataconfig=cfg.dataset_config,
        start=cfg.start,
        end=cfg.end,
        stride=cfg.stride,
        basedir=cfg.dataset_root,
        sequence=cfg.scene_id,
        desired_height=cfg.image_height,
        desired_width=cfg.image_width,
        device="cpu",
        dtype=torch.float,
    )
    # cam_K = dataset.get_cam_K()
    
    classes, class_colors = create_or_load_colors(cfg, cfg.color_file_name)

    objects = MapObjectList(device=cfg.device)
    
    if not cfg.skip_bg:
        # Handle the background detection separately 
        # Each class of them are fused into the map as a single object
        bg_objects = {
            c: None for c in BG_CLASSES
        }
    else:
        bg_objects = None
        
    # For visualization
    if cfg.vis_render:
        view_param = o3d.io.read_pinhole_camera_parameters(cfg.render_camera_path)
            
        obj_renderer = OnlineObjectRenderer(
            view_param = view_param,
            base_objects = None, 
            gray_map = False,
        )
        frames = []
        
    if cfg.save_objects_all_frames:
        save_all_folder = cfg.dataset_root \
            / cfg.scene_id / "objects_all_frames" / f"{cfg.gsa_variant}_{cfg.save_suffix}"
        os.makedirs(save_all_folder, exist_ok=True)

    for idx in trange(len(dataset)):
        # get color image
        color_path = dataset.color_paths[idx]
        image_original_pil = Image.open(color_path)

        color_tensor, depth_tensor, intrinsics, *_ = dataset[idx]
        # image_rgb = cv2.cvtColor(cv2.imread(color_path), cv2.COLOR_BGR2RGB)
        # Get the RGB image
        color_np = color_tensor.cpu().numpy() # (H, W, 3)
        image_rgb = (color_np).astype(np.uint8) # (H, W, 3)
        assert image_rgb.max() > 1, "Image is not in range [0, 255]"
        
        # Get the depth image
        depth_tensor = depth_tensor[..., 0]
        depth_array = depth_tensor.cpu().numpy()

        # Get the intrinsics matrix
        cam_K = intrinsics.cpu().numpy()[:3, :3]
        
        # load grounded SAM detections
        gobs = None # stands for grounded SAM observations

        color_path = Path(color_path)
        detections_path = color_path.parent.parent / cfg.detection_folder_name / color_path.name
        detections_path = detections_path.with_suffix(".pkl.gz")
        color_path = str(color_path)
        detections_path = str(detections_path)
        
        with gzip.open(detections_path, "rb") as f:
            gobs = pickle.load(f)

        
        # depth_image = Image.open(depth_path)
        # depth_array = np.array(depth_image) / dataset.png_depth_scale
        # depth_tensor = torch.from_numpy(depth_array).float().to(cfg['device']).T

        # get pose, this is the untrasformed pose.
        unt_pose = dataset.poses[idx]
        unt_pose = unt_pose.cpu().numpy()
        
        # Don't apply any transformation otherwise
        adjusted_pose = unt_pose
            
        # if idx == 71:
        #     fg_detection_list, bg_detection_list = gobs_to_detection_list(
        #         cfg = cfg,
        #         image = image_rgb,
        #         depth_array = depth_array,
        #         cam_K = cam_K,
        #         idx = idx,
        #         gobs = gobs,
        #         trans_pose = adjusted_pose,
        #         class_names = classes,
        #         BG_CLASSES = BG_CLASSES,
        #         color_path = color_path,
        #     )
        #     for det in fg_detection_list:
        #         o3d.visualization.draw_geometries([det['pcd']])
                
        #     exit()
        # else:
        #     continue
        
        fg_detection_list, bg_detection_list = gobs_to_detection_list(
            cfg = cfg,
            image = image_rgb,
            depth_array = depth_array,
            cam_K = cam_K,
            idx = idx,
            gobs = gobs,
            trans_pose = adjusted_pose,
            class_names = classes,
            BG_CLASSES = BG_CLASSES,
            color_path = color_path,
        )
        
        if len(bg_detection_list) > 0:
            for detected_object in bg_detection_list:
                class_name = detected_object['class_name'][0]
                if bg_objects[class_name] is None:
                    bg_objects[class_name] = detected_object
                else:
                    matched_obj = bg_objects[class_name]
                    matched_det = detected_object
                    bg_objects[class_name] = merge_obj2_into_obj1(cfg, matched_obj, matched_det, run_dbscan=False)
            
        if len(fg_detection_list) == 0:
            continue
            
        if cfg.use_contain_number:
            xyxy = fg_detection_list.get_stacked_values_torch('xyxy', 0)
            contain_numbers = compute_2d_box_contained_batch(xyxy, cfg.contain_area_thresh)
            for i in range(len(fg_detection_list)):
                fg_detection_list[i]['contain_number'] = [contain_numbers[i]]
            
        if len(objects) == 0:
            # Add all detections to the map
            for i in range(len(fg_detection_list)):
                objects.append(fg_detection_list[i])

            # Skip the similarity computation 
            continue
                
        spatial_sim = compute_spatial_similarities(cfg, fg_detection_list, objects)
        visual_sim = compute_visual_similarities(cfg, fg_detection_list, objects)
        agg_sim = aggregate_similarities(cfg, spatial_sim, visual_sim)
        
        # Compute the contain numbers for each detection
        if cfg.use_contain_number:
            # Get the contain numbers for all objects
            contain_numbers_objects = torch.Tensor([obj['contain_number'][0] for obj in objects])
            detection_contained = contain_numbers > 0 # (M,)
            object_contained = contain_numbers_objects > 0 # (N,)
            detection_contained = detection_contained.unsqueeze(1) # (M, 1)
            object_contained = object_contained.unsqueeze(0) # (1, N)                

            # Get the non-matching entries, penalize their similarities
            xor = detection_contained ^ object_contained
            agg_sim[xor] = agg_sim[xor] - cfg.contain_mismatch_penalty
        
        # Threshold sims according to cfg. Set to negative infinity if below threshold
        agg_sim[agg_sim < cfg.sim_threshold] = float('-inf')
        
        objects = merge_detections_to_objects(cfg, fg_detection_list, objects, agg_sim)
        
        # Perform post-processing periodically if told so
        if cfg.denoise_interval > 0 and (idx+1) % cfg.denoise_interval == 0:
            objects = denoise_objects(cfg, objects)
        if cfg.filter_interval > 0 and (idx+1) % cfg.filter_interval == 0:
            objects = filter_objects(cfg, objects)
        if cfg.merge_interval > 0 and (idx+1) % cfg.merge_interval == 0:
            objects = merge_objects(cfg, objects)
            
        if cfg.save_objects_all_frames:
            save_all_path = save_all_folder / f"{idx:06d}.pkl.gz"
            objects_to_save = MapObjectList([
                _ for _ in objects if _['num_detections'] >= cfg.obj_min_detections
            ])
            
            objects_to_save = prepare_objects_save_vis(objects_to_save)
            
            if not cfg.skip_bg:
                bg_objects_to_save = MapObjectList([_ for _ in bg_objects.values() if _ is not None])
                bg_objects_to_save = prepare_objects_save_vis(bg_objects_to_save)
            else:
                bg_objects_to_save = None
            
            result = {
                "camera_pose": adjusted_pose,
                "objects": objects_to_save,
                "bg_objects": bg_objects_to_save,
            }
            with gzip.open(save_all_path, 'wb') as f:
                pickle.dump(result, f)
        
        if cfg.vis_render:
            objects_vis = MapObjectList([
                copy.deepcopy(_) for _ in objects if _['num_detections'] >= cfg.obj_min_detections
            ])
            
            if cfg.class_agnostic:
                objects_vis.color_by_instance()
            else:
                objects_vis.color_by_most_common_classes(class_colors)
            
            rendered_image, vis = obj_renderer.step(
                image = image_original_pil,
                gt_pose = adjusted_pose,
                new_objects = objects_vis,
                paint_new_objects=False,
                return_vis_handle = cfg.debug_render,
            )

            if cfg.debug_render:
                vis.run()
                del vis
            
            # Convert to uint8
            if rendered_image is not None:
                rendered_image = (rendered_image * 255).astype(np.uint8)
                frames.append(rendered_image)
            
        # print(
        #     f"Finished image {idx} of {len(dataset)}", 
        #     f"Now we have {len(objects)} objects.",
        #     f"Effective objects {len([_ for _ in objects if _['num_detections'] >= cfg.obj_min_detections])}"
        # )
        
    if bg_objects is not None:
        bg_objects = MapObjectList([_ for _ in bg_objects.values() if _ is not None])
        bg_objects = denoise_objects(cfg, bg_objects)
        
    objects = denoise_objects(cfg, objects)
    
    # Save the full point cloud before post-processing
    if cfg.save_pcd:
        ts = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        results = {
            'objects': objects.to_serializable(),
            'bg_objects': None if bg_objects is None else bg_objects.to_serializable(),
            'cfg': cfg,
            'class_names': classes,
            'class_colors': class_colors,
        }

        pcd_save_path = cfg.dataset_root / \
            cfg.scene_id / 'pcd_saves' / f"full_pcd_{cfg.gsa_variant}_{cfg.save_suffix}.pkl.gz"
        # make the directory if it doesn't exist
        pcd_save_path.parent.mkdir(parents=True, exist_ok=True)
        pcd_save_path = str(pcd_save_path)
        
        with gzip.open(pcd_save_path, "wb") as f:
            pickle.dump(results, f)
        print(f"Saved full point cloud to {pcd_save_path}")
    
    objects = filter_objects(cfg, objects)
    objects = merge_objects(cfg, objects)
    
    # Save again the full point cloud after the post-processing
    if cfg.save_pcd:
        results['objects'] = objects.to_serializable()
        pcd_save_path = pcd_save_path[:-7] + "_post.pkl.gz"
        with gzip.open(pcd_save_path, "wb") as f:
            pickle.dump(results, f)
        print(f"Saved full point cloud after post-processing to {pcd_save_path}")
        
    if cfg.save_objects_all_frames:
        save_meta_path = save_all_folder / f"meta.pkl.gz"
        with gzip.open(save_meta_path, "wb") as f:
            pickle.dump({
                'cfg': cfg,
                'class_names': classes,
                'class_colors': class_colors,
            }, f)
        
    if cfg.vis_render:
        # Still render a frame after the post-processing
        objects_vis = MapObjectList([
            _ for _ in objects if _['num_detections'] >= cfg.obj_min_detections
        ])

        if cfg.class_agnostic:
            objects_vis.color_by_instance()
        else:
            objects_vis.color_by_most_common_classes(class_colors)
        
        rendered_image, vis = obj_renderer.step(
            image = image_original_pil,
            gt_pose = adjusted_pose,
            new_objects = objects_vis,
            paint_new_objects=False,
            return_vis_handle = False,
        )
        
        # Convert to uint8
        rendered_image = (rendered_image * 255).astype(np.uint8)
        frames.append(rendered_image)
        
        # Save frames as a mp4 video
        frames = np.stack(frames)
        video_save_path = (
            cfg.dataset_root
            / cfg.scene_id
            / ("objects_mapping-%s-%s.mp4" % (cfg.gsa_variant, cfg.save_suffix))
        )
        imageio.mimwrite(video_save_path, frames, fps=10)
        print("Save video to %s" % video_save_path)
        
if __name__ == "__main__":
    main()
conceptgraph/slam/cfslam_pipeline_batch.py
concept-graphs/concept-graphs
import torch
import torch.nn.functional as F

from conceptgraph.slam.slam_classes import MapObjectList, DetectionList
from conceptgraph.utils.general_utils import Timer
from conceptgraph.utils.ious import (
    compute_iou_batch, 
    compute_giou_batch, 
    compute_3d_iou_accuracte_batch, 
    compute_3d_giou_accurate_batch,
)
from conceptgraph.slam.utils import (
    merge_obj2_into_obj1, 
    compute_overlap_matrix_2set
)

def compute_spatial_similarities(cfg, detection_list: DetectionList, objects: MapObjectList) -> torch.Tensor:
    '''
    Compute the spatial similarities between the detections and the objects
    
    Args:
        detection_list: a list of M detections
        objects: a list of N objects in the map
    Returns:
        A MxN tensor of spatial similarities
    '''
    det_bboxes = detection_list.get_stacked_values_torch('bbox')
    obj_bboxes = objects.get_stacked_values_torch('bbox')

    if cfg.spatial_sim_type == "iou":
        spatial_sim = compute_iou_batch(det_bboxes, obj_bboxes)
    elif cfg.spatial_sim_type == "giou":
        spatial_sim = compute_giou_batch(det_bboxes, obj_bboxes)
    elif cfg.spatial_sim_type == "iou_accurate":
        spatial_sim = compute_3d_iou_accuracte_batch(det_bboxes, obj_bboxes)
    elif cfg.spatial_sim_type == "giou_accurate":
        spatial_sim = compute_3d_giou_accurate_batch(det_bboxes, obj_bboxes)
    elif cfg.spatial_sim_type == "overlap":
        spatial_sim = compute_overlap_matrix_2set(cfg, objects, detection_list)
        spatial_sim = torch.from_numpy(spatial_sim).T
    else:
        raise ValueError(f"Invalid spatial similarity type: {cfg.spatial_sim_type}")
    
    return spatial_sim

def compute_visual_similarities(cfg, detection_list: DetectionList, objects: MapObjectList) -> torch.Tensor:
    '''
    Compute the visual similarities between the detections and the objects
    
    Args:
        detection_list: a list of M detections
        objects: a list of N objects in the map
    Returns:
        A MxN tensor of visual similarities
    '''
    det_fts = detection_list.get_stacked_values_torch('clip_ft') # (M, D)
    obj_fts = objects.get_stacked_values_torch('clip_ft') # (N, D)

    det_fts = det_fts.unsqueeze(-1) # (M, D, 1)
    obj_fts = obj_fts.T.unsqueeze(0) # (1, D, N)
    
    visual_sim = F.cosine_similarity(det_fts, obj_fts, dim=1) # (M, N)
    
    return visual_sim

def aggregate_similarities(cfg, spatial_sim: torch.Tensor, visual_sim: torch.Tensor) -> torch.Tensor:
    '''
    Aggregate spatial and visual similarities into a single similarity score
    
    Args:
        spatial_sim: a MxN tensor of spatial similarities
        visual_sim: a MxN tensor of visual similarities
    Returns:
        A MxN tensor of aggregated similarities
    '''
    if cfg.match_method == "sim_sum":
        sims = (1 + cfg.phys_bias) * spatial_sim + (1 - cfg.phys_bias) * visual_sim # (M, N)
    else:
        raise ValueError(f"Unknown matching method: {cfg.match_method}")
    
    return sims

def merge_detections_to_objects(
    cfg, 
    detection_list: DetectionList, 
    objects: MapObjectList, 
    agg_sim: torch.Tensor
) -> MapObjectList:
    # Iterate through all detections and merge them into objects
    for i in range(agg_sim.shape[0]):
        # If not matched to any object, add it as a new object
        if agg_sim[i].max() == float('-inf'):
            objects.append(detection_list[i])
        # Merge with most similar existing object
        else:
            j = agg_sim[i].argmax()
            matched_det = detection_list[i]
            matched_obj = objects[j]
            merged_obj = merge_obj2_into_obj1(cfg, matched_obj, matched_det, run_dbscan=False)
            objects[j] = merged_obj
            
    return objects
conceptgraph/slam/mapping.py
concept-graphs/concept-graphs

from collections.abc import Iterable
import copy
import matplotlib
import torch
import torch.nn.functional as F
import numpy as np
import open3d as o3d

from conceptgraph.utils.general_utils import to_numpy, to_tensor

class DetectionList(list):
    def get_values(self, key, idx:int=None):
        if idx is None:
            return [detection[key] for detection in self]
        else:
            return [detection[key][idx] for detection in self]
    
    def get_stacked_values_torch(self, key, idx:int=None):
        values = []
        for detection in self:
            v = detection[key]
            if idx is not None:
                v = v[idx]
            if isinstance(v, o3d.geometry.OrientedBoundingBox) or \
                isinstance(v, o3d.geometry.AxisAlignedBoundingBox):
                v = np.asarray(v.get_box_points())
            if isinstance(v, np.ndarray):
                v = torch.from_numpy(v)
            values.append(v)
        return torch.stack(values, dim=0)
    
    def get_stacked_values_numpy(self, key, idx:int=None):
        values = self.get_stacked_values_torch(key, idx)
        return to_numpy(values)
    
    def __add__(self, other):
        new_list = copy.deepcopy(self)
        new_list.extend(other)
        return new_list
    
    def __iadd__(self, other):
        self.extend(other)
        return self
    
    def slice_by_indices(self, index: Iterable[int]):
        '''
        Return a sublist of the current list by indexing
        '''
        new_self = type(self)()
        for i in index:
            new_self.append(self[i])
        return new_self
    
    def slice_by_mask(self, mask: Iterable[bool]):
        '''
        Return a sublist of the current list by masking
        '''
        new_self = type(self)()
        for i, m in enumerate(mask):
            if m:
                new_self.append(self[i])
        return new_self
    
    def get_most_common_class(self) -> list[int]:
        classes = []
        for d in self:
            values, counts = np.unique(np.asarray(d['class_id']), return_counts=True)
            most_common_class = values[np.argmax(counts)]
            classes.append(most_common_class)
        return classes
    
    def color_by_most_common_classes(self, colors_dict: dict[str, list[float]], color_bbox: bool=True):
        '''
        Color the point cloud of each detection by the most common class
        '''
        classes = self.get_most_common_class()
        for d, c in zip(self, classes):
            color = colors_dict[str(c)]
            d['pcd'].paint_uniform_color(color)
            if color_bbox:
                d['bbox'].color = color
                
    def color_by_instance(self):
        if len(self) == 0:
            # Do nothing
            return
        
        if "inst_color" in self[0]:
            for d in self:
                d['pcd'].paint_uniform_color(d['inst_color'])
                d['bbox'].color = d['inst_color']
        else:
            cmap = matplotlib.colormaps.get_cmap("turbo")
            instance_colors = cmap(np.linspace(0, 1, len(self)))
            instance_colors = instance_colors[:, :3]
            for i in range(len(self)):
                self[i]['pcd'].paint_uniform_color(instance_colors[i])
                self[i]['bbox'].color = instance_colors[i]
            
    
class MapObjectList(DetectionList):
    def compute_similarities(self, new_clip_ft):
        '''
        The input feature should be of shape (D, ), a one-row vector
        This is mostly for backward compatibility
        '''
        # if it is a numpy array, make it a tensor 
        new_clip_ft = to_tensor(new_clip_ft)
        
        # assuming cosine similarity for features
        clip_fts = self.get_stacked_values_torch('clip_ft')

        similarities = F.cosine_similarity(new_clip_ft.unsqueeze(0), clip_fts)
        # return similarities.squeeze()
        return similarities
    
    def to_serializable(self):
        s_obj_list = []
        for obj in self:
            s_obj_dict = copy.deepcopy(obj)
            
            s_obj_dict['clip_ft'] = to_numpy(s_obj_dict['clip_ft'])
            s_obj_dict['text_ft'] = to_numpy(s_obj_dict['text_ft'])
            
            s_obj_dict['pcd_np'] = np.asarray(s_obj_dict['pcd'].points)
            s_obj_dict['bbox_np'] = np.asarray(s_obj_dict['bbox'].get_box_points())
            s_obj_dict['pcd_color_np'] = np.asarray(s_obj_dict['pcd'].colors)
            
            del s_obj_dict['pcd']
            del s_obj_dict['bbox']
            
            s_obj_list.append(s_obj_dict)
            
        return s_obj_list
    
    def load_serializable(self, s_obj_list):
        assert len(self) == 0, 'MapObjectList should be empty when loading'
        for s_obj_dict in s_obj_list:
            new_obj = copy.deepcopy(s_obj_dict)
            
            new_obj['clip_ft'] = to_tensor(new_obj['clip_ft'])
            new_obj['text_ft'] = to_tensor(new_obj['text_ft'])
            
            new_obj['pcd'] = o3d.geometry.PointCloud()
            new_obj['pcd'].points = o3d.utility.Vector3dVector(new_obj['pcd_np'])
            new_obj['bbox'] = o3d.geometry.OrientedBoundingBox.create_from_points(
                o3d.utility.Vector3dVector(new_obj['bbox_np']))
            new_obj['bbox'].color = new_obj['pcd_color_np'][0]
            new_obj['pcd'].colors = o3d.utility.Vector3dVector(new_obj['pcd_color_np'])
            
            del new_obj['pcd_np']
            del new_obj['bbox_np']
            del new_obj['pcd_color_np']
            
            self.append(new_obj)
conceptgraph/slam/slam_classes.py
concept-graphs/concept-graphs
from collections import Counter
import copy
import json
import cv2

import numpy as np
from omegaconf import DictConfig
import open3d as o3d
import torch

import torch.nn.functional as F

import faiss

from conceptgraph.utils.general_utils import to_tensor, to_numpy, Timer
from conceptgraph.slam.slam_classes import MapObjectList, DetectionList

from conceptgraph.utils.ious import compute_3d_iou, compute_3d_iou_accuracte_batch, mask_subtract_contained, compute_iou_batch
from conceptgraph.dataset.datasets_common import from_intrinsics_matrix

def get_classes_colors(classes):
    class_colors = {}

    # Generate a random color for each class
    for class_idx, class_name in enumerate(classes):
        # Generate random RGB values between 0 and 255
        r = np.random.randint(0, 256)/255.0
        g = np.random.randint(0, 256)/255.0
        b = np.random.randint(0, 256)/255.0

        # Assign the RGB values as a tuple to the class in the dictionary
        class_colors[class_idx] = (r, g, b)

    class_colors[-1] = (0, 0, 0)

    return class_colors

def create_or_load_colors(cfg, filename="gsa_classes_tag2text"):
    
    # get the classes, should be saved when making the dataset
    classes_fp = cfg['dataset_root'] / cfg['scene_id'] / f"{filename}.json"
    classes  = None
    with open(classes_fp, "r") as f:
        classes = json.load(f)
    
    # create the class colors, or load them if they exist
    class_colors  = None
    class_colors_fp = cfg['dataset_root'] / cfg['scene_id'] / f"{filename}_colors.json"
    if class_colors_fp.exists():
        with open(class_colors_fp, "r") as f:
            class_colors = json.load(f)
        print("Loaded class colors from ", class_colors_fp)
    else:
        class_colors = get_classes_colors(classes)
        class_colors = {str(k): v for k, v in class_colors.items()}
        with open(class_colors_fp, "w") as f:
            json.dump(class_colors, f)
        print("Saved class colors to ", class_colors_fp)
    return classes, class_colors

def create_object_pcd(depth_array, mask, cam_K, image, obj_color=None) -> o3d.geometry.PointCloud:
    fx, fy, cx, cy = from_intrinsics_matrix(cam_K)
    
    # Also remove points with invalid depth values
    mask = np.logical_and(mask, depth_array > 0)

    if mask.sum() == 0:
        pcd = o3d.geometry.PointCloud()
        return pcd
        
    height, width = depth_array.shape
    x = np.arange(0, width, 1.0)
    y = np.arange(0, height, 1.0)
    u, v = np.meshgrid(x, y)
    
    # Apply the mask, and unprojection is done only on the valid points
    masked_depth = depth_array[mask] # (N, )
    u = u[mask] # (N, )
    v = v[mask] # (N, )

    # Convert to 3D coordinates
    x = (u - cx) * masked_depth / fx
    y = (v - cy) * masked_depth / fy
    z = masked_depth

    # Stack x, y, z coordinates into a 3D point cloud
    points = np.stack((x, y, z), axis=-1)
    points = points.reshape(-1, 3)
    
    # Perturb the points a bit to avoid colinearity
    points += np.random.normal(0, 4e-3, points.shape)

    if obj_color is None: # color using RGB
        # # Apply mask to image
        colors = image[mask] / 255.0
    else: # color using group ID
        # Use the assigned obj_color for all points
        colors = np.full(points.shape, obj_color)
    
    if points.shape[0] == 0:
        import pdb; pdb.set_trace()

    # Create an Open3D PointCloud object
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points)
    pcd.colors = o3d.utility.Vector3dVector(colors)
    
    return pcd

def pcd_denoise_dbscan(pcd: o3d.geometry.PointCloud, eps=0.02, min_points=10) -> o3d.geometry.PointCloud:
    ### Remove noise via clustering
    pcd_clusters = pcd.cluster_dbscan(
        eps=eps,
        min_points=min_points,
    )
    
    # Convert to numpy arrays
    obj_points = np.asarray(pcd.points)
    obj_colors = np.asarray(pcd.colors)
    pcd_clusters = np.array(pcd_clusters)

    # Count all labels in the cluster
    counter = Counter(pcd_clusters)

    # Remove the noise label
    if counter and (-1 in counter):
        del counter[-1]

    if counter:
        # Find the label of the largest cluster
        most_common_label, _ = counter.most_common(1)[0]
        
        # Create mask for points in the largest cluster
        largest_mask = pcd_clusters == most_common_label

        # Apply mask
        largest_cluster_points = obj_points[largest_mask]
        largest_cluster_colors = obj_colors[largest_mask]
        
        # If the largest cluster is too small, return the original point cloud
        if len(largest_cluster_points) < 5:
            return pcd

        # Create a new PointCloud object
        largest_cluster_pcd = o3d.geometry.PointCloud()
        largest_cluster_pcd.points = o3d.utility.Vector3dVector(largest_cluster_points)
        largest_cluster_pcd.colors = o3d.utility.Vector3dVector(largest_cluster_colors)
        
        pcd = largest_cluster_pcd
        
    return pcd

def process_pcd(pcd, cfg, run_dbscan=True):
    
    pcd = pcd.voxel_down_sample(voxel_size=cfg.downsample_voxel_size)
        
    if cfg.dbscan_remove_noise and run_dbscan:
        # print("Before dbscan:", len(pcd.points))
        pcd = pcd_denoise_dbscan(
            pcd, 
            eps=cfg.dbscan_eps, 
            min_points=cfg.dbscan_min_points
        )
        # print("After dbscan:", len(pcd.points))
        
    return pcd

def get_bounding_box(cfg, pcd):
    if ("accurate" in cfg.spatial_sim_type or "overlap" in cfg.spatial_sim_type) and len(pcd.points) >= 4:
        try:
            return pcd.get_oriented_bounding_box(robust=True)
        except RuntimeError as e:
            print(f"Met {e}, use axis aligned bounding box instead")
            return pcd.get_axis_aligned_bounding_box()
    else:
        return pcd.get_axis_aligned_bounding_box()

def merge_obj2_into_obj1(cfg, obj1, obj2, run_dbscan=True):
    '''
    Merge the new object to the old object
    This operation is done in-place
    '''
    n_obj1_det = obj1['num_detections']
    n_obj2_det = obj2['num_detections']
    
    for k in obj1.keys():
        if k in ['caption']:
            # Here we need to merge two dictionaries and adjust the key of the second one
            for k2, v2 in obj2['caption'].items():
                obj1['caption'][k2 + n_obj1_det] = v2
        elif k not in ['pcd', 'bbox', 'clip_ft', "text_ft"]:
            if isinstance(obj1[k], list) or isinstance(obj1[k], int):
                obj1[k] += obj2[k]
            elif k == "inst_color":
                obj1[k] = obj1[k] # Keep the initial instance color
            else:
                # TODO: handle other types if needed in the future
                raise NotImplementedError
        else: # pcd, bbox, clip_ft, text_ft are handled below
            continue

    # merge pcd and bbox
    obj1['pcd'] += obj2['pcd']
    obj1['pcd'] = process_pcd(obj1['pcd'], cfg, run_dbscan=run_dbscan)
    obj1['bbox'] = get_bounding_box(cfg, obj1['pcd'])
    obj1['bbox'].color = [0,1,0]
    
    # merge clip ft
    obj1['clip_ft'] = (obj1['clip_ft'] * n_obj1_det +
                       obj2['clip_ft'] * n_obj2_det) / (
                       n_obj1_det + n_obj2_det)
    obj1['clip_ft'] = F.normalize(obj1['clip_ft'], dim=0)

    # merge text_ft
    obj2['text_ft'] = to_tensor(obj2['text_ft'], cfg.device)
    obj1['text_ft'] = to_tensor(obj1['text_ft'], cfg.device)
    obj1['text_ft'] = (obj1['text_ft'] * n_obj1_det +
                       obj2['text_ft'] * n_obj2_det) / (
                       n_obj1_det + n_obj2_det)
    obj1['text_ft'] = F.normalize(obj1['text_ft'], dim=0)
    
    return obj1

def compute_overlap_matrix(cfg, objects: MapObjectList):
    '''
    compute pairwise overlapping between objects in terms of point nearest neighbor. 
    Suppose we have a list of n point cloud, each of which is a o3d.geometry.PointCloud object. 
    Now we want to construct a matrix of size n x n, where the (i, j) entry is the ratio of points in point cloud i 
    that are within a distance threshold of any point in point cloud j. 
    '''
    n = len(objects)
    overlap_matrix = np.zeros((n, n))
    
    # Convert the point clouds into numpy arrays and then into FAISS indices for efficient search
    point_arrays = [np.asarray(obj['pcd'].points, dtype=np.float32) for obj in objects]
    indices = [faiss.IndexFlatL2(arr.shape[1]) for arr in point_arrays]
    
    # Add the points from the numpy arrays to the corresponding FAISS indices
    for index, arr in zip(indices, point_arrays):
        index.add(arr)

    # Compute the pairwise overlaps
    for i in range(n):
        for j in range(n):
            if i != j:  # Skip diagonal elements
                box_i = objects[i]['bbox']
                box_j = objects[j]['bbox']
                
                # Skip if the boxes do not overlap at all (saves computation)
                iou = compute_3d_iou(box_i, box_j)
                if iou == 0:
                    continue
                
                # # Use range_search to find points within the threshold
                # _, I = indices[j].range_search(point_arrays[i], threshold ** 2)
                D, I = indices[j].search(point_arrays[i], 1)

                # # If any points are found within the threshold, increase overlap count
                # overlap += sum([len(i) for i in I])

                overlap = (D < cfg.downsample_voxel_size ** 2).sum() # D is the squared distance

                # Calculate the ratio of points within the threshold
                overlap_matrix[i, j] = overlap / len(point_arrays[i])

    return overlap_matrix

def compute_overlap_matrix_2set(cfg, objects_map: MapObjectList, objects_new: DetectionList) -> np.ndarray:
    '''
    compute pairwise overlapping between two set of objects in terms of point nearest neighbor. 
    objects_map is the existing objects in the map, objects_new is the new objects to be added to the map
    Suppose len(objects_map) = m, len(objects_new) = n
    Then we want to construct a matrix of size m x n, where the (i, j) entry is the ratio of points 
    in point cloud i that are within a distance threshold of any point in point cloud j.
    '''
    m = len(objects_map)
    n = len(objects_new)
    overlap_matrix = np.zeros((m, n))
    
    # Convert the point clouds into numpy arrays and then into FAISS indices for efficient search
    points_map = [np.asarray(obj['pcd'].points, dtype=np.float32) for obj in objects_map] # m arrays
    indices = [faiss.IndexFlatL2(arr.shape[1]) for arr in points_map] # m indices
    
    # Add the points from the numpy arrays to the corresponding FAISS indices
    for index, arr in zip(indices, points_map):
        index.add(arr)
        
    points_new = [np.asarray(obj['pcd'].points, dtype=np.float32) for obj in objects_new] # n arrays
        
    bbox_map = objects_map.get_stacked_values_torch('bbox')
    bbox_new = objects_new.get_stacked_values_torch('bbox')
    try:
        iou = compute_3d_iou_accuracte_batch(bbox_map, bbox_new) # (m, n)
    except ValueError:
        print("Met `Plane vertices are not coplanar` error, use axis aligned bounding box instead")
        bbox_map = []
        bbox_new = []
        for pcd in objects_map.get_values('pcd'):
            bbox_map.append(np.asarray(
                pcd.get_axis_aligned_bounding_box().get_box_points()))
        for pcd in objects_new.get_values('pcd'):
            bbox_new.append(np.asarray(
                pcd.get_axis_aligned_bounding_box().get_box_points()))
        bbox_map = torch.from_numpy(np.stack(bbox_map))
        bbox_new = torch.from_numpy(np.stack(bbox_new))
        
        iou = compute_iou_batch(bbox_map, bbox_new) # (m, n)
            

    # Compute the pairwise overlaps
    for i in range(m):
        for j in range(n):
            if iou[i,j] < 1e-6:
                continue
            
            D, I = indices[i].search(points_new[j], 1) # search new object j in map object i

            overlap = (D < cfg.downsample_voxel_size ** 2).sum() # D is the squared distance

            # Calculate the ratio of points within the threshold
            overlap_matrix[i, j] = overlap / len(points_new[j])

    return overlap_matrix

def merge_overlap_objects(cfg, objects: MapObjectList, overlap_matrix: np.ndarray):
    x, y = overlap_matrix.nonzero()
    overlap_ratio = overlap_matrix[x, y]

    sort = np.argsort(overlap_ratio)[::-1]
    x = x[sort]
    y = y[sort]
    overlap_ratio = overlap_ratio[sort]

    kept_objects = np.ones(len(objects), dtype=bool)
    for i, j, ratio in zip(x, y, overlap_ratio):
        visual_sim = F.cosine_similarity(
            to_tensor(objects[i]['clip_ft']),
            to_tensor(objects[j]['clip_ft']),
            dim=0
        )
        text_sim = F.cosine_similarity(
            to_tensor(objects[i]['text_ft']),
            to_tensor(objects[j]['text_ft']),
            dim=0
        )
        if ratio > cfg.merge_overlap_thresh:
            if visual_sim > cfg.merge_visual_sim_thresh and \
                text_sim > cfg.merge_text_sim_thresh:
                if kept_objects[j]:
                    # Then merge object i into object j
                    objects[j] = merge_obj2_into_obj1(cfg, objects[j], objects[i], run_dbscan=True)
                    kept_objects[i] = False
        else:
            break
    
    # Remove the objects that have been merged
    new_objects = [obj for obj, keep in zip(objects, kept_objects) if keep]
    objects = MapObjectList(new_objects)
    
    return objects

def denoise_objects(cfg, objects: MapObjectList):
    for i in range(len(objects)):
        og_object_pcd = objects[i]['pcd']
        objects[i]['pcd'] = process_pcd(objects[i]['pcd'], cfg, run_dbscan=True)
        if len(objects[i]['pcd'].points) < 4:
            objects[i]['pcd'] = og_object_pcd
            continue
        objects[i]['bbox'] = get_bounding_box(cfg, objects[i]['pcd'])
        objects[i]['bbox'].color = [0,1,0]
        
    return objects

def filter_objects(cfg, objects: MapObjectList):
    # Remove the object that has very few points or viewed too few times
    print("Before filtering:", len(objects))
    objects_to_keep = []
    for obj in objects:
        if len(obj['pcd'].points) >= cfg.obj_min_points and obj['num_detections'] >= cfg.obj_min_detections:
            objects_to_keep.append(obj)
    objects = MapObjectList(objects_to_keep)
    print("After filtering:", len(objects))
    
    return objects

def merge_objects(cfg, objects: MapObjectList):
    if cfg.merge_overlap_thresh > 0:
        # Merge one object into another if the former is contained in the latter
        overlap_matrix = compute_overlap_matrix(cfg, objects)
        print("Before merging:", len(objects))
        objects = merge_overlap_objects(cfg, objects, overlap_matrix)
        print("After merging:", len(objects))
    
    return objects

def filter_gobs(
    cfg: DictConfig,
    gobs: dict,
    image: np.ndarray,
    BG_CLASSES = ["wall", "floor", "ceiling"],
):
    # If no detection at all
    if len(gobs['xyxy']) == 0:
        return gobs
    
    # Filter out the objects based on various criteria
    idx_to_keep = []
    for mask_idx in range(len(gobs['xyxy'])):
        local_class_id = gobs['class_id'][mask_idx]
        class_name = gobs['classes'][local_class_id]
        
        # SKip masks that are too small
        if gobs['mask'][mask_idx].sum() < max(cfg.mask_area_threshold, 10):
            continue
        
        # Skip the BG classes
        if cfg.skip_bg and class_name in BG_CLASSES:
            continue
        
        # Skip the non-background boxes that are too large
        if class_name not in BG_CLASSES:
            x1, y1, x2, y2 = gobs['xyxy'][mask_idx]
            bbox_area = (x2 - x1) * (y2 - y1)
            image_area = image.shape[0] * image.shape[1]
            if bbox_area > cfg.max_bbox_area_ratio * image_area:
                # print(f"Skipping {class_name} with area {bbox_area} > {cfg.max_bbox_area_ratio} * {image_area}")
                continue
            
        # Skip masks with low confidence
        if gobs['confidence'][mask_idx] < cfg.mask_conf_threshold:
            continue
        
        idx_to_keep.append(mask_idx)
    
    for k in gobs.keys():
        if isinstance(gobs[k], str) or k == "classes": # Captions
            continue
        elif isinstance(gobs[k], list):
            gobs[k] = [gobs[k][i] for i in idx_to_keep]
        elif isinstance(gobs[k], np.ndarray):
            gobs[k] = gobs[k][idx_to_keep]
        else:
            raise NotImplementedError(f"Unhandled type {type(gobs[k])}")
    
    return gobs

def resize_gobs(
    gobs,
    image
):
    n_masks = len(gobs['xyxy'])

    new_mask = []
    
    for mask_idx in range(n_masks):
        # TODO: rewrite using interpolation/resize in numpy or torch rather than cv2
        mask = gobs['mask'][mask_idx]
        if mask.shape != image.shape[:2]:
            # Rescale the xyxy coordinates to the image shape
            x1, y1, x2, y2 = gobs['xyxy'][mask_idx]
            x1 = round(x1 * image.shape[1] / mask.shape[1])
            y1 = round(y1 * image.shape[0] / mask.shape[0])
            x2 = round(x2 * image.shape[1] / mask.shape[1])
            y2 = round(y2 * image.shape[0] / mask.shape[0])
            gobs['xyxy'][mask_idx] = [x1, y1, x2, y2]
            
            # Reshape the mask to the image shape
            mask = cv2.resize(mask.astype(np.uint8), image.shape[:2][::-1], interpolation=cv2.INTER_NEAREST)
            mask = mask.astype(bool)
            new_mask.append(mask)

    if len(new_mask) > 0:
        gobs['mask'] = np.asarray(new_mask)
        
    return gobs

def gobs_to_detection_list(
    cfg, 
    image, 
    depth_array,
    cam_K, 
    idx, 
    gobs, 
    trans_pose = None,
    class_names = None,
    BG_CLASSES = ["wall", "floor", "ceiling"],
    color_path = None,
):
    '''
    Return a DetectionList object from the gobs
    All object are still in the camera frame. 
    '''
    fg_detection_list = DetectionList()
    bg_detection_list = DetectionList()
    
    gobs = resize_gobs(gobs, image)
    gobs = filter_gobs(cfg, gobs, image, BG_CLASSES)
    
    if len(gobs['xyxy']) == 0:
        return fg_detection_list, bg_detection_list
    
    # Compute the containing relationship among all detections and subtract fg from bg objects
    xyxy = gobs['xyxy']
    mask = gobs['mask']
    gobs['mask'] = mask_subtract_contained(xyxy, mask)
    
    n_masks = len(gobs['xyxy'])
    for mask_idx in range(n_masks):
        local_class_id = gobs['class_id'][mask_idx]
        mask = gobs['mask'][mask_idx]
        class_name = gobs['classes'][local_class_id]
        global_class_id = -1 if class_names is None else class_names.index(class_name)
        
        # make the pcd and color it
        camera_object_pcd = create_object_pcd(
            depth_array,
            mask,
            cam_K,
            image,
            obj_color = None
        )
        
        # It at least contains 5 points
        if len(camera_object_pcd.points) < max(cfg.min_points_threshold, 5): 
            continue
        
        if trans_pose is not None:
            global_object_pcd = camera_object_pcd.transform(trans_pose)
        else:
            global_object_pcd = camera_object_pcd
        
        # get largest cluster, filter out noise 
        global_object_pcd = process_pcd(global_object_pcd, cfg)
        
        pcd_bbox = get_bounding_box(cfg, global_object_pcd)
        pcd_bbox.color = [0,1,0]
        
        if pcd_bbox.volume() < 1e-6:
            continue
        
        # Treat the detection in the same way as a 3D object
        # Store information that is enough to recover the detection
        detected_object = {
            'image_idx' : [idx],                             # idx of the image
            'mask_idx' : [mask_idx],                         # idx of the mask/detection
            'color_path' : [color_path],                     # path to the RGB image
            'class_name' : [class_name],                         # global class id for this detection
            'class_id' : [global_class_id],                         # global class id for this detection
            'num_detections' : 1,                            # number of detections in this object
            'mask': [mask],
            'xyxy': [gobs['xyxy'][mask_idx]],
            'conf': [gobs['confidence'][mask_idx]],
            'n_points': [len(global_object_pcd.points)],
            'pixel_area': [mask.sum()],
            'contain_number': [None],                          # This will be computed later
            "inst_color": np.random.rand(3),                 # A random color used for this segment instance
            'is_background': class_name in BG_CLASSES,
            
            # These are for the entire 3D object
            'pcd': global_object_pcd,
            'bbox': pcd_bbox,
            'clip_ft': to_tensor(gobs['image_feats'][mask_idx]),
            'text_ft': to_tensor(gobs['text_feats'][mask_idx]),
        }
        
        if class_name in BG_CLASSES:
            bg_detection_list.append(detected_object)
        else:
            fg_detection_list.append(detected_object)
    
    return fg_detection_list, bg_detection_list

def transform_detection_list(
    detection_list: DetectionList,
    transform: torch.Tensor,
    deepcopy = False,
):
    '''
    Transform the detection list by the given transform
    
    Args:
        detection_list: DetectionList
        transform: 4x4 torch.Tensor
        
    Returns:
        transformed_detection_list: DetectionList
    '''
    transform = to_numpy(transform)
    
    if deepcopy:
        detection_list = copy.deepcopy(detection_list)
    
    for i in range(len(detection_list)):
        detection_list[i]['pcd'] = detection_list[i]['pcd'].transform(transform)
        detection_list[i]['bbox'] = detection_list[i]['bbox'].rotate(transform[:3, :3], center=(0, 0, 0))
        detection_list[i]['bbox'] = detection_list[i]['bbox'].translate(transform[:3, 3])
        # detection_list[i]['bbox'] = detection_list[i]['pcd'].get_oriented_bounding_box(robust=True)
    
    return detection_list
conceptgraph/slam/utils.py
concept-graphs/concept-graphs
import numpy as np
import torch


def compute_pred_gt_associations(pred, gt):
    # pred: predicted pointcloud
    # gt: GT pointcloud
    from chamferdist.chamfer import knn_points

    # pred = pointclouds.points_padded.cuda().contiguous()
    # gt = pts_gt.unsqueeze(0).cuda().contiguous()
    b, l, d = pred.shape
    lengths_src = torch.ones(b, dtype=torch.long, device=pred.device) * l
    b, l, d = gt.shape
    lengths_tgt = torch.ones(b, dtype=torch.long, device=pred.device) * l
    src_nn = knn_points(
        pred,
        gt,
        lengths1=lengths_src,
        lengths2=lengths_tgt,
        return_nn=True,
        return_sorted=True,
        K=1,
    )
    idx_pred_to_gt = src_nn.idx.squeeze(0).squeeze(-1)
    tgt_nn = knn_points(
        gt,
        pred,
        lengths1=lengths_tgt,
        lengths2=lengths_src,
        return_nn=True,
        return_sorted=True,
        K=1,
    )
    idx_gt_to_pred = tgt_nn.idx.squeeze(0).squeeze(-1)

    return idx_pred_to_gt, idx_gt_to_pred

def compute_confmatrix(
    labels_pred, labels_gt, idx_pred_to_gt, idx_gt_to_pred, class_names
):
    labels_gt = labels_gt[idx_pred_to_gt]
    # num_classes = labels_gt.max().item() + 1
    # print(num_classes)
    num_classes = len(class_names)
    print(num_classes)
    confmatrix = torch.zeros(num_classes, num_classes, device=labels_pred.device)
    for class_gt_int in range(num_classes):
        tensor_gt_class = torch.eq(labels_gt, class_gt_int).long()
        for class_pred_int in range(num_classes):
            tensor_pred_class = torch.eq(labels_pred, class_pred_int).long()
            tensor_pred_class = torch.mul(tensor_gt_class, tensor_pred_class)
            count = torch.sum(tensor_pred_class)
            confmatrix[class_gt_int, class_pred_int] += count

    return confmatrix


def compute_metrics(confmatrix, class_names):
    if isinstance(confmatrix, torch.Tensor):
        confmatrix = confmatrix.cpu().numpy()

    num_classes = len(class_names)
    ious = np.zeros((num_classes))
    precision = np.zeros((num_classes))
    recall = np.zeros((num_classes))
    f1score = np.zeros((num_classes))

    for _idx in range(num_classes):
        ious[_idx] = confmatrix[_idx, _idx] / (
            max(
                1,
                confmatrix[_idx, :].sum()
                + confmatrix[:, _idx].sum()
                - confmatrix[_idx, _idx],
            )
        )
        recall[_idx] = confmatrix[_idx, _idx] / max(1, confmatrix[_idx, :].sum())
        precision[_idx] = confmatrix[_idx, _idx] / max(1, confmatrix[:, _idx].sum())
        f1score[_idx] = (
            2 * precision[_idx] * recall[_idx] / max(1, precision[_idx] + recall[_idx])
        )

    fmiou = (ious * confmatrix.sum(1) / confmatrix.sum()).sum()
    # print(f"iou: {ious}")
    # print(f"miou: {ious.mean()}")
    # print(f"Acc>0.15: {(ious > 0.15).sum()}")
    # print(f"Acc>0.25: {(ious > 0.25).sum()}")
    # print(f"Acc>0.50: {(ious > 0.50).sum()}")
    # print(f"Acc>0.75: {(ious > 0.75).sum()}")
    # print(f"precision: {precision}")
    # print(f"recall: {recall}")
    # print(f"f1score: {f1score}")

    mdict = {}
    mdict["iou"] = ious.tolist()
    mdict["miou"] = ious.mean().item()
    mdict["fmiou"] = fmiou.item()
    mdict["num_classes"] = num_classes
    mdict["acc0.15"] = (ious > 0.15).sum().item()
    mdict["acc0.25"] = (ious > 0.25).sum().item()
    mdict["acc0.50"] = (ious > 0.50).sum().item()
    mdict["acc0.75"] = (ious > 0.75).sum().item()
    mdict["precision"] = precision.tolist()
    mdict["recall"] = recall.tolist()
    mdict["f1score"] = f1score.tolist()

    return mdict
conceptgraph/utils/eval.py
concept-graphs/concept-graphs


import torch
import numpy as np
import time

class Timer:
    def __init__(self, heading = "", verbose = True):
        self.verbose = verbose
        if not self.verbose:
            return
        self.heading = heading

    def __enter__(self):
        if not self.verbose:
            return self
        self.start = time.time()
        return self

    def __exit__(self, *args):
        if not self.verbose:
            return
        self.end = time.time()
        self.interval = self.end - self.start
        print(self.heading, self.interval)

def to_numpy(tensor):
    if isinstance(tensor, np.ndarray):
        return tensor
    return tensor.detach().cpu().numpy()

def to_tensor(numpy_array, device=None):
    if isinstance(numpy_array, torch.Tensor):
        return numpy_array
    if device is None:
        return torch.from_numpy(numpy_array)
    else:
        return torch.from_numpy(numpy_array).to(device)
    
def to_scalar(d: np.ndarray | torch.Tensor | float) -> int | float:
    '''
    Convert the d to a scalar
    '''
    if isinstance(d, float):
        return d
    
    elif "numpy" in str(type(d)):
        assert d.size == 1
        return d.item()
    
    elif isinstance(d, torch.Tensor):
        assert d.numel() == 1
        return d.item()
    
    else:
        raise TypeError(f"Invalid type for conversion: {type(d)}")

conceptgraph/utils/general_utils.py
concept-graphs/concept-graphs
from PIL import Image

import torch

def transform_points_batch(poses: torch.Tensor, points: torch.Tensor) -> torch.Tensor:
    """
    poses: (M, 4, 4)
    points: (N, 3)
    
    return: (M, N, 3)
    """
    N = points.shape[0]
    M = poses.shape[0]
    
    # Convert points to homogeneous coordinates (N, 4)
    points_homogeneous = torch.cat((points, torch.ones(N, 1).to(points.device)), dim=-1) # (N, 4)
    
    # Repeat points M times along a new dimension
    points_homogeneous = points_homogeneous.unsqueeze(0).repeat(M, 1, 1) # (M, N, 4)
    
    # Apply transformation: for each pose, do a matrix multiplication with points
    points_transformed = torch.bmm(poses, points_homogeneous.transpose(1,2)) # (M, 4, N)
    
    # Convert back to cartesian coordinates by dividing by the last coordinate
    points_transformed = points_transformed[:, :3, :] / points_transformed[:, 3:4, :] # (M, 3, N)
    
    return points_transformed.transpose(1,2) # (M, N, 3)

def project_points_camera2plane(points: torch.Tensor, K: torch.Tensor) -> torch.Tensor:
    """
    Project a set of points (in camera coordinates) to the image plane according to the camera intrinsics
    
    points: (N, 3)
    K: (3, 3)
    
    return: 
        points_proj: (N, 2) the projected points in the image plane
        points_depth: (N,) the depth of the points in the camera coordinates
    """
    # multiply points by the camera intrinsics
    points_proj = torch.matmul(K, points.t())  # result size is (3, N)
    
    # convert to Euclidean coordinates
    points_depth = points_proj[2, :]  # (N)
    points_coord = points_proj[:2, :] / points_proj[2, :].unsqueeze(0)  # divide x and y by z
    
    return points_coord.t(), points_depth # transpose back to (N, 2)

def project_points_world2plane(points: torch.Tensor, poses: torch.Tensor, K: torch.Tensor) -> torch.Tensor:
    """Project a set of points (in the world coordinates) to the image plane according to the camera pose and intrinsics

    Args:
        points (torch.Tensor): (N, 3)
        poses (torch.Tensor): (M, 4, 4) the camera poses in the world coordinates
        K (torch.Tensor): (3, 3)

    Returns:
        points_coord (torch.Tensor): (M, N, 2) the projected points in the image plane
        points_depth (torch.Tensor): (M, N) the depth of the points in the camera coordinates
    """
    poses_inv = torch.inverse(poses)  # (M, 4, 4)
    
    # Transform points from world coordinates to camera coordinates
    points_camera = transform_points_batch(poses_inv, points)  # (M, N, 3)

    M, N, _ = points_camera.shape

    points_camera = points_camera.reshape(-1, 3)  # (M*N, 3)
    points_coord, points_depth = project_points_camera2plane(points_camera, K)  # (M*N, 2), (M, N)
    
    points_coord = points_coord.reshape(M, N, 2)  # (M, N, 2)
    points_depth = points_depth.reshape(M, N)  # (M, N)

    return points_coord, points_depth  # (M, N, 2), (M, N)
    
def check_proj_points(
    points: torch.Tensor, 
    depth_tensor: torch.Tensor,
    K: torch.Tensor, 
    pose: torch.Tensor,
    depth_margin: float = 0.05,
) -> torch.Tensor:
    '''
    Project points to the image plane and perform visibility checks
    
    Args:
        points (torch.Tensor): (N, 3), points in the world coordinates
        depth_tensor (torch.Tensor): (H, W)
        K (torch.Tensor): (3, 3)
        pose (torch.Tensor): (4, 4)
        depth_margin (float, optional): depth margin. Defaults to 0.05.
        
    Returns:
        proj_xyz (torch.Tensor): (N, 3) the projected points
        proj_front (torch.Tensor): (N,) True if the point is projected in front of the depth map
        proj_align (torch.Tensor): (N,) True if the point is projected within the margin the depth map
        proj_behind (torch.Tensor): (N,) True if the point is projected behind the depth map
    '''
    proj_xy, proj_z = project_points_world2plane(
        points, # (N, 3)
        pose.unsqueeze(0), # (1, 4, 4) 
        K, # (3, 3)
    ) # (1, N, 2), (1, N)
    proj_xy = proj_xy.reshape((-1, 2)) # (N, 2)
    proj_z = proj_z.reshape((-1)) # (N, )

    proj_xyz = torch.cat([
        proj_xy, proj_z.unsqueeze(-1)
    ], dim=-1) # (N, 3)
    
    proj_xy_within = torch.stack([
        0 <= proj_xy[:, 0], proj_xy[:, 0] < depth_tensor.shape[1], 
        0 <= proj_xy[:, 1], proj_xy[:, 1] < depth_tensor.shape[0]
    ], dim=-1) # (N, 2)
    proj_xy_within = torch.all(proj_xy_within, dim=-1) # (N, )
    proj_z_within = proj_z > 0
    proj_within = torch.logical_and(
        proj_xy_within, proj_z_within
    ) # (N, )
    
    # Initialize the projection masks
    proj_front = torch.full_like(proj_within, fill_value=False) # (N, )
    proj_align = torch.full_like(proj_within, fill_value=False) # (N, )
    proj_behind = torch.full_like(proj_within, fill_value=False) # (N, )
    
    if proj_within.any():
        depth_within = depth_tensor[
            proj_xy[proj_within][:, 1].long(),
            proj_xy[proj_within][:, 0].long()
        ] # (N, )
        proj_z_within = proj_z[proj_within] # (N, )
        
        proj_front[proj_within] = proj_z_within < depth_within - depth_margin # (N, )
        proj_align[proj_within] = torch.abs(proj_z_within - depth_within) <= depth_margin # (N, )
        proj_behind[proj_within] = proj_z_within > depth_within + depth_margin # (N, )
    
    return proj_xyz, proj_front, proj_align, proj_behind
conceptgraph/utils/geometry.py
concept-graphs/concept-graphs
from PIL import Image

def crop_image_pil(image: Image, x1:int, y1:int, x2:int, y2:int, padding:int=0) -> Image:
    '''
    Crop the image with some padding
    
    Args:
        image: PIL image
        x1, y1, x2, y2: bounding box coordinates
        padding: padding around the bounding box
         
    Returns:
        image_crop: PIL image
    '''
    image_width, image_height = image.size
    x1 = max(0, x1 - padding)
    y1 = max(0, y1 - padding)
    x2 = min(image_width, x2 + padding)
    y2 = min(image_height, y2 + padding)
    
    image_crop = image.crop((x1, y1, x2, y2))
    return image_crop


conceptgraph/utils/image.py
concept-graphs/concept-graphs
import numpy as np
import torch

import open3d as o3d

def compute_3d_iou(bbox1, bbox2, padding=0, use_iou=True):
    # Get the coordinates of the first bounding box
    bbox1_min = np.asarray(bbox1.get_min_bound()) - padding
    bbox1_max = np.asarray(bbox1.get_max_bound()) + padding

    # Get the coordinates of the second bounding box
    bbox2_min = np.asarray(bbox2.get_min_bound()) - padding
    bbox2_max = np.asarray(bbox2.get_max_bound()) + padding

    # Compute the overlap between the two bounding boxes
    overlap_min = np.maximum(bbox1_min, bbox2_min)
    overlap_max = np.minimum(bbox1_max, bbox2_max)
    overlap_size = np.maximum(overlap_max - overlap_min, 0.0)

    overlap_volume = np.prod(overlap_size)
    bbox1_volume = np.prod(bbox1_max - bbox1_min)
    bbox2_volume = np.prod(bbox2_max - bbox2_min)
    
    obj_1_overlap = overlap_volume / bbox1_volume
    obj_2_overlap = overlap_volume / bbox2_volume
    max_overlap = max(obj_1_overlap, obj_2_overlap)

    iou = overlap_volume / (bbox1_volume + bbox2_volume - overlap_volume)

    if use_iou:
        return iou
    else:
        return max_overlap

def compute_iou_batch(bbox1: torch.Tensor, bbox2: torch.Tensor) -> torch.Tensor:
    '''
    Compute IoU between two sets of axis-aligned 3D bounding boxes.
    
    bbox1: (M, V, D), e.g. (M, 8, 3)
    bbox2: (N, V, D), e.g. (N, 8, 3)
    
    returns: (M, N)
    '''
    # Compute min and max for each box
    bbox1_min, _ = bbox1.min(dim=1) # Shape: (M, 3)
    bbox1_max, _ = bbox1.max(dim=1) # Shape: (M, 3)
    bbox2_min, _ = bbox2.min(dim=1) # Shape: (N, 3)
    bbox2_max, _ = bbox2.max(dim=1) # Shape: (N, 3)

    # Expand dimensions for broadcasting
    bbox1_min = bbox1_min.unsqueeze(1)  # Shape: (M, 1, 3)
    bbox1_max = bbox1_max.unsqueeze(1)  # Shape: (M, 1, 3)
    bbox2_min = bbox2_min.unsqueeze(0)  # Shape: (1, N, 3)
    bbox2_max = bbox2_max.unsqueeze(0)  # Shape: (1, N, 3)

    # Compute max of min values and min of max values
    # to obtain the coordinates of intersection box.
    inter_min = torch.max(bbox1_min, bbox2_min)  # Shape: (M, N, 3)
    inter_max = torch.min(bbox1_max, bbox2_max)  # Shape: (M, N, 3)

    # Compute volume of intersection box
    inter_vol = torch.prod(torch.clamp(inter_max - inter_min, min=0), dim=2)  # Shape: (M, N)

    # Compute volumes of the two sets of boxes
    bbox1_vol = torch.prod(bbox1_max - bbox1_min, dim=2)  # Shape: (M, 1)
    bbox2_vol = torch.prod(bbox2_max - bbox2_min, dim=2)  # Shape: (1, N)

    # Compute IoU, handling the special case where there is no intersection
    # by setting the intersection volume to 0.
    iou = inter_vol / (bbox1_vol + bbox2_vol - inter_vol + 1e-10)

    return iou
    
def compute_3d_giou(bbox1, bbox2):
    # Get the coordinates of the first bounding box
    bbox1_min = np.asarray(bbox1.get_min_bound())
    bbox1_max = np.asarray(bbox1.get_max_bound())

    # Get the coordinates of the second bounding box
    bbox2_min = np.asarray(bbox2.get_min_bound())
    bbox2_max = np.asarray(bbox2.get_max_bound())
    
    # Intersection
    intersec_min = np.maximum(bbox1_min, bbox2_min)
    intersec_max = np.minimum(bbox1_max, bbox2_max)
    intersec_size = np.maximum(intersec_max - intersec_min, 0.0)
    intersec_volume = np.prod(intersec_size)

    # Union
    bbox1_volume = np.prod(bbox1_max - bbox1_min)
    bbox2_volume = np.prod(bbox2_max - bbox2_min)
    union_volume = bbox1_volume + bbox2_volume - intersec_volume
    
    iou = intersec_volume / union_volume
    
    # Enclosing box
    enclosing_min = np.minimum(bbox1_min, bbox2_min)
    enclosing_max = np.maximum(bbox1_max, bbox2_max)
    enclosing_size = np.maximum(enclosing_max - enclosing_min, 0.0)
    enclosing_volume = np.prod(enclosing_size)
    
    giou = iou - (enclosing_volume - union_volume) / enclosing_volume
    
    return giou

def compute_giou_batch(bbox1: torch.Tensor, bbox2: torch.Tensor) -> torch.Tensor:
    '''
    Compute the generalized IoU between two sets of axis-aligned 3D bounding boxes.
    
    bbox1: (M, V, D), e.g. (M, 8, 3)
    bbox2: (N, V, D), e.g. (N, 8, 3)
    
    returns: (M, N)
    '''
    # Compute min and max for each box
    bbox1_min, _ = bbox1.min(dim=1) # Shape: (M, D)
    bbox1_max, _ = bbox1.max(dim=1) # Shape: (M, D)
    bbox2_min, _ = bbox2.min(dim=1) # Shape: (N, D)
    bbox2_max, _ = bbox2.max(dim=1) # Shape: (N, D)

    # Expand dimensions for broadcasting
    bbox1_min = bbox1_min.unsqueeze(1)  # Shape: (M, 1, D)
    bbox1_max = bbox1_max.unsqueeze(1)  # Shape: (M, 1, D)
    bbox2_min = bbox2_min.unsqueeze(0)  # Shape: (1, N, D)
    bbox2_max = bbox2_max.unsqueeze(0)  # Shape: (1, N, D)

    # Compute max of min values and min of max values
    # to obtain the coordinates of intersection box.
    inter_min = torch.max(bbox1_min, bbox2_min)  # Shape: (M, N, D)
    inter_max = torch.min(bbox1_max, bbox2_max)  # Shape: (M, N, D)
    
    # to obtain the coordinates of enclosing box
    enclosing_min = torch.min(bbox1_min, bbox2_min)  # Shape: (M, N, D)
    enclosing_max = torch.max(bbox1_max, bbox2_max)  # Shape: (M, N, D)

    # Compute volume of intersection box
    inter_vol = torch.prod(torch.clamp(inter_max - inter_min, min=0), dim=2)  # Shape: (M, N)
    enclosing_vol = torch.prod(enclosing_max - enclosing_min, dim=2)  # Shape: (M, N)

    # Compute volumes of the two sets of boxes
    bbox1_vol = torch.prod(bbox1_max - bbox1_min, dim=2)  # Shape: (M, 1)
    bbox2_vol = torch.prod(bbox2_max - bbox2_min, dim=2)  # Shape: (1, N)
    union_vol = bbox1_vol + bbox2_vol - inter_vol

    # Compute IoU, handling the special case where there is no intersection
    # by setting the intersection volume to 0.
    iou = inter_vol / (union_vol + 1e-10)
    giou = iou - (enclosing_vol - union_vol) / (enclosing_vol + 1e-10)

    return giou

def compute_3d_iou_accuracte_batch(bbox1, bbox2):
    '''
    Compute IoU between two sets of oriented (or axis-aligned) 3D bounding boxes.
    
    bbox1: (M, 8, D), e.g. (M, 8, 3)
    bbox2: (N, 8, D), e.g. (N, 8, 3)
    
    returns: (M, N)
    '''
    # Must expend the box beforehand, otherwise it may results overestimated results
    bbox1 = expand_3d_box(bbox1, 0.02)
    bbox2 = expand_3d_box(bbox2, 0.02)
    
    import pytorch3d.ops as ops

    bbox1 = bbox1[:, [0, 2, 5, 3, 1, 7, 4, 6]]
    bbox2 = bbox2[:, [0, 2, 5, 3, 1, 7, 4, 6]]
    
    inter_vol, iou = ops.box3d_overlap(bbox1.float(), bbox2.float())
    
    return iou

def compute_3d_giou_accurate(obj1, obj2):
    '''
    Compute the 3D GIoU in a more accurate way. 
    '''
    import pytorch3d.ops as ops
    
    # This is too slow
    # bbox1 = pcd1.get_minimal_oriented_bounding_box()
    # bbox2 = pcd2.get_minimal_oriented_bounding_box()
    
    # This is still slow ... 
    # Moved it outside of this function so that it is computed less times
    # bbox1 = pcd1.get_oriented_bounding_box()
    # bbox2 = pcd2.get_oriented_bounding_box()
    
    bbox1 = obj1['bbox']
    bbox2 = obj2['bbox']
    pcd1 = obj1['pcd']
    pcd2 = obj2['pcd']
    
    # Get the coordinates of the bounding boxes
    box_points1 = np.asarray(bbox1.get_box_points())
    box_points2 = np.asarray(bbox2.get_box_points())
    
    # Re-order the points to fit the format required in pytorch3d
    # xyz should be [---, -+-, -++, --+,    +--, ++-, +++, +-+]
    box_points1 = box_points1[[0, 2, 5, 3, 1, 7, 4, 6]]
    box_points2 = box_points2[[0, 2, 5, 3, 1, 7, 4, 6]]
    
    # Computet the intersection of the two boxes
    try:
        vols, ious = ops.box3d_overlap(
            torch.from_numpy(box_points1).unsqueeze(0).float(), 
            torch.from_numpy(box_points2).unsqueeze(0).float()
        )
    except ValueError as e: # This indicates colinear
        union_volume = 0.0
        iou = 0.0
    else:
        union_volume = vols[0,0].item()
        iou = ious[0,0].item()
    
    # Join the two point cloud
    pcd_union = pcd1 + pcd2

    # compute_convex_hull() somehow cannot provide watertight mesh
    # union_hull_mesh, union_hull_point_list = pcd_union.compute_convex_hull()
    # enclosing_volume = union_hull_mesh.get_volume()
    
    # enclosing_box = pcd_union.get_minimal_oriented_bounding_box()
    enclosing_box = pcd_union.get_oriented_bounding_box()
    enclosing_volume = enclosing_box.volume()
    
    giou = iou - (enclosing_volume - union_volume) / enclosing_volume
    
    # print(vols, ious)
    # print(enclosing_volume, union_volume, iou, giou)
    
    # o3d.visualization.draw_geometries([
    #     pcd1, pcd2, bbox1, bbox2
    # ])
    
    return giou

def compute_3d_box_volume_batch(bbox: torch.Tensor) -> torch.Tensor:
    '''
    Compute the volume of a set of rectangular boxes.
    This assumes bbox corner order follows the open3d convention, which is:
    ---, +--, -+-, --+, +++, -++, +-+, ++-
    See https://github.com/isl-org/Open3D/blob/47f4ee936841ae9f8a9c4ce5d9162bd5b3e0279f/cpp/open3d/geometry/BoundingVolume.cpp#L92
    
    bbox: (N, 8, D)
    
    returns: (N,)
    '''
    a = torch.linalg.vector_norm(bbox[:, 0, :] - bbox[:, 1, :], ord=2, dim=1)
    b = torch.linalg.vector_norm(bbox[:, 0, :] - bbox[:, 2, :], ord=2, dim=1)
    c = torch.linalg.vector_norm(bbox[:, 0, :] - bbox[:, 3, :], ord=2, dim=1)
    
    vol = a * b * c
    return vol
    
def expand_3d_box(bbox: torch.Tensor, eps=0.02) -> torch.Tensor:
    '''
    Expand the side of 3D boxes such that each side has at least eps length.
    Assumes the bbox cornder order in open3d convention. 
    
    bbox: (N, 8, D)
    
    returns: (N, 8, D)
    '''
    center = bbox.mean(dim=1)  # shape: (N, D)

    va = bbox[:, 1, :] - bbox[:, 0, :]  # shape: (N, D)
    vb = bbox[:, 2, :] - bbox[:, 0, :]  # shape: (N, D)
    vc = bbox[:, 3, :] - bbox[:, 0, :]  # shape: (N, D)
    
    a = torch.linalg.vector_norm(va, ord=2, dim=1, keepdim=True)  # shape: (N, 1)
    b = torch.linalg.vector_norm(vb, ord=2, dim=1, keepdim=True)  # shape: (N, 1)
    c = torch.linalg.vector_norm(vc, ord=2, dim=1, keepdim=True)  # shape: (N, 1)
    
    va = torch.where(a < eps, va / a * eps, va)  # shape: (N, D)
    vb = torch.where(b < eps, vb / b * eps, vb)  # shape: (N, D)
    vc = torch.where(c < eps, vc / c * eps, vc)  # shape: (N, D)
    
    new_bbox = torch.stack([
        center - va/2.0 - vb/2.0 - vc/2.0,
        center + va/2.0 - vb/2.0 - vc/2.0,
        center - va/2.0 + vb/2.0 - vc/2.0,
        center - va/2.0 - vb/2.0 + vc/2.0,
        center + va/2.0 + vb/2.0 + vc/2.0,
        center - va/2.0 + vb/2.0 + vc/2.0,
        center + va/2.0 - vb/2.0 + vc/2.0,
        center + va/2.0 + vb/2.0 - vc/2.0,
    ], dim=1) # shape: (N, 8, D)
    
    new_bbox = new_bbox.to(bbox.device)
    new_bbox = new_bbox.type(bbox.dtype)
    
    return new_bbox
    
def compute_enclosing_vol(bbox1: torch.Tensor, bbox2: torch.Tensor) -> torch.Tensor:
    '''
    Compute the enclosing volume between every pair of boxes in bbox1 and bbox2.
    This is an accurate but slow version using convex hull
    
    bbox1: (M, 8, D)
    bbox2: (N, 8, D)
    
    returns: (M, N)
    '''
    M = bbox1.shape[0]
    N = bbox2.shape[0]
    
    enclosing_vol = torch.zeros((M, N), dtype=bbox1.dtype, device=bbox1.device)
    for i in range(bbox1.shape[0]):
        for j in range(bbox2.shape[0]):
            pcd_union = o3d.geometry.PointCloud()
            bbox_points_union = torch.cat([bbox1[i], bbox2[j]], dim=0) # (16, 3)
            pcd_union.points = o3d.utility.Vector3dVector(bbox_points_union.cpu().numpy())
            enclosing_mesh, _ = pcd_union.compute_convex_hull(joggle_inputs=True)
            try:
                enclosing_vol[i, j] = enclosing_mesh.get_volume()
            except:
                # This occurs commonly when the enclosing_mesh is not watertight.
                enclosing_mesh = pcd_union.get_axis_aligned_bounding_box()
                enclosing_vol[i, j] = enclosing_mesh.volume()
                
    return enclosing_vol
    
def compute_enclosing_vol_fast(bbox1: torch.Tensor, bbox2: torch.Tensor) -> torch.Tensor:
    '''
    Compute the enclosing volume between every pair of boxes in bbox1 and bbox2.
    This is fast but approximate version using axis-aligned bounding box
    
    bbox1: (M, 8, 3)
    bbox2: (N, 8, 3)
    
    returns: (M, N)
    '''
    M = bbox1.shape[0]
    N = bbox2.shape[0]
    
    # Expand dimensions to compute the pairwise enclosing box
    bbox1 = bbox1.unsqueeze(1).expand(-1, N, -1, -1) # (M, N, 8, 3)
    bbox2 = bbox2.unsqueeze(0).expand(M, -1, -1, -1) # (M, N, 8, 3)
    
    # Compute the minimum and maximum coordinates for each pair of boxes
    min_coords = torch.minimum(bbox1, bbox2).amin(dim=2) # (M, N, 3)
    max_coords = torch.maximum(bbox1, bbox2).amax(dim=2) # (M, N, 3)

    # Compute the dimensions of the enclosing boxes
    enclosing_dims = max_coords - min_coords # (M, N, 3)
    
    # Clamp the dimensions to be non-negative (in case there's no overlap)
    enclosing_dims = torch.clamp(enclosing_dims, min=0) # (M, N, 3)
    
    # Compute the volume of the enclosing boxes
    vol = enclosing_dims[:, :, 0] * enclosing_dims[:, :, 1] * enclosing_dims[:, :, 2] # (M, N)

    return vol

def compute_3d_giou_accurate_batch(bbox1: torch.Tensor, bbox2: torch.Tensor) -> torch.Tensor:
    '''
    Compute Generalized IoU between two sets of oriented (or axis-aligned) 3D bounding boxes.
    
    bbox1: (M, 8, D), e.g. (M, 8, 3)
    bbox2: (N, 8, D), e.g. (N, 8, 3)
    
    returns: (M, N)
    '''
    # Must expend the box beforehand, otherwise it may results overestimated results
    bbox1 = expand_3d_box(bbox1, 0.02)
    bbox2 = expand_3d_box(bbox2, 0.02)
    
    bbox1_vol = compute_3d_box_volume_batch(bbox1)
    bbox2_vol = compute_3d_box_volume_batch(bbox2)
    
    import pytorch3d.ops as ops

    inter_vol, iou = ops.box3d_overlap(
        bbox1[:, [0, 2, 5, 3, 1, 7, 4, 6]].float(), 
        bbox2[:, [0, 2, 5, 3, 1, 7, 4, 6]].float()
    )
    union_vol = bbox1_vol.unsqueeze(1) + bbox2_vol.unsqueeze(0) - inter_vol
    
    enclosing_vol = compute_enclosing_vol(bbox1, bbox2)
    # enclosing_vol = compute_enclosing_vol_fast(bbox1, bbox2)
    
    giou = iou - (enclosing_vol - union_vol) / enclosing_vol
    
    return giou

def compute_3d_contain_ratio_accurate_batch(bbox1: torch.Tensor, bbox2: torch.Tensor) -> torch.Tensor:
    '''
    Compute for i-th box in bbox1, how much of it is contained in j-th box in bbox2.
    
    bbox1: (M, 8, D), e.g. (M, 8, 3)
    bbox2: (N, 8, D), e.g. (N, 8, 3)
    
    returns: (M, N)
    '''
    # Must expend the box beforehand, otherwise it may results overestimated results
    bbox1 = expand_3d_box(bbox1)
    bbox2 = expand_3d_box(bbox2)
    
    bbox1_vol = compute_3d_box_volume_batch(bbox1) # (M,)
    bbox2_vol = compute_3d_box_volume_batch(bbox2) # (M,)
    
    import pytorch3d.ops as ops
    
    inter_vol, iou = ops.box3d_overlap(
        bbox1[:, [0, 2, 5, 3, 1, 7, 4, 6]].float(), 
        bbox2[:, [0, 2, 5, 3, 1, 7, 4, 6]].float()
    ) # (M, N), (M, N)
    
    contain_ratio = inter_vol / bbox1_vol.unsqueeze(1) # (M, N)
    
    # Seems the following bug is unavoidable but happens very rarely
    # print((contain_ratio > 1.001).sum() / contain_ratio.numel())
    # if contain_ratio.amax() > 1.1:
    #     print('contain_ratio > 1.0')
    #     import pdb; pdb.set_trace()
    
    # Therefore we manually clamp it to [0, 1]
    contain_ratio = contain_ratio.clamp(min=0, max=1)
    
    return contain_ratio, iou

def compute_2d_box_contained_batch(bbox: torch.Tensor, thresh:float=0.95) -> torch.Tensor:
    '''
    For each bbox, compute how many other bboxes are containing it. 
    First compute the area of the intersection between each pair of bboxes. 
    Then for each bbox, count how many bboxes have the intersection area larger than thresh of its own area.
    
    bbox: (N, 4), in (x1, y1, x2, y2) format
    
    returns: (N,)
    '''
    N = bbox.shape[0]

    # Get areas of each bbox
    areas = (bbox[:, 2] - bbox[:, 0]) * (bbox[:, 3] - bbox[:, 1])

    # Compute intersection boxes
    lt = torch.max(bbox[:, None, :2], bbox[:, :2])  # left-top points
    rb = torch.min(bbox[:, None, 2:], bbox[:, 2:])  # right-bottom points

    inter = (rb - lt).clamp(min=0)  # intersection sizes (dx, dy), if no overlap, clamp to zero

    # Compute areas of intersection boxes
    inter_areas = inter[:, :, 0] * inter[:, :, 1]

    # Count how many boxes have intersection area larger than thresh of its own area
    mask = inter_areas > (areas * thresh).unsqueeze(1)
    count = mask.sum(dim=1) - 1  # exclude itself

    return count

def mask_subtract_contained(xyxy: np.ndarray, mask: np.ndarray, th1=0.8, th2=0.7):
    '''
    Compute the containing relationship between all pair of bounding boxes.
    For each mask, subtract the mask of bounding boxes that are contained by it.
     
    Args:
        xyxy: (N, 4), in (x1, y1, x2, y2) format
        mask: (N, H, W), binary mask
        th1: float, threshold for computing intersection over box1
        th2: float, threshold for computing intersection over box2
        
    Returns:
        mask_sub: (N, H, W), binary mask
    '''
    N = xyxy.shape[0] # number of boxes

    # Get areas of each xyxy
    areas = (xyxy[:, 2] - xyxy[:, 0]) * (xyxy[:, 3] - xyxy[:, 1]) # (N,)

    # Compute intersection boxes
    lt = np.maximum(xyxy[:, None, :2], xyxy[None, :, :2])  # left-top points (N, N, 2)
    rb = np.minimum(xyxy[:, None, 2:], xyxy[None, :, 2:])  # right-bottom points (N, N, 2)
    
    inter = (rb - lt).clip(min=0)  # intersection sizes (dx, dy), if no overlap, clamp to zero (N, N, 2)

    # Compute areas of intersection boxes
    inter_areas = inter[:, :, 0] * inter[:, :, 1] # (N, N)
    
    inter_over_box1 = inter_areas / areas[:, None] # (N, N)
    # inter_over_box2 = inter_areas / areas[None, :] # (N, N)
    inter_over_box2 = inter_over_box1.T # (N, N)
    
    # if the intersection area is smaller than th2 of the area of box1, 
    # and the intersection area is larger than th1 of the area of box2,
    # then box2 is considered contained by box1
    contained = (inter_over_box1 < th2) & (inter_over_box2 > th1) # (N, N)
    contained_idx = contained.nonzero() # (num_contained, 2)

    mask_sub = mask.copy() # (N, H, W)
    # mask_sub[contained_idx[0]] = mask_sub[contained_idx[0]] & (~mask_sub[contained_idx[1]])
    for i in range(len(contained_idx[0])):
        mask_sub[contained_idx[0][i]] = mask_sub[contained_idx[0][i]] & (~mask_sub[contained_idx[1][i]])

    return mask_sub
conceptgraph/utils/ious.py
concept-graphs/concept-graphs
import copy
from typing import Iterable
import dataclasses
from PIL import Image

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import torch
import open3d as o3d

import supervision as sv
from supervision.draw.color import Color, ColorPalette
from conceptgraph.slam.slam_classes import MapObjectList

class OnlineObjectRenderer():
    '''
    Refactor of the open3d visualization code to make it more modular
    '''
    def __init__(
        self, 
        view_param: str | dict,
        base_objects: MapObjectList | None = None,
        gray_map: bool = False
    ) -> None:
        # If the base objects are provided, we will visualize them
        if base_objects is not None:
            self.n_base_objects = len(base_objects)

            base_pcds_vis = copy.deepcopy(base_objects.get_values("pcd"))
            base_bboxes_vis = copy.deepcopy(base_objects.get_values("bbox"))
            for i in range(self.n_base_objects):
                base_pcds_vis[i] = base_pcds_vis[i].voxel_down_sample(voxel_size=0.08)
                if gray_map:
                    base_pcds_vis[i].paint_uniform_color([0.5, 0.5, 0.5])
            for i in range(self.n_base_objects):
                base_bboxes_vis[i].color = [0.5, 0.5, 0.5]
            
            self.base_pcds_vis = base_pcds_vis
            self.base_bboxes_vis = base_bboxes_vis
        else:
            self.n_base_objects = 0
        
        self.est_traj = []
        self.gt_traj = []
        
        self.cmap = matplotlib.colormaps.get_cmap("turbo")

        if isinstance(view_param, str):
            self.view_param = o3d.io.read_pinhole_camera_parameters(view_param)
        else:
            self.view_param = view_param
            
        self.window_height = self.view_param.intrinsic.height
        self.window_width = self.view_param.intrinsic.width
        
        self.vis = o3d.visualization.Visualizer()
        self.vis.create_window(
            width = self.window_width,
            height = self.window_height,
        )
        
        self.vis_ctrl = self.vis.get_view_control()
        self.vis_ctrl.convert_from_pinhole_camera_parameters(self.view_param)
        
    def filter_base_by_mask(self, mask: Iterable[bool]):
        assert len(mask) == self.n_base_objects
        self.base_pcds_vis = [pcd for pcd, m in zip(self.base_pcds_vis, mask) if m]
        self.base_bboxes_vis = [bbox for bbox, m in zip(self.base_bboxes_vis, mask) if m]
        self.n_base_objects = len(self.base_pcds_vis)
    
    def step(
        self,
        image: Image.Image,
        pcds: list[o3d.geometry.PointCloud] | None = None,
        pcd_colors: np.ndarray | None = None,
        est_pose: np.ndarray | None = None,
        gt_pose: np.ndarray | None = None,
        base_objects_color: dict | None = None,
        new_objects: MapObjectList = None,
        paint_new_objects: bool = True,
        return_vis_handle: bool = False,
    ):
        # Remove all the geometries
        self.vis.clear_geometries()
        
        # Add the pose cameras and trajectories
        if est_pose is not None:
            self.est_traj.append(est_pose)
            est_camera_frustum = better_camera_frustum(
                est_pose, image.height, image.width, scale=0.5, color=[1., 0, 0]
            )
            self.vis.add_geometry(est_camera_frustum)
            if len(self.est_traj) > 1:
                est_traj_lineset = poses2lineset(np.stack(self.est_traj), color=[1., 0, 0])
                self.vis.add_geometry(est_traj_lineset)
            
        if gt_pose is not None:
            self.gt_traj.append(gt_pose)
            gt_camera_frustum = better_camera_frustum(
                gt_pose, image.height, image.width, scale=0.5, color=[0, 1., 0]
            )
            self.vis.add_geometry(gt_camera_frustum)
            if len(self.gt_traj) > 1:
                gt_traj_lineset = poses2lineset(np.stack(self.gt_traj), color=[0, 1., 0])
                self.vis.add_geometry(gt_traj_lineset)
    
        # Add the base objects
        if self.n_base_objects > 0:
            if base_objects_color is not None:
                for obj_id in range(self.n_base_objects):
                    color = base_objects_color[obj_id]
                    self.base_pcds_vis[obj_id].paint_uniform_color(color)
                    self.base_bboxes_vis[obj_id].color = color
            
            for geom in self.base_pcds_vis + self.base_bboxes_vis:
                self.vis.add_geometry(geom)
            
        # Show the extra pcds to visualize
        if pcds is not None:
            for i in range(len(pcds)):
                pcds[i].transform(est_pose)
                if pcd_colors is not None:
                    pcds[i].paint_uniform_color(pcd_colors[i][:3])
                self.vis.add_geometry(pcds[i])
            
        # Show the extra new objects
        if new_objects is not None:
            for obj in new_objects:
                pcd = copy.deepcopy(obj['pcd'])
                bbox = copy.deepcopy(obj['bbox'])
                bbox.color = [0.0, 0.0, 1.0]
                if paint_new_objects:
                    pcd.paint_uniform_color([0.0, 1.0, 0.0])
                    bbox.color = [0.0, 1.0, 0.0]
                
                self.vis.add_geometry(pcd)
                self.vis.add_geometry(bbox)
        
        self.vis_ctrl.convert_from_pinhole_camera_parameters(self.view_param)
        
        self.vis.poll_events()
        self.vis.update_renderer()
        
        rendered_image = self.vis.capture_screen_float_buffer(False)
        rendered_image = np.asarray(rendered_image)
        
        if return_vis_handle:
            return rendered_image, self.vis
        else:
            return rendered_image, None

def get_random_colors(num_colors):
    '''
    Generate random colors for visualization
    
    Args:
        num_colors (int): number of colors to generate
        
    Returns:
        colors (np.ndarray): (num_colors, 3) array of colors, in RGB, [0, 1]
    '''
    colors = []
    for i in range(num_colors):
        colors.append(np.random.rand(3))
    colors = np.array(colors)
    return colors

def show_mask(mask, ax, random_color=False):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([30/255, 144/255, 255/255, 0.6])
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)
    
def show_points(coords, labels, ax, marker_size=375):
    pos_points = coords[labels==1]
    neg_points = coords[labels==0]
    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)   
    
def show_box(box, ax, label=None):
    x0, y0 = box[0], box[1]
    w, h = box[2] - box[0], box[3] - box[1]
    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0,0,0,0), lw=2))    
    
    if label is not None:
        ax.text(x0, y0, label)
        
def vis_result_fast(
    image: np.ndarray, 
    detections: sv.Detections, 
    classes: list[str], 
    color: Color | ColorPalette = ColorPalette.default(), 
    instance_random_color: bool = False,
    draw_bbox: bool = True,
) -> np.ndarray:
    '''
    Annotate the image with the detection results. 
    This is fast but of the same resolution of the input image, thus can be blurry. 
    '''
    # annotate image with detections
    box_annotator = sv.BoxAnnotator(
        color = color,
        text_scale=0.3,
        text_thickness=1,
        text_padding=2,
    )
    mask_annotator = sv.MaskAnnotator(
        color = color
    )
    labels = [
        f"{classes[class_id]} {confidence:0.2f}" 
        for _, _, confidence, class_id, _ 
        in detections]
    
    if instance_random_color:
        # generate random colors for each segmentation
        # First create a shallow copy of the input detections
        detections = dataclasses.replace(detections)
        detections.class_id = np.arange(len(detections))
        
    annotated_image = mask_annotator.annotate(scene=image.copy(), detections=detections)
    
    if draw_bbox:
        annotated_image = box_annotator.annotate(scene=annotated_image, detections=detections, labels=labels)
    return annotated_image, labels

def vis_result_slow_caption(image, masks, boxes_filt, pred_phrases, caption, text_prompt):
    '''
    Annotate the image with detection results, together with captions and text prompts.
    This function is very slow, but the output is more readable.
    '''
    plt.figure(figsize=(10, 10))
    plt.imshow(image)
    for mask in masks:
        show_mask(mask, plt.gca(), random_color=True)
    for box, label in zip(boxes_filt, pred_phrases):
        show_box(box, plt.gca(), label)

    plt.title('Tagging-Caption: ' + caption + '\n' + 'Tagging-classes: ' + text_prompt + '\n')
    plt.axis('off')
    
    # Convert the fig to a numpy array
    fig = plt.gcf()
    fig.tight_layout(pad=0)
    fig.canvas.draw()
    vis_image = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
    vis_image = vis_image.reshape(fig.canvas.get_width_height()[::-1] + (3,))
    plt.close()
    
    return vis_image

def vis_sam_mask(anns):
    sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)

    img = np.ones((sorted_anns[0]['segmentation'].shape[0], sorted_anns[0]['segmentation'].shape[1], 4))
    img[:,:,3] = 0
    for ann in sorted_anns:
        m = ann['segmentation']
        color_mask = np.concatenate([np.random.random(3), [0.35]])
        img[m] = color_mask
        
    return img

def poses2lineset(poses, color=[0, 0, 1]):
    '''
    Create a open3d line set from a batch of poses

    poses: (N, 4, 4)
    color: (3,)
    '''
    N = poses.shape[0]
    lineset = o3d.geometry.LineSet()
    lineset.points = o3d.utility.Vector3dVector(poses[:, :3, 3])
    lineset.lines = o3d.utility.Vector2iVector(
        np.array([[i, i + 1] for i in range(N - 1)])
    )
    lineset.colors = o3d.utility.Vector3dVector([color for _ in range(len(lineset.lines))])
    return lineset

def create_camera_frustum(
    camera_pose, width=1, height=1, z_near=0.5, z_far=1, color=[0, 0, 1]
):
    K = np.array([[z_near, 0, 0], [0, z_near, 0], [0, 0, z_near + z_far]])
    points = np.array(
        [
            [-width / 2, -height / 2, z_near],
            [width / 2, -height / 2, z_near],
            [width / 2, height / 2, z_near],
            [-width / 2, height / 2, z_near],
            [0, 0, 0],
        ]
    )
    points_transformed = camera_pose[:3, :3] @ (K @ points.T) + camera_pose[:3, 3:4]
    points_transformed = points_transformed.T
    frustum = o3d.geometry.LineSet()
    frustum.points = o3d.utility.Vector3dVector(points_transformed)
    lines = [[0, 1], [1, 2], [2, 3], [3, 0], [4, 0], [4, 1], [4, 2], [4, 3]]
    frustum.lines = o3d.utility.Vector2iVector(lines)
    frustum.colors = o3d.utility.Vector3dVector([color for _ in range(len(lines))])
    return frustum


def better_camera_frustum(camera_pose, img_h, img_w, scale=3.0, color=[0, 0, 1]):
    # Convert camera pose tensor to numpy array
    if isinstance(camera_pose, torch.Tensor):
        camera_pose = camera_pose.numpy()
    
    # Define near and far distance (adjust these as needed)
    near = scale * 0.1
    far = scale * 1.0
    
    # Define frustum dimensions at the near plane (replace with appropriate values)
    frustum_h = near
    frustum_w = frustum_h * img_w / img_h  # Set frustum width based on its height and the image aspect ratio
    
    # Compute the 8 points that define the frustum
    points = []
    for x in [-1, 1]:
        for y in [-1, 1]:
            for z in [-1, 1]:
                u = x * (frustum_w // 2 if z == -1 else frustum_w * far / near)
                v = y * (frustum_h // 2 if z == -1 else frustum_h * far / near)
                d = near if z == -1 else far # Negate depth here
                # d = -near if z == -1 else -far # Negate depth here
                point = np.array([u, v, d, 1]).reshape(-1, 1)
                transformed_point = (camera_pose @ point).ravel()[:3]
                # transformed_point[0] *= -1  # Flip X-coordinate
                points.append(transformed_point) # Using camera pose directly
                # points.append((camera_pose_np @ point).ravel()[:3]) # Using camera pose directly
    
    # Create lines that connect the 8 points
    lines = [[0, 1], [1, 3], [3, 2], [2, 0], [4, 5], [5, 7], [7, 6], [6, 4], [0, 4], [1, 5], [3, 7], [2, 6]]
    
    frustum = o3d.geometry.LineSet()
    frustum.points = o3d.utility.Vector3dVector(points)
    frustum.lines = o3d.utility.Vector2iVector(lines)
    frustum.colors = o3d.utility.Vector3dVector([color for _ in range(len(lines))])

    return frustum


# Copied from https://github.com/isl-org/Open3D/pull/738
def align_vector_to_another(a=np.array([0, 0, 1]), b=np.array([1, 0, 0])):
    """
    Aligns vector a to vector b with axis angle rotation
    """
    if np.array_equal(a, b):
        return None, None
    axis_ = np.cross(a, b)
    axis_ = axis_ / np.linalg.norm(axis_)
    angle = np.arccos(np.dot(a, b))

    return axis_, angle


def normalized(a, axis=-1, order=2):
    """Normalizes a numpy array of points"""
    l2 = np.atleast_1d(np.linalg.norm(a, order, axis))
    l2[l2 == 0] = 1
    return a / np.expand_dims(l2, axis), l2


class LineMesh(object):
    def __init__(self, points, lines=None, colors=[0, 1, 0], radius=0.15):
        """Creates a line represented as sequence of cylinder triangular meshes

        Arguments:
            points {ndarray} -- Numpy array of ponts Nx3.

        Keyword Arguments:
            lines {list[list] or None} -- List of point index pairs denoting line segments. If None, implicit lines from ordered pairwise points. (default: {None})
            colors {list} -- list of colors, or single color of the line (default: {[0, 1, 0]})
            radius {float} -- radius of cylinder (default: {0.15})
        """
        self.points = np.array(points)
        self.lines = np.array(
            lines) if lines is not None else self.lines_from_ordered_points(self.points)
        self.colors = np.array(colors)
        self.radius = radius
        self.cylinder_segments = []

        self.create_line_mesh()

    @staticmethod
    def lines_from_ordered_points(points):
        lines = [[i, i + 1] for i in range(0, points.shape[0] - 1, 1)]
        return np.array(lines)

    def create_line_mesh(self):
        first_points = self.points[self.lines[:, 0], :]
        second_points = self.points[self.lines[:, 1], :]
        line_segments = second_points - first_points
        line_segments_unit, line_lengths = normalized(line_segments)

        z_axis = np.array([0, 0, 1])
        # Create triangular mesh cylinder segments of line
        for i in range(line_segments_unit.shape[0]):
            line_segment = line_segments_unit[i, :]
            line_length = line_lengths[i]
            # get axis angle rotation to allign cylinder with line segment
            axis, angle = align_vector_to_another(z_axis, line_segment)
            # Get translation vector
            translation = first_points[i, :] + line_segment * line_length * 0.5
            # create cylinder and apply transformations
            cylinder_segment = o3d.geometry.TriangleMesh.create_cylinder(
                self.radius, line_length)
            cylinder_segment = cylinder_segment.translate(
                translation, relative=False)
            if axis is not None:
                axis_a = axis * angle
                # cylinder_segment = cylinder_segment.rotate(
                #     R=o3d.geometry.get_rotation_matrix_from_axis_angle(axis_a))
                # cylinder_segment = cylinder_segment.rotate(
                #   axis_a, center=True, type=o3d.geometry.RotationType.AxisAngle)
                cylinder_segment = cylinder_segment.rotate(
                    R=o3d.geometry.get_rotation_matrix_from_axis_angle(axis_a))
            # color cylinder
            color = self.colors if self.colors.ndim == 1 else self.colors[i, :]
            cylinder_segment.paint_uniform_color(color)

            self.cylinder_segments.append(cylinder_segment)

    def add_line(self, vis):
        """Adds this line to the visualizer"""
        for cylinder in self.cylinder_segments:
            vis.add_geometry(cylinder, reset_bounding_box=False)

    def remove_line(self, vis):
        """Removes this line from the visualizer"""
        for cylinder in self.cylinder_segments:
            vis.remove_geometry(cylinder, reset_bounding_box=False)
conceptgraph/utils/vis.py
concept-graphs/concept-graphs
from setuptools import setup, find_packages

setup(
    name='conceptgraph',
    version='1.0.0',
    description='ConceptGraphs: Open-Vocabulary 3D Scene Graphs for Perception and Planning',
    author='See https://concept-graphs.github.io/',
    packages=find_packages(),
)
setup.py
concept-graphs/concept-graphs