designer.py

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)