title point cloud segmentation with pointnet soumik rakshithttps github comsoumik12345 sayak paulhttps github comsayakpaul date created 20201023 last modified 20201024 description implementation of a pointnetbased model for segmenting point clouds accelerator gpu introduction a point cloud is an important type of data structure for storing geometric shape data due to its irregular format it s often transformed into regular 3d voxel grids or collections of images before being used in deep learning applications a step which makes the data unnecessarily large the pointnet family of models solves this problem by directly consuming point clouds respecting the permutationinvariance property of the point data the pointnet family of models provides a simple unified architecture for applications ranging from object classification part segmentation to scene semantic parsing in this example we demonstrate the implementation of the pointnet architecture for shape segmentation references pointnet deep learning on point sets for 3d classification and segmentationhttps arxiv orgabs1612 00593 point cloud classification with pointnethttps keras ioexamplesvisionpointnet spatial transformer networkshttps arxiv orgabs1506 02025 imports downloading dataset the shapenet datasethttps shapenet org is an ongoing effort to establish a richlyannotated largescale dataset of 3d shapes shapenetcore is a subset of the full shapenet dataset with clean single 3d models and manually verified category and alignment annotations it covers 55 common object categories with about 51 300 unique 3d models for this example we use one of the 12 object categories of pascal 3dhttp cvgl stanford eduprojectspascal3d html included as part of the shapenetcore dataset loading the dataset we parse the dataset metadata in order to easily map model categories to their respective directories and segmentation classes to colors for the purpose of visualization in this example we train pointnet to segment the parts of an airplane model structuring the dataset we generate the following inmemory data structures from the airplane point clouds and their labels pointclouds is a list of np array objects that represent the point cloud data in the form of x y and z coordinates axis 0 represents the number of points in the point cloud while axis 1 represents the coordinates alllabels is the list that represents the label of each coordinate as a string needed mainly for visualization purposes testpointclouds is in the same format as pointclouds but doesn t have corresponding the labels of the point clouds alllabels is a list of np array objects that represent the point cloud labels for each coordinate corresponding to the pointclouds list pointcloudlabels is a list of np array objects that represent the point cloud labels for each coordinate in onehot encoded form corresponding to the pointclouds list get the fileid of the current point cloud for parsing its labels point clouds having labels will be our training samples apply onehot encoding to the dense label representation next we take a look at some samples from the inmemory arrays we just generated now let s visualize some of the point clouds along with their labels preprocessing note that all the point clouds that we have loaded consist of a variable number of points which makes it difficult for us to batch them together in order to overcome this problem we randomly sample a fixed number of points from each point cloud we also normalize the point clouds in order to make the data scaleinvariant randomly sampling respective indices sampling points corresponding to sampled indices sampling corresponding onehot encoded labels sampling corresponding labels for visualization normalizing sampled point cloud let s visualize the sampled and normalized point clouds along with their corresponding labels creating tensorflow datasets we create tf data dataset objects for the training and validation data we also augment the training point clouds by applying random jitter to them pointnet model the figure below depicts the internals of the pointnet model family https i imgur comqflnw5l png given that pointnet is meant to consume an unordered set of coordinates as its input data its architecture needs to match the following characteristic properties of point cloud data permutation invariance given the unstructured nature of point cloud data a scan made up of n points has n permutations the subsequent data processing must be invariant to the different representations in order to make pointnet invariant to input permutations we use a symmetric function such as maxpooling once the n input points are mapped to higherdimensional space the result is a global feature vector that aims to capture an aggregate signature of the n input points the global feature vector is used alongside local point features for segmentation https i imgur com0mrvvjb png transformation invariance segmentation outputs should be unchanged if the object undergoes certain transformations such as translation or scaling for a given input point cloud we apply an appropriate rigid or affine transformation to achieve pose normalization because each of the n input points are represented as a vector and are mapped to the embedding spaces independently applying a geometric transformation simply amounts to matrix multiplying each point with a transformation matrix this is motivated by the concept of spatial transformer networkshttps arxiv orgabs1506 02025 the operations comprising the tnet are motivated by the higherlevel architecture of pointnet mlps or fullyconnected layers are used to map the input points independently and identically to a higherdimensional space maxpooling is used to encode a global feature vector whose dimensionality is then reduced with fullyconnected layers the inputdependent features at the final fullyconnected layer are then combined with globally trainable weights and biases resulting in a 3by3 transformation matrix https i imgur comaej3gyi png point interactions the interaction between neighboring points often carries useful information i e a single point should not be treated in isolation whereas classification need only make use of global features segmentation must be able to leverage local point features along with global point features note the figures presented in this section have been taken from the original paperhttps arxiv orgabs1612 00593 now that we know the pieces that compose the pointnet model we can implement the model we start by implementing the basic blocks i e the convolutional block and the multilayer perceptron block we implement a regularizer taken from this examplehttps keras ioexamplesvisionpointnetbuildamodel to enforce orthogonality in the feature space this is needed to ensure that the magnitudes of the transformed features do not vary too much reference https keras ioexamplesvisionpointnetbuildamodel def initself numfeatures l2reg0 001 self numfeatures numfeatures self l2reg l2reg self identity keras ops eyenumfeatures def callself x x keras ops reshapex 1 self numfeatures self numfeatures xxt keras ops tensordotx x axes2 2 xxt keras ops reshapexxt 1 self numfeatures self numfeatures return keras ops sumself l2reg keras ops squarexxt self identity def getconfigself config super getconfig config updatenumfeatures self numfeatures l2regstrength self l2reg return config def transformationnetinputs numfeatures name x convblockinputs filters64 namefname1 x convblockx filters128 namefname2 x convblockx filters1024 namefname3 x layers globalmaxpooling1dx x mlpblockx filters512 namefname11 x mlpblockx filters256 namefname21 return layers dense numfeatures numfeatures kernelinitializerzeros biasinitializerkeras initializers constantnp eyenumfeatures flatten activityregularizerorthogonalregularizernumfeatures namefnamefinal x def transformationblockinputs numfeatures name transformedfeatures transformationnetinputs numfeatures namename transformedfeatures layers reshapenumfeatures numfeatures transformedfeatures return layers dotaxes2 1 namefnamemminputs transformedfeatures def getshapesegmentationmodelnumpoints numclasses inputpoints keras inputshapenone 3 pointnet classification network transformedinputs transformationblock inputpoints numfeatures3 nameinputtransformationblock features64 convblocktransformedinputs filters64 namefeatures64 features1281 convblockfeatures64 filters128 namefeatures1281 features1282 convblockfeatures1281 filters128 namefeatures1282 transformedfeatures transformationblock features1282 numfeatures128 nametransformedfeatures features512 convblocktransformedfeatures filters512 namefeatures512 features2048 convblockfeatures512 filters2048 namepremaxpoolblock globalfeatures layers maxpool1dpoolsizenumpoints nameglobalfeatures features2048 globalfeatures keras ops tileglobalfeatures 1 numpoints 1 segmentation head segmentationinput layers concatenatenamesegmentationinput features64 features1281 features1282 transformedfeatures features512 globalfeatures segmentationfeatures convblock segmentationinput filters128 namesegmentationfeatures outputs layers conv1d numclasses kernelsize1 activationsoftmax namesegmentationhead segmentationfeatures return keras modelinputpoints outputs instantiate the model x y nextitertraindataset numpoints x shape1 numclasses y shape1 segmentationmodel getshapesegmentationmodelnumpoints numclasses segmentationmodel summary training stepsperepoch totaltrainingexamples batchsize totaltrainingsteps stepsperepoch epochs printfsteps per epoch stepsperepoch printftotal training steps totaltrainingsteps lrschedule keras optimizers schedules exponentialdecay initiallearningrate0 003 decaystepsstepsperepoch 5 decayrate0 5 staircasetrue steps rangetotaltrainingsteps lrs lrschedulestep for step in steps plt plotlrs plt xlabelsteps plt ylabellearning rate plt show def runexperimentepochs segmentationmodel getshapesegmentationmodelnumpoints numclasses segmentationmodel compile optimizerkeras optimizers adamlearningratelrschedule losskeras losses categoricalcrossentropy metricsaccuracy checkpointfilepath checkpoint weights h5 checkpointcallback keras callbacks modelcheckpoint checkpointfilepath monitorvalloss savebestonlytrue saveweightsonlytrue history segmentationmodel fit traindataset validationdatavaldataset epochsepochs callbackscheckpointcallback segmentationmodel loadweightscheckpointfilepath return segmentationmodel history segmentationmodel history runexperimentepochsepochs visualize the training landscape def plotresultitem plt plothistory historyitem labelitem plt plothistory historyval item labelval item plt xlabelepochs plt ylabelitem plt titletrain and validation over epochs formatitem fontsize14 plt legend plt grid plt show plotresultloss plotresultaccuracy inference validationbatch nextitervaldataset valpredictions segmentationmodel predictvalidationbatch0 printfvalidation prediction shape valpredictions shape def visualizesinglepointcloudpointclouds labelclouds idx labelmap labels none pointcloud pointcloudsidx labelcloud labelcloudsidx visualizedatapointcloud labelmapnp argmaxlabel for label in labelcloud idx np random choicelenvalidationbatch0 printfindex selected idx plotting with groundtruth visualizesinglepointcloudvalidationbatch0 validationbatch1 idx plotting with predicted labels visualizesinglepointcloudvalidationbatch0 valpredictions idx final notes introduction a point cloud is an important type of data structure for storing geometric shape data due to its irregular format it s often transformed into regular 3d voxel grids or collections of images before being used in deep learning applications a step which makes the data unnecessarily large the pointnet family of models solves this problem by directly consuming point clouds respecting the permutation invariance property of the point data the pointnet family of models provides a simple unified architecture for applications ranging from object classification part segmentation to scene semantic parsing in this example we demonstrate the implementation of the pointnet architecture for shape segmentation references pointnet deep learning on point sets for 3d classification and segmentation https arxiv org abs 1612 00593 point cloud classification with pointnet https keras io examples vision pointnet spatial transformer networks https arxiv org abs 1506 02025 imports for tf data downloading dataset the shapenet dataset https shapenet org is an ongoing effort to establish a richly annotated large scale dataset of 3d shapes shapenetcore is a subset of the full shapenet dataset with clean single 3d models and manually verified category and alignment annotations it covers 55 common object categories with about 51 300 unique 3d models for this example we use one of the 12 object categories of pascal 3d http cvgl stanford edu projects pascal3d html included as part of the shapenetcore dataset loading the dataset we parse the dataset metadata in order to easily map model categories to their respective directories and segmentation classes to colors for the purpose of visualization in this example we train pointnet to segment the parts of an airplane model structuring the dataset we generate the following in memory data structures from the airplane point clouds and their labels point_clouds is a list of np array objects that represent the point cloud data in the form of x y and z coordinates axis 0 represents the number of points in the point cloud while axis 1 represents the coordinates all_labels is the list that represents the label of each coordinate as a string needed mainly for visualization purposes test_point_clouds is in the same format as point_clouds but doesn t have corresponding the labels of the point clouds all_labels is a list of np array objects that represent the point cloud labels for each coordinate corresponding to the point_clouds list point_cloud_labels is a list of np array objects that represent the point cloud labels for each coordinate in one hot encoded form corresponding to the point_clouds list get the file id of the current point cloud for parsing its labels point clouds having labels will be our training samples apply one hot encoding to the dense label representation next we take a look at some samples from the in memory arrays we just generated now let s visualize some of the point clouds along with their labels preprocessing note that all the point clouds that we have loaded consist of a variable number of points which makes it difficult for us to batch them together in order to overcome this problem we randomly sample a fixed number of points from each point cloud we also normalize the point clouds in order to make the data scale invariant randomly sampling respective indices sampling points corresponding to sampled indices sampling corresponding one hot encoded labels sampling corresponding labels for visualization normalizing sampled point cloud let s visualize the sampled and normalized point clouds along with their corresponding labels creating tensorflow datasets we create tf data dataset objects for the training and validation data we also augment the training point clouds by applying random jitter to them pointnet model the figure below depicts the internals of the pointnet model family https i imgur com qflnw5l png given that pointnet is meant to consume an unordered set of coordinates as its input data its architecture needs to match the following characteristic properties of point cloud data permutation invariance given the unstructured nature of point cloud data a scan made up of n points has n permutations the subsequent data processing must be invariant to the different representations in order to make pointnet invariant to input permutations we use a symmetric function such as max pooling once the n input points are mapped to higher dimensional space the result is a global feature vector that aims to capture an aggregate signature of the n input points the global feature vector is used alongside local point features for segmentation https i imgur com 0mrvvjb png transformation invariance segmentation outputs should be unchanged if the object undergoes certain transformations such as translation or scaling for a given input point cloud we apply an appropriate rigid or affine transformation to achieve pose normalization because each of the n input points are represented as a vector and are mapped to the embedding spaces independently applying a geometric transformation simply amounts to matrix multiplying each point with a transformation matrix this is motivated by the concept of spatial transformer networks https arxiv org abs 1506 02025 the operations comprising the t net are motivated by the higher level architecture of pointnet mlps or fully connected layers are used to map the input points independently and identically to a higher dimensional space max pooling is used to encode a global feature vector whose dimensionality is then reduced with fully connected layers the input dependent features at the final fully connected layer are then combined with globally trainable weights and biases resulting in a 3 by 3 transformation matrix https i imgur com aej3gyi png point interactions the interaction between neighboring points often carries useful information i e a single point should not be treated in isolation whereas classification need only make use of global features segmentation must be able to leverage local point features along with global point features note the figures presented in this section have been taken from the original paper https arxiv org abs 1612 00593 now that we know the pieces that compose the pointnet model we can implement the model we start by implementing the basic blocks i e the convolutional block and the multi layer perceptron block we implement a regularizer taken from this example https keras io examples vision pointnet build a model to enforce orthogonality in the feature space this is needed to ensure that the magnitudes of the transformed features do not vary too much reference https keras io examples vision pointnet build a model the next piece is the transformation network which we explained earlier reference https keras io examples vision pointnet build a model the filters values come from the original paper https arxiv org abs 1612 00593 finally we piece the above blocks together and implement the segmentation model pointnet classification network segmentation head instantiate the model training for the training the s recommend using a learning rate schedule that decays the initial learning rate by half every 20 epochs in this example we use 5 epochs finally we implement a utility for running our experiments and launch model training visualize the training landscape inference plotting with ground truth plotting with predicted labels final notes if you are interested in learning more about this topic you may find this repository https github com soumik12345 point cloud segmentation useful
import os
import json
import random
import numpy as np
import pandas as pd
from tqdm import tqdm
from glob import glob

import tensorflow as tf  
import keras
from keras import layers

import matplotlib.pyplot as plt



dataset_url = "https://git.io/JiY4i"

dataset_path = keras.utils.get_file(
    fname="shapenet.zip",
    origin=dataset_url,
    cache_subdir="datasets",
    hash_algorithm="auto",
    extract=True,
    archive_format="auto",
    cache_dir="datasets",
)



with open("/tmp/.keras/datasets/PartAnnotation/metadata.json") as json_file:
    metadata = json.load(json_file)

print(metadata)



points_dir = "/tmp/.keras/datasets/PartAnnotation/{}/points".format(
    metadata["Airplane"]["directory"]
)
labels_dir = "/tmp/.keras/datasets/PartAnnotation/{}/points_label".format(
    metadata["Airplane"]["directory"]
)
LABELS = metadata["Airplane"]["lables"]
COLORS = metadata["Airplane"]["colors"]

VAL_SPLIT = 0.2
NUM_SAMPLE_POINTS = 1024
BATCH_SIZE = 32
EPOCHS = 60
INITIAL_LR = 1e-3



point_clouds, test_point_clouds = [], []
point_cloud_labels, all_labels = [], []

points_files = glob(os.path.join(points_dir, "*.pts"))
for point_file in tqdm(points_files):
    point_cloud = np.loadtxt(point_file)
    if point_cloud.shape[0] < NUM_SAMPLE_POINTS:
        continue

    
    
    file_id = point_file.split("/")[-1].split(".")[0]
    label_data, num_labels = {}, 0
    for label in LABELS:
        label_file = os.path.join(labels_dir, label, file_id + ".seg")
        if os.path.exists(label_file):
            label_data[label] = np.loadtxt(label_file).astype("float32")
            num_labels = len(label_data[label])

    
    try:
        label_map = ["none"] * num_labels
        for label in LABELS:
            for i, data in enumerate(label_data[label]):
                label_map[i] = label if data == 1 else label_map[i]
        label_data = [
            LABELS.index(label) if label != "none" else len(LABELS)
            for label in label_map
        ]
        
        label_data = keras.utils.to_categorical(label_data, num_classes=len(LABELS) + 1)

        point_clouds.append(point_cloud)
        point_cloud_labels.append(label_data)
        all_labels.append(label_map)
    except KeyError:
        test_point_clouds.append(point_cloud)



for _ in range(5):
    i = random.randint(0, len(point_clouds) - 1)
    print(f"point_clouds[{i}].shape:", point_clouds[0].shape)
    print(f"point_cloud_labels[{i}].shape:", point_cloud_labels[0].shape)
    for j in range(5):
        print(
            f"all_labels[{i}][{j}]:",
            all_labels[i][j],
            f"\tpoint_cloud_labels[{i}][{j}]:",
            point_cloud_labels[i][j],
            "\n",
        )




def visualize_data(point_cloud, labels):
    df = pd.DataFrame(
        data={
            "x": point_cloud[:, 0],
            "y": point_cloud[:, 1],
            "z": point_cloud[:, 2],
            "label": labels,
        }
    )
    fig = plt.figure(figsize=(15, 10))
    ax = plt.axes(projection="3d")
    for index, label in enumerate(LABELS):
        c_df = df[df["label"] == label]
        try:
            ax.scatter(
                c_df["x"], c_df["y"], c_df["z"], label=label, alpha=0.5, c=COLORS[index]
            )
        except IndexError:
            pass
    ax.legend()
    plt.show()


visualize_data(point_clouds[0], all_labels[0])
visualize_data(point_clouds[300], all_labels[300])



for index in tqdm(range(len(point_clouds))):
    current_point_cloud = point_clouds[index]
    current_label_cloud = point_cloud_labels[index]
    current_labels = all_labels[index]
    num_points = len(current_point_cloud)
    
    sampled_indices = random.sample(list(range(num_points)), NUM_SAMPLE_POINTS)
    
    sampled_point_cloud = np.array([current_point_cloud[i] for i in sampled_indices])
    
    sampled_label_cloud = np.array([current_label_cloud[i] for i in sampled_indices])
    
    sampled_labels = np.array([current_labels[i] for i in sampled_indices])
    
    norm_point_cloud = sampled_point_cloud - np.mean(sampled_point_cloud, axis=0)
    norm_point_cloud /= np.max(np.linalg.norm(norm_point_cloud, axis=1))
    point_clouds[index] = norm_point_cloud
    point_cloud_labels[index] = sampled_label_cloud
    all_labels[index] = sampled_labels



visualize_data(point_clouds[0], all_labels[0])
visualize_data(point_clouds[300], all_labels[300])




def load_data(point_cloud_batch, label_cloud_batch):
    point_cloud_batch.set_shape([NUM_SAMPLE_POINTS, 3])
    label_cloud_batch.set_shape([NUM_SAMPLE_POINTS, len(LABELS) + 1])
    return point_cloud_batch, label_cloud_batch


def augment(point_cloud_batch, label_cloud_batch):
    noise = tf.random.uniform(
        tf.shape(label_cloud_batch), -0.001, 0.001, dtype=tf.float64
    )
    point_cloud_batch += noise[:, :, :3]
    return point_cloud_batch, label_cloud_batch


def generate_dataset(point_clouds, label_clouds, is_training=True):
    dataset = tf.data.Dataset.from_tensor_slices((point_clouds, label_clouds))
    dataset = dataset.shuffle(BATCH_SIZE * 100) if is_training else dataset
    dataset = dataset.map(load_data, num_parallel_calls=tf.data.AUTOTUNE)
    dataset = dataset.batch(batch_size=BATCH_SIZE)
    dataset = (
        dataset.map(augment, num_parallel_calls=tf.data.AUTOTUNE)
        if is_training
        else dataset
    )
    return dataset


split_index = int(len(point_clouds) * (1 - VAL_SPLIT))
train_point_clouds = point_clouds[:split_index]
train_label_cloud = point_cloud_labels[:split_index]
total_training_examples = len(train_point_clouds)

val_point_clouds = point_clouds[split_index:]
val_label_cloud = point_cloud_labels[split_index:]

print("Num train point clouds:", len(train_point_clouds))
print("Num train point cloud labels:", len(train_label_cloud))
print("Num val point clouds:", len(val_point_clouds))
print("Num val point cloud labels:", len(val_label_cloud))

train_dataset = generate_dataset(train_point_clouds, train_label_cloud)
val_dataset = generate_dataset(val_point_clouds, val_label_cloud, is_training=False)

print("Train Dataset:", train_dataset)
print("Validation Dataset:", val_dataset)






def conv_block(x, filters, name):
    x = layers.Conv1D(filters, kernel_size=1, padding="valid", name=f"{name}_conv")(x)
    x = layers.BatchNormalization(name=f"{name}_batch_norm")(x)
    return layers.Activation("relu", name=f"{name}_relu")(x)


def mlp_block(x, filters, name):
    x = layers.Dense(filters, name=f"{name}_dense")(x)
    x = layers.BatchNormalization(name=f"{name}_batch_norm")(x)
    return layers.Activation("relu", name=f"{name}_relu")(x)





class OrthogonalRegularizer(keras.regularizers.Regularizer):
    

    def __init__(self, num_features, l2reg=0.001):
        self.num_features = num_features
        self.l2reg = l2reg
        self.identity = keras.ops.eye(num_features)

    def __call__(self, x):
        x = keras.ops.reshape(x, (-1, self.num_features, self.num_features))
        xxt = keras.ops.tensordot(x, x, axes=(2, 2))
        xxt = keras.ops.reshape(xxt, (-1, self.num_features, self.num_features))
        return keras.ops.sum(self.l2reg * keras.ops.square(xxt - self.identity))

    def get_config(self):
        config = super().get_config()
        config.update({"num_features": self.num_features, "l2reg_strength": self.l2reg})
        return config





def transformation_net(inputs, num_features, name):
    
    x = conv_block(inputs, filters=64, name=f"{name}_1")
    x = conv_block(x, filters=128, name=f"{name}_2")
    x = conv_block(x, filters=1024, name=f"{name}_3")
    x = layers.GlobalMaxPooling1D()(x)
    x = mlp_block(x, filters=512, name=f"{name}_1_1")
    x = mlp_block(x, filters=256, name=f"{name}_2_1")
    return layers.Dense(
        num_features * num_features,
        kernel_initializer="zeros",
        bias_initializer=keras.initializers.Constant(np.eye(num_features).flatten()),
        activity_regularizer=OrthogonalRegularizer(num_features),
        name=f"{name}_final",
    )(x)


def transformation_block(inputs, num_features, name):
    transformed_features = transformation_net(inputs, num_features, name=name)
    transformed_features = layers.Reshape((num_features, num_features))(
        transformed_features
    )
    return layers.Dot(axes=(2, 1), name=f"{name}_mm")([inputs, transformed_features])





def get_shape_segmentation_model(num_points, num_classes):
    input_points = keras.Input(shape=(None, 3))

    
    transformed_inputs = transformation_block(
        input_points, num_features=3, name="input_transformation_block"
    )
    features_64 = conv_block(transformed_inputs, filters=64, name="features_64")
    features_128_1 = conv_block(features_64, filters=128, name="features_128_1")
    features_128_2 = conv_block(features_128_1, filters=128, name="features_128_2")
    transformed_features = transformation_block(
        features_128_2, num_features=128, name="transformed_features"
    )
    features_512 = conv_block(transformed_features, filters=512, name="features_512")
    features_2048 = conv_block(features_512, filters=2048, name="pre_maxpool_block")
    global_features = layers.MaxPool1D(pool_size=num_points, name="global_features")(
        features_2048
    )
    global_features = keras.ops.tile(global_features, [1, num_points, 1])

    
    segmentation_input = layers.Concatenate(name="segmentation_input")(
        [
            features_64,
            features_128_1,
            features_128_2,
            transformed_features,
            features_512,
            global_features,
        ]
    )
    segmentation_features = conv_block(
        segmentation_input, filters=128, name="segmentation_features"
    )
    outputs = layers.Conv1D(
        num_classes, kernel_size=1, activation="softmax", name="segmentation_head"
    )(segmentation_features)
    return keras.Model(input_points, outputs)




x, y = next(iter(train_dataset))

num_points = x.shape[1]
num_classes = y.shape[-1]

segmentation_model = get_shape_segmentation_model(num_points, num_classes)
segmentation_model.summary()



steps_per_epoch = total_training_examples // BATCH_SIZE
total_training_steps = steps_per_epoch * EPOCHS
print(f"Steps per epoch: {steps_per_epoch}.")
print(f"Total training steps: {total_training_steps}.")

lr_schedule = keras.optimizers.schedules.ExponentialDecay(
    initial_learning_rate=0.003, decay_steps=steps_per_epoch * 5, decay_rate=0.5, staircase=True
)

steps = range(total_training_steps)
lrs = [lr_schedule(step) for step in steps]

plt.plot(lrs)
plt.xlabel("Steps")
plt.ylabel("Learning Rate")
plt.show()




def run_experiment(epochs):
    segmentation_model = get_shape_segmentation_model(num_points, num_classes)
    segmentation_model.compile(
        optimizer=keras.optimizers.Adam(learning_rate=lr_schedule),
        loss=keras.losses.CategoricalCrossentropy(),
        metrics=["accuracy"],
    )

    checkpoint_filepath = "checkpoint.weights.h5"
    checkpoint_callback = keras.callbacks.ModelCheckpoint(
        checkpoint_filepath,
        monitor="val_loss",
        save_best_only=True,
        save_weights_only=True,
    )

    history = segmentation_model.fit(
        train_dataset,
        validation_data=val_dataset,
        epochs=epochs,
        callbacks=[checkpoint_callback],
    )

    segmentation_model.load_weights(checkpoint_filepath)
    return segmentation_model, history


segmentation_model, history = run_experiment(epochs=EPOCHS)




def plot_result(item):
    plt.plot(history.history[item], label=item)
    plt.plot(history.history["val_" + item], label="val_" + item)
    plt.xlabel("Epochs")
    plt.ylabel(item)
    plt.title("Train and Validation {} Over Epochs".format(item), fontsize=14)
    plt.legend()
    plt.grid()
    plt.show()


plot_result("loss")
plot_result("accuracy")



validation_batch = next(iter(val_dataset))
val_predictions = segmentation_model.predict(validation_batch[0])
print(f"Validation prediction shape: {val_predictions.shape}")


def visualize_single_point_cloud(point_clouds, label_clouds, idx):
    label_map = LABELS + ["none"]
    point_cloud = point_clouds[idx]
    label_cloud = label_clouds[idx]
    visualize_data(point_cloud, [label_map[np.argmax(label)] for label in label_cloud])


idx = np.random.choice(len(validation_batch[0]))
print(f"Index selected: {idx}")


visualize_single_point_cloud(validation_batch[0], validation_batch[1], idx)


visualize_single_point_cloud(validation_batch[0], val_predictions, idx)