gradio demo added

app.py

@@ -0,0 +1,46 @@
+import gradio as gr
+from sketch_nn.designer import NeuralNetworkDesigner
+import tempfile
+import os
+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()

requirements.txt

@@ -0,0 +1,5 @@
+gradio
+torch
+opencv-python
+numpy
+pytesseract

sketch_ nn/__init__.py

@@ -0,0 +1,11 @@
+# Import the main class from the designer module
+from .designer import NeuralNetworkDesigner
+
+# Import any utility functions you want to make directly accessible
+from .utils import save_uploaded_file  # assuming you have this function in utils.py
+
+# You can also define the version of your package here
+__version__ = "0.1.0"
+
+# If you want to control what gets imported with "from sketch_nn import *"
+__all__ = ['NeuralNetworkDesigner', 'save_uploaded_file']

sketch_ nn/designer.py

@@ -0,0 +1,166 @@
+import cv2
+import numpy as np
+import pytesseract
+import torch
+import torch.nn as nn
+
+class NeuralNetworkDesigner:
+    def __init__(self):
+        self.layer_maps = {}
+
+    def process_image(self, image_path):
+        # Read the image
+        image = cv2.imread(image_path)
+        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+        
+        # Threshold the image
+        _, binary = cv2.threshold(gray, 225, 255, cv2.THRESH_BINARY_INV)
+        
+        # Find contours
+        contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        
+        # Sort contours from top to bottom
+        contours = sorted(contours, key=lambda c: cv2.boundingRect(c)[1])
+        
+        for i, contour in enumerate(contours):
+            x, y, w, h = cv2.boundingRect(contour)
+            roi = gray[y:y+h, x:x+w]
+            
+            # Perform OCR on the ROI
+            text = pytesseract.image_to_string(roi).strip()
+            self.parse_layer_info(i, text)
+
+    def parse_layer_info(self, layer_index, text):
+        lines = text.split('\n')
+        layer_info = {'type': 'Unknown', 'text': text}
+        
+        try:
+            if 'Input' in lines[0]:
+                layer_info['type'] = 'Input'
+                layer_info['channels'] = int(lines[1]) if len(lines) > 1 else None
+            elif 'Conv' in lines[0]:
+                layer_info['type'] = 'Conv2d'
+                layer_info['out_channels'] = int(lines[1]) if len(lines) > 1 else None
+                layer_info['kernel_size'] = int(lines[2]) if len(lines) > 2 else None
+            elif any(x in lines[0] for x in ['MaxPool', 'AvgPool']):
+                layer_info['type'] = 'MaxPool2d' if 'Max' in lines[0] else 'AvgPool2d'
+                layer_info['kernel_size'] = int(lines[1]) if len(lines) > 1 else None
+            elif 'Linear' in lines[0]:
+                layer_info['type'] = 'Linear'
+                if len(lines) > 1 and '*' in lines[1]:
+                    layer_info['in_features'] = lines[1]
+                layer_info['out_features'] = int(lines[-1]) if lines[-1].isdigit() else None
+            elif 'BatchNorm' in lines[0]:
+                layer_info['type'] = 'BatchNorm2d'
+                layer_info['num_features'] = int(lines[1]) if len(lines) > 1 else None
+            elif any(x in lines[0] for x in ['ReLU', 'LeakyReLU', 'Sigmoid', 'Tanh']):
+                layer_info['type'] = lines[0]
+            elif 'Dropout' in lines[0]:
+                layer_info['type'] = 'Dropout'
+                layer_info['p'] = float(lines[1]) if len(lines) > 1 else 0.5
+            elif 'Transformer' in lines[0]:
+                layer_info['type'] = 'Transformer'
+                layer_info['d_model'] = int(lines[1]) if len(lines) > 1 else 512
+                layer_info['nhead'] = int(lines[2]) if len(lines) > 2 else 8
+            elif 'Attention' in lines[0]:
+                layer_info['type'] = 'MultiheadAttention'
+                layer_info['embed_dim'] = int(lines[1]) if len(lines) > 1 else 512
+                layer_info['num_heads'] = int(lines[2]) if len(lines) > 2 else 8
+            elif 'LSTM' in lines[0] or 'GRU' in lines[0]:
+                layer_info['type'] = lines[0]
+                layer_info['hidden_size'] = int(lines[1]) if len(lines) > 1 else 256
+                layer_info['num_layers'] = int(lines[2]) if len(lines) > 2 else 1
+        except ValueError as e:
+            print(f"Error parsing layer {layer_index}: {e}")
+
+        self.layer_maps[layer_index] = layer_info
+        print(f"Parsed layer {layer_index}: {layer_info}")  # Debug print
+
+    def generate_pytorch_code(self):
+        code = "import torch\nimport torch.nn as nn\n\n"
+        code += "class CustomNN(nn.Module):\n"
+        code += "    def __init__(self):\n"
+        code += "        super(CustomNN, self).__init__()\n"
+        
+        forward_code = "    def forward(self, x):\n"
+        
+        in_channels = None
+        for i, layer_info in sorted(self.layer_maps.items()):
+            if layer_info['type'] == 'Input':
+                in_channels = layer_info.get('channels', 3)
+                continue
+            
+            if layer_info['type'] == 'Conv2d':
+                out_channels = layer_info.get('out_channels', 64)
+                kernel_size = layer_info.get('kernel_size', 3)
+                code += f"        self.conv{i} = nn.Conv2d({in_channels}, {out_channels}, kernel_size={kernel_size}, padding=1)\n"
+                forward_code += f"        x = self.conv{i}(x)\n"
+                in_channels = out_channels
+            
+            elif layer_info['type'] in ['MaxPool2d', 'AvgPool2d']:
+                kernel_size = layer_info.get('kernel_size', 2)
+                code += f"        self.pool{i} = nn.{layer_info['type']}(kernel_size={kernel_size})\n"
+                forward_code += f"        x = self.pool{i}(x)\n"
+            
+            elif layer_info['type'] == 'Linear':
+                out_features = layer_info.get('out_features')
+                if i == 1 or (i > 1 and self.layer_maps[i-1]['type'] not in ['Linear', 'Flatten']):
+                    code += f"        self.flatten = nn.Flatten()\n"
+                    forward_code += f"        x = self.flatten(x)\n"
+                    in_features = layer_info.get('in_features', 'x.shape[1]')
+                else:
+                    in_features = self.layer_maps[i-1].get('out_features', 64)
+                code += f"        self.fc{i} = nn.Linear({in_features}, {out_features})\n"
+                forward_code += f"        x = self.fc{i}(x)\n"
+            
+            elif layer_info['type'] == 'BatchNorm2d':
+                num_features = layer_info.get('num_features', in_channels)
+                code += f"        self.bn{i} = nn.BatchNorm2d({num_features})\n"
+                forward_code += f"        x = self.bn{i}(x)\n"
+            
+            elif layer_info['type'] in ['ReLU', 'LeakyReLU', 'Sigmoid', 'Tanh']:
+                code += f"        self.act{i} = nn.{layer_info['type']}()\n"
+                forward_code += f"        x = self.act{i}(x)\n"
+            
+            elif layer_info['type'] == 'Dropout':
+                p = layer_info.get('p', 0.5)
+                code += f"        self.dropout{i} = nn.Dropout(p={p})\n"
+                forward_code += f"        x = self.dropout{i}(x)\n"
+            
+            elif layer_info['type'] == 'Transformer':
+                d_model = layer_info.get('d_model', 512)
+                nhead = layer_info.get('nhead', 8)
+                code += f"        self.transformer{i} = nn.Transformer(d_model={d_model}, nhead={nhead})\n"
+                forward_code += f"        x = self.transformer{i}(x)\n"
+            
+            elif layer_info['type'] == 'MultiheadAttention':
+                embed_dim = layer_info.get('embed_dim', 512)
+                num_heads = layer_info.get('num_heads', 8)
+                code += f"        self.attention{i} = nn.MultiheadAttention(embed_dim={embed_dim}, num_heads={num_heads})\n"
+                forward_code += f"        x, _ = self.attention{i}(x, x, x)\n"
+            
+            elif layer_info['type'] in ['LSTM', 'GRU']:
+                hidden_size = layer_info.get('hidden_size', 256)
+                num_layers = layer_info.get('num_layers', 1)
+                code += f"        self.rnn{i} = nn.{layer_info['type']}(input_size={in_channels}, hidden_size={hidden_size}, num_layers={num_layers}, batch_first=True)\n"
+                forward_code += f"        x, _ = self.rnn{i}(x)\n"
+            
+            elif layer_info['type'] == 'Unknown':
+                print(f"Warning: Unknown layer type at index {i}. Layer info: {layer_info}")
+        
+        code += "\n" + forward_code
+        code += "        return x\n"
+        
+        return code
+
+    def write_to_file(self, code, filename):
+        with open(filename, 'w') as f:
+            f.write(code)
+
+    def design_network(self, image_path, output_file):
+        self.process_image(image_path)
+        pytorch_code = self.generate_pytorch_code()
+        self.write_to_file(pytorch_code, output_file)
+        # print(f"Neural network code has been generated and saved to '{output_file}'")
+        # print("\nGenerated PyTorch Code:")
+        # print(pytorch_code)

sketch_ nn/utils.py

@@ -0,0 +1,8 @@
+import os
+import tempfile
+
+def save_uploaded_file(uploaded_file):
+    temp_dir = tempfile.mkdtemp()
+    temp_path = os.path.join(temp_dir, "uploaded_image.png")
+    uploaded_file.save(temp_path)
+    return temp_path