okay here is the scannetdatasetwholescene class. i want to implement this num_points option to my code which is not working on that way. I want this option to reduce points on the scene like test_semseg.py file that i provide you here is my code also. 

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)


my code:
import torch
import numpy as np
import importlib
import os
import sys
import open3d as o3d
import argparse
import logging
from datetime import datetime

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

NUM_CLASSES = 13
BATCH_SIZE = 32
NUM_POINT = 4096  # Adjusted to match the original code

# 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 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 (TXT or PCD)')
    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('--model_name', type=str, default='log.sem_seg.pointnet2_sem_seg.pointnet2_sem_seg', help='Model script')
    parser.add_argument('--num_votes', type=int, default=3, help='Number of votes for segmentation')
    parser.add_argument('--block_size', type=float, default=1.0, help='Block size for block splitting')
    parser.add_argument('--stride', type=float, default=0.5, help='Stride for block splitting')
    return parser.parse_args()

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

def read_point_cloud(file_path):
    """
    Read XYZ and RGB data from a TXT or PCD file.
    """
    file_extension = os.path.splitext(file_path)[1].lower()

    if file_extension == '.txt':
        return read_txt_point_cloud(file_path)
    elif file_extension == '.pcd':
        return read_pcd_point_cloud(file_path)
    else:
        raise ValueError(f"Unsupported file format: {file_extension}")

def read_txt_point_cloud(file_path):
    """
    Read XYZ and RGB data from a text file.
    """
    data = np.loadtxt(file_path, delimiter=' ')
    xyz = data[:, :3]
    rgb = data[:, 3:6]
    return xyz, rgb

def read_pcd_point_cloud(file_path):
    """
    Read XYZ and RGB data from a PCD file.
    """
    pcd = o3d.io.read_point_cloud(file_path)
    xyz = np.asarray(pcd.points)
    rgb = np.asarray(pcd.colors)
    
    # If RGB values are normalized (0-1), scale them to 0-255
    if rgb.max() <= 1.0:
        rgb *= 255
    
    return xyz, rgb

def split_point_cloud_into_blocks(xyz, rgb, num_point, block_size=1.0, stride=0.5):
    """
    Split the point cloud into blocks.
    """
    # Determine the min and max coordinates
    coord_min = np.min(xyz, axis=0)
    coord_max = np.max(xyz, axis=0)
    
    # Calculate the number of blocks in each dimension
    grid_x = int(np.ceil((coord_max[0] - coord_min[0] - block_size) / stride)) + 1
    grid_y = int(np.ceil((coord_max[1] - coord_min[1] - block_size) / stride)) + 1
    
    blocks = []
    
    for idx_y in range(grid_y):
        for idx_x in range(grid_x):
            s_x = coord_min[0] + idx_x * stride
            s_y = coord_min[1] + idx_y * stride
            e_x = s_x + block_size
            e_y = s_y + block_size
            
            # Find points within this block
            block_point_indices = np.where(
                (xyz[:, 0] >= s_x) & (xyz[:, 0] <= e_x) &
                (xyz[:, 1] >= s_y) & (xyz[:, 1] <= e_y)
            )[0]
            
            if block_point_indices.size == 0:
                continue
            
            block_xyz = xyz[block_point_indices]
            block_rgb = rgb[block_point_indices]
            
            # Adjust coordinates within the block
            block_xyz_centered = block_xyz - np.array([s_x + block_size / 2.0, s_y + block_size / 2.0, 0])
            
            # Normalize XYZ coordinates
            block_xyz_normalized = block_xyz / coord_max
            
            # Normalize RGB
            block_rgb_normalized = block_rgb / 255.0
            
            # Combine features
            block_points = np.hstack((block_xyz_centered, block_rgb_normalized, block_xyz_normalized))
            
            # Pad or sample to num_point
            if block_points.shape[0] >= num_point:
                # Randomly sample num_point points
                idx = np.random.choice(block_points.shape[0], num_point, replace=False)
                block_points = block_points[idx]
                block_point_indices = block_point_indices[idx]
            else:
                # Pad with duplicated points
                idx = np.random.choice(block_points.shape[0], num_point - block_points.shape[0], replace=True)
                block_points = np.concatenate((block_points, block_points[idx]), axis=0)
                block_point_indices = np.concatenate((block_point_indices, block_point_indices[idx]), axis=0)
            
            blocks.append((block_points, block_point_indices))
    
    return blocks

def segment_point_cloud_with_voting(classifier, xyz, rgb, num_point, num_votes=3, block_size=1.0, stride=0.5):
    """
    Perform segmentation with voting.
    """
    num_classes = NUM_CLASSES
    num_points = xyz.shape[0]
    vote_label_pool = np.zeros((num_points, num_classes))
    
    for vote in range(num_votes):
        # Introduce randomness in point cloud for each vote
        xyz_vote = xyz + np.random.uniform(-0.01, 0.01, xyz.shape)
        # Split the point cloud into blocks
        blocks = split_point_cloud_into_blocks(xyz_vote, rgb, num_point, block_size, stride)
        
        for block_points, block_point_indices in blocks:
            # Prepare data for the model
            batch_points_tensor = torch.from_numpy(block_points).float().unsqueeze(0)
            batch_points_tensor = batch_points_tensor.transpose(2, 1)
            batch_points_tensor = batch_points_tensor.to(next(classifier.parameters()).device)
            
            # Perform segmentation
            with torch.no_grad():
                seg_pred, _ = classifier(batch_points_tensor)
                batch_pred_label = seg_pred.contiguous().cpu().data.max(2)[1].numpy()
            
            # Update vote label pool
            for idx, pred_label in zip(block_point_indices, batch_pred_label[0]):
                vote_label_pool[idx, pred_label] += 1
    
    # Get final predicted labels
    final_pred_labels = np.argmax(vote_label_pool, axis=1)
    return final_pred_labels

def save_as_txt(points, labels, output_file):
    """
    Save the point cloud with segmentation labels as a TXT file.
    """
    with open(output_file, 'w') as f:
        for i in range(points.shape[0]):
            xyz = points[i, :3]
            label_color = COLOR_MAP[LABEL_TO_NAMES[labels[i]]]
            f.write(f"{xyz[0]:.6f} {xyz[1]:.6f} {xyz[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 output file path
    output_file_txt = os.path.join(args.output_folder, os.path.splitext(os.path.basename(args.input_file))[0] + "_segmented.txt")

    # Load your model
    logger.info("Loading model")
    MODEL = importlib.import_module(args.model_name)
    classifier = MODEL.get_model(NUM_CLASSES)

    try:
        # Load model weights
        state_dict = torch.load(args.model_path, map_location=torch.device('cuda'))['model_state_dict']
        new_state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}
        classifier.load_state_dict(new_state_dict)
        classifier.eval()
    except Exception as e:
        logger.error(f"Failed to load model: {str(e)}")
        return

    # Read point cloud data
    logger.info(f"Reading point cloud from {args.input_file}")
    xyz, rgb = read_point_cloud(args.input_file)

    # Perform segmentation with voting
    logger.info("Starting segmentation with voting")
    segmentation_labels = segment_point_cloud_with_voting(
        classifier, xyz, rgb, args.num_point, num_votes=args.num_votes, block_size=args.block_size, stride=args.stride
    )

    # Save results as TXT file
    logger.info(f"Saving results to {output_file_txt}")
    points = np.hstack((xyz, rgb))
    save_as_txt(points, segmentation_labels, output_file_txt)

    # Print and log predicted labels
    label_counts = {}
    for label in segmentation_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(segmentation_labels)} points. Results saved to {output_file_txt}")

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

test_semseg.py code from original repo


"""
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)