I will provide you with two pieces of code and their external definitions. The first one is my code, which I am using for inference with PointNet for semantic segmentation. The second one is the inference code from the original PyTorch implementation. While my code performs correctly in some areas, it doesn't produce predictions as accurate as the original code. I am using the same pretrained models and working on the same dataset. My data contains only XYZ and RGB information—nothing more.

What is missing in my code, and what changes should I make to improve it? Please analyze both my code and the original one to provide me with a solution.

MY CODE
import torch
import numpy as np
import importlib
import os
import sys
import logging
from torch.utils.data import Dataset, DataLoader
import argparse
from datetime import datetime
import open3d as o3d

# Add the directory containing your model to the Python path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

def parse_args():
    parser = argparse.ArgumentParser(description='Point Cloud Segmentation')
    parser.add_argument('--input_file', type=str, required=True, help='Path to the input point cloud file (PCD or TXT)')
    parser.add_argument('--output_folder', type=str, required=True, help='Path to the output folder')
    parser.add_argument('--batch_size', type=int, default=32, help='Batch size for processing')
    parser.add_argument('--num_point', type=int, default=4096, help='Number of points per sample')
    parser.add_argument('--model_path', type=str, default='log/sem_seg/pointnet2_sem_seg/checkpoints/best_model.pth', help='Path to the trained model')
    parser.add_argument('--num_votes', type=int, default=5, help='Number of voting iterations')
    parser.add_argument('--model_name', type=str, default='log.sem_seg.pointnet2_sem_seg.pointnet2_sem_seg', help='Model script')
    parser.add_argument('--evaluation_mode', type=str, choices=['whole_scene', 'defined_points'], default='whole_scene', help='Evaluation mode: whole scene or defined number of points')
    return parser.parse_args()

NUM_CLASSES = 13
BATCH_SIZE = 32
NUM_POINT = 4096

# Define color map
COLOR_MAP = {
    'ceiling': [0, 255, 0],
    'floor': [0, 0, 255],
    'wall': [0, 255, 255],
    'beam': [255, 255, 0],
    'column': [255, 0, 255],
    'window': [100, 100, 255],
    'door': [200, 200, 100],
    'table': [170, 120, 200],
    'chair': [255, 0, 0],
    'sofa': [200, 100, 100],
    'bookcase': [10, 200, 100],
    'board': [200, 200, 200],
    'clutter': [50, 50, 50]
}

LABEL_TO_NAMES = {i: name for i, name in enumerate(COLOR_MAP.keys())}

def setup_logger(log_file):
    logger = logging.getLogger(__name__)
    logger.setLevel(logging.INFO)
    
    file_handler = logging.FileHandler(log_file)
    file_handler.setLevel(logging.INFO)
    
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)
    
    formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    file_handler.setFormatter(formatter)
    console_handler.setFormatter(formatter)
    
    logger.addHandler(file_handler)
    logger.addHandler(console_handler)
    
    return logger

class ImprovedPointCloudDataset(Dataset):
    def __init__(self, file_path, num_point=4096, block_size=1.0, stride=0.5, padding=0.001, evaluation_mode='whole_scene'):
        self.num_point = num_point
        self.block_size = block_size
        self.stride = stride
        self.padding = padding
        self.evaluation_mode = evaluation_mode
        
        # Load data
        self.points = self.load_point_cloud(file_path)
        self.coord_min = np.amin(self.points[:, :3], axis=0)
        self.coord_max = np.amax(self.points[:, :3], axis=0)
        
        if evaluation_mode == 'whole_scene':
            # Calculate grid size
            self.grid_x = int(np.ceil(float(self.coord_max[0] - self.coord_min[0] - self.block_size) / self.stride) + 1)
            self.grid_y = int(np.ceil(float(self.coord_max[1] - self.coord_min[1] - self.block_size) / self.stride) + 1)
        else:
            self.num_samples = len(self.points) // self.num_point
    
    def load_point_cloud(self, file_path):
        file_extension = os.path.splitext(file_path)[1].lower()
        
        if file_extension == '.pcd':
            pcd = o3d.io.read_point_cloud(file_path)
            points = np.asarray(pcd.points)
            colors = np.asarray(pcd.colors)
            if colors.shape[0] == 0:  # If no colors, use default values
                colors = np.ones_like(points) * 0.5
            return np.hstack((points, colors * 255))
        elif file_extension == '.txt':
            return np.loadtxt(file_path, delimiter=' ')
        else:
            raise ValueError(f"Unsupported file format: {file_extension}")

    def __getitem__(self, index):
        if self.evaluation_mode == 'whole_scene':
            return self._get_whole_scene_item(index)
        else:
            return self._get_defined_points_item(index)
    
    def _get_whole_scene_item(self, index):
        index_y = index // self.grid_x
        index_x = index % self.grid_x
        
        s_x = self.coord_min[0] + index_x * self.stride
        e_x = min(s_x + self.block_size, self.coord_max[0])
        s_x = e_x - self.block_size
        s_y = self.coord_min[1] + index_y * self.stride
        e_y = min(s_y + self.block_size, self.coord_max[1])
        s_y = e_y - self.block_size
        
        point_idxs = np.where(
            (self.points[:, 0] >= s_x - self.padding) & (self.points[:, 0] <= e_x + self.padding) &
            (self.points[:, 1] >= s_y - self.padding) & (self.points[:, 1] <= e_y + self.padding))[0]
        
        if point_idxs.size == 0:
            point_idxs = np.random.choice(len(self.points), self.num_point, replace=True)
        elif point_idxs.size < self.num_point:
            point_idxs = np.random.choice(point_idxs, self.num_point, replace=True)
        else:
            point_idxs = np.random.choice(point_idxs, self.num_point, replace=False)
        
        data_batch = self.points[point_idxs, :].copy()
        return self._process_data(data_batch), point_idxs
    
    def _get_defined_points_item(self, index):
        start_idx = index * self.num_point
        end_idx = start_idx + self.num_point
        data_batch = self.points[start_idx:end_idx, :].copy()
        return self._process_data(data_batch), np.arange(start_idx, end_idx)
    
    def _process_data(self, data_batch):
        # Normalize coordinates
        coord_mean = np.mean(data_batch[:, :3], axis=0)
        data_batch[:, :3] -= coord_mean
        coord_max = np.max(np.linalg.norm(data_batch[:, :3], axis=1))
        data_batch[:, :3] /= coord_max
        
        # Normalize color values
        data_batch[:, 3:6] /= 255.0
        
        # Compute additional features
        centroid = np.mean(data_batch[:, :3], axis=0)
        
        # Combine all features to make exactly 9 channels
        features = np.column_stack((
            data_batch,  # XYZ and RGB (6 channels)
            data_batch[:, :3] - centroid,  # Relative position (3 channels)*
        ))
        
        return torch.from_numpy(features.astype(np.float32))
    
    def __len__(self):
        if self.evaluation_mode == 'whole_scene':
            return self.grid_x * self.grid_y
        else:
            return self.num_samples

def inference(classifier, dataloader, device, num_votes=3):
    classifier.eval()
    all_pred_labels = []
    all_point_idxs = []

    with torch.no_grad():
        for batch_idx, (points, point_idxs) in enumerate(dataloader):
            points = points.to(device)
            points = points.transpose(2, 1).contiguous()  # Shape [B, C, N]
            vote_pool = torch.zeros(points.size(0), NUM_CLASSES, points.size(2)).to(device)

            for _ in range(num_votes):
                seg_pred, _ = classifier(points)
                if seg_pred.size(1) != NUM_CLASSES:
                    seg_pred = seg_pred.transpose(1, 2)  # [B, NUM_CLASSES, N]
                vote_pool += seg_pred

            pred_labels = vote_pool.max(1)[1].cpu().numpy()  # [B, N]
            all_pred_labels.extend(pred_labels)
            all_point_idxs.extend(point_idxs.numpy())

    return np.concatenate(all_pred_labels), np.concatenate(all_point_idxs)

def save_results(points, labels, output_file):
    if points.shape[0] != len(labels):
        raise ValueError(f"Mismatch between number of points ({points.shape[0]}) and labels ({len(labels)})")
    
    with open(output_file, 'w') as f:
        for i in range(points.shape[0]):
            label = int(labels[i])  # Ensure labels are integers
            label_color = COLOR_MAP[LABEL_TO_NAMES[label]]
            f.write(f"{points[i, 0]:.6f} {points[i, 1]:.6f} {points[i, 2]:.6f} {label_color[0]} {label_color[1]} {label_color[2]}\n")

def main(args):
    # Set up logging
    os.makedirs(args.output_folder, exist_ok=True)
    log_file = os.path.join(args.output_folder, f"segmentation_log_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt")
    logger = setup_logger(log_file)
    
    # Log arguments
    logger.info("Script arguments:")
    for arg, value in vars(args).items():
        logger.info(f"{arg}: {value}")
    
    # Set device
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    logger.info(f"Using device: {device}")
    
    # Load dataset
    logger.info(f"Loading dataset from {args.input_file}")
    dataset = ImprovedPointCloudDataset(args.input_file, num_point=args.num_point, evaluation_mode=args.evaluation_mode)
    dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, num_workers=4)
    
    # Load model
    logger.info("Loading model")
    MODEL = importlib.import_module(args.model_name)
    classifier = MODEL.get_model(NUM_CLASSES).to(device)
    
    try:
        checkpoint = torch.load(args.model_path, map_location=device)
        classifier.load_state_dict(checkpoint['model_state_dict'])
    except Exception as e:
        logger.error(f"Failed to load model: {str(e)}")
        return
    
    # Perform inference
    logger.info("Starting inference")
    pred_labels, point_idxs = inference(classifier, dataloader, device, num_votes=args.num_votes)
    
    # Reshape predictions to match original point cloud
    full_pred_labels = np.zeros(len(dataset.points), dtype=int)
    full_pred_labels[point_idxs] = pred_labels

    # Save results
    input_filename = os.path.splitext(os.path.basename(args.input_file))[0]
    output_file = os.path.join(args.output_folder, f"{input_filename}_segmented.txt")
    logger.info(f"Saving results to {output_file}")
    save_results(dataset.points, full_pred_labels, output_file)
    
    # Print and log statistics
    label_counts = {}
    for label in full_pred_labels:
        label_name = LABEL_TO_NAMES[label]
        if label_name not in label_counts:
            label_counts[label_name] = 0
        label_counts[label_name] += 1
    
    logger.info("Predicted label counts:")
    for label_name, count in label_counts.items():
        logger.info(f"{label_name}: {count}")
    
    logger.info(f"Segmented {len(full_pred_labels)} points. Results saved to {output_file}")

if __name__ == '__main__':
    args = parse_args()
    main(args)
ORIGINAL INFERENCE CODE

"""
Author: Benny
Date: Nov 2019
"""
import argparse
import os
from data_utils.S3DISDataLoader import ScannetDatasetWholeScene
from data_utils.indoor3d_util import g_label2color
import torch
import logging
from pathlib import Path
import sys
import importlib
from tqdm import tqdm
import provider
import numpy as np

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT_DIR = BASE_DIR
sys.path.append(os.path.join(ROOT_DIR, 'models'))

classes = ['ceiling', 'floor', 'wall', 'beam', 'column', 'window', 'door', 'table', 'chair', 'sofa', 'bookcase',
           'board', 'clutter']
class2label = {cls: i for i, cls in enumerate(classes)}
seg_classes = class2label
seg_label_to_cat = {}
for i, cat in enumerate(seg_classes.keys()):
    seg_label_to_cat[i] = cat


def parse_args():
    '''PARAMETERS'''
    parser = argparse.ArgumentParser('Model')
    parser.add_argument('--batch_size', type=int, default=32, help='batch size in testing [default: 32]')
    parser.add_argument('--gpu', type=str, default='0', help='specify gpu device')
    parser.add_argument('--num_point', type=int, default=4096, help='point number [default: 4096]')
    parser.add_argument('--log_dir', type=str, required=True, help='experiment root')
    parser.add_argument('--visual', action='store_true', default=False, help='visualize result [default: False]')
    parser.add_argument('--test_area', type=int, default=5, help='area for testing, option: 1-6 [default: 5]')
    parser.add_argument('--num_votes', type=int, default=3, help='aggregate segmentation scores with voting [default: 5]')
    return parser.parse_args()


def add_vote(vote_label_pool, point_idx, pred_label, weight):
    B = pred_label.shape[0]
    N = pred_label.shape[1]
    for b in range(B):
        for n in range(N):
            if weight[b, n] != 0 and not np.isinf(weight[b, n]):
                vote_label_pool[int(point_idx[b, n]), int(pred_label[b, n])] += 1
    return vote_label_pool


def main(args):
    def log_string(str):
        logger.info(str)
        print(str)

    '''HYPER PARAMETER'''
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
    experiment_dir = 'log/sem_seg/' + args.log_dir
    visual_dir = experiment_dir + '/visual/'
    visual_dir = Path(visual_dir)
    visual_dir.mkdir(exist_ok=True)

    '''LOG'''
    args = parse_args()
    logger = logging.getLogger("Model")
    logger.setLevel(logging.INFO)
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    file_handler = logging.FileHandler('%s/eval.txt' % experiment_dir)
    file_handler.setLevel(logging.INFO)
    file_handler.setFormatter(formatter)
    logger.addHandler(file_handler)
    log_string('PARAMETER ...')
    log_string(args)

    NUM_CLASSES = 13
    BATCH_SIZE = args.batch_size
    NUM_POINT = args.num_point

    root = 'data/s3dis/stanford_indoor3d/'

    TEST_DATASET_WHOLE_SCENE = ScannetDatasetWholeScene(root, split='test', test_area=args.test_area, block_points=NUM_POINT)
    log_string("The number of test data is: %d" % len(TEST_DATASET_WHOLE_SCENE))

    '''MODEL LOADING'''
    model_name = os.listdir(experiment_dir + '/logs')[0].split('.')[0]
    MODEL = importlib.import_module(model_name)
    classifier = MODEL.get_model(NUM_CLASSES).cuda()
    checkpoint = torch.load(str(experiment_dir) + '/checkpoints/best_model.pth')
    classifier.load_state_dict(checkpoint['model_state_dict'])
    classifier = classifier.eval()

    with torch.no_grad():
        scene_id = TEST_DATASET_WHOLE_SCENE.file_list
        scene_id = [x[:-4] for x in scene_id]
        num_batches = len(TEST_DATASET_WHOLE_SCENE)

        total_seen_class = [0 for _ in range(NUM_CLASSES)]
        total_correct_class = [0 for _ in range(NUM_CLASSES)]
        total_iou_deno_class = [0 for _ in range(NUM_CLASSES)]

        log_string('---- EVALUATION WHOLE SCENE----')

        for batch_idx in range(num_batches):
            print("Inference [%d/%d] %s ..." % (batch_idx + 1, num_batches, scene_id[batch_idx]))
            total_seen_class_tmp = [0 for _ in range(NUM_CLASSES)]
            total_correct_class_tmp = [0 for _ in range(NUM_CLASSES)]
            total_iou_deno_class_tmp = [0 for _ in range(NUM_CLASSES)]
            if args.visual:
                fout = open(os.path.join(visual_dir, scene_id[batch_idx] + '_pred.obj'), 'w')
                fout_gt = open(os.path.join(visual_dir, scene_id[batch_idx] + '_gt.obj'), 'w')

            whole_scene_data = TEST_DATASET_WHOLE_SCENE.scene_points_list[batch_idx]
            whole_scene_label = TEST_DATASET_WHOLE_SCENE.semantic_labels_list[batch_idx]
            vote_label_pool = np.zeros((whole_scene_label.shape[0], NUM_CLASSES))
            for _ in tqdm(range(args.num_votes), total=args.num_votes):
                scene_data, scene_label, scene_smpw, scene_point_index = TEST_DATASET_WHOLE_SCENE[batch_idx]
                num_blocks = scene_data.shape[0]
                s_batch_num = (num_blocks + BATCH_SIZE - 1) // BATCH_SIZE
                batch_data = np.zeros((BATCH_SIZE, NUM_POINT, 9))

                batch_label = np.zeros((BATCH_SIZE, NUM_POINT))
                batch_point_index = np.zeros((BATCH_SIZE, NUM_POINT))
                batch_smpw = np.zeros((BATCH_SIZE, NUM_POINT))

                for sbatch in range(s_batch_num):
                    start_idx = sbatch * BATCH_SIZE
                    end_idx = min((sbatch + 1) * BATCH_SIZE, num_blocks)
                    real_batch_size = end_idx - start_idx
                    batch_data[0:real_batch_size, ...] = scene_data[start_idx:end_idx, ...]
                    batch_label[0:real_batch_size, ...] = scene_label[start_idx:end_idx, ...]
                    batch_point_index[0:real_batch_size, ...] = scene_point_index[start_idx:end_idx, ...]
                    batch_smpw[0:real_batch_size, ...] = scene_smpw[start_idx:end_idx, ...]
                    batch_data[:, :, 3:6] /= 1.0

                    torch_data = torch.Tensor(batch_data)
                    torch_data = torch_data.float().cuda()
                    torch_data = torch_data.transpose(2, 1)
                    seg_pred, _ = classifier(torch_data)
                    batch_pred_label = seg_pred.contiguous().cpu().data.max(2)[1].numpy()

                    vote_label_pool = add_vote(vote_label_pool, batch_point_index[0:real_batch_size, ...],
                                               batch_pred_label[0:real_batch_size, ...],
                                               batch_smpw[0:real_batch_size, ...])

            pred_label = np.argmax(vote_label_pool, 1)

            for l in range(NUM_CLASSES):
                total_seen_class_tmp[l] += np.sum((whole_scene_label == l))
                total_correct_class_tmp[l] += np.sum((pred_label == l) & (whole_scene_label == l))
                total_iou_deno_class_tmp[l] += np.sum(((pred_label == l) | (whole_scene_label == l)))
                total_seen_class[l] += total_seen_class_tmp[l]
                total_correct_class[l] += total_correct_class_tmp[l]
                total_iou_deno_class[l] += total_iou_deno_class_tmp[l]

            iou_map = np.array(total_correct_class_tmp) / (np.array(total_iou_deno_class_tmp, dtype=np.float) + 1e-6)
            print(iou_map)
            arr = np.array(total_seen_class_tmp)
            tmp_iou = np.mean(iou_map[arr != 0])
            log_string('Mean IoU of %s: %.4f' % (scene_id[batch_idx], tmp_iou))
            print('----------------------------')

            filename = os.path.join(visual_dir, scene_id[batch_idx] + '.txt')
            with open(filename, 'w') as pl_save:
                for i in pred_label:
                    pl_save.write(str(int(i)) + '\n')
                pl_save.close()
            for i in range(whole_scene_label.shape[0]):
                color = g_label2color[pred_label[i]]
                color_gt = g_label2color[whole_scene_label[i]]
                if args.visual:
                    fout.write('v %f %f %f %d %d %d\n' % (
                        whole_scene_data[i, 0], whole_scene_data[i, 1], whole_scene_data[i, 2], color[0], color[1],
                        color[2]))
                    fout_gt.write(
                        'v %f %f %f %d %d %d\n' % (
                            whole_scene_data[i, 0], whole_scene_data[i, 1], whole_scene_data[i, 2], color_gt[0],
                            color_gt[1], color_gt[2]))
            if args.visual:
                fout.close()
                fout_gt.close()

        IoU = np.array(total_correct_class) / (np.array(total_iou_deno_class, dtype=np.float) + 1e-6)
        iou_per_class_str = '------- IoU --------\n'
        for l in range(NUM_CLASSES):
            iou_per_class_str += 'class %s, IoU: %.3f \n' % (
                seg_label_to_cat[l] + ' ' * (14 - len(seg_label_to_cat[l])),
                total_correct_class[l] / float(total_iou_deno_class[l]))
        log_string(iou_per_class_str)
        log_string('eval point avg class IoU: %f' % np.mean(IoU))
        log_string('eval whole scene point avg class acc: %f' % (
            np.mean(np.array(total_correct_class) / (np.array(total_seen_class, dtype=np.float) + 1e-6))))
        log_string('eval whole scene point accuracy: %f' % (
                np.sum(total_correct_class) / float(np.sum(total_seen_class) + 1e-6)))

        print("Done!")


if __name__ == '__main__':
    args = parse_args()
    main(args)



sacnnetdatasetwholescene function

class ScannetDatasetWholeScene():
    # prepare to give prediction on each points
    def __init__(self, root, block_points=4096, split='test', test_area=5, stride=0.5, block_size=1.0, padding=0.001):
        self.block_points = block_points
        self.block_size = block_size
        self.padding = padding
        self.root = root
        self.split = split
        self.stride = stride
        self.scene_points_num = []
        assert split in ['train', 'test']
        if self.split == 'train':
            self.file_list = [d for d in os.listdir(root) if d.find('Area_%d' % test_area) is -1]
        else:
            self.file_list = [d for d in os.listdir(root) if d.find('Area_%d' % test_area) is not -1]
        self.scene_points_list = []
        self.semantic_labels_list = []
        self.room_coord_min, self.room_coord_max = [], []
        for file in self.file_list:
            data = np.load(root + file)
            points = data[:, :3]
            self.scene_points_list.append(data[:, :6])
            self.semantic_labels_list.append(data[:, 6])
            coord_min, coord_max = np.amin(points, axis=0)[:3], np.amax(points, axis=0)[:3]
            self.room_coord_min.append(coord_min), self.room_coord_max.append(coord_max)
        assert len(self.scene_points_list) == len(self.semantic_labels_list)

        labelweights = np.zeros(13)
        for seg in self.semantic_labels_list:
            tmp, _ = np.histogram(seg, range(14))
            self.scene_points_num.append(seg.shape[0])
            labelweights += tmp
        labelweights = labelweights.astype(np.float32)
        labelweights = labelweights / np.sum(labelweights)
        self.labelweights = np.power(np.amax(labelweights) / labelweights, 1 / 3.0)

    def __getitem__(self, index):
        point_set_ini = self.scene_points_list[index]
        points = point_set_ini[:,:6]
        labels = self.semantic_labels_list[index]
        coord_min, coord_max = np.amin(points, axis=0)[:3], np.amax(points, axis=0)[:3]
        grid_x = int(np.ceil(float(coord_max[0] - coord_min[0] - self.block_size) / self.stride) + 1)
        grid_y = int(np.ceil(float(coord_max[1] - coord_min[1] - self.block_size) / self.stride) + 1)
        data_room, label_room, sample_weight, index_room = np.array([]), np.array([]), np.array([]),  np.array([])
        for index_y in range(0, grid_y):
            for index_x in range(0, grid_x):
                s_x = coord_min[0] + index_x * self.stride
                e_x = min(s_x + self.block_size, coord_max[0])
                s_x = e_x - self.block_size
                s_y = coord_min[1] + index_y * self.stride
                e_y = min(s_y + self.block_size, coord_max[1])
                s_y = e_y - self.block_size
                point_idxs = np.where(
                    (points[:, 0] >= s_x - self.padding) & (points[:, 0] <= e_x + self.padding) & (points[:, 1] >= s_y - self.padding) & (
                                points[:, 1] <= e_y + self.padding))[0]
                if point_idxs.size == 0:
                    continue
                num_batch = int(np.ceil(point_idxs.size / self.block_points))
                point_size = int(num_batch * self.block_points)
                replace = False if (point_size - point_idxs.size <= point_idxs.size) else True
                point_idxs_repeat = np.random.choice(point_idxs, point_size - point_idxs.size, replace=replace)
                point_idxs = np.concatenate((point_idxs, point_idxs_repeat))
                np.random.shuffle(point_idxs)
                data_batch = points[point_idxs, :]
                normlized_xyz = np.zeros((point_size, 3))
                normlized_xyz[:, 0] = data_batch[:, 0] / coord_max[0]
                normlized_xyz[:, 1] = data_batch[:, 1] / coord_max[1]
                normlized_xyz[:, 2] = data_batch[:, 2] / coord_max[2]
                data_batch[:, 0] = data_batch[:, 0] - (s_x + self.block_size / 2.0)
                data_batch[:, 1] = data_batch[:, 1] - (s_y + self.block_size / 2.0)
                data_batch[:, 3:6] /= 255.0
                data_batch = np.concatenate((data_batch, normlized_xyz), axis=1)
                label_batch = labels[point_idxs].astype(int)
                batch_weight = self.labelweights[label_batch]

                data_room = np.vstack([data_room, data_batch]) if data_room.size else data_batch
                label_room = np.hstack([label_room, label_batch]) if label_room.size else label_batch
                sample_weight = np.hstack([sample_weight, batch_weight]) if label_room.size else batch_weight
                index_room = np.hstack([index_room, point_idxs]) if index_room.size else point_idxs
        data_room = data_room.reshape((-1, self.block_points, data_room.shape[1]))
        label_room = label_room.reshape((-1, self.block_points))
        sample_weight = sample_weight.reshape((-1, self.block_points))
        index_room = index_room.reshape((-1, self.block_points))
        return data_room, label_room, sample_weight, index_room

    def __len__(self):
        return len(self.scene_points_list)

g_label2color variable:

g_classes = [x.rstrip() for x in open(os.path.join(BASE_DIR, 'meta/class_names.txt'))]
g_class2label = {cls: i for i,cls in enumerate(g_classes)}
g_class2color = {'ceiling':	[0,255,0],
                 'floor':	[0,0,255],
                 'wall':	[0,255,255],
                 'beam':        [255,255,0],
                 'column':      [255,0,255],
                 'window':      [100,100,255],
                 'door':        [200,200,100],
                 'table':       [170,120,200],
                 'chair':       [255,0,0],
                 'sofa':        [200,100,100],
                 'bookcase':    [10,200,100],
                 'board':       [200,200,200],
                 'clutter':     [50,50,50]} 
g_easy_view_labels = [7,8,9,10,11,1]
g_label2color = {g_classes.index(cls): g_class2color[cls] for cls in g_classes}