Create app.py

app.py

@@ -0,0 +1,97 @@
+import streamlit as st
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+import open3d as o3d
+import numpy as np
+import cadquery as cq
+
+# Load the tokenizer from Qwen2-1.5B and model weights from filapro/cad-recode
+tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-1.5B", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained("filapro/cad-recode", trust_remote_code=True)
+
+# Set device (GPU if available, CPU otherwise)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+print(f"Model loaded on {device}")
+
+@st.cache(allow_output_mutation=True)
+def load_point_cloud(file):
+    """Loads a point cloud from a uploaded file."""
+    if not file:
+        return None
+
+    if file.type not in ("application/octet-stream", "text/plain"):
+        st.error("Please upload a point cloud file (.pcd, .xyz, etc.)")
+        return None
+
+    try:
+        point_cloud = o3d.io.read_point_cloud(file)
+    except Exception as e:
+        st.error(f"Error loading point cloud: {e}")
+        return None
+
+    return point_cloud
+
+def prepare_input_data(point_cloud):
+    """Prepares point cloud data for model input."""
+    if not point_cloud:
+        return None
+
+    point_cloud_array = np.asarray(point_cloud.points).flatten()
+    input_text = " ".join(map(str, point_cloud_array))
+    return input_text
+
+def generate_cad_code(input_text):
+    """Runs inference and decodes generated output."""
+    if not input_text:
+        return None
+
+    inputs = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=512)
+    inputs = {key: val.to(device) for key, val in inputs.items()}
+
+    with torch.no_grad():
+        outputs = model.generate(**inputs, max_new_tokens=256, pad_token_id=tokenizer.eos_token_id)
+
+    cad_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return cad_code
+
+def generate_cad_model(cad_code):
+    """Generates a CAD model from the provided code."""
+    if not cad_code:
+        return None
+
+    try:
+        # Execute CAD code using CadQuery library
+        exec(cad_code)
+        cad_model = cq.Workplane("XY").val()
+    except Exception as e:
+        st.error(f"Error generating CAD model: {e}")
+        return None
+
+    return cad_model
+
+def main():
+    """Streamlit app for point cloud to CAD code conversion."""
+    st.title("Point Cloud to CAD Code Converter")
+    st.write("This app uses the filapro/cad-recode model to generate Python code for a 3D CAD model from your point cloud data.")
+
+    uploaded_file = st.file_uploader("Upload Point Cloud File")
+    point_cloud = load_point_cloud(uploaded_file)
+
+    if point_cloud:
+        input_text = prepare_input_data(point_cloud)
+        cad_code = generate_cad_code(input_text)
+
+        if cad_code:
+            st.success("Generated Python CAD Code:")
+            st.code(cad_code)
+
+            cad_model = generate_cad_model(cad_code)
+            if cad_model:
+                # Optionally, use a 3D visualization library like trimesh
+                # to display the generated CAD model (not included)
+                st.success("Generated CAD Model (Visualization not yet implemented)")
+                # st.write(cad_model)  # Replace with visualization code
+
+if __name__ == "__main__":
+    main()