Update app.py

app.py

@@ -28,7 +28,6 @@ def generate_die(length, width, thickness):
         return f"Error generating die: {str(e)}"
 
 
-# Function to visualize die in 3D
 def visualize_die(length, width, thickness):
     try:
         output_dir = "outputs"
@@ -43,7 +42,7 @@ def visualize_die(length, width, thickness):
         stl_file = os.path.join(output_dir, "progressive_die.stl")
         cq.exporters.exportShape(die.val(), "STL", stl_file)
 
-        # Generate 3D visualization
+        # 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)
@@ -51,11 +50,14 @@ def visualize_die(length, width, thickness):
         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)}"
 
-
 # Function for Stress Analysis (including thermal stress and fatigue strength)
 def stress_analysis(force, die_width, die_height, material_strength, temperature_change=50, alpha=1e-5, elastic_modulus=200000, fatigue_strength=150):
     try: