From 3ba4828278b587f2dad0171b643675a162882379 Mon Sep 17 00:00:00 2001
From: Tai-Wang <tab_wang@outlook.com>
Date: Tue, 26 Dec 2023 15:11:52 +0800
Subject: [PATCH] Support adjusting the line thickness of visualized boxes and
 update readme

---
 README.md                               |  77 +++++++++++-
 embodiedscan/explorer.py                |  56 +++------
 embodiedscan/tutorial.ipynb             |   8 ++
 embodiedscan/utils/continuous_drawer.py |  10 +-
 embodiedscan/utils/line_mesh.py         | 150 ++++++++++++++++++++++++
 embodiedscan/utils/utils.py             |  12 ++
 6 files changed, 268 insertions(+), 45 deletions(-)
 create mode 100644 embodiedscan/utils/line_mesh.py

diff --git a/README.md b/README.md
index bb47bc9..c82ce7d 100644
--- a/README.md
+++ b/README.md
@@ -44,6 +44,19 @@
 
 [![demo](assets/demo_fig.png "demo")](https://tai-wang.github.io/embodiedscan)
 
+<!-- contents with emoji -->
+
+## 📋 Contents
+
+1. [About](#-about)
+2. [News](#-news)
+3. [Getting Started](#-getting-started)
+4. [Model and Benchmark](#-model-and-benchmark)
+5. [TODO List](#-todo-list)
+6. [Citation](#-citation)
+7. [License](#-license)
+8. [Acknowledgements](#-acknowledgements)
+
 ## 🏠 About
 
 <!-- ![Teaser](assets/teaser.jpg) -->
@@ -62,9 +75,63 @@ Building upon this database, we introduce a baseline framework named <b>Embodied
 
 - \[2023-12\] We release the [paper](./assets/EmbodiedScan.pdf) of EmbodiedScan. Please check the [webpage](https://tai-wang.github.io/embodiedscan) and view our demos!
 
-## 🔍 Overview
+## 📚 Getting Started
+
+### Installation
 
-### Model
+We test our codes under the following environment:
+
+- Ubuntu 20.04
+- NVIDIA Driver: 525.147.05
+- CUDA 12.0
+- Python 3.8.18
+- PyTorch 1.11.0+cu113
+- PyTorch3D 0.7.2
+
+1. Clone this repository.
+
+```bash
+git clone https://github.com/OpenRobotLab/EmbodiedScan.git
+cd EmbodiedScan
+```
+
+2. Install [PyTorch3D](https://github.com/facebookresearch/pytorch3d/blob/main/INSTALL.md)
+
+```bash
+conda create -n embodiedscan python=3.8 -y  # pytorch3d needs python>3.7
+conda activate embodiedscan
+# We recommend installing pytorch3d with pre-compiled packages
+# For example, to install for Python 3.8, PyTorch 1.11.0 and CUDA 11.3
+# For more information, please refer to https://github.com/facebookresearch/pytorch3d/blob/main/INSTALL.md#2-install-wheels-for-linux
+pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py38_cu113_pyt1110/download.html
+```
+
+3. Install EmbodiedScan
+
+```bash
+# We plan to make EmbodiedScan easier to install by "pip install EmbodiedScan".
+# Please stay tuned for the future official release.
+# Make sure you are under ./EmbodiedScan/
+pip install -e .
+```
+
+### Data Preparation
+
+Please download ScanNet, 3RScan and matterport3d from its official website.
+
+We will release the demo data, re-organized file structure, post-processing script and annotation files in the near future.
+
+Please stay tuned.
+
+### Tutorial
+
+We provide a simple [tutorial](https://github.com/OpenRobotLab/EmbodiedScan/blob/main/embodiedscan/tutorial.ipynb) here as a guideline for the basic analysis and visualization of our dataset. Welcome to try and post your suggestions!
+
+## 📦 Model and Benchmark
+
+We will release the code for model training and benchmark with pretrained checkpoints in the 2024 Q1.
+
+### Model Overview
 
 <p align="center">
   <img src="assets/framework.png" align="center" width="100%">
@@ -91,6 +158,10 @@ Embodied Perceptron accepts RGB-D sequence with any number of views along with t
 <video src="assets/scannet_two_bed_demo.mp4" controls>
 </video> -->
 
+### Benchmark
+
+Please see the [paper](./assets/EmbodiedScan.pdf) for details of our two benchmarks, fundamental 3D perception and language-grounded benchmarks. This dataset is still scaling up and the benchmark is being polished and extended. Please stay tuned for our recent updates.
+
 ## 📝 TODO List
 
 - \[x\] Paper and partial code release.
@@ -153,7 +224,7 @@ This work is under the <a rel="license" href="http://creativecommons.org/license
 
 ## 👏 Acknowledgements
 
-- [OpenMMLab](https://github.com/open-mmlab): Our dataset code uses [MMEngine](https://github.com/open-mmlab/mmengine) and our model is built upon of [MMDetection3D](https://github.com/open-mmlab/mmdetection3d).
+- [OpenMMLab](https://github.com/open-mmlab): Our dataset code uses [MMEngine](https://github.com/open-mmlab/mmengine) and our model is built upon [MMDetection3D](https://github.com/open-mmlab/mmdetection3d).
 - [PyTorch3D](https://github.com/facebookresearch/pytorch3d): We use some functions supported in PyTorch3D for efficient computations on fundamental 3D data structures.
 - [ScanNet](https://github.com/ScanNet/ScanNet), [3RScan](https://github.com/WaldJohannaU/3RScan), [Matterport3D](https://github.com/niessner/Matterport): Our dataset uses the raw data from these datasets.
 - [ReferIt3D](https://github.com/referit3d/referit3d): We refer to the SR3D's approach to obtaining the language prompt annotations.
diff --git a/embodiedscan/explorer.py b/embodiedscan/explorer.py
index bf8ceee..c5d2e11 100644
--- a/embodiedscan/explorer.py
+++ b/embodiedscan/explorer.py
@@ -9,6 +9,7 @@
 from embodiedscan.utils.continuous_drawer import (ContinuousDrawer,
                                                   ContinuousOccupancyDrawer)
 from embodiedscan.utils.img_drawer import ImageDrawer
+from embodiedscan.utils.utils import _9dof_to_box, _box_add_thickness
 
 DATASETS = ['scannet', '3rscan', 'matterport3d']
 
@@ -27,13 +28,12 @@ class EmbodiedScanExplorer:
             Defaults to None.
     """
 
-    def __init__(
-        self,
-        data_root: Union[dict, List],
-        ann_file: Union[dict, List, str],
-        verbose: bool = False,
-        color_setting: str = None,
-    ):
+    def __init__(self,
+                 data_root: Union[dict, List],
+                 ann_file: Union[dict, List, str],
+                 verbose: bool = False,
+                 color_setting: str = None,
+                 thickness: float = 0.01):
 
         if isinstance(ann_file, dict):
             ann_file = list(ann_file.values())
@@ -56,6 +56,7 @@ def __init__(
             self.data_root = data_root
 
         self.verbose = verbose
+        self.thickness = thickness
 
         if self.verbose:
             print('Dataset root')
@@ -239,9 +240,10 @@ def render_scene(self, scene_name, render_box=False):
             if self.verbose:
                 print('Rendering box')
             for instance in select['instances']:
-                box = self._9dof_to_box(instance['bbox_3d'],
-                                        instance['bbox_label_3d'])
-                boxes.append(box)
+                box = _9dof_to_box(instance['bbox_3d'],
+                                   self.classes[instance['bbox_label_3d'] - 1],
+                                   self.color_selector)
+                boxes += _box_add_thickness(box, self.thickness)
             if self.verbose:
                 print('Rendering complete')
         o3d.visualization.draw_geometries([mesh, frame] + boxes)
@@ -297,7 +299,7 @@ def render_continuous_scene(self,
         drawer = ContinuousDrawer(dataset, self.data_root[dataset],
                                   selected_scene, self.classes,
                                   self.color_selector, start_idx,
-                                  pcd_downsample)
+                                  pcd_downsample, self.thickness)
         drawer.begin()
 
     def render_continuous_occupancy(self, scene_name, start_cam=None):
@@ -443,8 +445,10 @@ def show_image(self, scene_name, camera_name, render_box=False):
                         print('Rendering box')
                     for i in camera['visible_instance_ids']:
                         instance = select['instances'][i]
-                        box = self._9dof_to_box(instance['bbox_3d'],
-                                                instance['bbox_label_3d'])
+                        box = _9dof_to_box(
+                            instance['bbox_3d'],
+                            self.classes[instance['bbox_label_3d'] - 1],
+                            self.color_selector)
                         label = self.classes[instance['bbox_label_3d'] - 1]
                         color = self.color_selector.get_color(label)
                         img_drawer.draw_box3d(box,
@@ -461,32 +465,6 @@ def show_image(self, scene_name, camera_name, render_box=False):
         print('No such camera')
         return
 
-    def _9dof_to_box(self, box, label_id):
-        """Convert 9-DoF box annotations to open3d oriented boxes.
-
-        Args:
-            box (list | np.ndarray): Original box annotations.
-            label_id (int): Category ID.
-
-        Returns:
-            open3d.geometry.OrientedBoundingBox: Converted boxes for the
-            subsequent processing and visualization.
-        """
-        if isinstance(box, list):
-            box = np.array(box)
-        center = box[:3].reshape(3, 1)
-        scale = box[3:6].reshape(3, 1)
-        rot = box[6:].reshape(3, 1)
-        rot_mat = \
-            o3d.geometry.OrientedBoundingBox.get_rotation_matrix_from_zxy(rot)
-        geo = o3d.geometry.OrientedBoundingBox(center, rot_mat, scale)
-
-        label = self.classes[label_id - 1]
-        color = self.color_selector.get_color(label)
-        color = [x / 255.0 for x in color]
-        geo.color = color
-        return geo
-
 
 if __name__ == '__main__':
     explorer = EmbodiedScanExplorer(
diff --git a/embodiedscan/tutorial.ipynb b/embodiedscan/tutorial.ipynb
index d9ebfb0..1f3dea8 100644
--- a/embodiedscan/tutorial.ipynb
+++ b/embodiedscan/tutorial.ipynb
@@ -1105,6 +1105,14 @@
       "Rendering box\n",
       "Rendering complete\n"
      ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/data/EmbodiedScan/release/EmbodiedScan/embodiedscan/utils/line_mesh.py:25: RuntimeWarning: invalid value encountered in divide\n",
+      "  axis_ = axis_ / np.linalg.norm(axis_)\n"
+     ]
     }
    ],
    "source": [
diff --git a/embodiedscan/utils/continuous_drawer.py b/embodiedscan/utils/continuous_drawer.py
index ce38c88..242b4b3 100644
--- a/embodiedscan/utils/continuous_drawer.py
+++ b/embodiedscan/utils/continuous_drawer.py
@@ -5,7 +5,8 @@
 import numpy as np
 import open3d as o3d
 
-from .utils import _9dof_to_box, draw_camera, from_depth_to_point
+from .utils import (_9dof_to_box, _box_add_thickness, draw_camera,
+                    from_depth_to_point)
 
 
 class ContinuousDrawer:
@@ -26,7 +27,7 @@ class ContinuousDrawer:
     """
 
     def __init__(self, dataset, dir, scene, classes, color_selector, start_idx,
-                 pcd_downsample):
+                 pcd_downsample, thickness):
         self.dir = dir
         self.dataset = dataset
         self.scene = scene
@@ -34,6 +35,7 @@ def __init__(self, dataset, dir, scene, classes, color_selector, start_idx,
         self.color_selector = color_selector
         self.idx = start_idx
         self.downsample = pcd_downsample
+        self.thickness = thickness
         self.camera = None
         self.occupied = np.zeros((len(self.scene['instances']), ), dtype=bool)
         self.vis = o3d.visualization.VisualizerWithKeyCallback()
@@ -134,7 +136,9 @@ def draw_next(self, vis):
             box = _9dof_to_box(instance['bbox_3d'],
                                self.classes[instance['bbox_label_3d'] - 1],
                                self.color_selector)
-            vis.add_geometry(box)
+            box = _box_add_thickness(box, self.thickness)
+            for item in box:
+                vis.add_geometry(item)
 
         self.idx += 1
         ctr = vis.get_view_control()
diff --git a/embodiedscan/utils/line_mesh.py b/embodiedscan/utils/line_mesh.py
new file mode 100644
index 0000000..1d47a5b
--- /dev/null
+++ b/embodiedscan/utils/line_mesh.py
@@ -0,0 +1,150 @@
+"""Adapted from https://github.com/isl-
+org/Open3D/pull/738#issuecomment-564785941.
+
+Module which creates mesh lines from a line set
+Open3D relies upon using glLineWidth to set line width on a LineSet
+However, this method is now deprecated and not fully supporeted in
+newer OpenGL versions.
+See:
+    Open3D Github Pull Request:
+        https://github.com/intel-isl/Open3D/pull/738
+    Other Framework Issues:
+        https://github.com/openframeworks/openFrameworks/issues/3460
+
+This module aims to solve this by converting a line into a triangular mesh
+(which has thickness). The basic idea is to create a cylinder for each line
+segment, translate it, and then rotate it.
+
+License: MIT
+"""
+import numpy as np
+import open3d as o3d
+
+
+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))  # center=True)
+                # cylinder_segment = cylinder_segment.rotate(
+                #   axis_a, center=True,
+                #   type=o3d.geometry.RotationType.AxisAngle)
+            # 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)
+
+    def remove_line(self, vis):
+        """Removes this line from the visualizer."""
+        for cylinder in self.cylinder_segments:
+            vis.remove_geometry(cylinder)
+
+
+if __name__ == '__main__':
+
+    def main():
+        print('Demonstrating LineMesh vs LineSet')
+        # Create Line Set
+        points = [[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], [0, 0, 1],
+                  [1, 0, 1], [0, 1, 1], [1, 1, 1]]
+        lines = [[0, 1], [0, 2], [1, 3], [2, 3], [4, 5], [4, 6], [5, 7],
+                 [6, 7], [0, 4], [1, 5], [2, 6], [3, 7]]
+        colors = [[1, 0, 0] for i in range(len(lines))]
+
+        line_set = o3d.geometry.LineSet()
+        line_set.points = o3d.utility.Vector3dVector(points)
+        line_set.lines = o3d.utility.Vector2iVector(lines)
+        line_set.colors = o3d.utility.Vector3dVector(colors)
+
+        # Create Line Mesh 1
+        points = np.array(points) + [0, 0, 2]
+        line_mesh1 = LineMesh(points, lines, colors, radius=0.02)
+        line_mesh1_geoms = line_mesh1.cylinder_segments
+
+        # Create Line Mesh 1
+        points = np.array(points) + [0, 2, 0]
+        line_mesh2 = LineMesh(points, radius=0.03)
+        line_mesh2_geoms = line_mesh2.cylinder_segments
+
+        o3d.visualization.draw_geometries(
+            [line_set, *line_mesh1_geoms, *line_mesh2_geoms])
+
+    main()
diff --git a/embodiedscan/utils/utils.py b/embodiedscan/utils/utils.py
index dc22fad..eab1a08 100644
--- a/embodiedscan/utils/utils.py
+++ b/embodiedscan/utils/utils.py
@@ -2,6 +2,8 @@
 import numpy as np
 import open3d as o3d
 
+from .line_mesh import LineMesh
+
 
 def from_depth_to_point(rgb, depth, mask, intrinsic, depth_intrinsic,
                         extrinsic):
@@ -27,6 +29,16 @@ def from_depth_to_point(rgb, depth, mask, intrinsic, depth_intrinsic,
     return pc[:, :3], color
 
 
+def _box_add_thickness(box, thickness):
+    bbox_lines = o3d.geometry.LineSet.create_from_oriented_bounding_box(box)
+    bbox_lines_width = LineMesh(points=bbox_lines.points,
+                                lines=bbox_lines.lines,
+                                colors=box.color,
+                                radius=thickness)
+    results = bbox_lines_width.cylinder_segments
+    return results
+
+
 def _9dof_to_box(box, label, color_selector):
     if isinstance(box, list):
         box = np.array(box)