Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from transformers import pipeline | |
| import tempfile | |
| import os | |
| # Load zero-shot classification model | |
| classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli") | |
| # Define candidate screw sizes and head types | |
| screw_sizes = ["M3 screw", "M4 screw", "M6 screw"] | |
| head_types = ["hex head", "flat head", "pan head", "round head", "socket head"] | |
| def parse_prompt_to_scad(prompt: str): | |
| # Detect screw size | |
| screw_result = classifier(prompt, screw_sizes) | |
| screw_type = screw_result["labels"][0] | |
| # Detect head type | |
| head_result = classifier(prompt, head_types) | |
| head_type = head_result["labels"][0] | |
| # Define dimensions (approximate values) | |
| if screw_type == "M3 screw": | |
| d = 3 | |
| h = 15 | |
| pitch = 0.5 | |
| elif screw_type == "M4 screw": | |
| d = 4 | |
| h = 20 | |
| pitch = 0.7 | |
| elif screw_type == "M6 screw": | |
| d = 6 | |
| h = 30 | |
| pitch = 1.0 | |
| else: | |
| d = 5 | |
| h = 25 | |
| pitch = 0.8 | |
| # Head geometry (basic representation) | |
| head_code = "" | |
| if head_type == "hex head": | |
| head_code = f"cylinder(d=10, h=4); translate([0,0,4])" | |
| elif head_type == "flat head": | |
| head_code = f"cylinder(d1=10, d2={d}, h=4); translate([0,0,4])" | |
| elif head_type == "pan head": | |
| head_code = f"sphere(r=5); translate([0,0,5])" | |
| elif head_type == "round head": | |
| head_code = f"sphere(r=5); translate([0,0,5])" | |
| elif head_type == "socket head": | |
| head_code = f"cylinder(d=8, h=5); translate([0,0,5])" | |
| # Thread simulation (simplified) | |
| threads = f""" | |
| // Simulated threads using rotate_extrude | |
| module thread() {{ | |
| translate([d/2, 0, 0]) | |
| rotate_extrude(angle=360, $fn=100) | |
| polygon(points=[[0,0],[pitch/2,0],[pitch/2,1],[0,1]]); | |
| }} | |
| translate([0,0,0]) | |
| for(i = [0:{int(h/pitch)-1}]) | |
| translate([0,0,i*pitch]) | |
| thread(); | |
| """ | |
| # Final OpenSCAD script | |
| scad_script = f""" | |
| // Auto-generated screw with thread and {head_type} | |
| d = {d}; | |
| h = {h}; | |
| pitch = {pitch}; | |
| // Head | |
| {head_code} | |
| // Shaft | |
| cylinder(d=d, h=h, center=false); | |
| // Threads | |
| {threads} | |
| """ | |
| return scad_script.strip(), screw_type.replace(" ", "_"), head_type.replace(" ", "_"), d, h | |
| def generate_scad_file(prompt: str): | |
| scad_script, screw_type, head_type, d, h = parse_prompt_to_scad(prompt) | |
| filename = f"{screw_type}_{head_type}_d{d}_h{h}.scad" | |
| filepath = os.path.join(tempfile.gettempdir(), filename) | |
| with open(filepath, "w") as f: | |
| f.write(scad_script) | |
| return filepath | |
| iface = gr.Interface( | |
| fn=generate_scad_file, | |
| inputs=gr.Textbox(label="Enter prompt (e.g. 'Make an M6 hex head screw with threads')"), | |
| outputs=gr.File(label="Download .scad file"), | |
| title="Threaded Screw Generator", | |
| description="Generates OpenSCAD scripts for threaded screws with customizable head types." | |
| ) | |
| if __name__ == "__main__": | |
| iface.launch() | |