File size: 1,690 Bytes
6ce527a ad1cf6d 6ce527a 8462fc9 29b11c2 8462fc9 873d368 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import gradio as gr
from designer import NeuralNetworkDesigner
import tempfile
import cv2
import numpy as np
def generate_nn_code(image):
designer = NeuralNetworkDesigner()
temp_dir = tempfile.mkdtemp()
temp_path = os.path.join(temp_dir, "input_image.png")
if isinstance(image, np.ndarray):
# If it's a numpy array (captured image or uploaded image)
cv2.imwrite(temp_path, cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
else:
# If it's a file path (should not happen with current setup, but just in case)
os.rename(image, temp_path)
output_file = os.path.join(temp_dir, "custom_nn.py")
designer.design_network(temp_path, output_file)
with open(output_file, 'r') as f:
code = f.read()
return output_file, code
# Check if the version of Gradio supports the 'source' parameter
try:
image_input = gr.Image(source=["upload", "webcam"], type="numpy", label="Upload or Capture Flowchart")
except TypeError:
# Fallback for older Gradio versions
image_input = gr.Image(type="numpy", label="Upload Flowchart")
iface = gr.Interface(
fn=generate_nn_code,
inputs=[image_input],
outputs=[
gr.File(label="Download Generated Code"),
gr.Code(language="python", label="Generated PyTorch Code")
],
title="Sketch NN: Neural Network Designer",
description="Upload a flowchart image or capture one using your webcam to generate PyTorch code for your neural network architecture."
)
if __name__ == "__main__":
iface.launch(share=True)
else:
iface.launch(inline=False) |