Update app.py

app.py

@@ -4,7 +4,6 @@ import gradio as gr
 import cadquery as cq
 from cadquery import exporters
 import trimesh
-import pyvista as pv
 from io import BytesIO
 import json
 import atexit
@@ -22,17 +21,17 @@ def cleanup_temp_files():
 atexit.register(cleanup_temp_files)
 
 # Define 2D and 3D view generation functions with error handling
-def generate_2d_view(step_file):
+def generate_2d_view(file):
     try:
-        model = cq.importers.importStep(step_file.name)
+        model = process_cad_file(file)
         two_d_image = model.toSvg()
         return two_d_image
     except Exception as e:
         return f"Error generating 2D view: {str(e)}"
 
-def generate_3d_view(step_file):
+def generate_3d_view(file):
     try:
-        model = cq.importers.importStep(step_file.name)
+        model = process_cad_file(file)
         mesh = trimesh.Trimesh(vertices=model.val().vertices(), faces=model.val().faces())
         three_d_image = mesh.show()
         return three_d_image
@@ -40,14 +39,29 @@ def generate_3d_view(step_file):
         return f"Error generating 3D view: {str(e)}"
 
 # Wrapper function to generate views
-def generate_views(step_file):
+def generate_views(file):
     try:
-        two_d_view = generate_2d_view(step_file)
-        three_d_view = generate_3d_view(step_file)
+        two_d_view = generate_2d_view(file)
+        three_d_view = generate_3d_view(file)
         return two_d_view, three_d_view
     except Exception as e:
         return f"Error: {str(e)}", None
 
+# Function to process different CAD file formats
+def process_cad_file(file):
+    file_extension = os.path.splitext(file.name)[-1].lower()
+    if file_extension in [".stp", ".step"]:
+        model = cq.importers.importStep(file.name)
+    elif file_extension in [".iges", ".igs"]:
+        model = cq.importers.importStep(file.name)  # Same as STEP for cadquery
+    elif file_extension == ".sldprt":
+        raise NotImplementedError("SLDPRT support is not yet implemented.")
+    elif file_extension == ".dwg":
+        raise NotImplementedError("DWG support is not yet implemented.")
+    else:
+        raise ValueError(f"Unsupported file format: {file_extension}")
+    return model
+
 # APDL script generation
 def generate_apdl_script(file, press_force, material_json):
     try:
@@ -57,7 +71,7 @@ def generate_apdl_script(file, press_force, material_json):
 
         # Use temporary file path for processing
         step_file_path = file.name
-        model = cq.importers.importStep(step_file_path)
+        model = process_cad_file(file)
 
         stl_file = step_file_path.replace(".stp", ".stl")
         exporters.export(model, stl_file)
@@ -100,7 +114,10 @@ def main():
         gr.Markdown("# Press Tool Design and APDL Script Generator")
 
         with gr.Row():
-            step_input = gr.File(label="Upload STEP File")
+            step_input = gr.File(
+                label="Upload CAD File",
+                file_types=[".stp", ".step", ".iges", ".igs", ".sldprt", ".dwg"]
+            )
 
         with gr.Row():
             two_d_output = gr.Textbox(label="2D View/Error Message", interactive=False)