following code can not segment point cloud compeletly true. examine the method of inference using on the original pointnet semantic segmentation pytorch implementation github repository and make necessary changes to create inference correctly. give me the whole code with changes

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 = 16384

# 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')
    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=16384, 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('--use_full_cloud', action='store_true', help='Use full point cloud for segmentation')
    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 text file.
    """
    data = np.loadtxt(file_path, delimiter=' ')
    xyz = data[:, :3]
    rgb = data[:, 3:6]
    return xyz, rgb

def sample_point_cloud(xyz, rgb, num_points):
    """
    Randomly sample a subset of points from the point cloud.
    """
    if len(xyz) <= num_points:
        return xyz, rgb
    
    indices = np.random.choice(len(xyz), num_points, replace=False)
    return xyz[indices], rgb[indices]

def process_point_cloud(xyz, rgb, num_neighbors=30):
    """
    Process XYZ and RGB data using Open3D.
    """
    # Create Open3D PointCloud object
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(xyz)
    pcd.colors = o3d.utility.Vector3dVector(rgb)

    # Normalize coordinates
    pcd.normalize_normals()
    
    xyz_normalized = np.asarray(pcd.points)
    
    # Estimate normals
    pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamKNN(knn=num_neighbors))
    pcd.orient_normals_consistent_tangent_plane(k=num_neighbors)
    normals = np.asarray(pcd.normals)
    
    # Combine XYZ, RGB, and normals (9 channels in total)
    points = np.hstack((xyz_normalized, rgb , normals))
    points[:, 3:6] /= 255.0
    return points

def segment_point_cloud(classifier, xyz, rgb, use_full_cloud=True, num_point=NUM_POINT):
    """
    Perform segmentation on input XYZ and RGB data.
    """
    if not use_full_cloud:
        xyz, rgb = sample_point_cloud(xyz, rgb, num_point)
    
    points = process_point_cloud(xyz, rgb)

    # Implement batching
    num_points = points.shape[0]
    num_batches = (num_points + num_point - 1) // num_point
    pred_labels = []

    for i in range(num_batches):
        start_idx = i * num_point
        end_idx = min((i + 1) * num_point, num_points)
        batch_points = points[start_idx:end_idx]

        # Pad the batch if necessary
        if batch_points.shape[0] < num_point:
            pad_size = num_point - batch_points.shape[0]
            batch_points = np.pad(batch_points, ((0, pad_size), (0, 0)), mode='constant')

        # Prepare data for the model
        batch_points_tensor = torch.from_numpy(batch_points).float().unsqueeze(0)
        batch_points_tensor = batch_points_tensor.transpose(2, 1)

        # Perform segmentation
        with torch.no_grad():
            seg_pred, _ = classifier(batch_points_tensor)

        batch_pred_label = seg_pred.contiguous().data.max(2)[1].numpy()
        pred_labels.append(batch_pred_label[0][:end_idx-start_idx])  # Remove padding

    pred_labels = np.concatenate(pred_labels)
    return pred_labels, points

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.basename(args.input_file))

    # 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('cpu'))['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
    logger.info("Starting segmentation")
    segmentation_labels, processed_points = segment_point_cloud(classifier, xyz, rgb, args.use_full_cloud, args.num_point)

    # Save results as TXT file
    logger.info(f"Saving results to {output_file_txt}")
    save_as_txt(processed_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)