Spaces:
No application file
No application file
File size: 3,810 Bytes
a2fddfe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | import os
import gradio as gr
import torch
import trimesh
from PIL import Image
from rembg import remove
from tsr.system import TSR
from tsr.utils import remove_background, resize_foreground
# Check for hardware acceleration on the server host
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Loading Generative 3D Transformer Model on: {device.upper()}...")
# Initialize the model weights globally once when the web server starts up
model = TSR.from_pretrained(
"stabilityai/TripoSR",
config_name="config.yaml",
weight_name="model.ckpt"
)
model.to(device)
model.eval()
def process_image_to_stl(input_image):
if input_image is None:
return None, "Error: No image uploaded."
try:
# Step 1: Strip image backgrounds natively on the server
print("Executing background removal layers...")
no_bg_image = remove(input_image)
# Step 2: Clear artifacts and scale to the neural network's bounding box
processed_img = remove_background(no_bg_image, "white")
processed_img = resize_foreground(processed_img, 0.85)
# Step 3: Run the Large Reconstruction Model to infer 3D spatial values
print("Processing 3D tensor field reconstruction...")
with torch.no_grad():
scene_codes = model([processed_img], device=device)
# Use Marching Cubes algorithm at a standard 256^3 resolution grid
meshes = model.extract_mesh(scene_codes, resolution=256)
ai_mesh = meshes
# Step 4: Extract the vertex mathematical arrays
vertices = ai_mesh.vertices.cpu().numpy()
faces = ai_mesh.faces.cpu().numpy()
# Step 5: Convert vertex points from local coordinates to a true 3D printable bed layout
vertices[:, [1, 2]] = vertices[:, [2, 1]] # Swap Y and Z axes
vertices[:, 1] *= -1 # Correct face-up inversion
# Create a solid geometry object
mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
mesh.process(validate=True) # Remove overlapping nodes
# Snap the absolute bottom boundary of the 3D mesh flat to Z=0 coordinate
z_min = mesh.bounds[0][2]
mesh.apply_translation([0, 0, -z_min])
# Step 6: Write out a local binary file on the server partition
output_filename = "generated_model.stl"
mesh.export(output_filename, file_type='stl')
status_msg = f"Success! Polygon Count: {len(mesh.faces)} | Solid Manifold: {mesh.is_watertight}"
return output_filename, status_msg
except Exception as e:
return None, f"An algorithmic pipeline error occurred: {str(e)}"
# Define the HTML/CSS user portal via Gradio framework
with gr.Blocks(theme=gr.themes.Soft()) as demo:
gr.Markdown("# Local-Engine AI Image-to-STL Converter")
gr.Markdown("Upload any image (objects, shapes, drawings) to synthesize a watertight, 3D-printable solid model without external API subscriptions.")
with gr.Row():
with gr.Column(scale=1):
input_img_slot = gr.Image(type="pil", label="Step 1: Upload Source Image")
submit_btn = gr.Button("Generate 3D STL Mesh", variant="primary")
with gr.Column(scale=1):
output_file_slot = gr.File(label="Step 2: Download Ready-to-Print STL File")
execution_log = gr.Textbox(label="System Pipeline Output Logs", interactive=False)
# Bind elements to backend trigger functions
submit_btn.click(
fn=process_image_to_stl,
inputs=[input_img_slot],
outputs=[output_file_slot, execution_log]
)
# Fire up the local webserver link on port 7860
if __name__ == "__main__":
demo.launch(server_name="0.0.0.0", server_port=7860)
|