Update pages/01_🦷 Segment.py

pages/01_🦷 Segment.py

@@ -1,27 +1,25 @@
+from streamlit import session_state as session
 import shutil
-import os
 
+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
-
-import pyvista as pv
 from stpyvista import stpyvista
-
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 from torch.autograd import Variable
-
+import torch.nn.functional as F
 import streamlit as st
+import pyvista as pv
 
-from streamlit import session_state as session
 from PIL import Image
 
 class TeethApp:
@@ -896,6 +894,43 @@ class Segment(TeethApp):
             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":