voxel_to_SDF_to_STL.py

import numpy as np
from scipy.ndimage import morphology
from skimage import measure
from stl import mesh
import pyvista as pv
import streamlit as st
from stpyvista import stpyvista
import tempfile
import os


########################################################################################################################
def voxel_to_sdf(voxel, voxel_threshold, pad_size=5):
    """
    :param voxel: - a continuous or binary 3D array of values that range [0,1]
    :param voxel_threshold: - the value that is used to threshold the array and convert it into a binary array,
    a larger value will restrict the array to more concrete values but will produce a rougher shape.
    :param pad_size: - the size of padding to encompass the array to ensure the body is internal to the SDF
    :return: sdf - a signed distance array, where -'s are the internal components and +'s and the external components
    """
    # Convert a 3D voxel array to a 3D SDF
    # Where voxel is a continuous array [0,1]
    # Pad the voxel array with extra border elements
    padded_voxel = np.pad(voxel, pad_size, mode='constant', constant_values=0)

    # The values are switched, since SDF considers +'s to be outside the body and -'s internal to the body
    binary_voxel_data = np.where(np.array(padded_voxel) > voxel_threshold, 0, 1)

    # Calculate the Signed Distance Field (SDF)
    sdf = morphology.distance_transform_edt(binary_voxel_data) - 0.5

    # Plot Binary Data to see how it looks
    # voxel_plot = pv.wrap(np.array(voxel))
    # voxel_plot.plot(volume=True)

    # Plot the SDF to see how it looks before proceeding

    # HOW TO PLOT THE SDF IN STREAMLIT
    # # Create a PyVista grid from the SDF
    # grid = pv.wrap(sdf)
    # # # Plot the grid as a volume rendering
    # # grid.plot(volume=True)
    #
    # # Set up plotter
    # plotter = pv.Plotter(window_size=[600,600])
    # plotter.add_volume(grid, opacity='linear')
    # plotter.view_zy()
    #
    # # Pass the plotter (not the mesh) to stpyvista
    # stpyvista(plotter)

    return sdf


########################################################################################################################
def single_body_check(sdf, threshold_divisor):
    """
    :param sdf: - a signed distance array, where -'s are the internal components and +'s and the external components
    :return: boolean - indicates whether the SDF is a single body, which can be converted into an STL file
    """
    # This function is intended to be used to determine whether the sdf contains more than one body prior to converting to 3MF

    # Determine the SDF range (minimum and maximum values)
    sdf_min = np.min(sdf)
    sdf_max = np.max(sdf)

    # Set the surface level within the SDF range (or close to zero for the surface)
    surface_level = (sdf_min + sdf_max) / threshold_divisor  # For example, set it to the middle of the range

    # Threshold the SDF to identify regions close to the surface (threshold_value may vary based on your data)
    threshold_value = 0.05
    surface_regions = (np.abs(sdf) <= surface_level)

    # Perform connected component analysis on the thresholded SDF to label individual bodies
    connected_components = measure.label(surface_regions)

    # Count the number of connected components to determine the number of bodies
    num_labels = np.max(connected_components)

    # Count the number of connected components to determine the number of bodies
    if num_labels == 1:
        print("SDF contains single body, processing into STL")
        return True
    else:
        st.warning(f"ERROR: SDF was not a single body, contained {num_labels} bodies, therefore model requires further "
                   f"modification before proceeding. Recommended to start by changing threshold value for SDF (3), "
                   f"or restart with modifying the image threshold (2).")
        return False


########################################################################################################################
def sdf_to_stl(sdf, threshold_divsor):
    """
    :param sdf: - a signed distance array, where -'s are the internal components and +'s and the external components
    :param filename_stl: - the file name desired for the STL
    :param threshold_divsor: - this divisor determines the threshold for the SDF to be converted into the STL.
    A lower threshold will result in a lower resolution model with more spherical shapes,
    and a higher threshold will produce rugged shapes that better match the original voxel format
    :return: stl file - an STL file is saved based on the SDF provided, it is also displayed in the IDE
    """
    # Determine the SDF range (minimum and maximum values)
    sdf_min = np.min(sdf)
    sdf_max = np.max(sdf)

    # Set the surface level within the SDF range (or close to zero for the surface)
    surface_level = (sdf_min + abs(sdf_max)) / threshold_divsor  # For example, set it to the middle of the range

    # Generate a 3D surface mesh using marching cubes from the SDF
    vertices, triangles, _, _ = measure.marching_cubes(sdf, level=surface_level)

    '''
    # Display resulting triangular mesh using Matplotlib. This can also be done
    # with mayavi (see skimage.measure.marching_cubes docstring).
    fig = plt.figure(figsize=(10, 10))
    ax = fig.add_subplot(111, projection='3d')

    # Fancy indexing: `verts[faces]` to generate a collection of triangles
    mesh = Poly3DCollection(vertices[triangles])
    mesh.set_edgecolor('k')
    ax.add_collection3d(mesh)

    ax.set_xlim(0, 50)  # a = 6 (times two for 2nd ellipsoid)
    ax.set_ylim(0, 50)  # b = 10
    ax.set_zlim(0, 50)  # c = 16

    plt.tight_layout()
    plt.show()
    '''

    # Create an STL mesh object
    stl_mesh = mesh.Mesh(np.zeros(triangles.shape[0], dtype=mesh.Mesh.dtype))
    for i, triangle in enumerate(triangles):
        for j in range(3):
            stl_mesh.vectors[i][j] = vertices[triangle[j], :]


    # # Save the STL mesh to a file
    # stl_mesh.save(filename_stl)
    #
    # # Load the STL file using trimesh
    # stl = trimesh.load(filename_stl)
    #
    # # Show the mesh using trimesh's built-in viewer
    # stl.show()


    # # Set up plotter
    # plotter = pv.Plotter(window_size=[600, 600])
    # plotter.add_mesh(stl_mesh, color='teal')

    # Create a PyVista mesh from the stl_mesh
    pyvista_mesh = pv.PolyData(vertices)
    # pyvista_mesh = pv.PolyData(stl_mesh.vectors.copy().reshape(-1, 3))
    faces = np.hstack([np.full((len(triangles), 1), 3), triangles])
    pyvista_mesh.faces = faces

    # # Optionally, you can assign a color to the mesh
    # pyvista_mesh['color'] = [0.5, 0.5, 1.0]  # RGB color, e.g., teal

    # Create a plotter and add the PyVista mesh to it
    plotter = pv.Plotter(window_size=[600, 600])
    plotter.add_mesh(pyvista_mesh, style='wireframe', show_edges=True, lighting=True, edge_color='black')
    plotter.background_color = 'white'
    plotter.view_zy()  # if mesh_2D is on the xy plane.

    # # Show the plot
    # plotter.show()

    # Pass the plotter (not the mesh) to stpyvista
    stpyvista(plotter)

    # Create a temporary directory to store the STL file
    temp_dir = tempfile.mkdtemp()

    # Define the temporary STL file path
    temp_stl_file = os.path.join(temp_dir, 'interpolation.stl')

    # Save the STL data to the temporary file
    stl_mesh.save(temp_stl_file)

    return temp_stl_file


########################################################################################################################
# Creating a unit cell using two interpolation points
#
# def linear_interp_unit_cell(image_1, image_2, test_data, geometry_train, encoder, decoder, number_of_inputs, filename_stl):
#     image_size = np.shape(test_data[image_1])[0]
#     interpolation_length = int(image_size/2)
#
#     # 1. Select Latent Endpoints to Interpolate Between
#     sample_indices, [latent_point_1, latent_point_2] = encode_random_latent_endpoints(test_data, encoder,
#                                                                                       number_latent_points=2,
#                                                                                       sample_indices=[image_1, image_2],
#                                                                                       random_seed=0,
#                                                                                       num_inputs=number_of_inputs)
#
#     # 2. Create a list of all latent points in a linear interpolation between two endpoints
#     latent_interpolation_points = linear_interpolation_matrix(interpolation_length, latent_point_1, latent_point_2)
#
#     # 3. Plot the results of the linear interpolation
#     linear_predicted_geometry = linear_interpolation_reconstruction_plot(latent_interpolation_points,
#                                                                          geometry_train, decoder,
#                                                                          sample_indices[0], sample_indices[1])
#
#
#     # Mirror the array along the x-axis
#     mirrored_array = np.flip(linear_predicted_geometry, axis=0)
#
#     # Create the combined array with shape (50, 50, 50) by concatenating the mirrored array with the original array
#     single_axis_cell = np.concatenate((linear_predicted_geometry, mirrored_array), axis=0)
#
#     sdf = voxel_to_sdf(single_axis_cell, 0.1)
#
#     sdf_to_stl(sdf, filename_stl, 21.0)