use pointnet semantic segmentation pytorch implementation github repo and make sure that following code works. my input is x y z r g b. i get following error

2024-09-16 00:33:11,727 - INFO - Script arguments:
2024-09-16 00:33:11,727 - INFO - input_file: data/org_data/s3dis/office_3.txt
2024-09-16 00:33:11,727 - INFO - output_folder: data/segmented/s3dis/
2024-09-16 00:33:11,727 - INFO - batch_size: 32
2024-09-16 00:33:11,727 - INFO - num_point: 16384
2024-09-16 00:33:11,727 - INFO - model_path: log/sem_seg/pointnet2_sem_seg/checkpoints/best_model.pth
2024-09-16 00:33:11,727 - INFO - model_name: log.sem_seg.pointnet2_sem_seg.pointnet2_sem_seg
2024-09-16 00:33:11,727 - INFO - use_full_cloud: False
2024-09-16 00:33:11,727 - INFO - Using device: cpu
2024-09-16 00:33:11,728 - INFO - Loading model
2024-09-16 00:33:11,825 - INFO - Model loaded successfully
2024-09-16 00:33:11,825 - INFO - Reading point cloud from data/org_data/s3dis/office_3.txt
2024-09-16 00:33:12,317 - INFO - Starting segmentation
2024-09-16 00:33:24,781 - ERROR - Segmentation failed: Given groups=1, weight of size [32, 12, 1, 1], expected input[1, 16387, 32, 1024] to have 12 channels, but got 16387 channels instead


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
    file_handler = logging.FileHandler(log_file)
    file_handler.setLevel(logging.INFO)
    
    # Console handler
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)
    
    # Formatter
    formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    file_handler.setFormatter(formatter)
    console_handler.setFormatter(formatter)
    
    # Add handlers
    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]
    if data.shape[1] >=6:
        rgb = data[:, 3:6]
    else:
        rgb = None
    return xyz, rgb

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

def process_point_cloud(xyz, rgb, num_neighbors=30):
    """
    Process XYZ and RGB data using Open3D.
    Correct normalization and feature extraction.
    """
    # Create Open3D PointCloud object
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(xyz)
    if rgb is not None:
        pcd.colors = o3d.utility.Vector3dVector(rgb / 255.0)  # Normalize RGB to [0,1]
    
    # Normalize the point cloud: zero mean, unit scale
    points = np.asarray(pcd.points)
    mean = np.mean(points, axis=0)
    points -= mean
    scale = np.max(np.linalg.norm(points, axis=1))
    points /= scale
    pcd.points = o3d.utility.Vector3dVector(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 and normals (6 channels)
    points = np.hstack((points, normals))
    
    # If using RGB, include it
    if rgb is not None:
        points = np.hstack((points, np.asarray(pcd.colors)))
    
    return points

def segment_point_cloud(classifier, xyz, rgb, device, 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)
    
    # Move to device
    points_tensor = torch.from_numpy(points).float().to(device)
    
    # Depending on the model's expected input dimensions
    # The model likely expects [B, C, N]
    points_tensor = points_tensor.unsqueeze(0)  # [1, C, N]
    
    # Check if num of points matches, else pad or truncate
    current_num_points = points_tensor.shape[2]
    if current_num_points < num_point:
        pad_size = num_point - current_num_points
        padding = torch.zeros((1, points_tensor.shape[1], pad_size)).to(device)
        points_tensor = torch.cat([points_tensor, padding], dim=2)
    elif current_num_points > num_point:
        points_tensor = points_tensor[:, :, :num_point]
    
    # Perform segmentation
    with torch.no_grad():
        seg_pred, _ = classifier(points_tensor)
    
    # Get predicted labels
    batch_pred_label = seg_pred.contiguous().data.max(2)[1].cpu().numpy()
    pred_labels = batch_pred_label[0][:xyz.shape[0]]  # Get the labels for actual points
    
    return pred_labels, points

def save_as_txt(original_xyz, 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(original_xyz.shape[0]):
            xyz = original_xyz[i]
            label = labels[i]
            label_name = LABEL_TO_NAMES.get(label, 'unknown')
            label_color = COLOR_MAP.get(label_name, [0, 0, 0])
            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))
    
    # Device configuration
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    logger.info(f"Using device: {device}")
    
    # Load your model
    logger.info("Loading model")
    MODEL = importlib.import_module(args.model_name)
    classifier = MODEL.get_model(NUM_CLASSES)
    classifier.to(device)
    
    try:
        # Load model weights
        state_dict = torch.load(args.model_path, map_location=device)['model_state_dict']
        # Remove 'module.' prefix if present
        new_state_dict = {}
        for k, v in state_dict.items():
            if k.startswith('module.'):
                new_state_dict[k[7:]] = v
            else:
                new_state_dict[k] = v
        classifier.load_state_dict(new_state_dict)
        classifier.eval()
        logger.info("Model loaded successfully")
    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}")
    try:
        xyz, rgb = read_point_cloud(args.input_file)
    except Exception as e:
        logger.error(f"Failed to read point cloud: {str(e)}")
        return
    
    # Perform segmentation
    logger.info("Starting segmentation")
    try:
        pred_labels, processed_points = segment_point_cloud(classifier, xyz, rgb, device, args.use_full_cloud, args.num_point)
    except Exception as e:
        logger.error(f"Segmentation failed: {str(e)}")
        return
    
    # Save results as TXT file
    logger.info(f"Saving results to {output_file_txt}")
    try:
        save_as_txt(xyz, pred_labels, output_file_txt)
    except Exception as e:
        logger.error(f"Failed to save segmented point cloud: {str(e)}")
        return
    
    # Print and log predicted labels
    label_counts = {}
    for label in pred_labels:
        label_name = LABEL_TO_NAMES.get(label, 'unknown')
        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(pred_labels)} points. Results saved to {output_file_txt}")

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