Update pages/01_🦷 Segment.py

pages/01_🦷 Segment.py

@@ -1,27 +1,29 @@
-from streamlit import session_state as session
+import os
 import shutil
+import json
 
-import os
 import numpy as np
-from sklearn import neighbors
 from scipy.spatial import distance_matrix
+from sklearn import neighbors
 from pygco import cut_from_graph
 import open3d as o3d
 import matplotlib.pyplot as plt
 import matplotlib.colors as mcolors
-from stqdm import stqdm
-import json
-from stpyvista import stpyvista
 import torch
 import torch.nn as nn
 from torch.autograd import Variable
 import torch.nn.functional as F
 import streamlit as st
-import pyvista as pv
-
+from streamlit import session_state as session
+from stpyvista import stpyvista
+from stqdm import stqdm
 from PIL import Image
 
+# Configure Streamlit page
 class TeethApp:
+    """
+    Base class for Streamlit app
+    """
     def __init__(self):
         # Font
         with open("utils/style.css") as css:
@@ -48,44 +50,6 @@ class TeethApp:
             unsafe_allow_html=True,
         )
 
-
-class STN3d(nn.Module):
-    def __init__(self, channel):
-        super(STN3d, self).__init__()
-        self.conv1 = torch.nn.Conv1d(channel, 64, 1)
-        self.conv2 = torch.nn.Conv1d(64, 128, 1)
-        self.conv3 = torch.nn.Conv1d(128, 1024, 1)
-        self.fc1 = nn.Linear(1024, 512)
-        self.fc2 = nn.Linear(512, 256)
-        self.fc3 = nn.Linear(256, 9)
-        self.relu = nn.ReLU()
-
-        self.bn1 = nn.BatchNorm1d(64)
-        self.bn2 = nn.BatchNorm1d(128)
-        self.bn3 = nn.BatchNorm1d(1024)
-        self.bn4 = nn.BatchNorm1d(512)
-        self.bn5 = nn.BatchNorm1d(256)
-
-    def forward(self, x):
-        batchsize = x.size()[0]
-        x = F.relu(self.bn1(self.conv1(x)))
-        x = F.relu(self.bn2(self.conv2(x)))
-        x = F.relu(self.bn3(self.conv3(x)))
-        x = torch.max(x, 2, keepdim=True)[0]
-        x = x.view(-1, 1024)
-
-        x = F.relu(self.bn4(self.fc1(x)))
-        x = F.relu(self.bn5(self.fc2(x)))
-        x = self.fc3(x)
-
-        iden = Variable(torch.from_numpy(np.array([1, 0, 0, 0, 1, 0, 0, 0, 1]).astype(np.float32))).view(1, 9).repeat(
-            batchsize, 1)
-        if x.is_cuda:
-            iden = iden.to(x.get_device())
-        x = x + iden
-        x = x.view(-1, 3, 3)
-        return x
-
 class STNkd(nn.Module):
     def __init__(self, k=64):
         super(STNkd, self).__init__()
@@ -133,13 +97,15 @@ class MeshSegNet(nn.Module):
         self.with_dropout = with_dropout
         self.dropout_p = dropout_p
 
-        # MLP-1 [64, 64]
+        # MLP-1 -shape: [64, 64]
         self.mlp1_conv1 = torch.nn.Conv1d(self.num_channels, 64, 1)
         self.mlp1_conv2 = torch.nn.Conv1d(64, 64, 1)
         self.mlp1_bn1 = nn.BatchNorm1d(64)
         self.mlp1_bn2 = nn.BatchNorm1d(64)
+
         # FTM (feature-transformer module)
         self.fstn = STNkd(k=64)
+
         # GLM-1 (graph-contrained learning modulus)
         self.glm1_conv1_1 = torch.nn.Conv1d(64, 32, 1)
         self.glm1_conv1_2 = torch.nn.Conv1d(64, 32, 1)
@@ -147,6 +113,7 @@ class MeshSegNet(nn.Module):
         self.glm1_bn1_2 = nn.BatchNorm1d(32)
         self.glm1_conv2 = torch.nn.Conv1d(32+32, 64, 1)
         self.glm1_bn2 = nn.BatchNorm1d(64)
+
         # MLP-2
         self.mlp2_conv1 = torch.nn.Conv1d(64, 64, 1)
         self.mlp2_bn1 = nn.BatchNorm1d(64)
@@ -154,6 +121,7 @@ class MeshSegNet(nn.Module):
         self.mlp2_bn2 = nn.BatchNorm1d(128)
         self.mlp2_conv3 = torch.nn.Conv1d(128, 512, 1)
         self.mlp2_bn3 = nn.BatchNorm1d(512)
+
         # GLM-2 (graph-contrained learning modulus)
         self.glm2_conv1_1 = torch.nn.Conv1d(512, 128, 1)
         self.glm2_conv1_2 = torch.nn.Conv1d(512, 128, 1)
@@ -163,6 +131,7 @@ class MeshSegNet(nn.Module):
         self.glm2_bn1_3 = nn.BatchNorm1d(128)
         self.glm2_conv2 = torch.nn.Conv1d(128*3, 512, 1)
         self.glm2_bn2 = nn.BatchNorm1d(512)
+
         # MLP-3
         self.mlp3_conv1 = torch.nn.Conv1d(64+512+512+512, 256, 1)
         self.mlp3_conv2 = torch.nn.Conv1d(256, 256, 1)
@@ -172,7 +141,8 @@ class MeshSegNet(nn.Module):
         self.mlp3_conv4 = torch.nn.Conv1d(128, 128, 1)
         self.mlp3_bn2_1 = nn.BatchNorm1d(128)
         self.mlp3_bn2_2 = nn.BatchNorm1d(128)
-        # output
+
+        # Output
         self.output_conv = torch.nn.Conv1d(128, self.num_classes, 1)
         if self.with_dropout:
             self.dropout = nn.Dropout(p=self.dropout_p)
@@ -180,13 +150,16 @@ class MeshSegNet(nn.Module):
     def forward(self, x, a_s, a_l):
         batchsize = x.size()[0]
         n_pts = x.size()[2]
+
         # MLP-1
         x = F.relu(self.mlp1_bn1(self.mlp1_conv1(x)))
         x = F.relu(self.mlp1_bn2(self.mlp1_conv2(x)))
+
         # FTM
         trans_feat = self.fstn(x)
         x = x.transpose(2, 1)
         x_ftm = torch.bmm(x, trans_feat)
+
         # GLM-1
         sap = torch.bmm(a_s, x_ftm)
         sap = sap.transpose(2, 1)
@@ -195,12 +168,14 @@ class MeshSegNet(nn.Module):
         glm_1_sap = F.relu(self.glm1_bn1_2(self.glm1_conv1_2(sap)))
         x = torch.cat([x, glm_1_sap], dim=1)
         x = F.relu(self.glm1_bn2(self.glm1_conv2(x)))
+
         # MLP-2
         x = F.relu(self.mlp2_bn1(self.mlp2_conv1(x)))
         x = F.relu(self.mlp2_bn2(self.mlp2_conv2(x)))
         x_mlp2 = F.relu(self.mlp2_bn3(self.mlp2_conv3(x)))
         if self.with_dropout:
             x_mlp2 = self.dropout(x_mlp2)
+
         # GLM-2
         x_mlp2 = x_mlp2.transpose(2, 1)
         sap_1 = torch.bmm(a_s, x_mlp2)
@@ -213,12 +188,16 @@ class MeshSegNet(nn.Module):
         glm_2_sap_2 = F.relu(self.glm2_bn1_3(self.glm2_conv1_3(sap_2)))
         x = torch.cat([x, glm_2_sap_1, glm_2_sap_2], dim=1)
         x_glm2 = F.relu(self.glm2_bn2(self.glm2_conv2(x)))
+
         # GMP
         x = torch.max(x_glm2, 2, keepdim=True)[0]
+
         # Upsample
         x = torch.nn.Upsample(n_pts)(x)
+
         # Dense fusion
         x = torch.cat([x, x_ftm, x_mlp2, x_glm2], dim=1)
+
         # MLP-3
         x = F.relu(self.mlp3_bn1_1(self.mlp3_conv1(x)))
         x = F.relu(self.mlp3_bn1_2(self.mlp3_conv2(x)))
@@ -226,6 +205,7 @@ class MeshSegNet(nn.Module):
         if self.with_dropout:
             x = self.dropout(x)
         x = F.relu(self.mlp3_bn2_2(self.mlp3_conv4(x)))
+
         # output
         x = self.output_conv(x)
         x = x.transpose(2,1).contiguous()
@@ -235,20 +215,25 @@ class MeshSegNet(nn.Module):
         return x
 
 def clone_runoob(li1):
+    """
+    copy list
+    """
     li_copy = li1[:]
+
     return li_copy
 
-# 对离群点重新进行分类
+# Reclassify outliers
 def class_inlier_outlier(label_list, mean_points,cloud, ind, label_index, points, labels):
     label_change = clone_runoob(labels)
     outlier_index = clone_runoob(label_index)
     ind_reverse = clone_runoob(ind)
-    # 得到离群点的label下标
+
+    # Get the label subscript of the outlier point
     ind_reverse.reverse()
     for i in ind_reverse:
         outlier_index.pop(i)
 
-    # 获取离群点
+    # Get outliers
     inlier_cloud = cloud.select_by_index(ind)
     outlier_cloud = cloud.select_by_index(ind, invert=True)
     outlier_points = np.array(outlier_cloud.points)
@@ -256,29 +241,27 @@ def class_inlier_outlier(label_list, mean_points,cloud, ind, label_index, points
     for i in range(len(outlier_points)):
         distance = []
         for j in range(len(mean_points)):
-            dis = np.linalg.norm(outlier_points[i] - mean_points[j], ord=2)  # 计算tooth和GT质心之间的距离
+            dis = np.linalg.norm(outlier_points[i] - mean_points[j], ord=2)  # Compute the distance between tooth and GT centroid
             distance.append(dis)
-        min_index = distance.index(min(distance))  # 获取和离群点质心最近label的index
-        outlier_label = label_list[min_index]  # 获取离群点应该的label
+        min_index = distance.index(min(distance))  # Get the index of the label closest to the centroid of the outlier point
+        outlier_label = label_list[min_index]  # Get the label of the outlier point
         index = outlier_index[i]
         label_change[index] = outlier_label
 
     return label_change
 
-# 利用knn算法消除离群点
+# Use knn algorithm to eliminate outliers
 def remove_outlier(points, labels):
-    # points = np.array(point_cloud_o3d_orign.points)
-    # global label_list
     same_label_points = {}
 
     same_label_index = {}
 
-    mean_points = []  # 所有label种类对应点云的质心坐标
+    mean_points = [] # All label types correspond to the centroid coordinates of the point cloud.
 
     label_list = []
     for i in range(len(labels)):
         label_list.append(labels[i])
-    label_list = list(set(label_list))  # 去重获从小到大排序取GT_label=[0, 11, 12, 13, 14, 15, 16, 17, 21, 22, 23, 24, 25, 26, 27]
+    label_list = list(set(label_list)) # To retrieve the order from small to large, take GT_label=[0, 11, 12, 13, 14, 15, 16, 17, 21, 22, 23, 24, 25, 26, 27]
     label_list.sort()
     label_list = label_list[1:]
 
@@ -289,7 +272,7 @@ def remove_outlier(points, labels):
         for j in range(len(labels)):
             if labels[j] == i:
                 points_list.append(points[j].tolist())
-                all_label_index.append(j)  # 得到label为 i 的点对应的label的下标
+                all_label_index.append(j) # Get the subscript of the label corresponding to the point with label i
         same_label_points[key] = points_list
         same_label_index[key] = all_label_index
 
@@ -299,106 +282,102 @@ def remove_outlier(points, labels):
 
     for i in label_list:
         points_array = same_label_points[i]
-        # 建立一个o3d的点云对象
+        # Build one o3d object
         pcd = o3d.geometry.PointCloud()
-        # 使用Vector3dVector方法转换
+        # UseVector3dVector conversion method
         pcd.points = o3d.utility.Vector3dVector(points_array)
 
-        # 对label i 对应的点云进行统计离群值去除，找出离群点并显示
-        # 统计式离群点移除
+        # Perform statistical outlier removal on the point cloud corresponding to label i, find outliers and display them
+        # Statistical outlier removal
         cl, ind = pcd.remove_statistical_outlier(nb_neighbors=200, std_ratio=2.0)  # cl是选中的点，ind是选中点index
-        # 可视化
-        # display_inlier_outlier(pcd, ind)
 
-        # 对分出来的离群点重新分类
+        # Reclassify the separated outliers
         label_index = same_label_index[i]
         labels = class_inlier_outlier(label_list, mean_points, pcd, ind, label_index, points, labels)
         # print(f"label_change{labels[4400]}")
 
     return labels
 
-
-# 消除离群点，保存最后的输出
+# Eliminate outliers and save the final output
 def remove_outlier_main(jaw, pcd_points, labels, instances_labels):
-    # point_cloud_o3d_orign = o3d.io.read_point_cloud('E:/tooth/data/MeshSegNet-master/test_upsample_15/upsample_01K17AN8_upper_refined.pcd')
-    # 原始点
+    # original point
     points = pcd_points.copy()
     label = remove_outlier(points, labels)
 
-    # 保存json文件
+    # Save json file
     label_dict = {}
     label_dict["id_patient"] = ""
     label_dict["jaw"] = jaw
     label_dict["labels"] = label.tolist()
     label_dict["instances"] = instances_labels.tolist()
+
     b = json.dumps(label_dict)
     with open('dental-labels4' + '.json', 'w') as f_obj:
         f_obj.write(b)
     f_obj.close()
 
-
 same_points_list = {}
 
-
-# 体素下采样
+# voxel downsampling
 def voxel_filter(point_cloud, leaf_size):
     same_points_list = {}
     filtered_points = []
-    # step1 计算边界点
+
+    # step1 Calculate boundary points
     x_max, y_max, z_max = np.amax(point_cloud, axis=0)  # 计算 x,y,z三个维度的最值
     x_min, y_min, z_min = np.amin(point_cloud, axis=0)
 
-    # step2 确定体素的尺寸
+    # step2 Determine the size of the voxel
     size_r = leaf_size
 
-    # step3 计算每个 volex的维度 voxel grid
+    # step3 Calculate the dimensions of each volex voxel grid
     Dx = (x_max - x_min) // size_r + 1
     Dy = (y_max - y_min) // size_r + 1
     Dz = (z_max - z_min) // size_r + 1
 
     # print("Dx x Dy x Dz is {} x {} x {}".format(Dx, Dy, Dz))
 
-    # step4 计算每个点在volex grid内每一个维度的值
-    h = list()  # h 为保存索引的列表
+    # step4 Calculate the value of each point in each dimension in the volex grid
+    h = list()  # h is a list of saved indexes
     for i in range(len(point_cloud)):
         hx = np.floor((point_cloud[i][0] - x_min) // size_r)
         hy = np.floor((point_cloud[i][1] - y_min) // size_r)
         hz = np.floor((point_cloud[i][2] - z_min) // size_r)
         h.append(hx + hy * Dx + hz * Dx * Dy)
-    # print(h[60581])
 
-    # step5 对h值进行排序
+    # step5 Sort h values
     h = np.array(h)
-    h_indice = np.argsort(h)  # 提取索引,返回h里面的元素按从小到大排序的  索引
-    h_sorted = h[h_indice]  # 升序
-    count = 0  # 用于维度的累计
+    h_indice = np.argsort(h)  # Extract the index and return the index of the elements in h sorted from small to large.
+    h_sorted = h[h_indice]  # Ascending order
+    count = 0  # used for accumulation of dimensions
     step = 20
-    # 将h值相同的点放入到同一个grid中，并进行筛选
-    for i in range(1, len(h_sorted)):  # 0-19999个数据点
-        # if i == len(h_sorted)-1:
-        #     print("aaa")
+
+    # Put points with the same h value into the same grid and filter them
+    for i in range(1, len(h_sorted)):  # 0-19999 data points
         if h_sorted[i] == h_sorted[i - 1] and (i != len(h_sorted) - 1):
             continue
+
         elif h_sorted[i] == h_sorted[i - 1] and (i == len(h_sorted) - 1):
             point_idx = h_indice[count:]
             key = h_sorted[i - 1]
             same_points_list[key] = point_idx
-            _G = np.mean(point_cloud[point_idx], axis=0)  # 所有点的重心
-            _d = np.linalg.norm(point_cloud[point_idx] - _G, axis=1, ord=2)  # 计算到重心的距离
+            _G = np.mean(point_cloud[point_idx], axis=0)  # center of gravity of all points
+            _d = np.linalg.norm(point_cloud[point_idx] - _G, axis=1, ord=2)  # Calculate distance to center of gravity
             _d.sort()
-            inx = [j for j in range(0, len(_d), step)]  # 获取指定间隔元素下标
+            inx = [j for j in range(0, len(_d), step)]  # Get the index of the specified interval element
             for j in inx:
                 index = point_idx[j]
                 filtered_points.append(point_cloud[index])
             count = i
+
         elif h_sorted[i] != h_sorted[i - 1] and (i == len(h_sorted) - 1):
             point_idx1 = h_indice[count:i]
             key1 = h_sorted[i - 1]
             same_points_list[key1] = point_idx1
-            _G = np.mean(point_cloud[point_idx1], axis=0)  # 所有点的重心
-            _d = np.linalg.norm(point_cloud[point_idx1] - _G, axis=1, ord=2)  # 计算到重心的距离
+            _G = np.mean(point_cloud[point_idx1], axis=0)  # center of gravity of all points
+            _d = np.linalg.norm(point_cloud[point_idx1] - _G, axis=1, ord=2)  # Calculate distance to center of gravity
             _d.sort()
-            inx = [j for j in range(0, len(_d), step)]  # 获取指定间隔元素下标
+            inx = [j for j in range(0, len(_d), step)]  # Get the index of the specified interval element
             for j in inx:
                 index = point_idx1[j]
                 filtered_points.append(point_cloud[index])
@@ -406,10 +385,10 @@ def voxel_filter(point_cloud, leaf_size):
             point_idx2 = h_indice[i:]
             key2 = h_sorted[i]
             same_points_list[key2] = point_idx2
-            _G = np.mean(point_cloud[point_idx2], axis=0)  # 所有点的重心
-            _d = np.linalg.norm(point_cloud[point_idx2] - _G, axis=1, ord=2)  # 计算到重心的距离
+            _G = np.mean(point_cloud[point_idx2], axis=0)  # center of gravity of all points
+            _d = np.linalg.norm(point_cloud[point_idx2] - _G, axis=1, ord=2)  # Calculate distance to center of gravity
             _d.sort()
-            inx = [j for j in range(0, len(_d), step)]  # 获取指定间隔元素下标
+            inx = [j for j in range(0, len(_d), step)]  # Get the index of the specified interval element
             for j in inx:
                 index = point_idx2[j]
                 filtered_points.append(point_cloud[index])
@@ -419,38 +398,42 @@ def voxel_filter(point_cloud, leaf_size):
             point_idx = h_indice[count: i]
             key = h_sorted[i - 1]
             same_points_list[key] = point_idx
-            _G = np.mean(point_cloud[point_idx], axis=0)  # 所有点的重心
-            _d = np.linalg.norm(point_cloud[point_idx] - _G, axis=1, ord=2)  # 计算到重心的距离
+            _G = np.mean(point_cloud[point_idx], axis=0)  # center of gravity of all points
+            _d = np.linalg.norm(point_cloud[point_idx] - _G, axis=1, ord=2)  # Calculate distance to center of gravity
             _d.sort()
-            inx = [j for j in range(0, len(_d), step)]  # 获取指定间隔元素下标
+            inx = [j for j in range(0, len(_d), step)]  # Get the index of the specified interval element
             for j in inx:
                 index = point_idx[j]
                 filtered_points.append(point_cloud[index])
             count = i
 
-    # 把点云格式改成array，并对外返回
+    # Change the point cloud format to array and return it externally
     # print(f'filtered_points[0]为{filtered_points[0]}')
     filtered_points = np.array(filtered_points, dtype=np.float64)
+
     return filtered_points,same_points_list
 
 
-# 体素上采样
+# voxel upsampling
 def voxel_upsample(same_points_list, point_cloud, filtered_points, filter_labels, leaf_size):
     upsample_label = []
     upsample_point = []
     upsample_index = []
-    # step1 计算边界点
-    x_max, y_max, z_max = np.amax(point_cloud, axis=0)  # 计算 x,y,z三个维度的最值
+
+    # step1 Calculate boundary points
+    x_max, y_max, z_max = np.amax(point_cloud, axis=0) # Calculate the maximum value of the three dimensions x, y, z
     x_min, y_min, z_min = np.amin(point_cloud, axis=0)
-    # step2 确定体素的尺寸
+
+    # step2 Determine the size of the voxel
     size_r = leaf_size
-    # step3 计算每个 volex的维度 voxel grid
+
+    # step3 Calculate the dimensions of each volex voxel grid
     Dx = (x_max - x_min) // size_r + 1
     Dy = (y_max - y_min) // size_r + 1
     Dz = (z_max - z_min) // size_r + 1
     print("Dx x Dy x Dz is {} x {} x {}".format(Dx, Dy, Dz))
 
-    # step4 计算每个点（采样后的点）在volex grid内每一个维度的值
+    # step4 Calculate the value of each point (sampled point) in each dimension within the volex grid
     h = list()
     for i in range(len(filtered_points)):
         hx = np.floor((filtered_points[i][0] - x_min) // size_r)
@@ -458,30 +441,33 @@ def voxel_upsample(same_points_list, point_cloud, filtered_points, filter_labels
         hz = np.floor((filtered_points[i][2] - z_min) // size_r)
         h.append(hx + hy * Dx + hz * Dx * Dy)
 
-    # step5 根据h值查询字典same_points_list
+    # step5 Query the dictionary same_points_list based on the h value
     h = np.array(h)
     count = 0
     for i in range(1, len(h)):
         if h[i] == h[i - 1] and i != (len(h) - 1):
             continue
+
         elif h[i] == h[i - 1] and i == (len(h) - 1):
             label = filter_labels[count:]
             key = h[i - 1]
             count = i
-            # 累计label次数，classcount：{‘A’:2,'B':1}
+
+            # Cumulative number of labels, classcount: {‘A’: 2, ‘B’: 1}
             classcount = {}
             for i in range(len(label)):
                 vote = label[i]
                 classcount[vote] = classcount.get(vote, 0) + 1
-            # 对map的value排序
+
+            # Sort map values
             sortedclass = sorted(classcount.items(), key=lambda x: (x[1]), reverse=True)
-            # key = h[i-1]
-            point_index = same_points_list[key]  # h对应的point index列表
+            point_index = same_points_list[key]  # Point index list corresponding to h
             for j in range(len(point_index)):
                 upsample_label.append(sortedclass[0][0])
                 index = point_index[j]
                 upsample_point.append(point_cloud[index])
                 upsample_index.append(index)
+
         elif h[i] != h[i - 1] and (i == len(h) - 1):
             label1 = filter_labels[count:i]
             key1 = h[i - 1]
@@ -493,8 +479,8 @@ def voxel_upsample(same_points_list, point_cloud, filtered_points, filter_labels
             for i in range(len(label1)):
                 vote = label1[i]
                 classcount[vote] = classcount.get(vote, 0) + 1
+
             sortedclass = sorted(classcount.items(), key=lambda x: (x[1]), reverse=True)
-            # key1 = h[i-1]
             point_index = same_points_list[key1]
             for j in range(len(point_index)):
                 upsample_label.append(sortedclass[0][0])
@@ -502,13 +488,12 @@ def voxel_upsample(same_points_list, point_cloud, filtered_points, filter_labels
                 upsample_point.append(point_cloud[index])
                 upsample_index.append(index)
 
-            # label2 = filter_labels[i:]
             classcount = {}
             for i in range(len(label2)):
                 vote = label2[i]
                 classcount[vote] = classcount.get(vote, 0) + 1
+
             sortedclass = sorted(classcount.items(), key=lambda x: (x[1]), reverse=True)
-            # key2 = h[i]
             point_index = same_points_list[key2]
             for j in range(len(point_index)):
                 upsample_label.append(sortedclass[0][0])
@@ -523,58 +508,51 @@ def voxel_upsample(same_points_list, point_cloud, filtered_points, filter_labels
             for i in range(len(label)):
                 vote = label[i]
                 classcount[vote] = classcount.get(vote, 0) + 1
+
             sortedclass = sorted(classcount.items(), key=lambda x: (x[1]), reverse=True)
-            # key = h[i-1]
             point_index = same_points_list[key]  # h对应的point index列表
             for j in range(len(point_index)):
                 upsample_label.append(sortedclass[0][0])
                 index = point_index[j]
                 upsample_point.append(point_cloud[index])
                 upsample_index.append(index)
-            # count = i
 
-    # 恢复原始顺序
-    # print(f'upsample_index[0]的值为{upsample_index[0]}')
-    # print(f'upsample_index的总长度为{len(upsample_index)}')
-
-    # 恢复index原始顺序
+    # Restore the original order of index
     upsample_index = np.array(upsample_index)
-    upsample_index_indice = np.argsort(upsample_index)  # 提取索引,返回h里面的元素按从小到大排序的  索引
+    upsample_index_indice = np.argsort(upsample_index) # Extract the index and return the index of the elements in h sorted from small to large.
     upsample_index_sorted = upsample_index[upsample_index_indice]
 
     upsample_point = np.array(upsample_point)
     upsample_label = np.array(upsample_label)
-    # 恢复point和label的原始顺序
+
+    # Restore the original order of points and labels
     upsample_point_sorted = upsample_point[upsample_index_indice]
     upsample_label_sorted = upsample_label[upsample_index_indice]
 
     return upsample_point_sorted, upsample_label_sorted
 
-
-# 利用knn算法上采样
+# Upsampling using knn algorithm
 def KNN_sklearn_Load_data(voxel_points, center_points, labels):
-    # 载入数据
-    # x_train, x_test, y_train, y_test = train_test_split(center_points, labels, test_size=0.1)
-    # 构建模型
+    # Build model
     model = neighbors.KNeighborsClassifier(n_neighbors=3)
     model.fit(center_points, labels)
     prediction = model.predict(voxel_points.reshape(1, -1))
-    # meshtopoints_labels = classification_report(voxel_points, prediction)
-    return prediction[0]
 
+    return prediction[0]
 
-# 加载点进行knn上采样
+# Loading points for knn upsampling
 def Load_data(voxel_points, center_points, labels):
     meshtopoints_labels = []
-    # meshtopoints_labels.append(SVC_sklearn_Load_data(voxel_points[i], center_points, labels))
     for i in range(0, voxel_points.shape[0]):
         meshtopoints_labels.append(KNN_sklearn_Load_data(voxel_points[i], center_points, labels))
+
     return np.array(meshtopoints_labels)
 
-# 将三角网格数据上采样回原始点云数据
+# Upsample triangular mesh data back to original point cloud data
 def mesh_to_points_main(jaw, pcd_points, center_points, labels):
     points = pcd_points.copy()
-    # 下采样
+
+    # Downsampling
     voxel_points, same_points_list = voxel_filter(points, 0.6)
 
     after_labels = Load_data(voxel_points, center_points, labels)
@@ -584,8 +562,8 @@ def mesh_to_points_main(jaw, pcd_points, center_points, labels):
     new_pcd = o3d.geometry.PointCloud()
     new_pcd.points = o3d.utility.Vector3dVector(upsample_point)
     instances_labels = upsample_label.copy()
-    # '''
-    # o3d.io.write_point_cloud(os.path.join(save_path, 'upsample_' + name + '.pcd'), new_pcd, write_ascii=True)
+
+    # Reclassify the label of the upper and lower jaws
     for i in stqdm(range(0, upsample_label.shape[0])):
         if jaw == 'upper':
             if (upsample_label[i] >= 1) and (upsample_label[i] <= 8):
@@ -597,13 +575,14 @@ def mesh_to_points_main(jaw, pcd_points, center_points, labels):
                 upsample_label[i] = upsample_label[i] + 30
             elif (upsample_label[i] >= 9) and (upsample_label[i] <= 16):
                 upsample_label[i] = upsample_label[i] + 32
+
     remove_outlier_main(jaw, pcd_points, upsample_label, instances_labels)
 
 
-# 将原始点云数据转换为三角网格
+# Convert raw point cloud data to triangular mesh
 def mesh_grid(pcd_points):
     new_pcd,_ = voxel_filter(pcd_points, 0.6)
-    # pcd需要有法向量
+    # pcd needs to have a normal vector
 
     # estimate radius for rolling ball
     pcd_new = o3d.geometry.PointCloud()
@@ -615,12 +594,10 @@ def mesh_grid(pcd_points):
     mesh = o3d.geometry.TriangleMesh.create_from_point_cloud_ball_pivoting(
         pcd_new,
         o3d.utility.DoubleVector([radius, radius * 2]))
-    # o3d.io.write_triangle_mesh("./tooth date/test.ply", mesh)
 
     return mesh
 
-
-# 读取obj文件内容
+# Read the contents of obj file
 def read_obj(obj_path):
     jaw = None
     with open(obj_path) as file:
@@ -642,14 +619,12 @@ def read_obj(obj_path):
 
     points = np.array(points)
     faces = np.array(faces)
-
     if jaw is None:
         raise ValueError("Jaw type not found in OBJ file")
 
     return points, faces, jaw
 
-
-# obj文件转为pcd文件
+# Convert obj file to pcd file
 def obj2pcd(obj_path):
     if os.path.exists(obj_path):
         print('yes')
@@ -661,13 +636,14 @@ def obj2pcd(obj_path):
         pcd_list.append(new_line.split())
 
     pcd_points = np.array(pcd_list).astype(np.float64)
-    return pcd_points, jaw
 
+    return pcd_points, jaw
 
+# Main function for segment
 def segmentation_main(obj_path):
     upsampling_method = 'KNN'
 
-    model_path = 'Mesh_Segementation_MeshSegNet_17_classes_60samples_best.tar'
+    model_path = 'model.tar'
     num_classes = 17
     num_channels = 15
 
@@ -737,6 +713,7 @@ def segmentation_main(obj_path):
             nmeans = normals.mean(axis=0)
             nstds = normals.std(axis=0)
 
+            # normalization
             for i in range(3):
                 cells[:, i] = (cells[:, i] - means[i]) / stds[i]  # point 1
                 cells[:, i + 3] = (cells[:, i + 3] - means[i]) / stds[i]  # point 2
@@ -744,6 +721,7 @@ def segmentation_main(obj_path):
                 barycenters[:, i] = (barycenters[:, i] - mins[i]) / (maxs[i] - mins[i])
                 normals[:, i] = (normals[:, i] - nmeans[i]) / nstds[i]
 
+            # concatenate
             X = np.column_stack((cells, barycenters, normals))
 
             # computing A_S and A_L
@@ -794,6 +772,7 @@ def segmentation_main(obj_path):
                     if i_node < i_nei:
                         cos_theta = np.dot(normals[i_node, 0:3], normals[i_nei, 0:3]) / np.linalg.norm(
                             normals[i_node, 0:3]) / np.linalg.norm(normals[i_nei, 0:3])
+
                         if cos_theta >= 1.0:
                             cos_theta = 0.9999
                         theta = np.arccos(cos_theta)
@@ -806,6 +785,7 @@ def segmentation_main(obj_path):
                             edges = np.concatenate(
                                 (edges, np.array([i_node, i_nei, -beta * np.log10(theta / np.pi) * phi]).reshape(1, 3)),
                                 axis=0)
+
             edges = np.delete(edges, 0, 0)
             edges[:, 2] *= lambda_c * round_factor
             edges = edges.astype(np.int32)
@@ -913,9 +893,9 @@ class Segment(TeethApp):
                 # Create a pyvista plotter
                 plotter = pv.Plotter()
 
-                cmap = plt.cm.get_cmap('jet', 27)  # Using a colormap with sufficient distinct colors
+                cmap = plt.cm.get_cmap('jet', 27) # Using a colormap with sufficient distinct colors
 
-                colors = cmap(np.linspace(0, 1, 27))  # Generate colors
+                colors = cmap(np.linspace(0, 1, 27)) # Generate colors
 
                 # Convert colors to a format acceptable by PyVista
                 colormap = mcolors.ListedColormap(colors)
@@ -930,8 +910,6 @@ class Segment(TeethApp):
                 with st.expander("Ground Truth - scroll for zoom", expanded=False):
                     stpyvista(plotter)
 
-            
-
         elif inputs == "Upload Scan":
             file = st.file_uploader("Please upload an OBJ Object file", type=["OBJ"])
             st.markdown("Expected time per prediction: 7-10 min.")
@@ -939,7 +917,7 @@ class Segment(TeethApp):
                 # save the uploaded file to disk
                 with open("file.obj", "wb") as buffer:
                     shutil.copyfileobj(file, buffer)
-                # 复制数据
+
                 obj_path = "file.obj"
 
                 mesh = pv.read(obj_path)
@@ -957,9 +935,5 @@ class Segment(TeethApp):
                 if segment:
                     segmentation_main(obj_path)
 
-
-                
-                    
-
 if __name__ == "__main__":
     app = Segment()