Update app.py

app.py

@@ -1,28 +1,64 @@
 import gradio as gr
 import cadquery as cq
 from cadquery import exporters
+import trimesh
+import pyvista as pv
+from io import BytesIO
+import matplotlib.pyplot as plt
 import json
 import os
 
-# Generate APDL Script
+# 2D View Function
+def generate_2d_view(step_file):
+    try:
+        mesh = trimesh.load(step_file.name)
+        projection = mesh.section(plane_origin=[0, 0, 0], plane_normal=[0, 0, 1])
+        projection_img = projection.to_image() if projection else None
+
+        fig, ax = plt.subplots()
+        ax.imshow(projection_img, cmap="gray")
+        ax.axis("off")
+
+        buffer = BytesIO()
+        plt.savefig(buffer, format="png")
+        buffer.seek(0)
+        return buffer
+    except Exception as e:
+        return f"Error in generating 2D view: {e}"
+
+# 3D View Function
+def generate_3d_view(step_file):
+    try:
+        model = cq.importers.importStep(step_file.name)
+        stl_file = "temp_model.stl"
+        exporters.export(model, stl_file)
+
+        mesh = pv.read(stl_file)
+        plotter = pv.Plotter(off_screen=True)
+        plotter.add_mesh(mesh, color="lightblue")
+        screenshot = plotter.screenshot(asarray=True)
+
+        buffer = BytesIO()
+        plt.imshow(screenshot)
+        plt.axis("off")
+        plt.savefig(buffer, format="png")
+        buffer.seek(0)
+        return buffer
+    except Exception as e:
+        return f"Error in generating 3D view: {e}"
+
+# APDL Script Generation
 def generate_apdl_script(file, press_force, material_json):
     try:
-        # Parse material properties from JSON
         material = json.loads(material_json)
         elastic_modulus = material.get("elastic_modulus", 2e11)
         poisson = material.get("poisson", 0.3)
 
-        # Get the uploaded file path
         step_file_path = file.name
-
-        # Load STEP file with CadQuery
         model = cq.importers.importStep(step_file_path)
-
-        # Export as STL for meshing
         stl_file = step_file_path.replace(".stp", ".stl")
         exporters.export(model, stl_file)
 
-        # Generate APDL script
         apdl_script = f"""\
 /PREP7
 ! Importing Geometry
@@ -39,12 +75,10 @@ SOLVE
 *GET, stress, NODE, 0, S, MAX
 /EXIT, SAVE
 """
-        # Save APDL script to file
         apdl_file_path = step_file_path.replace(".stp", ".inp")
         with open(apdl_file_path, "w") as apdl_file:
             apdl_file.write(apdl_script)
 
-        # Return script content and file path for download
         return apdl_script, apdl_file_path
     except Exception as e:
         return f"Error: {str(e)}", None
@@ -52,43 +86,45 @@ SOLVE
 # Gradio App
 def main():
     with gr.Blocks() as app:
-        gr.Markdown("# APDL Script Generator for Press Tool Design")
+        gr.Markdown("# Press Tool Design and APDL Script Generator")
 
-        # Input Row
         with gr.Row():
             step_input = gr.File(label="Upload STEP File")
+
+        with gr.Row():
+            two_d_output = gr.Image(label="2D View")
+            three_d_output = gr.Image(label="3D View")
+            generate_views_btn = gr.Button("Generate 2D & 3D Views")
+
+        with gr.Row():
             force_input = gr.Number(label="Press Force (N)", value=1000)
             material_input = gr.Textbox(
                 label="Material Properties (JSON)",
                 value='{"elastic_modulus": 2e11, "poisson": 0.3}'
             )
 
-        # Output Row
         with gr.Row():
             script_output = gr.Textbox(label="Generated APDL Script", interactive=False)
             download_button = gr.File(label="Download Script")
+            generate_apdl_btn = gr.Button("Generate APDL Script")
 
-        # Button Row
-        with gr.Row():
-            generate_button = gr.Button("Generate APDL Script")
-            submit_button = gr.Button("Submit")
+        # Events for Visualization
+        generate_views_btn.click(
+            fn=[generate_2d_view, generate_3d_view],
+            inputs=step_input,
+            outputs=[two_d_output, three_d_output]
+        )
 
-        # Events
+        # Events for APDL Script
         def handle_download(apdl_script, apdl_file_path):
             if apdl_file_path and os.path.exists(apdl_file_path):
                 return apdl_file_path
             return None
 
-        generate_button.click(
+        generate_apdl_btn.click(
             fn=generate_apdl_script,
             inputs=[step_input, force_input, material_input],
-            outputs=[script_output, download_button],
-        )
-        
-        submit_button.click(
-            fn=lambda: "Submission successful!", 
-            inputs=None, 
-            outputs=gr.Textbox(label="Submission Status")
+            outputs=[script_output, download_button]
         )
 
     app.launch()