Spaces:
Running
Running
| import gradio as gr | |
| import trimesh | |
| import numpy as np | |
| import tempfile | |
| import os | |
| import traceback | |
| def fit_mesh(mesh_file): | |
| """ | |
| SuperFit Algorithm: | |
| Replaces each isolated part of the input mesh with a best-fit geometric primitive | |
| (an oriented bounding box). | |
| """ | |
| if mesh_file is None: | |
| return None | |
| try: | |
| # Load the mesh/scene | |
| scene = trimesh.load(mesh_file.name, force='scene') | |
| primitives = [] | |
| def process_geometry(geom): | |
| # Generate the Oriented Bounding Box (OBB) for the geometry | |
| obb = geom.bounding_box_oriented | |
| # Transfer the color if available | |
| try: | |
| color = None | |
| if hasattr(geom.visual, 'face_colors') and len(geom.visual.face_colors) > 0: | |
| color = geom.visual.face_colors[0] | |
| elif hasattr(geom.visual, 'main_color'): | |
| color = geom.visual.main_color | |
| if color is not None: | |
| obb.visual.face_colors = color | |
| except Exception: | |
| pass | |
| return obb | |
| # Process each geometry in the scene | |
| if isinstance(scene, trimesh.Scene): | |
| for name, geom in scene.geometry.items(): | |
| if isinstance(geom, trimesh.Trimesh): | |
| # Sometimes a single geometry node contains multiple disconnected parts | |
| components = geom.split(only_watertight=False) | |
| for comp in components: | |
| if len(comp.vertices) > 4: # Ignore tiny dust/noise | |
| primitives.append(process_geometry(comp)) | |
| elif isinstance(scene, trimesh.Trimesh): | |
| components = scene.split(only_watertight=False) | |
| for comp in components: | |
| if len(comp.vertices) > 4: | |
| primitives.append(process_geometry(comp)) | |
| # If we failed to extract any primitives, return the original | |
| if not primitives: | |
| return mesh_file.name | |
| # Create a new scene with the bounding box primitives | |
| out_scene = trimesh.Scene(primitives) | |
| # Export to a temporary GLB file | |
| out_path = tempfile.mktemp(suffix=".glb") | |
| out_scene.export(out_path) | |
| return out_path | |
| except Exception as e: | |
| print("Error in SuperFit algorithm:") | |
| traceback.print_exc() | |
| # Fallback to the original mesh if the algorithm fails | |
| return mesh_file.name | |
| def segment_mesh(mesh_file, num_clusters=5): | |
| """ | |
| Lightweight geometric segmentation algorithm to replace Stirnix/PartField. | |
| Uses K-Means clustering on face centroids to split the mesh into distinct parts. | |
| """ | |
| if mesh_file is None: | |
| return None | |
| try: | |
| from sklearn.cluster import KMeans | |
| # Force load as a single mesh | |
| mesh = trimesh.load(mesh_file.name, force='mesh') | |
| # Calculate face centroids | |
| centroids = mesh.triangles.mean(axis=1) | |
| # Cluster faces | |
| num_clusters = int(num_clusters) | |
| clustering = KMeans(n_clusters=num_clusters, random_state=42, n_init=10).fit(centroids) | |
| labels = clustering.labels_ | |
| # Split mesh into submeshes | |
| submeshes = [] | |
| for i in range(num_clusters): | |
| face_mask = (labels == i) | |
| sub_faces = mesh.faces[face_mask] | |
| if len(sub_faces) > 0: | |
| submesh = trimesh.Trimesh(vertices=mesh.vertices, faces=sub_faces, process=True) | |
| submeshes.append(submesh) | |
| out_scene = trimesh.Scene(submeshes) | |
| out_path = tempfile.mktemp(suffix=".glb") | |
| out_scene.export(out_path) | |
| return out_path | |
| except Exception as e: | |
| print("Error in segmentation algorithm:") | |
| traceback.print_exc() | |
| return mesh_file.name | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# SuperFit API (All-in-One)") | |
| gr.Markdown("Replaces complex meshes with fitted primitive shapes AND provides geometric segmentation.") | |
| with gr.Tab("SuperFit (Fit Primitives)"): | |
| with gr.Row(): | |
| mesh_input = gr.File(label="Input Mesh (GLB)") | |
| mesh_output = gr.File(label="Fitted Primitives (GLB)") | |
| btn = gr.Button("Fit Primitives") | |
| btn.click(fn=fit_mesh, inputs=[mesh_input], outputs=[mesh_output], api_name="fit_mesh") | |
| with gr.Tab("Segmentation (Stirnix Replacement)"): | |
| with gr.Row(): | |
| seg_input = gr.File(label="Input Mesh (GLB)") | |
| num_clusters = gr.Slider(minimum=1, maximum=20, value=5, step=1, label="Number of Parts") | |
| seg_output = gr.File(label="Segmented Mesh (GLB)") | |
| seg_btn = gr.Button("Segment Mesh") | |
| seg_btn.click(fn=segment_mesh, inputs=[seg_input, num_clusters], outputs=[seg_output], api_name="segment") | |
| if __name__ == "__main__": | |
| demo.launch() | |