Update app.py

app.py

@@ -1,80 +1,54 @@
 import streamlit as st
+import cadquery as cq
 from datasets import load_dataset
 import json
-from cadquery import Workplane, exporters
-import tempfile
-import os
 
-# Load CADBench dataset
-dataset = load_dataset("FreedomIntelligence/CADBench", split="train")
+# Load dataset
+dataset = load_dataset("FreedomIntelligence/CADBench")
 
-# Initialize the app
-st.title("Quick CAD Model Generator")
-st.sidebar.title("Select Case")
-case = st.sidebar.radio("Choose a functionality:", ["DFM Analysis", "Predefined Template Selection"])
+# Helper to generate Cargo Crane CAD
+def generate_cargo_crane(inputs):
+    # Parse inputs
+    boom_length = inputs.get("Boom Length (in mm)", 3000)
+    lifting_capacity = inputs.get("Lifting Capacity (in tons)", 10)
+    base_height = inputs.get("Base Height (in mm)", 500)
+    material_type = inputs.get("Material Type", "Steel")
+
+    # Create base
+    base = cq.Workplane("XY").box(200, 200, base_height)
 
-def analyze_dfm(file, file_format):
-    """Perform Design for Manufacturing (DFM) analysis."""
-    try:
-        st.write(f"Analyzing {file_format} file...")
-        # Placeholder: Replace with real DFM logic
-        return "The design is suitable for CNC manufacturing!"
-    except Exception as e:
-        return f"Error analyzing file: {e}"
+    # Create boom
+    boom = cq.Workplane("XY").box(50, boom_length, 50).translate((0, base_height, 0))
 
-def generate_cad_model(responses):
-    """Generate CAD model based on user inputs."""
-    try:
-        # Placeholder: CAD generation logic
-        result = Workplane("XY").box(1, 1, 1)  # Example model
-        return result
-    except Exception as e:
-        raise RuntimeError(f"Error generating CAD model: {e}")
+    # Combine parts
+    crane = base.union(boom)
 
-def export_cad(model, file_format):
-    """Export CAD model to selected format."""
-    with tempfile.NamedTemporaryFile(delete=False, suffix=f".{file_format}") as temp_file:
-        if file_format == "stl":
-            exporters.exportShape(model, exporters.ExportTypes.STL, temp_file.name)
-        elif file_format == "step":
-            exporters.exportShape(model, exporters.ExportTypes.STEP, temp_file.name)
-        elif file_format == "dwg":
-            # Placeholder: Add DWG export logic
-            pass
-        return temp_file.name
+    return crane
 
-if case == "DFM Analysis":
-    st.subheader("Case 1: DFM Analysis")
-    cad_file = st.file_uploader("Upload a CAD file (.stl, .step, .dwg)", type=["stl", "step", "dwg"])
-    if cad_file:
-        file_format = os.path.splitext(cad_file.name)[1][1:]
-        analysis_result = analyze_dfm(cad_file, file_format)
-        st.success(analysis_result)
+# Streamlit Interface
+st.title("Quick CAD Model Generator")
+case = st.radio("Select Case", ["DFM Analysis", "Predefined Template"])
 
-elif case == "Predefined Template Selection":
-    st.subheader("Case 2: Predefined Template Selection")
-    template_names = dataset["name"]
-    selected_template = st.selectbox("Select a template", template_names)
+if case == "Predefined Template":
+    st.subheader("Predefined Template Selection")
+    template_names = dataset["train"]["name"]
+    selected_template = st.selectbox("Select a Template", template_names)
 
+    # Retrieve criteria
     if selected_template:
-        st.write(f"Selected Template: {selected_template}")
-        template_data = dataset.filter(lambda x: x["name"] == selected_template).to_pandas()
-        criteria = template_data.iloc[0]["criteria"]  # Already a dict
+        template_data = dataset["train"].filter(lambda x: x["name"] == selected_template).to_pandas()
+        criteria = json.loads(template_data.iloc[0]["criteria"])
 
-        responses = {}
-        for criterion, prompt in criteria.items():
-            prefilled_text = f"Enter {criterion} ({prompt}):"
-            response = st.text_input(prefilled_text, key=criterion)
-            if response:
-                responses[criterion] = response
+        # Dynamic input collection
+        inputs = {}
+        for question in criteria.keys():
+            inputs[question] = st.text_input(f"{question}: ")
 
         if st.button("Generate CAD Model"):
-            try:
-                model = generate_cad_model(responses)
-                file_format = st.selectbox("Select file format for download", ["stl", "step", "dwg"])
-                file_path = export_cad(model, file_format)
-                st.success("CAD Model Generated Successfully!")
-                with open(file_path, "rb") as f:
-                    st.download_button("Download CAD File", data=f.read(), file_name=f"model.{file_format}")
-            except Exception as e:
-                st.error(f"Error generating CAD model: {e}")
+            if selected_template == "Cargo Crane":
+                cad_model = generate_cargo_crane(inputs)
+                st.success("CAD Model Generated!")
+                cq.exporters.export(cad_model, "cargo_crane.step")
+                cq.exporters.export(cad_model, "cargo_crane.stl")
+                st.download_button("Download STEP File", open("cargo_crane.step", "rb").read(), "cargo_crane.step")
+                st.download_button("Download STL File", open("cargo_crane.stl", "rb").read(), "cargo_crane.stl")