Annotations for the ScanNet dataset generated using scannotate and HOC-Search.
class ScanNetAnnotation(object):
    def __init__(self, scene_name, obj_annotation_list, inst_seg_3d, scene_type):
        self.scene_name = scene_name
        self.obj_annotation_list = obj_annotation_list
        self.inst_seg_3d = inst_seg_3d
        self.scene_type = scene_type


class ObjectAnnotation(object):
    def __init__(self, object_id, category_label, scannet_category_label, view_params, transform3d=None,
                 transform_dict=None, id_cad=None, catid_cad=None, cad_symmetry='__SYM_NONE', is_in_scan2cad=False,
                 scan2cad_annotation_dict=None):
        self.object_id = object_id
        self.category_label = category_label
        self.view_params = view_params
        self.scannet_category_label = scannet_category_label
        self.scan2cad_annotation_dict = scan2cad_annotation_dict

        self.transform3d = transform3d
        self.transform_dict = transform_dict
        self.id_cad = id_cad
        self.catid_cad = catid_cad
        self.cad_symmetry = cad_symmetry
        self.is_in_scan2cad = is_in_scan2cad

ScanNetAnnotation.py
stefan-ainetter/SCANnotateDataset
import os
import numpy as np


def load_axis_alignment_mat(meta_file_path):
    metaFile = meta_file_path  # includes axisAlignment info
    assert os.path.exists(metaFile), '%s' % metaFile
    axis_align_matrix = np.identity(4)
    if os.path.isfile(metaFile):
        with open(metaFile) as f:
            lines = f.readlines()

        for line in lines:
            if 'axisAlignment' in line:
                axis_align_matrix = [float(x) \
                                     for x in line.rstrip().strip('axisAlignment = ').split(' ')]
                break

    return axis_align_matrix

load_ScanNet_data.py
stefan-ainetter/SCANnotateDataset
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)
        # line_segments_unit += np.array([1e-2, 0, 0])

        z_axis = np.array([0, 0, 1])
        # Create triangular mesh cylinder segments of line
        # assert line_segments_unit.shape[0] == 12
        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# + np.array([1e-2, 0, 0])
            # create cylinder and apply transformations
            # assert line_length > 0
            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
                aa = o3d.geometry.get_rotation_matrix_from_axis_angle(axis_a)
                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)
misc/line_mesh.py
stefan-ainetter/SCANnotateDataset
import argparse
import os
import sys

parser = argparse.ArgumentParser(description="Data preprocessing for ShapeNetCore.v2")
parser.add_argument("--device", type=str, default="", help="device")

os.environ["CUDA_VISIBLE_DEVICES"] = parser.parse_args().device

current = os.path.dirname(os.path.realpath(__file__))
parent = os.path.dirname(current)
sys.path.append(parent)

import torch
import numpy as np
from pytorch3d.io import IO
from misc.utils import shapenet_category_dict
import trimesh
import open3d as o3d
from misc.utils_CAD_retrieval import normalize_mesh


def simplify_mesh(mesh_path, max_faces=np.inf):
    mesh = trimesh.load(mesh_path, force='mesh', file_type='obj', resolver=None)

    obj_mesh = o3d.geometry.TriangleMesh()
    obj_mesh.vertices = o3d.utility.Vector3dVector(mesh.vertices)
    obj_mesh.triangles = o3d.utility.Vector3iVector(mesh.faces)

    obj_mesh.remove_unreferenced_vertices()
    obj_mesh.compute_vertex_normals()

    if np.asarray(obj_mesh.triangles).shape[0] <= max_faces:
        obj_mesh.remove_unreferenced_vertices()
        obj_mesh.remove_degenerate_triangles()
        obj_mesh.remove_duplicated_triangles()
        obj_mesh.remove_duplicated_vertices()

        return obj_mesh

    else:
        voxel_size = max(obj_mesh.get_max_bound() - obj_mesh.get_min_bound()) / 64
        mesh_smp = obj_mesh.simplify_vertex_clustering(
            voxel_size=voxel_size,
            contraction=o3d.geometry.SimplificationContraction.Quadric)

        mesh_smp.remove_unreferenced_vertices()
        mesh_smp.remove_degenerate_triangles()
        mesh_smp.remove_duplicated_triangles()
        mesh_smp.remove_duplicated_vertices()

        return mesh_smp


def preprocess_shapenet(shapenet_core_path, synset_id_core, shapenet_out_path, device):
    all_obj_list = os.listdir(os.path.join(shapenet_core_path, synset_id_core))

    for obj_id in all_obj_list:
        obj_path = os.path.join(shapenet_core_path, synset_id_core, obj_id, 'models', 'model_normalized.obj')
        if not os.path.exists(obj_path):
            continue

        out_path = os.path.join(shapenet_out_path, synset_id_core, obj_id, 'models')
        if not os.path.exists(out_path):
            os.makedirs(out_path)
        out_path_obj = os.path.join(out_path, 'model_normalized.obj')

        if os.path.exists(out_path_obj):
            continue

        try:
            mesh_o3d = simplify_mesh(obj_path)
        except:
            continue

        verts = torch.tensor(np.asarray(mesh_o3d.vertices)).float().to(device)
        faces = torch.tensor(np.asarray(mesh_o3d.triangles)).to(device)
        mesh_normalized, _ = normalize_mesh(verts, faces, device)

        if not torch.isfinite(mesh_normalized.verts_packed()).all():
            print("Meshes contain nan or inf.")
        else:
            IO().save_mesh(data=mesh_normalized, path=out_path_obj, include_textures=False)

    return


def main(args):
    # Setup
    if args.device == '':
        if torch.cuda.is_available():
            device = torch.device("cuda:0")
            torch.cuda.set_device(device)
        else:
            device = torch.device("cpu")
    else:
        device = 'cuda:' + str(torch.cuda.current_device())

    shapenet_ScanNet_classes = ['trash bin', 'table', 'printer', 'basket', 'bookshelf', 'flowerpot',
                                'laptop', 'sofa', 'chair', 'file cabinet', 'display', 'bag', 'bathtub',
                                'cabinet', 'clock', 'bed', 'bench', 'stove', 'lamp', 'faucet', 'bowl',
                                'keyboard', 'piano', 'microwaves', 'dishwasher', 'washer', 'guitar', 'pillow',
                                'motorbike']

    shapenet_path = os.path.join(parent, 'data', 'ShapeNet')
    shapenet_core_path = os.path.join(shapenet_path, 'ShapeNetCore.v2')
    shapenet_out_path = os.path.join(shapenet_path, 'ShapeNet_preprocessed')

    if not os.path.exists(shapenet_out_path):
        os.makedirs(shapenet_out_path)

    for cls_name in shapenet_ScanNet_classes:
        print('Current Class = ' + str(cls_name))
        try:
            synset_id_core = shapenet_category_dict[cls_name]
        except:
            assert False

        preprocess_shapenet(shapenet_core_path, synset_id_core, shapenet_out_path, device)


if __name__ == "__main__":
    main(parser.parse_args())

misc/preprocess_ShapeNet.py
stefan-ainetter/SCANnotateDataset
import numpy as np
import open3d as o3d

COLOR_DETECTRON2 = np.array(
    [
        0.000, 0.447, 0.741,
        0.850, 0.325, 0.098,
        0.929, 0.694, 0.125,
        0.494, 0.184, 0.556,
        0.466, 0.674, 0.188,
        0.301, 0.745, 0.933,
        0.635, 0.078, 0.184,
        0.000, 0.667, 0.000,
        1.000, 0.000, 0.000,
        1.000, 0.500, 0.000,
        0.000, 0.000, 1.000,
        0.000, 1.000, 0.000,
        0.749, 0.749, 0.000,
        0.667, 0.000, 1.000,
        0.333, 0.333, 0.000,
        0.333, 0.667, 0.000,
        0.333, 1.000, 0.000,
        0.667, 0.333, 0.000,
        0.667, 0.667, 0.000,
        0.667, 1.000, 0.000,
        1.000, 0.333, 0.000,
        0.200, 0.733, 1.000,
        1.000, 1.000, 0.000,
        0.000, 0.333, 0.500,
        0.000, 0.667, 0.500,
        0.000, 1.000, 0.500,
        0.333, 0.000, 0.500,
        0.333, 0.333, 0.500,
        0.333, 0.667, 0.500,
        0.333, 1.000, 0.500,
        0.667, 0.000, 0.500,
        0.667, 0.333, 0.500,
        0.667, 0.667, 0.500,
        0.667, 1.000, 0.500,
        1.000, 0.000, 0.500,
        1.000, 0.333, 0.500,
        1.000, 0.667, 0.500,
        1.000, 1.000, 0.500,
        0.000, 0.333, 1.000,
        0.000, 0.667, 1.000,
        0.000, 1.000, 1.000,
        0.333, 0.000, 1.000,
        0.333, 0.333, 1.000,
        0.333, 0.667, 1.000,
        0.333, 1.000, 1.000,
        0.667, 0.000, 1.000,
        0.667, 0.333, 1.000,
        0.667, 0.667, 1.000,
        0.667, 1.000, 1.000,
        1.000, 0.000, 1.000,
        1.000, 0.333, 1.000,
        1.000, 0.667, 1.000,
        # 0.333, 0.000, 0.000,
        0.500, 0.000, 0.000,
        0.667, 0.000, 0.000,
        0.833, 0.000, 0.000,
        1.000, 0.000, 0.000,
        0.000, 0.167, 0.000,
        # 0.000, 0.333, 0.000,
        0.000, 0.500, 0.000,
        # 0.000, 0.667, 0.000,
        0.000, 0.833, 0.000,
        0.000, 1.000, 0.000,
        0.000, 0.000, 0.167,
        # 0.000, 0.000, 0.333,
        0.000, 0.000, 0.500,
        0.000, 0.000, 0.667,
        0.000, 0.000, 0.833,
        0.000, 0.000, 1.000,
        # 0.000, 0.000, 0.000,
        0.143, 0.143, 0.143,
        0.857, 0.857, 0.857,
        # 1.000, 1.000, 1.000
    ]).astype(np.float32).reshape(-1, 3)

SEMANTIC_IDXS = np.array(
    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
     30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40])

SEMANTIC_NAMES = np.array(['unannotated', 'wall', 'floor', 'cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 'window',
                           'bookshelf', 'picture', 'counter', 'blinds', 'desk', 'shelves', 'curtain', 'dresser',
                           'pillow',
                           'mirror', 'floor_mat', 'clothes', 'ceiling', 'book', 'refridgerator', 'television', 'paper',
                           'towel',
                           'shower curtain', 'box', 'whiteboard', 'person', 'night stand',
                           'toilet', 'sink', 'lamp', 'bathtub', 'bag', 'otherstructure', 'otherfurniture', 'otherprop'])

SEMANTIC_IDX2NAME = {0: 'unannotated', 1: 'wall', 2: 'floor', 3: 'cabinet', 4: 'bed', 5: 'chair', 6: 'sofa',
                     7: 'table', 8: 'door', 9: 'window', 10: 'bookshelf', 11: 'picture',
                     12: 'counter', 13: 'blinds', 14: 'desk', 15: 'shelves', 16: 'curtain', 17: 'dresser', 18: 'pillow',
                     19: 'mirror', 20: 'floor_mat', 21: 'clothes', 22: 'ceiling', 23: 'book', 24: 'refridgerator',
                     25: 'television', 26: 'paper', 27: 'towel', 28: 'shower curtain', 29: 'box', 30: 'whiteboard',
                     31: 'person', 32: 'night stand', 33: 'toilet', 34: 'sink', 35: 'lamp', 36: 'bathtub', 37: 'bag',
                     38: 'otherstructure', 39: 'otherfurniture', 40: 'otherprop'}

shapenet_category_dict = {'airplane': '02691156', 'trash bin': '02747177', 'bag': '02773838', 'basket': '02801938',
                          'bathtub': '02808440', 'bed': '02818832', 'bench': '02828884', 'birdhouse': '02843684',
                          'bookshelf': '02871439', 'bottle': '02876657', 'bowl': '02880940', 'bus': '02924116',
                          'cabinet': '02933112', 'camera': '02942699', 'can': '02946921', 'cap': '02954340',
                          'car': '02958343', 'cellphone': '02992529', 'chair': '03001627', 'clock': '03046257',
                          'keyboard': '03085013', 'dishwasher': '03207941', 'display': '03211117',
                          'earphone': '03261776', 'faucet': '03325088', 'file cabinet': '03337140',
                          'guitar': '03467517',
                          'helmet': '03513137', 'jar': '03593526', 'knife': '03624134', 'lamp': '03636649',
                          'laptop': '03642806', 'loudspeaker': '03691459', 'mailbox': '03710193',
                          'microphone': '03759954', 'microwaves': '03761084', 'motorbike': '03790512',
                          'mug': '03797390', 'piano': '03928116', 'pillow': '03938244', 'pistol': '03948459',
                          'flowerpot': '03991062', 'printer': '04004475', 'remote': '04074963', 'rifle': '04090263',
                          'rocket': '04099429', 'skateboard': '04225987', 'sofa': '04256520', 'stove': '04330267',
                          'table': '04379243', 'telephone': '04401088', 'tower': '04460130', 'train': '04468005',
                          'watercraft': '04530566', 'washer': '04554684', 'desk': '03179701', 'dresser': '03237340',
                          'bed cabinet': '20000008'}


def Ry(t):
    """Rotation about the y-axis."""
    c = np.cos(t)
    s = np.sin(t)
    return np.array([[c, 0, s],
                     [0, 1, 0],
                     [-s, 0, c]])


def Rx(t):
    """Rotation about the x-axis."""
    c = np.cos(t)
    s = np.sin(t)
    return np.array([[1, 0, 0],
                     [0, c, -s],
                     [0, s, c]])


def Rz(t):
    """Rotation about the z-axis."""
    c = np.cos(t)
    s = np.sin(t)
    return np.array([[c, -s, 0],
                     [s, c, 0],
                     [0, 0, 1]])


def alignPy3D_format(pclMesh, T_mat):
    verts = np.array(pclMesh.verts_list()[0])
    newVerts = np.ones((verts.shape[0], 4))
    newVerts[:, :3] = verts
    newVerts = newVerts.dot(T_mat.T)
    pclMesh.vertices = o3d.utility.Vector3dVector(newVerts[:, :3])
    return pclMesh


def transform_ScanNet_to_py3D():
    rot_tmp1 = Rx(np.deg2rad(-90))
    rot_tmp2 = Ry(np.deg2rad(-90))
    rot3 = np.asarray(np.dot(rot_tmp2, rot_tmp1))
    T = np.eye(4)
    T[:3, :3] = rot3
    return T


def transform_ARKIT_to_py3D():
    rot_tmp1 = Rx(np.deg2rad(-90))
    rot_tmp2 = Ry(np.deg2rad(-180))
    rot3 = np.asarray(np.dot(rot_tmp2, rot_tmp1))
    T = np.eye(4)
    T[:3, :3] = rot3
    return T


def alignPclMesh(pclMesh, axis_align_matrix=np.eye(4), T=np.eye(4)):
    if isinstance(pclMesh, o3d.geometry.TriangleMesh):
        verts = np.array(pclMesh.vertices)
        newVerts = np.ones((verts.shape[0], 4))
        newVerts[:, :3] = verts
        newVerts = newVerts.dot(axis_align_matrix.T)
        newVerts = newVerts.dot(T.T)
        pclMesh.vertices = o3d.utility.Vector3dVector(newVerts[:, :3])

    elif isinstance(pclMesh, o3d.geometry.PointCloud):
        points = np.array(pclMesh.points)
        newPoints = np.ones((points.shape[0], 4))
        newPoints[:, :3] = points
        newPoints = newPoints.dot(axis_align_matrix.T)
        newPoints = newPoints.dot(T.T)
        pclMesh.points = o3d.utility.Vector3dVector(newPoints[:, :3])

    elif isinstance(pclMesh, np.ndarray):
        points = pclMesh
        newPoints = np.ones((points.shape[0], 4))
        newPoints[:, :3] = points
        newPoints = newPoints.dot(axis_align_matrix.T)
        newPoints = newPoints.dot(T.T)
        return newPoints

    else:
        raise NotImplementedError

    return pclMesh

misc/utils.py
stefan-ainetter/SCANnotateDataset
import copy
import os

import numpy as np
import open3d as o3d
import torch
from pytorch3d.io import load_obj
from pytorch3d.structures import Meshes
from pytorch3d.transforms import Transform3d

from misc.utils import SEMANTIC_IDX2NAME, COLOR_DETECTRON2
from misc.line_mesh import LineMesh

def normalize_mesh(verts_init, faces, device):
    mesh = Meshes(
        verts=[verts_init],
        faces=[faces],
    )
    bbox = mesh.get_bounding_boxes().squeeze(dim=0)
    bbox = bbox.cpu().detach().numpy()

    center = torch.tensor(bbox.mean(1)).float().to(device)
    vector_x = np.array([bbox[0, 1] - bbox[0, 0], 0, 0])
    vector_y = np.array([0, bbox[1, 1] - bbox[1, 0], 0])
    vector_z = np.array([0, 0, bbox[2, 1] - bbox[2, 0]])

    coeff_x = np.linalg.norm(vector_x)
    coeff_y = np.linalg.norm(vector_y)
    coeff_z = np.linalg.norm(vector_z)

    mesh = mesh.offset_verts(-center)
    transform_func = Transform3d().scale(x=(1 / coeff_x), y=(1 / coeff_y), z=1 / coeff_z).to(device)
    tverts = transform_func.transform_points(mesh.verts_list()[0]).unsqueeze(dim=0)

    mesh = Meshes(
        verts=[tverts.squeeze(dim=0)],
        faces=[faces]
    )

    return mesh, tverts


def cut_meshes(mesh_o3d, indices_list):
    mesh_o3d_obj = copy.deepcopy(mesh_o3d)
    mesh_o3d_obj = mesh_o3d_obj.select_by_index(indices_list)
    mesh_o3d.remove_vertices_by_index(indices_list)

    face_list_bg = np.asarray(mesh_o3d.triangles)
    face_list_obj = np.asarray(mesh_o3d_obj.triangles)

    mesh_bg = Meshes(
        verts=[torch.tensor(np.asarray(mesh_o3d.vertices)).float()],
        faces=[torch.tensor(np.asarray(face_list_bg))]
    )

    mesh_obj = Meshes(
        verts=[torch.tensor(np.asarray(mesh_o3d_obj.vertices)).float()],
        faces=[torch.tensor(np.asarray(face_list_obj))]
    )
    return mesh_bg, mesh_obj


def load_textured_cad_model_prepro(model_path, cad_transform_base, cls_name, device='cpu'):
    """Load CAD models from ShapeNet_preprocessed (using run_shapenet_prepro.sh)."""

    try:
        sem_id = list(SEMANTIC_IDX2NAME.keys())[
            list(SEMANTIC_IDX2NAME.values()).index(cls_name)]
    except:
        sem_id = 0
    mesh_color = COLOR_DETECTRON2[sem_id]

    verts, faces_, _ = load_obj(model_path, load_textures=False, device=device)
    faces = faces_.verts_idx

    _, tverts = normalize_mesh(verts, faces, device)

    tverts_final = cad_transform_base.transform_points(tverts)

    if device == 'cpu':
        tverts_final_ary = tverts_final.squeeze(dim=0).detach().numpy()
        faces_ary = faces.detach().numpy().astype(np.int32)

    else:
        tverts_final_ary = tverts_final.squeeze(dim=0).cpu().detach().numpy()
        faces_ary = faces.cpu().detach().numpy().astype(np.int32)

    cad_model_o3d = o3d.geometry.TriangleMesh()
    cad_model_o3d.vertices = o3d.utility.Vector3dVector(tverts_final_ary)
    cad_model_o3d.triangles = o3d.utility.Vector3iVector(faces_ary)
    vertex_n = np.array(cad_model_o3d.vertices).shape[0]
    vertex_colors = np.ones((vertex_n, 3)) * mesh_color
    cad_model_o3d.vertex_colors = o3d.utility.Vector3dVector(vertex_colors)

    return cad_model_o3d


def load_textured_cad_model(model_path, cad_transform_base, cls_name, device='cpu'):
    """Load CAD models from raw ShapeNetCore.v2 and scale-normalize + center model."""

    try:
        sem_id = list(SEMANTIC_IDX2NAME.keys())[
            list(SEMANTIC_IDX2NAME.values()).index(cls_name)]
    except:
        sem_id = 0
    mesh_color = COLOR_DETECTRON2[sem_id]

    verts, faces_, _ = load_obj(model_path, load_textures=False, device=device)
    faces = faces_.verts_idx

    _, tverts = normalize_mesh(verts, faces, device)

    tverts_final = cad_transform_base.transform_points(tverts)

    if device == 'cpu':
        tverts_final_ary = tverts_final.squeeze(dim=0).detach().numpy()
        faces_ary = faces.detach().numpy().astype(np.int32)

    else:
        tverts_final_ary = tverts_final.squeeze(dim=0).cpu().detach().numpy()
        faces_ary = faces.cpu().detach().numpy().astype(np.int32)

    cad_model_o3d = o3d.geometry.TriangleMesh()
    cad_model_o3d.vertices = o3d.utility.Vector3dVector(tverts_final_ary)
    cad_model_o3d.triangles = o3d.utility.Vector3iVector(faces_ary)
    vertex_n = np.array(cad_model_o3d.vertices).shape[0]
    vertex_colors = np.ones((vertex_n, 3)) * mesh_color
    cad_model_o3d.vertex_colors = o3d.utility.Vector3dVector(vertex_colors)

    return cad_model_o3d


def drawOpen3dCylLines(bbListIn,col=None):

    lines = [[0, 1],[1,2],[2,3],[3,0],[4,5],[5,6],[6,7],[7,4],
             [0,4],[1,5],[2,6],[3,7]]

    line_sets = []

    for bb in bbListIn:
        points = bb
        if col is None:
            col = [0,0,1]
        colors = [col for i in range(len(lines))]

        line_mesh1 = LineMesh(points, lines, colors, radius=0.02)
        line_mesh1_geoms = line_mesh1.cylinder_segments
        line_sets = line_mesh1_geoms[0]
        for l in line_mesh1_geoms[1:]:
            line_sets = line_sets + l

    return line_sets
misc/utils_CAD_retrieval.py
stefan-ainetter/SCANnotateDataset
import argparse
import copy
import os
import sys

parser = argparse.ArgumentParser(description="Code for visualization of scannotate annotations for ScanNet")
parser.add_argument("--device", type=str, default="", help="device")

os.environ["CUDA_VISIBLE_DEVICES"] = parser.parse_args().device

current_dir = os.path.dirname(os.path.realpath(__file__))
sys.path.append(current_dir)

import pickle
import torch
import numpy as np
import open3d as o3d
from load_ScanNet_data import load_axis_alignment_mat
from misc.utils import transform_ScanNet_to_py3D, alignPclMesh
from misc.utils import COLOR_DETECTRON2
from misc.utils_CAD_retrieval import load_textured_cad_model_prepro, drawOpen3dCylLines
from ScanNetAnnotation import *


def load_scene_mesh(meta_file_path, scene_path):
    mesh_o3d = o3d.io.read_triangle_mesh(scene_path)

    # Transfer points to py3d coord system
    T_mat = transform_ScanNet_to_py3D()

    align_mat_Scannet = load_axis_alignment_mat(meta_file_path=meta_file_path)
    align_mat_Scannet = np.reshape(np.asarray(align_mat_Scannet), (4, 4))

    mesh_o3d = alignPclMesh(mesh_o3d, axis_align_matrix=align_mat_Scannet, T=T_mat)

    return mesh_o3d


def visualize_annotations(scene_obj, shapenet_preprocessed_path, device):
    obj_mesh_all = None
    lineSets_all = None
    for box_item in scene_obj.obj_annotation_list:

        model_path = os.path.join(shapenet_preprocessed_path, box_item.catid_cad, box_item.id_cad, 'models',
                                  'model_normalized.obj')
        cad_transform_base = box_item.transform3d.to(device)

        cad_model_o3d = load_textured_cad_model_prepro(model_path, cad_transform_base, box_item.category_label, device)

        if obj_mesh_all is None:
            obj_mesh_all = cad_model_o3d
        else:
            obj_mesh_all += cad_model_o3d

        x_shift = .5
        y_shift = .5
        z_shift = .5

        box = np.array([[x_shift, y_shift, z_shift], [-x_shift, y_shift, z_shift], [-x_shift, -y_shift, z_shift],
                        [x_shift, -y_shift, z_shift], [x_shift, y_shift, -z_shift], [-x_shift, y_shift, -z_shift],
                        [-x_shift, -y_shift, -z_shift], [x_shift, -y_shift, -z_shift]])

        box_tensor = torch.Tensor(box).to(device)

        box_transformed = cad_transform_base.transform_points(box_tensor)
        box_transformed = box_transformed.cpu().detach().numpy()
        if box_item.is_in_scan2cad:
            line_color = [0, 0, 1]
        else:
            line_color = [1, 0, 0]

        lineSets = drawOpen3dCylLines([box_transformed], line_color)

        if lineSets_all is None:
            lineSets_all = lineSets
        else:
            lineSets_all += lineSets

    return obj_mesh_all, lineSets_all


def main(args):
    # Setup
    if args.device == '':
        if torch.cuda.is_available():
            device = torch.device("cuda:0")
            torch.cuda.set_device(device)
        else:
            device = torch.device("cpu")
    else:
        device = 'cuda:' + str(torch.cuda.current_device())

    scannet_path = os.path.join(current_dir, 'data', 'ScanNet')
    annotations_path = os.path.join(scannet_path, 'annotations')
    scans_path = os.path.join(scannet_path, 'scans')
    shapenet_path = os.path.join(current_dir, 'data', 'ShapeNet', 'ShapeNet_preprocessed')

    scene_list = os.listdir(annotations_path)

    for scene_name in scene_list:

        meta_file_path = os.path.join(scans_path, scene_name,
                                      scene_name + '.txt')

        scene_path = os.path.join(scans_path, scene_name,
                                  scene_name + '_vh_clean_2.ply')

        annotation_file = os.path.join(annotations_path, scene_name, scene_name + '.pkl')

        if not os.path.exists(meta_file_path) or not os.path.exists(scene_path) or not os.path.exists(annotation_file):
            print('Data not available for scene ' + str(scene_name))
            continue

        scene_mesh = load_scene_mesh(meta_file_path, scene_path)

        pkl_file = open(annotation_file, 'rb')
        scene_obj = pickle.load(pkl_file)

        obj_mesh_all, lineSets_all = visualize_annotations(scene_obj, shapenet_path, device)

        final_inst_seg_3d = scene_obj.inst_seg_3d
        points_mesh = copy.deepcopy(np.asarray(scene_mesh.vertices))
        pcl_color_tmp = np.zeros((points_mesh.shape[0], 3))
        for color_id in np.unique(final_inst_seg_3d):
            if color_id < 1.:
                continue
            pcl_color_tmp[final_inst_seg_3d == int(color_id)] = COLOR_DETECTRON2[int(color_id), :]

        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(points_mesh)
        pcd.colors = o3d.utility.Vector3dVector(pcl_color_tmp)

        indices = np.where(final_inst_seg_3d > 0)[0].tolist()
        if not indices:
            continue

        mesh_o3d = copy.deepcopy(scene_mesh)
        mesh_o3d.remove_vertices_by_index(indices)
        bg_and_anno = obj_mesh_all + mesh_o3d + lineSets_all

        # Visualize
        vis = o3d.visualization.Visualizer()
        vis.create_window(window_name='Open3D', width=640 * 2, height=480 * 2, left=100, top=200,
                          visible=True)
        vis.get_render_option().mesh_show_back_face = False

        sceneMeshVert = np.array(scene_mesh.vertices)
        trans = np.zeros((3,)) * 1.2
        trans[0] = np.max(sceneMeshVert[:, 0]) - np.min(sceneMeshVert[:, 0])
        vis.add_geometry(bg_and_anno)

        vis.add_geometry(pcd.translate(trans * 1.2))
        scene_and_boxes = scene_mesh + lineSets_all
        vis.add_geometry(scene_and_boxes.translate(trans * 1.2 * 2))
        vis.run()


if __name__ == "__main__":
    main(parser.parse_args())

visualize_annotations.py
stefan-ainetter/SCANnotateDataset