Create apply.py

apply.py

@@ -0,0 +1,77 @@
+from util import UIDataset, Vocabulary
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import torch
+from torch.utils.data import DataLoader
+from model import *
+from torchvision import transforms
+from PIL import Image
+
+# Carrega o modelo treinado
+net = Pix2Code()
+net.load_state_dict(torch.load('./pix2code.weights'))
+net.cuda().eval()
+
+# Carrega o vocabulário
+vocab = Vocabulary('voc.pkl')
+
+# Define uma transformação para redimensionar e normalizar as imagens
+transform = transforms.Compose([
+    transforms.Resize((256, 256)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+])
+
+# Função que receberá a imagem e retornará o código GUI gerado
+def generate_gui(image):
+    # Aplica a transformação na imagem
+    image = transform(image).unsqueeze(0).cuda()
+    
+    # Cria um contexto inicial
+    context = torch.tensor([vocab.to_vec(' '), vocab.to_vec('<START>')]).unsqueeze(0).float().cuda()
+    
+    # Inicializa uma lista para armazenar o código gerado
+    code = []
+    
+    # Gera o código iterativamente até encontrar o token <END>
+    for i in range(200):
+        # Passa a imagem e o contexto para a rede neural e obtém o índice do token com maior probabilidade
+        index = torch.argmax(net(image, context), 2).squeeze()[-1:].squeeze()
+        
+        # Converte o índice para o token correspondente
+        token = vocab.to_vocab(int(index))
+        
+        # Se encontrar o token <END>, interrompe a geração do código
+        if token == '<END>':
+            break
+        
+        # Adiciona o token à lista de código gerado
+        code.append(token)
+        
+        # Atualiza o contexto com o token gerado
+        context = torch.cat([context, torch.tensor([vocab.to_vec(token)]).unsqueeze(0).float().cuda()], dim=1)
+    
+    # Retorna o código gerado como uma string
+    return ''.join(code)
+
+import gradio as gr
+
+# Define o componente de entrada
+image_input = gr.inputs.Image()
+
+# Define o componente de saída
+text_output = gr.outputs.Textbox()
+
+# Cria a interface Gradio
+iface = gr.Interface(
+    fn=generate_gui,
+    inputs=image_input,
+    outputs=text_output,
+    title='Pix2Code',
+    description='Gerador de código GUI a partir de imagens',
+    theme='default'
+)
+
+# Executa a interface
+iface.launch()