Update app.py

app.py

@@ -1,43 +1,77 @@
 import gradio as gr
-import os
+import numpy as np
+import matplotlib.pyplot as plt
+import plotly.graph_objects as go
+import io
 
 # Function to generate G-code based on user input
 def generate_gcode(length, diameter, feed_rate):
-    """
-    Generates a simple G-code string based on user inputs.
-    """
     gcode = []
     gcode.append(f"(Generated G-code for length={length}, diameter={diameter}, feed_rate={feed_rate})")
     gcode.append("G21   ; Set units to millimeters")
     gcode.append("G90   ; Absolute positioning mode")
-    gcode.append("G17   ; Select XY plane")
     gcode.append("M03 S1000   ; Start spindle at 1000 RPM")
     gcode.append("G00 Z5.0   ; Move to safe height")
-    gcode.append(f"G00 X0 Y0  ; Move to start position")
-    gcode.append(f"G01 Z-1.0 F{feed_rate}   ; Start cutting")
-    gcode.append(f"G01 X{length} F{feed_rate} ; Linear move along X-axis")
-    gcode.append(f"G01 Y{diameter} F{feed_rate} ; Linear move along Y-axis")
+    gcode.append(f"G01 X{length} Y0 F{feed_rate}   ; Linear move to length")
+    gcode.append(f"G01 X{length} Y{diameter} F{feed_rate}   ; Move to diameter")
     gcode.append("G00 Z5.0   ; Retract to safe height")
     gcode.append("M05   ; Stop spindle")
-    gcode.append("M30   ; End of program")
+    gcode.append("M30   ; End program")
     return "\n".join(gcode)
 
-# Gradio UI setup
+# Function to generate 2D visualization
+def generate_2d_plot(length, diameter):
+    fig, ax = plt.subplots()
+    ax.plot([0, length, length, 0, 0], [0, 0, diameter, diameter, 0], color='blue')
+    ax.set_title("2D Visualization of Toolpath")
+    ax.set_xlabel("Length (mm)")
+    ax.set_ylabel("Diameter (mm)")
+    ax.set_aspect('equal')
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png')
+    buf.seek(0)
+    return buf
+
+# Function to generate 3D visualization
+def generate_3d_plot(length, diameter):
+    X = np.linspace(0, length, 50)
+    Y = np.linspace(0, diameter, 50)
+    X, Y = np.meshgrid(X, Y)
+    Z = np.zeros_like(X)
+
+    fig = go.Figure(data=[go.Surface(z=Z, x=X, y=Y, colorscale="Blues")])
+    fig.update_layout(
+        title="3D Visualization of Toolpath",
+        scene=dict(
+            xaxis_title="Length (mm)",
+            yaxis_title="Diameter (mm)",
+            zaxis_title="Height"
+        )
+    )
+    return fig
+
+# Main interface function
 def main_interface(length, diameter, feed_rate):
     gcode = generate_gcode(length, diameter, feed_rate)
-    return gcode
+    plot_2d = generate_2d_plot(length, diameter)
+    plot_3d = generate_3d_plot(length, diameter)
+    return gcode, plot_2d, plot_3d
 
-# Create Gradio Interface
+# Gradio Interface
 inputs = [
     gr.Number(label="Length (mm)", value=50),
     gr.Number(label="Diameter (mm)", value=20),
     gr.Number(label="Feed Rate (mm/min)", value=100)
 ]
 
-outputs = gr.Textbox(label="Generated G-code", lines=20)
+outputs = [
+    gr.Textbox(label="Generated G-code", lines=15),
+    gr.Image(label="2D Visualization"),
+    gr.Plot(label="3D Visualization")
+]
 
-title = "CNC G-code Generator"
-description = "Enter the desired length, diameter, and feed rate to generate a basic G-code file."
+title = "CNC G-code Generator with Visualization"
+description = "Generate CNC G-code with 2D and 3D visualizations. Input the length, diameter, and feed rate to preview the toolpath."
 
 demo = gr.Interface(
     fn=main_interface,