Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import cadquery as cq | |
| from datasets import load_dataset | |
| import json | |
| # Load dataset | |
| dataset = load_dataset("FreedomIntelligence/CADBench") | |
| # 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) | |
| # Create boom | |
| boom = cq.Workplane("XY").box(50, boom_length, 50).translate((0, base_height, 0)) | |
| # Combine parts | |
| crane = base.union(boom) | |
| return crane | |
| # Streamlit Interface | |
| st.title("Quick CAD Model Generator") | |
| case = st.radio("Select Case", ["DFM Analysis", "Predefined Template"]) | |
| 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: | |
| template_data = dataset["train"].filter(lambda x: x["name"] == selected_template).to_pandas() | |
| criteria = template_data.iloc[0]["criteria"] | |
| # Dynamic input collection | |
| inputs = {} | |
| for question in criteria.keys(): | |
| inputs[question] = st.text_input(f"{question}: ") | |
| if st.button("Generate CAD Model"): | |
| 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") | |
| 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) | |
| 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 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 | |
| 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 | |
| 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}") | |