utils.py

import os
import numpy as np
import glob
import open3d as o3d
import math
from tqdm import tqdm
import tensorflow as tf
import h5py


def compose_model_args(dataset, dataset_dir, params):
    b, _ = load_block(block_dir=dataset_dir, name=0)
    n_classes = None
    if dataset == "S3DIS":
        n_classes = 14
    elif dataset == "PCG":
        n_classes = 12
    elif dataset == "Scannet":
        n_classes = 524
    model_args = {
        "name": params["model_name"],
        "n_classes": n_classes,
        "n_points": b.shape[0],
        "seed": params["seed"],
        "trainable": params["trainable"],
        "check_numerics": params["check_numerics"],
        "initializer": params["initializer"],
        "trainable_net": params["trainable_net"],
        "p_dim": params["p_dim"]
        }
    return model_args


def get_loss(seg_pred, seg, t=None, reg_f=1e-3, check_numerics=True):
    # loss calculation
    seg = seg.astype(np.int32)
    #print(seg.shape, seg_pred.shape)
    ce = tf.nn.sparse_softmax_cross_entropy_with_logits(
                    labels=seg, logits=seg_pred)
    if check_numerics:
        ce = tf.debugging.check_numerics(ce, "ce")

    per_instance_seg_loss = tf.reduce_mean(ce, axis=1)
    seg_loss = tf.reduce_mean(per_instance_seg_loss)

    total = seg_loss
    mat_diff_loss = None
    if t is not None:
        K = tf.shape(t)[1]
        mul = tf.matmul(t, tf.transpose(t, perm=[0,2,1]))
        mat_diff = mul - tf.constant(np.eye(K), dtype=tf.float32)
        mat_diff_loss = tf.nn.l2_loss(mat_diff)
        if check_numerics:
            mat_diff_loss = tf.debugging.check_numerics(mat_diff_loss, "mat_diff_loss")

        total = seg_loss + mat_diff_loss * reg_f
    return total, seg_loss, mat_diff_loss


def load_folds(dataset_dir, k_fold_dir, train_folds, test_fold):
    test_fold_fname = k_fold_dir + "/" + str(test_fold) + ".h5"
    train_folds_fnames = [k_fold_dir + "/" + str(train_fold) + ".h5" for train_fold in train_folds]

    hf_test = h5py.File(test_fold_fname, "r")
    test_files = list(hf_test["files"])
    try:
        test_files = [ta.decode() for ta in test_files]
    except:
        test_files = [ta for ta in test_files]
    hf_test.close()

    train_files = []
    for fname in train_folds_fnames:
        hf_train = h5py.File(fname, "r")
        train_f = list(hf_train["files"])
        try:
            train_f = [ta.decode() for ta in train_f]
        except:
            train_f = [ta for ta in train_f]
        hf_train.close()
        train_files.extend(train_f)
    return train_files, test_files


def load_block(block_dir, name, spatial_only=False):
    # load a block of the point cloud
    filename = block_dir + "/" + str(name) + ".npz"
    data = np.load(filename)
    block = data["block"]
    b_labels = data["labels"]

    # render_point_cloud(block)
    if spatial_only:
        block = block[:, :3]
    #if len(block.shape) == 2:
    #    block = np.expand_dims(block, axis=0)
    return block, b_labels


def load_batch(i, train_idxs, block_dir, blocks, labels, batch_size, apply_random_rotation=False, spatial_only=False):
    # load a batch of blocks and their corresponding labels
    j = i * batch_size
    idxs = train_idxs[j:j+batch_size]
    for k in range(idxs.shape[0]):
        name = idxs[k]
        block, b_labels = load_block(block_dir, name, spatial_only=spatial_only)
        if len(b_labels.shape) == 2:
            b_labels = np.squeeze(b_labels, -1)
        elif len(b_labels.shape) > 2:
            raise Exception("Unexpected shape of labels" + str(b_labels.shape))
        if apply_random_rotation:
            #rot = R.random().as_matrix()
            rotation_angle = np.random.uniform() * 2 * np.pi
            cosval = np.cos(rotation_angle)
            sinval = np.sin(rotation_angle)
            rot = np.array([
                    [cosval, 0, sinval],
                    [0, 1, 0],
                    [-sinval, 0, cosval]
                ])
            block[:, :3] = np.matmul(block[:, :3], rot)
            #block[:, 6:9] = np.matmul(block[:, 6:9], rot)
        blocks[k] = block
        labels[k] = b_labels
    return blocks, labels


def load_batch2(i, train_idxs, block_dir, blocks, labels, batch_size, apply_random_rotation=False, spatial_only=False):
    name = train_idxs[i]
    blocks, labels = load_block(block_dir, name, spatial_only=spatial_only)
    labels = np.squeeze(labels, -1)
    if apply_random_rotation:
        for j in range(blocks.shape[0]):
            block = blocks[j]
            #rot = R.random().as_matrix()
            rotation_angle = np.random.uniform() * 2 * np.pi
            cosval = np.cos(rotation_angle)
            sinval = np.sin(rotation_angle)
            rot = np.array([
                    [cosval, 0, sinval],
                    [0, 1, 0],
                    [-sinval, 0, cosval]
                ])
            blocks[j, :, :3] = np.matmul(blocks[j, :, :3], rot)
    #if len(blocks.shape) == 2:
    #blocks = np.expand_dims(blocks, axis=0)
    #labels = np.expand_dims(labels, axis=0)
    return blocks, labels


def mkdir(directory):
    """Method to create a new directory.

    Parameters
    ----------
    directory : str
        Relative or absolute path.
    """
    if not os.path.isdir(directory):
        os.makedirs(directory)


def file_exists(filepath):
    """Check if a file exists.

    Parameters
    ----------
    filepath : str
        Relative or absolute path to a file.

    Returns
    -------
    boolean
        True if the file exists.

    """
    return os.path.isfile(filepath)


def load_scene(dataset, scene):
    filename = "./Scenes/" + dataset + "/" + scene + "/P.npz"
    data = np.load(filename)
    P = data["P"]
    labels = data["labels"]
    return P, labels


def coordinate_system():
    """Returns a coordinate system.

    Returns
    -------
    o3d.geometry.LineSet
        The lines of a coordinate system that are colored in red, green
        and blue.

    """
    line_set = o3d.geometry.LineSet()
    points = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]])
    colors = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
    lines = np.array([[0, 1], [0, 2], [0, 3]]).astype(int)
    line_set.points = o3d.utility.Vector3dVector(points)
    line_set.colors = o3d.utility.Vector3dVector(colors)
    line_set.lines = o3d.utility.Vector2iVector(lines)
    return line_set


def render_pc(pcd, animate=False, x_speed=2.5, y_speed=0.0, width=1920, left=0):
    """Render a point cloud.

    Parameters
    ----------
    pcd : o3d.PointCloud
        Open3D point cloud.
    animate : boolean
        If True, the point cloud will be rotated with x_speed and y_speed. A
        simulation of dragging the mouse in standard rendering mode is
        simulated.
    x_speed : float
        Used if point cloud will be animated. Strength if the horizontal mouse
        drag.
    y_speed : float
        Used if point cloud will be animated. Strength if the vertical mouse
        drag.
    width : int
        Set the width of the open3D plot
    left : int
        How much should the open3D plot be dragged to the right.

    """
    if animate:
        def rotate_view(vis):
            ctr = vis.get_view_control()
            ctr.rotate(x_speed, y_speed)
            return False
        o3d.visualization.draw_geometries_with_animation_callback(
            [pcd, coordinate_system()], rotate_view, width=width, left=left)
    else:
        o3d.visualization.draw_geometries([pcd, coordinate_system()], width=width, left=left)


def render_point_cloud(
        P, partition_vec=None, classification=None, gt_partition=None, animate=False, x_speed=2.5, y_speed=0.0, width=1920, left=0):
    """Displays a point cloud.

    Parameters
    ----------
    P : np.ndarray
        Nx3 or Nx6 matrix. N is the number of points. A point should have at
        least 3 spatial coordinates and can have optionally 3 color values.
    partition_vec : np.ndarray
        The partition of the point cloud.
    classification : np.ndarray
        Match classification matrix.
    gt_partition : np.ndarray
        Ground truth partition.
    animate : boolean
        If True, the point cloud will be rotated with x_speed and y_speed. A
        simulation of dragging the mouse in standard rendering mode is
        simulated.
    x_speed : float
        Used if point cloud will be animated. Strength if the horizontal mouse
        drag.
    y_speed : float
        Used if point cloud will be animated. Strength if the vertical mouse
        drag.
    width : int
        Set the width of the open3D plot
    left : int
        How much should the open3D plot be dragged to the right.
    """
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(P[:, :3])
    if partition_vec is not None:
        col_mat = np.zeros((P.shape[0], 3))
        data = np.load("colors.npz")
        colors = data["colors"]
        superpoints = np.unique(partition_vec)
        n_superpoints = superpoints.shape[0]
        if classification is not None and gt_partition is not None:
            segment_val_to_color_idx = {}
            if n_superpoints != classification.shape[0]:
                raise Exception("Mismatch of number of superpoints in the classification and the partition.")
            for i in range(n_superpoints):
                superpoint_value = superpoints[i]
                idxs = np.where(classification[i, :] != 0)[0]
                if idxs.shape[0] > 1:
                    raise Exception("One-to-many assignment in classification")
                if idxs.shape[0] == 1:
                    color_idx = idxs[0]
                    segment_val_to_color_idx[superpoint_value] = color_idx
            color_idx_offset = classification.shape[1]
            for i in range(n_superpoints):
                superpoint_value = superpoints[i]
                idx = np.where(partition_vec == superpoint_value)[0]
                col_idx = color_idx_offset + i
                if superpoint_value in segment_val_to_color_idx:
                    col_idx = segment_val_to_color_idx[superpoint_value]
                col_mat[idx, :] = colors[col_idx, :] / 255
        else:
            for i in range(n_superpoints):
                superpoint_value = superpoints[i]
                idx = np.where(partition_vec == superpoint_value)[0]
                color = colors[i, :] / 255
                col_mat[idx, :] = color
        pcd.colors = o3d.utility.Vector3dVector(col_mat)
    else:
        try:
            # print(P[:5, 3:6] / 255.0)
            pcd.colors = o3d.utility.Vector3dVector(P[:, 3:6] / 255.0)
        except Exception as e:
            print(e)
    render_pc(pcd=pcd, animate=animate, x_speed=x_speed, y_speed=y_speed, width=width, left=left)
    return pcd


def collect_point_label(anno_path, out_filename, file_format='txt'):
    """ Convert original dataset files to data_label file (each line is XYZRGBL).
        We aggregated all the points from each instance in the room.
    Args:
        anno_path: path to annotations. e.g. Area_1/office_2/Annotations/
        out_filename: path to save collected points and labels (each line is XYZRGBL)
        file_format: txt or numpy, determines what file format to save.
    Returns:
        None
    Note:
        the points are shifted before save, the most negative point is now at origin.
    """
    points_list = []

    for f in glob.glob(os.path.join(anno_path, '*.txt')):
        cls = os.path.basename(f).split('_')[0]
        if cls not in g_classes: # note: in some room there is 'staris' class..
            cls = 'clutter'
        points = np.loadtxt(f)
        labels = np.ones((points.shape[0],1)) * g_class2label[cls]
        points_list.append(np.concatenate([points, labels], 1)) # Nx7

    data_label = np.concatenate(points_list, 0)
    xyz_min = np.amin(data_label, axis=0)[0:3]
    data_label[:, 0:3] -= xyz_min

    if file_format=='txt':
        fout = open(out_filename, 'w')
        for i in range(data_label.shape[0]):
            fout.write('%f %f %f %d %d %d %d\n' % \
                          (data_label[i,0], data_label[i,1], data_label[i,2],
                           data_label[i,3], data_label[i,4], data_label[i,5],
                           data_label[i,6]))
        fout.close()
    elif file_format=='numpy':
        np.save(out_filename, data_label)
    else:
        print('ERROR!! Unknown file format: %s, please use txt or numpy.' % \
            (file_format))
        exit()

def point_label_to_obj(input_filename, out_filename, label_color=True, easy_view=False, no_wall=False):
    """ For visualization of a room from data_label file,
    input_filename: each line is X Y Z R G B L
    out_filename: OBJ filename,
            visualize input file by coloring point with label color
        easy_view: only visualize furnitures and floor
    """
    data_label = np.loadtxt(input_filename)
    data = data_label[:, 0:6]
    label = data_label[:, -1].astype(int)
    fout = open(out_filename, 'w')
    for i in range(data.shape[0]):
        color = g_label2color[label[i]]
        if easy_view and (label[i] not in g_easy_view_labels):
            continue
        if no_wall and ((label[i] == 2) or (label[i]==0)):
            continue
        if label_color:
            fout.write('v %f %f %f %d %d %d\n' % \
                (data[i,0], data[i,1], data[i,2], color[0], color[1], color[2]))
        else:
            fout.write('v %f %f %f %d %d %d\n' % \
                (data[i,0], data[i,1], data[i,2], data[i,3], data[i,4], data[i,5]))
    fout.close()


def get_labelled_blocks(scene, dataset, num_points=4096):
    P, labels = load_scene(dataset=dataset, scene=scene)
    labels = labels.reshape(labels.shape[0], )
    #print(labels)
    #blocks, b_labels, sample_indices = room2blocks(data=P, label=labels, num_point=num_points)
    #blocks, b_labels, sample_indices = room2blocks_plus_normalized(data=P, label=labels, num_points=num_points)
    blocks, b_labels, sample_indices = room2samples_plus_normalized(data=P, label=labels, num_point=num_points)
    #print(labels.shape)
    #blocks, b_labels, sample_indices = room2samples(data=P, label=labels, sample_num_point=num_points)
    return blocks, b_labels, P, labels, sample_indices


def get_blocks(P, num_points=4096):
    #blocks, sample_indices = room2samples(data=P, label=None, sample_num_point=num_points)
    blocks, sample_indices = room2samples_plus_normalized(data=P, label=None, num_point=num_points)
    #blocks, sample_indices = room2blocks_plus_normalized(data=P, label=None, num_points=num_points)
    return blocks, sample_indices


def create_blocks(dataset, num_points=4096, batch_size=-1):
    block_dir = "./Blocks/" + dataset
    mkdir(block_dir)
    scenes = os.listdir("./Scenes/" + dataset)
    block_n = 0
    for i in tqdm(range(len(scenes)), desc="Blocks"):
        scene = scenes[i]
        #P, labels = load_scene(dataset=dataset, scene=scene)
        #blocks, b_labels = room2blocks(data=P, label=labels, num_point=num_points)
        blocks, b_labels, _, _, _ = get_labelled_blocks(scene=scene, dataset=dataset, num_points=num_points)
        if batch_size > 1:
            n_blocks = blocks.shape[0]
            n_batches = math.floor(n_blocks / batch_size)
            for k in range(n_batches):
                start = k*batch_size
                stop = start + batch_size
                batch = blocks[start:stop]
                batch_labels = b_labels[start:stop]
                ok = True
                for j in range(batch_size):
                    block = batch[j]
                    b_label = batch_labels[j]
                    if b_label.shape[0] != num_points or block.shape[0] != num_points:
                        ok = False
                        break
                if not ok:
                    continue
                np.savez(block_dir + "/" + str(block_n) + ".npz", block=batch, labels=batch_labels)
                block_n += 1
        else:
            for k in range(blocks.shape[0]):
                block = blocks[k]
                #print("min", np.min(block, axis=0))
                #print("max", np.max(block, axis=0))
                b_label = b_labels[k]
                if b_label.shape[0] != num_points or block.shape[0] != num_points:
                    continue
                np.savez(block_dir + "/" + str(block_n) + ".npz", block=block, labels=b_label)
                block_n += 1
    print(block_n, "Blocks saved.")


# -----------------------------------------------------------------------------
# PREPARE BLOCK DATA FOR DEEPNETS TRAINING/TESTING
# -----------------------------------------------------------------------------

def sample_data(data, num_sample):
    """ data is in N x ...
        we want to keep num_samplexC of them.
        if N > num_sample, we will randomly keep num_sample of them.
        if N < num_sample, we will randomly duplicate samples.
    """
    N = data.shape[0]
    if (N == num_sample):
        return data, range(N)
    elif (N > num_sample):
        sample = np.random.choice(N, num_sample)
        return data[sample, ...], sample
    else:
        sample = np.random.choice(N, num_sample-N)
        dup_data = data[sample, ...]
        return np.concatenate([data, dup_data], 0), list(range(N))+list(sample)

def sample_data_label(data, label, num_sample):
    new_data, sample_indices = sample_data(data, num_sample)
    new_label = label[sample_indices]
    return new_data, new_label, sample_indices

def sample_wo_label(data, num_sample):
    new_data, sample_indices = sample_data(data, num_sample)
    return new_data, sample_indices


def room2blocks(data, label, num_point, block_size=1.0, stride=1.0,
                random_sample=False, sample_num=None, sample_aug=1):
    """ Prepare block training data.
    Args:
        data: N x 6 numpy array, 012 are XYZ in meters, 345 are RGB in [0,1]
            assumes the data is shifted (min point is origin) and aligned
            (aligned with XYZ axis)
        label: N size uint8 numpy array from 0-n_classes
        num_point: int, how many points to sample in each block
        block_size: float, physical size of the block in meters
        stride: float, stride for block sweeping
        random_sample: bool, if True, we will randomly sample blocks in the room
        sample_num: int, if random sample, how many blocks to sample
            [default: room area]
        sample_aug: if random sample, how much aug
    Returns:
        block_datas: K x num_point x 6 np array of XYZRGB, RGB is in [0,1]
        block_labels: K x num_point x 1 np array of uint8 labels

    TODO: for this version, blocking is in fixed, non-overlapping pattern.
    """
    assert(stride<=block_size)

    limit = np.amax(data, 0)[0:3]

    wo_label = label is None

    # Get the corner location for our sampling blocks
    xbeg_list = []
    ybeg_list = []
    if not random_sample:
        num_block_x = int(np.ceil((limit[0] - block_size) / stride)) + 1
        num_block_y = int(np.ceil((limit[1] - block_size) / stride)) + 1
        for i in range(num_block_x):
            for j in range(num_block_y):
                xbeg_list.append(i*stride)
                ybeg_list.append(j*stride)
    else:
        num_block_x = int(np.ceil(limit[0] / block_size))
        num_block_y = int(np.ceil(limit[1] / block_size))
        if sample_num is None:
            sample_num = num_block_x * num_block_y * sample_aug
        for _ in range(sample_num):
            xbeg = np.random.uniform(-block_size, limit[0])
            ybeg = np.random.uniform(-block_size, limit[1])
            xbeg_list.append(xbeg)
            ybeg_list.append(ybeg)

    # Collect blocks
    block_data_list = []
    block_label_list = []
    sample_indices_list = []
    idx = 0
    for idx in range(len(xbeg_list)):
       xbeg = xbeg_list[idx]
       ybeg = ybeg_list[idx]
       xcond = (data[:,0]<=xbeg+block_size) & (data[:,0]>=xbeg)
       ycond = (data[:,1]<=ybeg+block_size) & (data[:,1]>=ybeg)
       cond = xcond & ycond
       if np.sum(cond) < 100: # discard block if there are less than 100 pts.
           continue

       block_data = data[cond, :]
       if wo_label:
           block_data_sampled, sample_indices = sample_wo_label(block_data, num_point)
       else: # with label
           block_label = label[cond]
           # randomly subsample data
           block_data_sampled, block_label_sampled, sample_indices = \
               sample_data_label(block_data, block_label, num_point)
           block_label_list.append(np.expand_dims(block_label_sampled, 0))
       block_data_list.append(np.expand_dims(block_data_sampled, 0))
       sample_indices_list.append(np.expand_dims(sample_indices, 0))
    if wo_label:
       return np.concatenate(block_data_list, 0), \
              np.concatenate(sample_indices_list, 0)    
    return np.concatenate(block_data_list, 0), \
           np.concatenate(block_label_list, 0), \
           np.concatenate(sample_indices_list, 0)


def room2blocks_plus(data_label, num_point, block_size, stride,
                     random_sample, sample_num, sample_aug):
    """ room2block with input filename and RGB preprocessing.
    """
    data = data_label[:,0:6]
    data[:,3:6] /= 255.0
    label = data_label[:,-1].astype(np.uint8)

    return room2blocks(data, label, num_point, block_size, stride,
                       random_sample, sample_num, sample_aug)

def room2blocks_wrapper(data_label_filename, num_point, block_size=1.0, stride=1.0,
                        random_sample=False, sample_num=None, sample_aug=1):
    if data_label_filename[-3:] == 'txt':
        data_label = np.loadtxt(data_label_filename)
    elif data_label_filename[-3:] == 'npy':
        data_label = np.load(data_label_filename)
    else:
        print('Unknown file type! exiting.')
        exit()
    return room2blocks_plus(data_label, num_point, block_size, stride,
                            random_sample, sample_num, sample_aug)

def room2blocks_plus_normalized(data, label, num_points, block_size=1.0, stride=1.0,
                                random_sample=False, sample_num=None, sample_aug=1):
    """ room2block, with input filename and RGB preprocessing.
        for each block centralize XYZ, add normalized XYZ as 678 channels
    """
    #data = data_label[:,0:6]
    #data[:,3:6] /= 255.0
    #label = data_label[:,-1].astype(np.uint8)

    wo_label = label is None

    max_room_x = max(data[:,0])
    max_room_y = max(data[:,1])
    max_room_z = max(data[:,2])

    if wo_label:
        data_batch, sample_indices_batch = room2blocks(data=data, label=label, num_point=num_points, block_size=block_size,
            stride=stride, random_sample=random_sample, sample_num=sample_num, sample_aug=sample_aug)
    else:
        data_batch, label_batch, sample_indices_batch = room2blocks(data=data, label=label, num_point=num_points, block_size=block_size,
            stride=stride, random_sample=random_sample, sample_num=sample_num, sample_aug=sample_aug)
    new_data_batch = np.zeros((data_batch.shape[0], num_points, 9))
    for b in range(data_batch.shape[0]):
        new_data_batch[b, :, 6] = data_batch[b, :, 0]/max_room_x
        new_data_batch[b, :, 7] = data_batch[b, :, 1]/max_room_y
        new_data_batch[b, :, 8] = data_batch[b, :, 2]/max_room_z
        minx = min(data_batch[b, :, 0])
        miny = min(data_batch[b, :, 1])
        data_batch[b, :, 0] -= (minx+block_size/2)
        data_batch[b, :, 1] -= (miny+block_size/2)
    new_data_batch[:, :, 0:6] = data_batch
    if wo_label:
        return new_data_batch, sample_indices_batch
    return new_data_batch, label_batch, sample_indices_batch


def room2blocks_wrapper_normalized(data_label_filename, num_point, block_size=1.0, stride=1.0,
                                   random_sample=False, sample_num=None, sample_aug=1):
    if data_label_filename[-3:] == 'txt':
        data_label = np.loadtxt(data_label_filename)
    elif data_label_filename[-3:] == 'npy':
        data_label = np.load(data_label_filename)
    else:
        print('Unknown file type! exiting.')
        exit()
    return room2blocks_plus_normalized(data_label, num_point, block_size, stride,
                                       random_sample, sample_num, sample_aug)

def room2samples(data, label, sample_num_point):
    """ Prepare whole room samples.
    Args:
        data: N x 6 numpy array, 012 are XYZ in meters, 345 are RGB in [0,1]
            assumes the data is shifted (min point is origin) and
            aligned (aligned with XYZ axis)
        label: N size uint8 numpy array from 0-12
        sample_num_point: int, how many points to sample in each sample
    Returns:
        sample_datas: K x sample_num_point x 9
                     numpy array of XYZRGBX'Y'Z', RGB is in [0,1]
        sample_labels: K x sample_num_point x 1 np array of uint8 labels
    """
    N = data.shape[0]
    order = np.arange(N)
    np.random.shuffle(order)
    data = data[order, :]
    wo_label = label is None
    if not wo_label:
        label = label[order]

    batch_num = int(np.ceil(N / float(sample_num_point)))
    sample_datas = np.zeros((batch_num, sample_num_point, 6))
    if not wo_label:
        sample_labels = np.zeros((batch_num, sample_num_point, 1))
    sample_indices_list = []
    for i in range(batch_num):
        beg_idx = i*sample_num_point
        end_idx = min((i+1)*sample_num_point, N)
        num = end_idx - beg_idx
        sample_indices = np.arange(start=beg_idx, stop=end_idx, dtype=np.int32)
        sample_indices_list.append(sample_indices)
        sample_datas[i,0:num,:] = data[beg_idx:end_idx, :]
        if not wo_label:
            sample_labels[i,0:num,0] = label[beg_idx:end_idx]
        if num < sample_num_point:
            makeup_indices = np.random.choice(N, sample_num_point - num)
            sample_datas[i,num:,:] = data[makeup_indices, :]
            if not wo_label:
                sample_labels[i,num:,0] = label[makeup_indices]
    if wo_label:
        return sample_datas, sample_indices_list
    return sample_datas, sample_labels, sample_indices_list

def room2samples_plus_normalized(data, label, num_point):
    """ room2sample, with input filename and RGB preprocessing.
        for each block centralize XYZ, add normalized XYZ as 678 channels
    """
    # data = data_label[:,0:6]
    # data[:,3:6] /= 255.0
    # label = data_label[:,-1].astype(np.uint8)
    wo_label = label is None
    max_room_x = max(data[:,0])
    max_room_y = max(data[:,1])
    max_room_z = max(data[:,2])
    #print(max_room_x, max_room_y, max_room_z)
    
    #blocks, sample_indices = room2samples(data=P, label=None, sample_num_point=num_points)
    if wo_label:
        data_batch, sample_indices = room2samples(data=data, label=None, sample_num_point=num_point)
    else:
        data_batch, label_batch, sample_indices = room2samples(
            data=data, label=label, sample_num_point=num_point)
    """
    new_data_batch = np.zeros((data_batch.shape[0], num_point, 9))
    for b in range(data_batch.shape[0]):
        new_data_batch[b, :, 6] = data_batch[b, :, 0]/max_room_x
        new_data_batch[b, :, 7] = data_batch[b, :, 1]/max_room_y
        new_data_batch[b, :, 8] = data_batch[b, :, 2]/max_room_z
        #minx = min(data_batch[b, :, 0])
        #miny = min(data_batch[b, :, 1])
        #data_batch[b, :, 0] -= (minx+block_size/2)
        #data_batch[b, :, 1] -= (miny+block_size/2)
    new_data_batch[:, :, 0:6] = data_batch
    """
    new_data_batch = np.zeros((data_batch.shape[0], num_point, 6))
    for b in range(data_batch.shape[0]):
        new_data_batch[b, :, 0] = data_batch[b, :, 0]/max_room_x
        new_data_batch[b, :, 1] = data_batch[b, :, 1]/max_room_y
        new_data_batch[b, :, 2] = data_batch[b, :, 2]/max_room_z
        #minx = min(data_batch[b, :, 0])
        #miny = min(data_batch[b, :, 1])
        #data_batch[b, :, 0] -= (minx+block_size/2)
        #data_batch[b, :, 1] -= (miny+block_size/2)
    new_data_batch[:, :, 3:6] = data_batch[:, :, 3:6]
    if wo_label:
        return new_data_batch, sample_indices
    return new_data_batch, label_batch, sample_indices


def room2samples_wrapper_normalized(data_label_filename, num_point):
    if data_label_filename[-3:] == 'txt':
        data_label = np.loadtxt(data_label_filename)
    elif data_label_filename[-3:] == 'npy':
        data_label = np.load(data_label_filename)
    else:
        print('Unknown file type! exiting.')
        exit()
    return room2samples_plus_normalized(data_label, num_point)


# -----------------------------------------------------------------------------
# EXTRACT INSTANCE BBOX FROM ORIGINAL DATA (for detection evaluation)
# -----------------------------------------------------------------------------

def collect_bounding_box(anno_path, out_filename):
    """ Compute bounding boxes from each instance in original dataset files on
        one room. **We assume the bbox is aligned with XYZ coordinate.**

    Args:
        anno_path: path to annotations. e.g. Area_1/office_2/Annotations/
        out_filename: path to save instance bounding boxes for that room.
            each line is x1 y1 z1 x2 y2 z2 label,
            where (x1,y1,z1) is the point on the diagonal closer to origin
    Returns:
        None
    Note:
        room points are shifted, the most negative point is now at origin.
    """
    bbox_label_list = []

    for f in glob.glob(os.path.join(anno_path, '*.txt')):
        cls = os.path.basename(f).split('_')[0]
        if cls not in g_classes: # note: in some room there is 'staris' class..
            cls = 'clutter'
        points = np.loadtxt(f)
        label = g_class2label[cls]
        # Compute tightest axis aligned bounding box
        xyz_min = np.amin(points[:, 0:3], axis=0)
        xyz_max = np.amax(points[:, 0:3], axis=0)
        ins_bbox_label = np.expand_dims(
            np.concatenate([xyz_min, xyz_max, np.array([label])], 0), 0)
        bbox_label_list.append(ins_bbox_label)

    bbox_label = np.concatenate(bbox_label_list, 0)
    room_xyz_min = np.amin(bbox_label[:, 0:3], axis=0)
    bbox_label[:, 0:3] -= room_xyz_min
    bbox_label[:, 3:6] -= room_xyz_min

    fout = open(out_filename, 'w')
    for i in range(bbox_label.shape[0]):
        fout.write('%f %f %f %f %f %f %d\n' % \
                      (bbox_label[i,0], bbox_label[i,1], bbox_label[i,2],
                       bbox_label[i,3], bbox_label[i,4], bbox_label[i,5],
                       bbox_label[i,6]))
    fout.close()

def bbox_label_to_obj(input_filename, out_filename_prefix, easy_view=False):
    """ Visualization of bounding boxes.

    Args:
        input_filename: each line is x1 y1 z1 x2 y2 z2 label
        out_filename_prefix: OBJ filename prefix,
            visualize object by g_label2color
        easy_view: if True, only visualize furniture and floor
    Returns:
        output a list of OBJ file and MTL files with the same prefix
    """
    bbox_label = np.loadtxt(input_filename)
    bbox = bbox_label[:, 0:6]
    label = bbox_label[:, -1].astype(int)
    v_cnt = 0 # count vertex
    ins_cnt = 0 # count instance
    for i in range(bbox.shape[0]):
        if easy_view and (label[i] not in g_easy_view_labels):
            continue
        obj_filename = out_filename_prefix+'_'+g_classes[label[i]]+'_'+str(ins_cnt)+'.obj'
        mtl_filename = out_filename_prefix+'_'+g_classes[label[i]]+'_'+str(ins_cnt)+'.mtl'
        fout_obj = open(obj_filename, 'w')
        fout_mtl = open(mtl_filename, 'w')
        fout_obj.write('mtllib %s\n' % (os.path.basename(mtl_filename)))

        length = bbox[i, 3:6] - bbox[i, 0:3]
        a = length[0]
        b = length[1]
        c = length[2]
        x = bbox[i, 0]
        y = bbox[i, 1]
        z = bbox[i, 2]
        color = np.array(g_label2color[label[i]], dtype=float) / 255.0

        material = 'material%d' % (ins_cnt)
        fout_obj.write('usemtl %s\n' % (material))
        fout_obj.write('v %f %f %f\n' % (x,y,z+c))
        fout_obj.write('v %f %f %f\n' % (x,y+b,z+c))
        fout_obj.write('v %f %f %f\n' % (x+a,y+b,z+c))
        fout_obj.write('v %f %f %f\n' % (x+a,y,z+c))
        fout_obj.write('v %f %f %f\n' % (x,y,z))
        fout_obj.write('v %f %f %f\n' % (x,y+b,z))
        fout_obj.write('v %f %f %f\n' % (x+a,y+b,z))
        fout_obj.write('v %f %f %f\n' % (x+a,y,z))
        fout_obj.write('g default\n')
        v_cnt = 0 # for individual box
        fout_obj.write('f %d %d %d %d\n' % (4+v_cnt, 3+v_cnt, 2+v_cnt, 1+v_cnt))
        fout_obj.write('f %d %d %d %d\n' % (1+v_cnt, 2+v_cnt, 6+v_cnt, 5+v_cnt))
        fout_obj.write('f %d %d %d %d\n' % (7+v_cnt, 6+v_cnt, 2+v_cnt, 3+v_cnt))
        fout_obj.write('f %d %d %d %d\n' % (4+v_cnt, 8+v_cnt, 7+v_cnt, 3+v_cnt))
        fout_obj.write('f %d %d %d %d\n' % (5+v_cnt, 8+v_cnt, 4+v_cnt, 1+v_cnt))
        fout_obj.write('f %d %d %d %d\n' % (5+v_cnt, 6+v_cnt, 7+v_cnt, 8+v_cnt))
        fout_obj.write('\n')

        fout_mtl.write('newmtl %s\n' % (material))
        fout_mtl.write('Kd %f %f %f\n' % (color[0], color[1], color[2]))
        fout_mtl.write('\n')
        fout_obj.close()
        fout_mtl.close()

        v_cnt += 8
        ins_cnt += 1

def bbox_label_to_obj_room(input_filename, out_filename_prefix, easy_view=False, permute=None, center=False, exclude_table=False):
    """ Visualization of bounding boxes.

    Args:
        input_filename: each line is x1 y1 z1 x2 y2 z2 label
        out_filename_prefix: OBJ filename prefix,
            visualize object by g_label2color
        easy_view: if True, only visualize furniture and floor
        permute: if not None, permute XYZ for rendering, e.g. [0 2 1]
        center: if True, move obj to have zero origin
    Returns:
        output a list of OBJ file and MTL files with the same prefix
    """
    bbox_label = np.loadtxt(input_filename)
    bbox = bbox_label[:, 0:6]
    if permute is not None:
        assert(len(permute)==3)
        permute = np.array(permute)
        bbox[:,0:3] = bbox[:,permute]
        bbox[:,3:6] = bbox[:,permute+3]
    if center:
        xyz_max = np.amax(bbox[:,3:6], 0)
        bbox[:,0:3] -= (xyz_max/2.0)
        bbox[:,3:6] -= (xyz_max/2.0)
        bbox /= np.max(xyz_max/2.0)
    label = bbox_label[:, -1].astype(int)
    obj_filename = out_filename_prefix+'.obj'
    mtl_filename = out_filename_prefix+'.mtl'

    fout_obj = open(obj_filename, 'w')
    fout_mtl = open(mtl_filename, 'w')
    fout_obj.write('mtllib %s\n' % (os.path.basename(mtl_filename)))
    v_cnt = 0 # count vertex
    ins_cnt = 0 # count instance
    for i in range(bbox.shape[0]):
        if easy_view and (label[i] not in g_easy_view_labels):
            continue
        if exclude_table and label[i] == g_classes.index('table'):
            continue

        length = bbox[i, 3:6] - bbox[i, 0:3]
        a = length[0]
        b = length[1]
        c = length[2]
        x = bbox[i, 0]
        y = bbox[i, 1]
        z = bbox[i, 2]
        color = np.array(g_label2color[label[i]], dtype=float) / 255.0

        material = 'material%d' % (ins_cnt)
        fout_obj.write('usemtl %s\n' % (material))
        fout_obj.write('v %f %f %f\n' % (x,y,z+c))
        fout_obj.write('v %f %f %f\n' % (x,y+b,z+c))
        fout_obj.write('v %f %f %f\n' % (x+a,y+b,z+c))
        fout_obj.write('v %f %f %f\n' % (x+a,y,z+c))
        fout_obj.write('v %f %f %f\n' % (x,y,z))
        fout_obj.write('v %f %f %f\n' % (x,y+b,z))
        fout_obj.write('v %f %f %f\n' % (x+a,y+b,z))
        fout_obj.write('v %f %f %f\n' % (x+a,y,z))
        fout_obj.write('g default\n')
        fout_obj.write('f %d %d %d %d\n' % (4+v_cnt, 3+v_cnt, 2+v_cnt, 1+v_cnt))
        fout_obj.write('f %d %d %d %d\n' % (1+v_cnt, 2+v_cnt, 6+v_cnt, 5+v_cnt))
        fout_obj.write('f %d %d %d %d\n' % (7+v_cnt, 6+v_cnt, 2+v_cnt, 3+v_cnt))
        fout_obj.write('f %d %d %d %d\n' % (4+v_cnt, 8+v_cnt, 7+v_cnt, 3+v_cnt))
        fout_obj.write('f %d %d %d %d\n' % (5+v_cnt, 8+v_cnt, 4+v_cnt, 1+v_cnt))
        fout_obj.write('f %d %d %d %d\n' % (5+v_cnt, 6+v_cnt, 7+v_cnt, 8+v_cnt))
        fout_obj.write('\n')

        fout_mtl.write('newmtl %s\n' % (material))
        fout_mtl.write('Kd %f %f %f\n' % (color[0], color[1], color[2]))
        fout_mtl.write('\n')

        v_cnt += 8
        ins_cnt += 1

    fout_obj.close()
    fout_mtl.close()


def collect_point_bounding_box(anno_path, out_filename, file_format):
    """ Compute bounding boxes from each instance in original dataset files on
        one room. **We assume the bbox is aligned with XYZ coordinate.**
        Save both the point XYZRGB and the bounding box for the point's
        parent element.

    Args:
        anno_path: path to annotations. e.g. Area_1/office_2/Annotations/
        out_filename: path to save instance bounding boxes for each point,
            plus the point's XYZRGBL
            each line is XYZRGBL offsetX offsetY offsetZ a b c,
            where cx = X+offsetX, cy=X+offsetY, cz=Z+offsetZ
            where (cx,cy,cz) is center of the box, a,b,c are distances from center
            to the surfaces of the box, i.e. x1 = cx-a, x2 = cx+a, y1=cy-b etc.
        file_format: output file format, txt or numpy
    Returns:
        None
    Note:
        room points are shifted, the most negative point is now at origin.
    """
    point_bbox_list = []

    for f in glob.glob(os.path.join(anno_path, '*.txt')):
        cls = os.path.basename(f).split('_')[0]
        if cls not in g_classes: # note: in some room there is 'staris' class..
            cls = 'clutter'
        points = np.loadtxt(f) # Nx6
        label = g_class2label[cls] # N,
        # Compute tightest axis aligned bounding box
        xyz_min = np.amin(points[:, 0:3], axis=0) # 3,
        xyz_max = np.amax(points[:, 0:3], axis=0) # 3,
        xyz_center = (xyz_min + xyz_max) / 2
        dimension = (xyz_max - xyz_min) / 2

        xyz_offsets = xyz_center - points[:,0:3] # Nx3
        dimensions = np.ones((points.shape[0],3)) * dimension # Nx3
        labels = np.ones((points.shape[0],1)) * label # N
        point_bbox_list.append(np.concatenate([points, labels,
                                           xyz_offsets, dimensions], 1)) # Nx13

    point_bbox = np.concatenate(point_bbox_list, 0) # KxNx13
    room_xyz_min = np.amin(point_bbox[:, 0:3], axis=0)
    point_bbox[:, 0:3] -= room_xyz_min

    if file_format == 'txt':
        fout = open(out_filename, 'w')
        for i in range(point_bbox.shape[0]):
            fout.write('%f %f %f %d %d %d %d %f %f %f %f %f %f\n' % \
                          (point_bbox[i,0], point_bbox[i,1], point_bbox[i,2],
                           point_bbox[i,3], point_bbox[i,4], point_bbox[i,5],
                           point_bbox[i,6],
                           point_bbox[i,7], point_bbox[i,8], point_bbox[i,9],
                           point_bbox[i,10], point_bbox[i,11], point_bbox[i,12]))

        fout.close()
    elif file_format == 'numpy':
        np.save(out_filename, point_bbox)
    else:
        print('ERROR!! Unknown file format: %s, please use txt or numpy.' % \
            (file_format))
        exit()