Update app.py

app.py

@@ -3,21 +3,19 @@ import re
 import math
 import streamlit as st
 import cadquery as cq
+from cadquery import exporters
 from PyPDF2 import PdfReader
 from groq import Groq
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d.art3d import Poly3DCollection
 
 # Initialize Groq API client
 client = Groq(api_key="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")
 
 # Function to parse profile description
 def parse_profile(description):
-    """
-    Parse text input for dimensions and angles.
-    Example input: "profile with 50mm then 90 degree bend..."
-    """
     pattern = r"(\d+)\s*mm|(\d+)\s*degree"
     matches = re.findall(pattern, description)
-
     profile = []
     for match in matches:
         if match[0]:
@@ -28,13 +26,9 @@ def parse_profile(description):
 
 # Function to generate CAD model using CadQuery
 def generate_cad_model(profile, is_3d=False):
-    """
-    Generate CAD model using CadQuery based on parsed profile data.
-    """
     model = cq.Workplane("XY")
     x, y = 0, 0
     angle = 0
-
     for segment in profile:
         if segment["type"] == "length":
             x += segment["value"] * math.cos(math.radians(angle))
@@ -42,27 +36,54 @@ def generate_cad_model(profile, is_3d=False):
             model = model.lineTo(x, y)
         elif segment["type"] == "angle":
             angle += segment["value"]
-
     if is_3d:
-        model = model.extrude(10)  # Example extrusion height for 3D models
+        model = model.extrude(10)
+    return model
 
+# Function to create gear profile
+def generate_gear_model(teeth, module, thickness):
+    pitch_diameter = teeth * module
+    model = cq.Workplane("XY").gear(teeth, pitch_diameter / teeth, thickness=thickness)
     return model
 
+# Function to preview 2D profile
+def preview_2d(profile):
+    x, y = [0], [0]
+    angle = 0
+    for segment in profile:
+        if segment["type"] == "length":
+            x.append(x[-1] + segment["value"] * math.cos(math.radians(angle)))
+            y.append(y[-1] + segment["value"] * math.sin(math.radians(angle)))
+        elif segment["type"] == "angle":
+            angle += segment["value"]
+    plt.figure(figsize=(6, 6))
+    plt.plot(x, y, marker="o")
+    plt.title("2D Profile Preview")
+    plt.xlabel("X (mm)")
+    plt.ylabel("Y (mm)")
+    plt.grid(True)
+    st.pyplot(plt)
+
+# Function to preview 3D model
+def preview_3d(model):
+    shape = model.val()
+    faces = shape.faces().toWires()
+    verts = [tuple(p.toTuple()) for p in shape.Vertices()]
+    fig = plt.figure()
+    ax = fig.add_subplot(111, projection="3d")
+    ax.add_collection3d(Poly3DCollection(faces))
+    ax.scatter(*zip(*verts))
+    plt.show()
+    st.pyplot(fig)
+
 # Function to export CAD model
 def export_model(model, output_format):
-    """
-    Export CAD model to the desired format.
-    """
     output_file = f"custom_profile.{output_format}"
-    shape = model.val()  # Convert Workplane to Shape
-    if output_format == "step":
-        cq.exporters.export(shape, output_file)
-    elif output_format == "stl":
-        cq.exporters.export(shape, output_file)
-    elif output_format == "svg":
-        cq.exporters.export(shape, output_file)
-    else:
-        raise ValueError(f"Unsupported format: {output_format}")
+    shape = model.val()
+    if output_format in ["step", "iges", "stl"]:
+        exporters.export(shape, output_file)
+    elif output_format == "prt":
+        raise NotImplementedError("CreoParametric .prt export requires additional tools.")
     return output_file
 
 # Streamlit UI
@@ -89,10 +110,21 @@ if st.button("Generate CAD Model") and query:
     try:
         profile = parse_profile(query)
         is_3d = st.checkbox("Generate 3D Model", value=False)
-        model = generate_cad_model(profile, is_3d=is_3d)
+        if "gear" in query.lower():
+            teeth = st.number_input("Enter number of teeth", min_value=10, value=20)
+            module = st.number_input("Enter module size", min_value=1, value=2)
+            thickness = st.number_input("Enter gear thickness (mm)", min_value=1, value=5)
+            model = generate_gear_model(teeth, module, thickness)
+        else:
+            model = generate_cad_model(profile, is_3d=is_3d)
 
         st.success("CAD Model Generated!")
-        output_format = st.selectbox("Select Download Format", ["step", "stl", "svg"])
+        if is_3d:
+            preview_3d(model)
+        else:
+            preview_2d(profile)
+
+        output_format = st.selectbox("Select Download Format", ["step", "iges", "stl"])
         output_file = export_model(model, output_format)
         with open(output_file, "rb") as f:
             st.download_button("Download CAD File", f.read(), file_name=output_file)