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

class TeethApp:
    def __init__(self):
        # 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)
        width, height = self.image.size
        scale = 12
        new_width, new_height = width / scale, height / scale
        self.image = self.image.resize((int(new_width), int(new_height)))

        # Streamlit side navigation bar
        st.sidebar.markdown("# AI ToothSeg")
        st.sidebar.markdown("Automatic teeth segmentation with Deep Learning")
        st.sidebar.markdown(" ")
        st.sidebar.image(self.image, use_column_width=False)
        st.markdown(
            """
                <style>
                .css-1bxukto {
                background-color: rgb(255, 255, 255) ;""",
            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** - scroll for zoom", expanded=False):
                    stpyvista(plotter)

# Configure Streamlit page
st.set_page_config(page_title="Teeth Segmentation", page_icon="🦷")

class Segment(TeethApp):
    def __init__(self):
        TeethApp.__init__(self)
        self.build_app()

    def build_app(self):

        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 - scroll for zoom", expanded=False):
                stpyvista(plotter)

            if segment:
                segmentation_main("ZOUIF2W4_upper.obj")

                # Load the JSON file
                with open('ZOUIF2W4_upper.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("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.")
            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 - scroll for zoom", expanded=False):
                    stpyvista(plotter)

                if segment:
                    segmentation_main(obj_path)


                
                    

if __name__ == "__main__":
    app = Segment()