Update pages/01_🦷 Segment.py

pages/01_🦷 Segment.py

@@ -9,6 +9,7 @@ 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
@@ -662,6 +663,201 @@ def obj2pcd(obj_path):
     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="🦷")
 
@@ -681,218 +877,45 @@ class Segment(TeethApp):
         )
         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='black', show_edges=True)
-            visualize = st.button("Segment")
-            if visualize:
+            plotter.add_mesh(mesh, color='white', show_edges=True)
+            segment = st.button("Apply 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"
-                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 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
+                mesh = pv.read(obj_path)
+                plotter = pv.Plotter()
 
-                    # Convert colors to a format acceptable by PyVista
-                    colormap = mcolors.ListedColormap(colors)
+                # Add the mesh to the plotter
+                plotter.add_mesh(mesh, color='white', show_edges=True)
+                segment = st.button("Apply Segmentation")
+                with st.expander("View Scan", expanded=False):
+                    stpyvista(plotter)
 
-                    # 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])
+                if segment:
+                    segmentation_main(obj_path)
 
-                    # Show the plot
-                    #plotter.show()
-                    ## Send to streamlit
-                    stpyvista(plotter)
+                
+                    
 
 if __name__ == "__main__":
     app = Segment()