Create Main

Main

@@ -0,0 +1,93 @@
+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)