Update design.py

design.py

@@ -2,41 +2,33 @@ import cadquery as cq
 import os
 import pyvista as pv
 
-def generate_die(length, width, thickness):
-    try:
-        output_dir = "outputs"
-        os.makedirs(output_dir, exist_ok=True)
-
-        plate = cq.Workplane("XY").box(length, width, thickness)
-        punch = cq.Workplane("XY").rect(10, 10).extrude(5).translate((length / 4, width / 4, thickness / 2))
-        die = plate.cut(punch)
-
-        filename = os.path.join(output_dir, "progressive_die.step")
-        cq.exporters.export(die, filename)
-
-        return filename
-    except Exception as e:
-        return f"Error generating die: {str(e)}"
-
 def visualize_die(length, width, thickness):
     try:
         output_dir = "outputs"
         os.makedirs(output_dir, exist_ok=True)
 
+        # Create die geometry
         plate = cq.Workplane("XY").box(length, width, thickness)
         punch = cq.Workplane("XY").rect(10, 10).extrude(5).translate((length / 4, width / 4, thickness / 2))
         die = plate.cut(punch)
 
+        # Export STL file for visualization
         stl_file = os.path.join(output_dir, "progressive_die.stl")
-        cq.exporters.exportShape(die.val(), "STL", stl_file)
+        cq.exporters.export(die, stl_file)  # Corrected line
 
-        pv.global_theme.off_screen = True
+        # Generate 3D visualization screenshot
+        pv.global_theme.off_screen = True  # Ensure off-screen rendering
         mesh = pv.read(stl_file)
         plotter = pv.Plotter(off_screen=True)
         plotter.add_mesh(mesh, color="blue")
         screenshot = os.path.join(output_dir, "progressive_die_visualization.png")
         plotter.screenshot(screenshot)
 
+        # Ensure the screenshot file exists
+        if not os.path.exists(screenshot):
+            raise FileNotFoundError("Screenshot file was not generated.")
+
         return screenshot
     except Exception as e:
         return f"Error visualizing die: {str(e)}"
+