Update ⓘ_Introduction.py

ⓘ_Introduction.py

@@ -1,5 +1,24 @@
-import streamlit as st
 from streamlit import session_state as session
+import shutil
+
+import os
+import numpy as np
+from sklearn import neighbors
+from scipy.spatial import distance_matrix
+from pygco import cut_from_graph
+import streamlit_ext as ste
+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 PIL import Image
 
@@ -8,7 +27,7 @@ class TeethApp:
         # Font
         with open("utils/style.css") as css:
             st.markdown(f"<style>{css.read()}</style>", unsafe_allow_html=True)
-    
+
         # Logo
         self.image_path = "utils/teeth-295404_1280.png"
         self.image = Image.open(self.image_path)
@@ -30,20 +49,881 @@ 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__()
+        self.conv1 = torch.nn.Conv1d(k, 64, 1)
+        self.conv2 = torch.nn.Conv1d(64, 128, 1)
+        self.conv3 = torch.nn.Conv1d(128, 512, 1)
+        self.fc1 = nn.Linear(512, 256)
+        self.fc2 = nn.Linear(256, 128)
+        self.fc3 = nn.Linear(128, k * k)
+        self.relu = nn.ReLU()
+
+        self.bn1 = nn.BatchNorm1d(64)
+        self.bn2 = nn.BatchNorm1d(128)
+        self.bn3 = nn.BatchNorm1d(512)
+        self.bn4 = nn.BatchNorm1d(256)
+        self.bn5 = nn.BatchNorm1d(128)
+
+        self.k = k
+
+    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, 512)
+
+        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.eye(self.k).flatten().astype(np.float32))).view(1, self.k * self.k).repeat(
+            batchsize, 1)
+        if x.is_cuda:
+            iden = iden.to(x.get_device())
+        x = x + iden
+        x = x.view(-1, self.k, self.k)
+        return x
+
+class MeshSegNet(nn.Module):
+    def __init__(self, num_classes=17, num_channels=15, with_dropout=True, dropout_p=0.5):
+        super(MeshSegNet, self).__init__()
+        self.num_classes = num_classes
+        self.num_channels = num_channels
+        self.with_dropout = with_dropout
+        self.dropout_p = dropout_p
+
+        # MLP-1 [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)
+        self.glm1_bn1_1 = nn.BatchNorm1d(32)
+        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)
+        self.mlp2_conv2 = torch.nn.Conv1d(64, 128, 1)
+        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)
+        self.glm2_conv1_3 = torch.nn.Conv1d(512, 128, 1)
+        self.glm2_bn1_1 = nn.BatchNorm1d(128)
+        self.glm2_bn1_2 = nn.BatchNorm1d(128)
+        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)
+        self.mlp3_bn1_1 = nn.BatchNorm1d(256)
+        self.mlp3_bn1_2 = nn.BatchNorm1d(256)
+        self.mlp3_conv3 = torch.nn.Conv1d(256, 128, 1)
+        self.mlp3_conv4 = torch.nn.Conv1d(128, 128, 1)
+        self.mlp3_bn2_1 = nn.BatchNorm1d(128)
+        self.mlp3_bn2_2 = nn.BatchNorm1d(128)
+        # output
+        self.output_conv = torch.nn.Conv1d(128, self.num_classes, 1)
+        if self.with_dropout:
+            self.dropout = nn.Dropout(p=self.dropout_p)
+
+    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)
+        x_ftm = x_ftm.transpose(2, 1)
+        x = F.relu(self.glm1_bn1_1(self.glm1_conv1_1(x_ftm)))
+        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)
+        sap_2 = torch.bmm(a_l, x_mlp2)
+        x_mlp2 = x_mlp2.transpose(2, 1)
+        sap_1 = sap_1.transpose(2, 1)
+        sap_2 = sap_2.transpose(2, 1)
+        x = F.relu(self.glm2_bn1_1(self.glm2_conv1_1(x_mlp2)))
+        glm_2_sap_1 = F.relu(self.glm2_bn1_2(self.glm2_conv1_2(sap_1)))
+        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)))
+        x = F.relu(self.mlp3_bn2_1(self.mlp3_conv3(x)))
+        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()
+        x = torch.nn.Softmax(dim=-1)(x.view(-1, self.num_classes))
+        x = x.view(batchsize, n_pts, self.num_classes)
+
+        return x
+
+def clone_runoob(li1):
+    li_copy = li1[:]
+    return li_copy
+
+# 对离群点重新进行分类
+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下标
+    ind_reverse.reverse()
+    for i in ind_reverse:
+        outlier_index.pop(i)
+
+    # 获取离群点
+    inlier_cloud = cloud.select_by_index(ind)
+    outlier_cloud = cloud.select_by_index(ind, invert=True)
+    outlier_points = np.array(outlier_cloud.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质心之间的距离
+            distance.append(dis)
+        min_index = distance.index(min(distance))  # 获取和离群点质心最近label的index
+        outlier_label = label_list[min_index]  # 获取离群点应该的label
+        index = outlier_index[i]
+        label_change[index] = outlier_label
+
+    return label_change
+
+# 利用knn算法消除离群点
+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种类对应点云的质心坐标
+
+    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.sort()
+    label_list = label_list[1:]
+
+    for i in label_list:
+        key = i
+        points_list = []
+        all_label_index = []
+        for j in range(len(labels)):
+            if labels[j] == i:
+                points_list.append(points[j].tolist())
+                all_label_index.append(j)  # 得到label为 i 的点对应的label的下标
+        same_label_points[key] = points_list
+        same_label_index[key] = all_label_index
+
+        tooth_mean = np.mean(points_list, axis=0)
+        mean_points.append(tooth_mean)
+        # print(mean_points)
+
+    for i in label_list:
+        points_array = same_label_points[i]
+        # 建立一个o3d的点云对象
+        pcd = o3d.geometry.PointCloud()
+        # 使用Vector3dVector方法转换
+        pcd.points = o3d.utility.Vector3dVector(points_array)
+
+        # 对label i 对应的点云进行统计离群值去除，找出离群点并显示
+        # 统计式离群点移除
+        cl, ind = pcd.remove_statistical_outlier(nb_neighbors=200, std_ratio=2.0)  # cl是选中的点，ind是选中点index
+        # 可视化
+        # display_inlier_outlier(pcd, ind)
+
+        # 对分出来的离群点重新分类
+        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
+
+
+# 消除离群点，保存最后的输出
+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')
+    # 原始点
+    points = pcd_points.copy()
+    label = remove_outlier(points, labels)
+
+    # 保存json文件
+    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 = {}
+
+
+# 体素下采样
+def voxel_filter(point_cloud, leaf_size):
+    same_points_list = {}
+    filtered_points = []
+    # step1 计算边界点
+    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 确定体素的尺寸
+    size_r = leaf_size
+
+    # step3 计算每个 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 为保存索引的列表
+    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值进行排序
+    h = np.array(h)
+    h_indice = np.argsort(h)  # 提取索引,返回h里面的元素按从小到大排序的  索引
+    h_sorted = h[h_indice]  # 升序
+    count = 0  # 用于维度的累计
+    step = 20
+    # 将h值相同的点放入到同一个grid中，并进行筛选
+    for i in range(1, len(h_sorted)):  # 0-19999个数据点
+        # if i == len(h_sorted)-1:
+        #     print("aaa")
+        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)  # 计算到重心的距离
+            _d.sort()
+            inx = [j for j in range(0, len(_d), step)]  # 获取指定间隔元素下标
+            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)  # 计算到重心的距离
+            _d.sort()
+            inx = [j for j in range(0, len(_d), step)]  # 获取指定间隔元素下标
+            for j in inx:
+                index = point_idx1[j]
+                filtered_points.append(point_cloud[index])
+
+            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)  # 计算到重心的距离
+            _d.sort()
+            inx = [j for j in range(0, len(_d), step)]  # 获取指定间隔元素下标
+            for j in inx:
+                index = point_idx2[j]
+                filtered_points.append(point_cloud[index])
+            count = i
+
+        else:
+            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)  # 计算到重心的距离
+            _d.sort()
+            inx = [j for j in range(0, len(_d), step)]  # 获取指定间隔元素下标
+            for j in inx:
+                index = point_idx[j]
+                filtered_points.append(point_cloud[index])
+            count = i
+
+    # 把点云格式改成array，并对外返回
+    # print(f'filtered_points[0]为{filtered_points[0]}')
+    filtered_points = np.array(filtered_points, dtype=np.float64)
+    return filtered_points,same_points_list
+
+
+# 体素上采样
+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三个维度的最值
+    x_min, y_min, z_min = np.amin(point_cloud, axis=0)
+    # step2 确定体素的尺寸
+    size_r = leaf_size
+    # step3 计算每个 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()
+    for i in range(len(filtered_points)):
+        hx = np.floor((filtered_points[i][0] - x_min) // size_r)
+        hy = np.floor((filtered_points[i][1] - y_min) // size_r)
+        hz = np.floor((filtered_points[i][2] - z_min) // size_r)
+        h.append(hx + hy * Dx + hz * Dx * Dy)
+
+    # step5 根据h值查询字典same_points_list
+    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}
+            classcount = {}
+            for i in range(len(label)):
+                vote = label[i]
+                classcount[vote] = classcount.get(vote, 0) + 1
+            # 对map的value排序
+            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)
+        elif h[i] != h[i - 1] and (i == len(h) - 1):
+            label1 = filter_labels[count:i]
+            key1 = h[i - 1]
+            label2 = filter_labels[i:]
+            key2 = h[i]
+            count = i
+
+            classcount = {}
+            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])
+                index = point_index[j]
+                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])
+                index = point_index[j]
+                upsample_point.append(point_cloud[index])
+                upsample_index.append(index)
+        else:
+            label = filter_labels[count:i]
+            key = h[i - 1]
+            count = i
+            classcount = {}
+            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原始顺序
+    upsample_index = np.array(upsample_index)
+    upsample_index_indice = np.argsort(upsample_index)  # 提取索引,返回h里面的元素按从小到大排序的  索引
+    upsample_index_sorted = upsample_index[upsample_index_indice]
+
+    upsample_point = np.array(upsample_point)
+    upsample_label = np.array(upsample_label)
+    # 恢复point和label的原始顺序
+    upsample_point_sorted = upsample_point[upsample_index_indice]
+    upsample_label_sorted = upsample_label[upsample_index_indice]
+
+    return upsample_point_sorted, upsample_label_sorted
+
+
+# 利用knn算法上采样
+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)
+    # 构建模型
+    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]
+
+
+# 加载点进行knn上采样
+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)
+
+# 将三角网格数据上采样回原始点云数据
+def mesh_to_points_main(jaw, pcd_points, center_points, labels):
+    points = pcd_points.copy()
+    # 下采样
+    voxel_points, same_points_list = voxel_filter(points, 0.6)
+
+    after_labels = Load_data(voxel_points, center_points, labels)
+
+    upsample_point, upsample_label = voxel_upsample(same_points_list, points, voxel_points, after_labels, 0.6)
+
+    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)
+    for i in stqdm(range(0, upsample_label.shape[0])):
+        if jaw == 'upper':
+            if (upsample_label[i] >= 1) and (upsample_label[i] <= 8):
+                upsample_label[i] = upsample_label[i] + 10
+            elif (upsample_label[i] >= 9) and (upsample_label[i] <= 16):
+                upsample_label[i] = upsample_label[i] + 12
+        else:
+            if (upsample_label[i] >= 1) and (upsample_label[i] <= 8):
+                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)
+
+
+# 将原始点云数据转换为三角网格
+def mesh_grid(pcd_points):
+    new_pcd,_ = voxel_filter(pcd_points, 0.6)
+    # pcd需要有法向量
+
+    # estimate radius for rolling ball
+    pcd_new = o3d.geometry.PointCloud()
+    pcd_new.points = o3d.utility.Vector3dVector(new_pcd)
+    pcd_new.estimate_normals()
+    distances = pcd_new.compute_nearest_neighbor_distance()
+    avg_dist = np.mean(distances)
+    radius = 6 * avg_dist
+    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文件内容
+def read_obj(obj_path):
+    jaw = None
+    with open(obj_path) as file:
+        points = []
+        faces = []
+        while 1:
+            line = file.readline()
+            if not line:
+                break
+            strs = line.split(" ")
+            if strs[0] == "v":
+                points.append((float(strs[1]), float(strs[2]), float(strs[3])))
+            elif strs[0] == "f":
+                faces.append((int(strs[1]), int(strs[2]), int(strs[3])))
+            elif strs[1][0:5] == 'lower':
+                jaw = 'lower'
+            elif strs[1][0:5] == 'upper':
+                jaw = 'upper'
+
+    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文件
+def obj2pcd(obj_path):
+    if os.path.exists(obj_path):
+        print('yes')
+    points, _, jaw = read_obj(obj_path)
+    pcd_list = []
+    num_points = np.shape(points)[0]
+    for i in range(num_points):
+        new_line = str(points[i, 0]) + ' ' + str(points[i, 1]) + ' ' + str(points[i, 2])
+        pcd_list.append(new_line.split())
+
+    pcd_points = np.array(pcd_list).astype(np.float64)
+    return pcd_points, jaw
+
+
+def segmentation_main(obj_path):
+    upsampling_method = 'KNN'
+
+    model_path = 'Mesh_Segementation_MeshSegNet_17_classes_60samples_best.tar'
+    num_classes = 17
+    num_channels = 15
+
+    # set model
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = MeshSegNet(num_classes=num_classes, num_channels=num_channels).to(device, dtype=torch.float)
+
+    # load trained model
+    # checkpoint = torch.load(os.path.join(model_path, model_name), map_location='cpu')
+    checkpoint = torch.load(model_path, map_location='cpu')
+    model.load_state_dict(checkpoint['model_state_dict'])
+    del checkpoint
+    model = model.to(device, dtype=torch.float)
+
+    # cudnn
+    torch.backends.cudnn.benchmark = True
+    torch.backends.cudnn.enabled = True
+
+    # Predicting
+    model.eval()
+    with torch.no_grad():
+        pcd_points, jaw = obj2pcd(obj_path)
+        mesh = mesh_grid(pcd_points)
+
+        # move mesh to origin
+        with st.spinner("Patience please, AI at work. Grab a coffee while you wait ☕."):
+            vertices_points = np.asarray(mesh.vertices)
+            triangles_points = np.asarray(mesh.triangles)
+            N = triangles_points.shape[0]
+            cells = np.zeros((triangles_points.shape[0], 9))
+            cells = vertices_points[triangles_points].reshape(triangles_points.shape[0], 9)
+
+            mean_cell_centers = mesh.get_center()
+            cells[:, 0:3] -= mean_cell_centers[0:3]
+            cells[:, 3:6] -= mean_cell_centers[0:3]
+            cells[:, 6:9] -= mean_cell_centers[0:3]
+
+            v1 = np.zeros([triangles_points.shape[0], 3], dtype='float32')
+            v2 = np.zeros([triangles_points.shape[0], 3], dtype='float32')
+            v1[:, 0] = cells[:, 0] - cells[:, 3]
+            v1[:, 1] = cells[:, 1] - cells[:, 4]
+            v1[:, 2] = cells[:, 2] - cells[:, 5]
+            v2[:, 0] = cells[:, 3] - cells[:, 6]
+            v2[:, 1] = cells[:, 4] - cells[:, 7]
+            v2[:, 2] = cells[:, 5] - cells[:, 8]
+            mesh_normals = np.cross(v1, v2)
+            mesh_normal_length = np.linalg.norm(mesh_normals, axis=1)
+            mesh_normals[:, 0] /= mesh_normal_length[:]
+            mesh_normals[:, 1] /= mesh_normal_length[:]
+            mesh_normals[:, 2] /= mesh_normal_length[:]
+
+            # prepare input
+            points = vertices_points.copy()
+            points[:, 0:3] -= mean_cell_centers[0:3]
+            normals = np.nan_to_num(mesh_normals).copy()
+            barycenters = np.zeros((triangles_points.shape[0], 3))
+            s = np.sum(vertices_points[triangles_points], 1)
+            barycenters = 1 / 3 * s
+            center_points = barycenters.copy()
+            barycenters -= mean_cell_centers[0:3]
+
+            # normalized data
+            maxs = points.max(axis=0)
+            mins = points.min(axis=0)
+            means = points.mean(axis=0)
+            stds = points.std(axis=0)
+            nmeans = normals.mean(axis=0)
+            nstds = normals.std(axis=0)
+
+            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
+                cells[:, i + 6] = (cells[:, i + 6] - means[i]) / stds[i]  # point 3
+                barycenters[:, i] = (barycenters[:, i] - mins[i]) / (maxs[i] - mins[i])
+                normals[:, i] = (normals[:, i] - nmeans[i]) / nstds[i]
+
+            X = np.column_stack((cells, barycenters, normals))
+
+            # computing A_S and A_L
+            A_S = np.zeros([X.shape[0], X.shape[0]], dtype='float32')
+            A_L = np.zeros([X.shape[0], X.shape[0]], dtype='float32')
+            D = distance_matrix(X[:, 9:12], X[:, 9:12])
+            A_S[D < 0.1] = 1.0
+            A_S = A_S / np.dot(np.sum(A_S, axis=1, keepdims=True), np.ones((1, X.shape[0])))
+
+            A_L[D < 0.2] = 1.0
+            A_L = A_L / np.dot(np.sum(A_L, axis=1, keepdims=True), np.ones((1, X.shape[0])))
+
+            # numpy -> torch.tensor
+            X = X.transpose(1, 0)
+            X = X.reshape([1, X.shape[0], X.shape[1]])
+            X = torch.from_numpy(X).to(device, dtype=torch.float)
+            A_S = A_S.reshape([1, A_S.shape[0], A_S.shape[1]])
+            A_L = A_L.reshape([1, A_L.shape[0], A_L.shape[1]])
+            A_S = torch.from_numpy(A_S).to(device, dtype=torch.float)
+            A_L = torch.from_numpy(A_L).to(device, dtype=torch.float)
+
+            tensor_prob_output = model(X, A_S, A_L).to(device, dtype=torch.float)
+            patch_prob_output = tensor_prob_output.cpu().numpy()
+
+        # refinement
+        with st.spinner("Refining..."):
+            round_factor = 100
+            patch_prob_output[patch_prob_output < 1.0e-6] = 1.0e-6
+
+            # unaries
+            unaries = -round_factor * np.log10(patch_prob_output)
+            unaries = unaries.astype(np.int32)
+            unaries = unaries.reshape(-1, num_classes)
+
+            # parawisex
+            pairwise = (1 - np.eye(num_classes, dtype=np.int32))
+
+            cells = cells.copy()
+
+            cell_ids = np.asarray(triangles_points)
+            lambda_c = 20
+            edges = np.empty([1, 3], order='C')
+            for i_node in stqdm(range(cells.shape[0])):
+                # Find neighbors
+                nei = np.sum(np.isin(cell_ids, cell_ids[i_node, :]), axis=1)
+                nei_id = np.where(nei == 2)
+                for i_nei in nei_id[0][:]:
+                    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)
+                        phi = np.linalg.norm(barycenters[i_node, :] - barycenters[i_nei, :])
+                        if theta > np.pi / 2.0:
+                            edges = np.concatenate(
+                                (edges, np.array([i_node, i_nei, -np.log10(theta / np.pi) * phi]).reshape(1, 3)), axis=0)
+                        else:
+                            beta = 1 + np.linalg.norm(np.dot(normals[i_node, 0:3], normals[i_nei, 0:3]))
+                            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)
+
+            refine_labels = cut_from_graph(edges, unaries, pairwise)
+            refine_labels = refine_labels.reshape([-1, 1])
+
+            predicted_labels_3 = refine_labels.reshape(refine_labels.shape[0])
+            mesh_to_points_main(jaw, pcd_points, center_points, predicted_labels_3)
+
+            import pyvista as pv
+            with st.spinner("Rendering..."):
+                # Load the .obj file
+                mesh = pv.read('file.obj')
+
+                # Load the JSON file
+                with open('dental-labels4.json', 'r') as file:
+                    labels_data = json.load(file)
+
+                # Assuming labels_data['labels'] is a list of labels
+                labels = labels_data['labels']
+
+                # Make sure the number of labels matches the number of vertices or faces
+                assert len(labels) == mesh.n_points or len(labels) == mesh.n_cells
+
+                # If labels correspond to vertices
+                if len(labels) == mesh.n_points:
+                    mesh.point_data['Labels'] = labels
+                # If labels correspond to faces
+                elif len(labels) == mesh.n_cells:
+                    mesh.cell_data['Labels'] = labels
+                
+                # Create a pyvista plotter
+                plotter = pv.Plotter()
+
+                cmap = plt.cm.get_cmap('jet', 27)  # Using a colormap with sufficient distinct colors
+
+                colors = cmap(np.linspace(0, 1, 27))  # Generate colors
+
+                # Convert colors to a format acceptable by PyVista
+                colormap = mcolors.ListedColormap(colors)
+
+                # Add the mesh to the plotter with labels as a scalar field
+                #plotter.add_mesh(mesh, scalars='Labels', show_scalar_bar=True, cmap='jet')
+                plotter.add_mesh(mesh, scalars='Labels', show_scalar_bar=True, cmap=colormap, clim=[0, 27])
+
+                # Show the plot
+                #plotter.show()
+                ## Send to streamlit
+                with st.expander("**View Segmentation Result** - ", expanded=False):
+                    stpyvista(plotter)
+
 # Configure Streamlit page
-st.set_page_config(page_title="Teeth Segmentation", page_icon="ⓘ")
+st.set_page_config(page_title="Teeth Segmentation", page_icon="🦷")
 
-class Intro(TeethApp):
+class Segment(TeethApp):
     def __init__(self):
         TeethApp.__init__(self)
         self.build_app()
 
     def build_app(self):
-        st.title("AI-assited Tooth Segmentation")
-        st.markdown("This app automatically segments intra-oral scans of teeth using machine learning.")
-        st.markdown("Head to the 'Segment' tab to try it out!")
-        st.markdown("**Example:**")
-        st.image("illu.png")
+
+        st.title("Segment Intra-oral Scans")
+        st.markdown("Identify and segment teeth. Segmentation is performed using MeshSegNet, a deep learning model trained on both upper and lower jaws.")
+
+        inputs = st.radio(
+            "Select scan for segmentation:",
+            ("Upload Scan", "Example Scan"),
+        )
+        import pyvista as pv
+        if inputs == "Example Scan":
+            st.markdown("Expected time per prediction: 7-10 min.")
+            mesh = pv.read("ZOUIF2W4_upper.obj")
+            plotter = pv.Plotter()
+
+            # Add the mesh to the plotter
+            plotter.add_mesh(mesh, color='white', show_edges=False)
+            segment = st.button(
+                "✔️ Submit",
+                help="Submit 3D scan for segmentation",
+            )
+            with st.expander("View Scan", expanded=False):
+                stpyvista(plotter)
+
+            if segment:
+                segmentation_main("ZOUIF2W4_upper.obj")
+
+            
+
+        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.")
+            if file is not None:
+                # 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)
+                plotter = pv.Plotter()
+
+                # Add the mesh to the plotter
+                plotter.add_mesh(mesh, color='white', show_edges=False)
+                segment = st.button(
+                "✔️ Submit",
+                help="Submit 3D scan for segmentation",
+            )
+                with st.expander("View Scan", expanded=False):
+                    stpyvista(plotter)
+
+                if segment:
+                    segmentation_main(obj_path)
+
+
+                
+                    
 
 if __name__ == "__main__":
-    app = Intro()
+    app = Segment()