go3

app.py

@@ -1,7 +1,3 @@
-# ============================================================================
-# SISTEMA DE CLASSIFICAÇÃO DE IMAGENS - HUGGING FACE SPACE
-# ============================================================================
-
 import os
 import shutil
 import gradio as gr
@@ -19,492 +15,392 @@ import tempfile
 import warnings
 warnings.filterwarnings("ignore")
 
-print("✅ Todas as bibliotecas carregadas com sucesso!")
-print(f"🔥 GPU disponível: {torch.cuda.is_available()}")
-print(f"🖥️ Device: {torch.device('cuda' if torch.cuda.is_available() else 'cpu')}")
-
-# ============================================================================
-# CONFIGURAÇÕES E VARIÁVEIS GLOBAIS
-# ============================================================================
+# Configuração do device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"🖥️ Usando device: {device}")
 
 # Modelos disponíveis
-model_dict = {
-    'AlexNet': models.alexnet,
+MODELS = {
     'ResNet18': models.resnet18,
-    'ResNet34': models.resnet34,
-    'ResNet50': models.resnet50,
     'MobileNetV2': models.mobilenet_v2
 }
 
-# Variáveis globais
-model = None
-train_loader = None
-val_loader = None
-test_loader = None
-dataset_path = tempfile.mkdtemp()  # Usar diretório temporário
-class_dirs = []
-class_labels = []
-test_dataset_path = tempfile.mkdtemp()
-test_class_dirs = []
-num_classes = 2
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-print("🎯 Configurações inicializadas!")
-
-# ============================================================================
-# FUNÇÕES PRINCIPAIS
-# ============================================================================
+# Estado global da aplicação
+class AppState:
+    def __init__(self):
+        self.model = None
+        self.train_loader = None
+        self.val_loader = None
+        self.test_loader = None
+        self.dataset_path = None
+        self.class_dirs = []
+        self.class_labels = []
+        self.num_classes = 2
+
+# Instância global do estado
+app_state = AppState()
 
 def setup_classes(num_classes_value):
     """Configura o número de classes e cria diretórios"""
-    global class_dirs, dataset_path, num_classes, class_labels
-
-    num_classes = int(num_classes_value)
-
-    # Limpar e criar diretórios
-    if os.path.exists(dataset_path):
-        shutil.rmtree(dataset_path)
-    dataset_path = tempfile.mkdtemp()
-
-    # Inicializar rótulos padrão
-    class_labels = [f'class_{i}' for i in range(num_classes)]
-
-    # Criar diretórios para cada classe
-    class_dirs = [os.path.join(dataset_path, f'class_{i}') for i in range(num_classes)]
-    for class_dir in class_dirs:
-        os.makedirs(class_dir)
-
-    return f"✅ Criados {num_classes} diretórios para classes"
+    try:
+        app_state.num_classes = int(num_classes_value)
+        
+        # Criar diretório temporário
+        app_state.dataset_path = tempfile.mkdtemp()
+        
+        # Inicializar rótulos padrão
+        app_state.class_labels = [f'classe_{i}' for i in range(app_state.num_classes)]
+        
+        # Criar diretórios para cada classe
+        app_state.class_dirs = []
+        for i in range(app_state.num_classes):
+            class_dir = os.path.join(app_state.dataset_path, f'classe_{i}')
+            os.makedirs(class_dir, exist_ok=True)
+            app_state.class_dirs.append(class_dir)
+        
+        choices = [(f"{i} - {app_state.class_labels[i]}", i) for i in range(app_state.num_classes)]
+        
+        return (
+            f"✅ Criados {app_state.num_classes} diretórios para classes",
+            gr.Dropdown(choices=choices, value=0)
+        )
+    except Exception as e:
+        return f"❌ Erro: {str(e)}", gr.Dropdown()
 
-def set_class_labels(*labels):
+def set_class_labels(label0, label1, label2, label3, label4):
     """Define rótulos personalizados para as classes"""
-    global class_labels
-
-    # Filtrar labels vazios
-    filtered_labels = [label.strip() for label in labels if label.strip()][:num_classes]
-
-    if len(filtered_labels) != num_classes:
-        return f"❌ Erro: Você deve fornecer exatamente {num_classes} rótulos. Fornecidos: {len(filtered_labels)}"
-
-    class_labels = filtered_labels
-    return f"✅ Rótulos definidos: {', '.join(class_labels)}"
+    try:
+        labels = [label0, label1, label2, label3, label4]
+        filtered_labels = [label.strip() for label in labels if label.strip()][:app_state.num_classes]
+        
+        if len(filtered_labels) != app_state.num_classes:
+            return f"❌ Erro: Forneça exatamente {app_state.num_classes} rótulos.", gr.Dropdown()
+        
+        app_state.class_labels = filtered_labels
+        choices = [(f"{i} - {app_state.class_labels[i]}", i) for i in range(app_state.num_classes)]
+        
+        return (
+            f"✅ Rótulos definidos: {', '.join(app_state.class_labels)}",
+            gr.Dropdown(choices=choices, value=0)
+        )
+    except Exception as e:
+        return f"❌ Erro: {str(e)}", gr.Dropdown()
 
 def upload_images(class_id, images):
     """Faz upload das imagens para a classe especificada"""
-    if int(class_id) >= len(class_dirs):
-        return f"❌ Erro: ID da classe {class_id} inválido."
-
-    if not images:
-        return "❌ Nenhuma imagem foi selecionada."
-
-    class_dir = class_dirs[int(class_id)]
-    count = 0
-
-    for image in images:
-        try:
-            shutil.copy(image, class_dir)
-            count += 1
-        except Exception as e:
-            return f"❌ Erro ao copiar imagem: {str(e)}"
-
-    class_name = class_labels[int(class_id)] if int(class_id) < len(class_labels) else f"class_{class_id}"
-    return f"✅ {count} imagens salvas na classe {class_id} ({class_name})"
+    try:
+        if not images:
+            return "❌ Nenhuma imagem selecionada."
+        
+        if int(class_id) >= len(app_state.class_dirs):
+            return f"❌ Classe {class_id} inválida."
+        
+        class_dir = app_state.class_dirs[int(class_id)]
+        count = 0
+        
+        for image in images:
+            if image is not None:
+                shutil.copy2(image, class_dir)
+                count += 1
+        
+        class_name = app_state.class_labels[int(class_id)]
+        return f"✅ {count} imagens salvas na classe {class_id} ({class_name})"
+    except Exception as e:
+        return f"❌ Erro: {str(e)}"
 
-def prepare_data(batch_size=32, resize_str="224,224"):
+def prepare_data(batch_size):
     """Prepara os dados para treinamento"""
-    global train_loader, val_loader, test_loader, num_classes
-
     try:
-        # Parse do resize
-        resize = tuple(map(int, resize_str.split(',')))
-
+        if not app_state.dataset_path or not os.path.exists(app_state.dataset_path):
+            return "❌ Configure as classes primeiro."
+        
         # Transformações
         transform = transforms.Compose([
-            transforms.Resize(resize),
+            transforms.Resize((224, 224)),
             transforms.ToTensor(),
             transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
         ])
-
-        dataset = datasets.ImageFolder(dataset_path, transform=transform)
-
+        
+        dataset = datasets.ImageFolder(app_state.dataset_path, transform=transform)
+        
         if len(dataset.classes) == 0:
             return "❌ Nenhuma classe encontrada. Faça upload das imagens primeiro."
-
-        if len(dataset.classes) != num_classes:
-            return f"❌ Erro: {len(dataset.classes)} classes detectadas, {num_classes} esperadas."
-
-        # Verificar se há imagens suficientes
-        if len(dataset) < 10:
-            return f"❌ Muito poucas imagens ({len(dataset)}). Adicione pelo menos 10 imagens por classe."
-
-        # Divisão dos dados: 70% treino, 20% validação, 10% teste
+        
+        if len(dataset) < 6:
+            return f"❌ Muito poucas imagens ({len(dataset)}). Adicione pelo menos 2 imagens por classe."
+        
+        # Divisão dos dados
         train_size = int(0.7 * len(dataset))
         val_size = int(0.2 * len(dataset))
         test_size = len(dataset) - train_size - val_size
-
+        
         train_dataset, val_dataset, test_dataset = random_split(
             dataset, [train_size, val_size, test_size],
             generator=torch.Generator().manual_seed(42)
         )
-
-        train_loader = DataLoader(train_dataset, batch_size=int(batch_size), shuffle=True)
-        val_loader = DataLoader(val_dataset, batch_size=int(batch_size), shuffle=False)
-        test_loader = DataLoader(test_dataset, batch_size=int(batch_size), shuffle=False)
-
+        
+        app_state.train_loader = DataLoader(train_dataset, batch_size=int(batch_size), shuffle=True)
+        app_state.val_loader = DataLoader(val_dataset, batch_size=int(batch_size), shuffle=False)
+        app_state.test_loader = DataLoader(test_dataset, batch_size=int(batch_size), shuffle=False)
+        
         return f"✅ Dados preparados: {train_size} treino, {val_size} validação, {test_size} teste"
-
+    
     except Exception as e:
         return f"❌ Erro na preparação: {str(e)}"
 
-def start_training(model_name, epochs, lr, progress=gr.Progress()):
+def start_training(model_name, epochs, lr):
     """Inicia o treinamento do modelo"""
-    global model, train_loader, val_loader, device
-
-    if train_loader is None or val_loader is None:
-        return "❌ Erro: Dados não preparados. Execute a preparação dos dados primeiro."
-
     try:
-        # Carregar modelo pré-treinado
-        model = model_dict[model_name](pretrained=True)
-
-        # Adaptar última camada para número de classes
-        if hasattr(model, 'fc'):
-            model.fc = nn.Linear(model.fc.in_features, num_classes)
-        elif hasattr(model, 'classifier'):
-            model.classifier[-1] = nn.Linear(model.classifier[-1].in_features, num_classes)
-
-        model = model.to(device)
-
+        if app_state.train_loader is None:
+            return "❌ Erro: Dados não preparados."
+        
+        # Carregar modelo
+        app_state.model = MODELS[model_name](pretrained=True)
+        
+        # Adaptar última camada
+        if hasattr(app_state.model, 'fc'):
+            app_state.model.fc = nn.Linear(app_state.model.fc.in_features, app_state.num_classes)
+        elif hasattr(app_state.model, 'classifier'):
+            if isinstance(app_state.model.classifier, nn.Sequential):
+                app_state.model.classifier[-1] = nn.Linear(app_state.model.classifier[-1].in_features, app_state.num_classes)
+            else:
+                app_state.model.classifier = nn.Linear(app_state.model.classifier.in_features, app_state.num_classes)
+        
+        app_state.model = app_state.model.to(device)
+        
         criterion = nn.CrossEntropyLoss()
-        optimizer = optim.Adam(model.parameters(), lr=float(lr))
-
-        results = []
-        results.append(f"🚀 Iniciando treinamento com {model_name}")
-        results.append(f"📊 Épocas: {epochs}, Learning Rate: {lr}")
-        results.append(f"🔥 Device: {device}")
-        results.append("-" * 50)
-
-        model.train()
-
+        optimizer = optim.Adam(app_state.model.parameters(), lr=float(lr))
+        
+        app_state.model.train()
+        
+        results = [f"🚀 Treinando {model_name} por {epochs} épocas"]
+        
         for epoch in range(int(epochs)):
             running_loss = 0.0
             correct = 0
             total = 0
-
-            progress(epoch / int(epochs), desc=f"Época {epoch+1}/{epochs}")
-
-            for batch_idx, (inputs, labels) in enumerate(train_loader):
+            
+            for inputs, labels in app_state.train_loader:
                 inputs, labels = inputs.to(device), labels.to(device)
-
+                
                 optimizer.zero_grad()
-                outputs = model(inputs)
+                outputs = app_state.model(inputs)
                 loss = criterion(outputs, labels)
                 loss.backward()
                 optimizer.step()
-
+                
                 running_loss += loss.item()
                 _, predicted = torch.max(outputs.data, 1)
                 total += labels.size(0)
                 correct += (predicted == labels).sum().item()
-
-            # Estatísticas da época
-            epoch_loss = running_loss / len(train_loader)
+            
+            epoch_loss = running_loss / len(app_state.train_loader)
             epoch_acc = 100. * correct / total
-            results.append(f"📈 Época {epoch+1}: Loss={epoch_loss:.4f}, Accuracy={epoch_acc:.2f}%")
-
-        # Salvar modelo
-        model_path = tempfile.NamedTemporaryFile(suffix='.pth', delete=False).name
-        torch.save(model.state_dict(), model_path)
-        results.append("-" * 50)
-        results.append(f"✅ Treinamento concluído! Modelo salvo temporariamente.")
-
+            results.append(f"Época {epoch+1}: Loss={epoch_loss:.4f}, Acc={epoch_acc:.2f}%")
+        
+        results.append("✅ Treinamento concluído!")
         return "\n".join(results)
-
+    
     except Exception as e:
         return f"❌ Erro durante treinamento: {str(e)}"
 
 def evaluate_model():
     """Avalia o modelo no conjunto de teste"""
-    global model, device, num_classes, class_labels, test_loader
-
-    if model is None:
-        return "❌ Erro: Modelo não treinado."
-
-    if test_loader is None:
-        return "❌ Erro: Conjunto de dados não preparado."
-
-    model.eval()
-    all_preds = []
-    all_labels = []
-
     try:
+        if app_state.model is None or app_state.test_loader is None:
+            return "❌ Modelo ou dados não disponíveis."
+        
+        app_state.model.eval()
+        all_preds = []
+        all_labels = []
+        
         with torch.no_grad():
-            for inputs, labels in test_loader:
+            for inputs, labels in app_state.test_loader:
                 inputs, labels = inputs.to(device), labels.to(device)
-                outputs = model(inputs)
+                outputs = app_state.model(inputs)
                 _, preds = torch.max(outputs, 1)
                 all_preds.extend(preds.cpu().numpy())
                 all_labels.extend(labels.cpu().numpy())
-
-        # Usar rótulos personalizados
-        target_names = class_labels if len(class_labels) == num_classes else [f"class_{i}" for i in range(num_classes)]
-        report = classification_report(all_preds, all_labels, target_names=target_names, zero_division=0)
-
+        
+        report = classification_report(all_labels, all_preds, target_names=app_state.class_labels, zero_division=0)
         return f"📊 RELATÓRIO DE CLASSIFICAÇÃO:\n\n{report}"
-
+    
     except Exception as e:
         return f"❌ Erro durante avaliação: {str(e)}"
 
-def show_confusion_matrix():
-    """Gera matriz de confusão"""
-    global model, device, num_classes, class_labels, test_loader
-
-    if model is None:
-        return None
-
-    if test_loader is None:
-        return None
-
-    model.eval()
-    all_preds = []
-    all_labels = []
-
-    with torch.no_grad():
-        for inputs, labels in test_loader:
-            inputs, labels = inputs.to(device), labels.to(device)
-            outputs = model(inputs)
-            _, preds = torch.max(outputs, 1)
-            all_preds.extend(preds.cpu().numpy())
-            all_labels.extend(labels.cpu().numpy())
-
-    cm = confusion_matrix(all_labels, all_preds)
-    labels_for_cm = class_labels if len(class_labels) == num_classes else [f"class_{i}" for i in range(num_classes)]
-
-    plt.figure(figsize=(8, 6))
-    sns.heatmap(cm, annot=True, fmt="d", cmap="Blues",
-                xticklabels=labels_for_cm,
-                yticklabels=labels_for_cm)
-    plt.xlabel('Predições')
-    plt.ylabel('Valores Reais')
-    plt.title('Matriz de Confusão')
-    plt.tight_layout()
-
-    # Salvar em arquivo temporário
-    temp_path = tempfile.NamedTemporaryFile(suffix='.png', delete=False).name
-    plt.savefig(temp_path, dpi=150, bbox_inches='tight')
-    plt.close()
-
-    return temp_path
-
 def predict_images(images):
     """Faz predições em novas imagens"""
-    global model, device, class_labels
-
-    if model is None:
-        return "❌ Erro: Modelo não treinado."
-
-    if not images:
-        return "❌ Nenhuma imagem selecionada."
-
-    transform = transforms.Compose([
-        transforms.Resize((224, 224)),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-    ])
-
-    model.eval()
-    results = []
-
-    for image_path in images:
-        try:
-            image = Image.open(image_path).convert('RGB')
-            img_tensor = transform(image).unsqueeze(0).to(device)
-
-            with torch.no_grad():
-                outputs = model(img_tensor)
-                probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
-                _, predicted = torch.max(outputs, 1)
-
-                predicted_class_id = predicted.item()
-                confidence = probabilities[predicted_class_id].item() * 100
-
-                if predicted_class_id < len(class_labels):
-                    predicted_class_name = class_labels[predicted_class_id]
-                else:
-                    predicted_class_name = f"class_{predicted_class_id}"
-
-                results.append(f"📸 {os.path.basename(image_path)}")
-                results.append(f"   🎯 Classe: {predicted_class_name}")
-                results.append(f"   📊 Confiança: {confidence:.2f}%")
-                results.append("-" * 40)
-
-        except Exception as e:
-            results.append(f"❌ Erro ao processar {os.path.basename(image_path)}: {str(e)}")
-
-    return "\n".join(results)
-
-# ============================================================================
-# INTERFACE GRADIO
-# ============================================================================
+    try:
+        if app_state.model is None:
+            return "❌ Modelo não treinado."
+        
+        if not images:
+            return "❌ Nenhuma imagem selecionada."
+        
+        transform = transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        ])
+        
+        app_state.model.eval()
+        results = []
+        
+        for image_path in images:
+            if image_path is not None:
+                image = Image.open(image_path).convert('RGB')
+                img_tensor = transform(image).unsqueeze(0).to(device)
+                
+                with torch.no_grad():
+                    outputs = app_state.model(img_tensor)
+                    probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
+                    _, predicted = torch.max(outputs, 1)
+                    
+                    predicted_class_id = predicted.item()
+                    confidence = probabilities[predicted_class_id].item() * 100
+                    predicted_class_name = app_state.class_labels[predicted_class_id]
+                    
+                    results.append(f"📸 {os.path.basename(image_path)}")
+                    results.append(f"   🎯 Classe: {predicted_class_name}")
+                    results.append(f"   📊 Confiança: {confidence:.2f}%")
+                    results.append("-" * 40)
+        
+        return "\n".join(results) if results else "❌ Nenhuma predição realizada."
+    
+    except Exception as e:
+        return f"❌ Erro: {str(e)}"
 
+# Interface Gradio
 def create_interface():
-    """Cria a interface Gradio"""
-
     with gr.Blocks(title="🖼️ Classificador de Imagens", theme=gr.themes.Soft()) as demo:
-
+        
         gr.Markdown("""
-        # 🖼️ Sistema de Classificação de Imagens Interativo
-        #### Por [Ramon Mayor Martins](https://rmayormartins.github.io/)
-
-        **Instruções:**
-        1. Configure o número de classes e defina os rótulos
-        2. Faça upload das imagens para cada classe
-        3. Prepare os dados e treine o modelo
-        4. Avalie o desempenho e faça predições!
+        # 🖼️ Sistema de Classificação de Imagens
         
-        ⚠️ **Nota**: Este é um ambiente temporário. Os dados e modelos serão perdidos ao reiniciar.
+        **Instruções:**
+        1. Configure as classes e rótulos
+        2. Faça upload das imagens
+        3. Prepare os dados e treine
+        4. Avalie e faça predições!
         """)
-
+        
         with gr.Tab("1️⃣ Configuração"):
-            gr.Markdown("### 🎯 Configurar Classes")
-
-            num_classes_input = gr.Number(
-                label="Número de Classes",
-                value=2,
-                precision=0,
-                minimum=2,
-                maximum=10
-            )
-            setup_button = gr.Button("🔧 Configurar Classes", variant="primary")
-            setup_output = gr.Textbox(label="📋 Status", lines=2)
-
-            gr.Markdown("### 🏷️ Definir Rótulos")
-
-            # Campos para rótulos dinâmicos
-            label_inputs = []
-            for i in range(10):
-                label_input = gr.Textbox(
-                    label=f"Rótulo da Classe {i}",
-                    placeholder=f"Ex: gato, cachorro, pássaro...",
-                    visible=(i < 2)
-                )
-                label_inputs.append(label_input)
-
-            set_labels_button = gr.Button("🏷️ Definir Rótulos", variant="secondary")
-            labels_output = gr.Textbox(label="📋 Status dos Rótulos")
-
-            # Atualizar visibilidade dos campos
-            def update_label_visibility(num_classes_value):
-                updates = []
-                for i in range(10):
-                    updates.append(gr.update(visible=(i < int(num_classes_value))))
-                return updates
-
-            # Conectar eventos
-            setup_button.click(setup_classes, inputs=num_classes_input, outputs=setup_output)
-            num_classes_input.change(update_label_visibility, inputs=num_classes_input, outputs=label_inputs)
-            set_labels_button.click(set_class_labels, inputs=label_inputs, outputs=labels_output)
-
-        with gr.Tab("2️⃣ Upload de Imagens"):
-            gr.Markdown("### 📤 Upload de Imagens por Classe")
-
             with gr.Row():
-                class_selector = gr.Dropdown(
-                    label="Selecionar Classe",
-                    choices=[(f"Classe 0", 0), (f"Classe 1", 1)],
-                    value=0
+                num_classes_input = gr.Number(
+                    label="Número de Classes", 
+                    value=2, 
+                    minimum=2, 
+                    maximum=5,
+                    precision=0
                 )
-                images_upload = gr.File(
-                    label="Selecionar Imagens",
-                    file_count="multiple",
-                    type="filepath",
-                    file_types=["image"]
-                )
-
-            upload_button = gr.Button("📤 Fazer Upload", variant="primary")
-            upload_output = gr.Textbox(label="📋 Status do Upload")
-
-            # Atualizar dropdown de classes
-            def update_class_dropdown(num_classes_value):
-                choices = []
-                for i in range(int(num_classes_value)):
-                    label = class_labels[i] if i < len(class_labels) else f"Classe {i}"
-                    choices.append((f"{i} - {label}", i))
-                return gr.update(choices=choices, value=0)
-
-            # Conectar eventos
-            upload_button.click(upload_images, inputs=[class_selector, images_upload], outputs=upload_output)
-            num_classes_input.change(update_class_dropdown, inputs=num_classes_input, outputs=class_selector)
-            set_labels_button.click(update_class_dropdown, inputs=num_classes_input, outputs=class_selector)
-
-        with gr.Tab("3️⃣ Preparação & Treinamento"):
-            gr.Markdown("### ⚙️ Configurar Parâmetros")
-
+                setup_button = gr.Button("🔧 Configurar Classes", variant="primary")
+            
+            setup_output = gr.Textbox(label="Status", lines=2)
+            
+            gr.Markdown("### Rótulos das Classes")
+            
             with gr.Row():
-                batch_size = gr.Number(label="Batch Size", value=16, minimum=1, maximum=64)
-                resize_input = gr.Textbox(label="Tamanho da Imagem (formato: 224,224)", value="224,224")
-
+                label0 = gr.Textbox(label="Classe 0", placeholder="Ex: gato")
+                label1 = gr.Textbox(label="Classe 1", placeholder="Ex: cachorro")
+            
+            with gr.Row():
+                label2 = gr.Textbox(label="Classe 2", placeholder="Ex: pássaro", visible=False)
+                label3 = gr.Textbox(label="Classe 3", placeholder="Ex: peixe", visible=False)
+                label4 = gr.Textbox(label="Classe 4", placeholder="Ex: hamster", visible=False)
+            
+            set_labels_button = gr.Button("🏷️ Definir Rótulos")
+            labels_output = gr.Textbox(label="Status dos Rótulos")
+            
+            # Dropdown que será atualizado
+            class_selector = gr.Dropdown(
+                label="Selecionar Classe",
+                choices=[(f"Classe 0", 0), (f"Classe 1", 1)],
+                value=0
+            )
+            
+        with gr.Tab("2️⃣ Upload"):
+            images_upload = gr.File(
+                label="Selecionar Imagens", 
+                file_count="multiple", 
+                file_types=["image"]
+            )
+            upload_button = gr.Button("📤 Fazer Upload", variant="primary")
+            upload_output = gr.Textbox(label="Status do Upload")
+            
+        with gr.Tab("3️⃣ Treinamento"):
+            batch_size = gr.Number(label="Batch Size", value=8, minimum=1, maximum=32)
             prepare_button = gr.Button("⚙️ Preparar Dados", variant="primary")
-            prepare_output = gr.Textbox(label="📋 Status da Preparação", lines=3)
-
-            gr.Markdown("### 🚀 Treinamento")
-
+            prepare_output = gr.Textbox(label="Status", lines=3)
+            
             with gr.Row():
                 model_name = gr.Dropdown(
-                    label="Modelo",
-                    choices=list(model_dict.keys()),
+                    label="Modelo", 
+                    choices=list(MODELS.keys()), 
                     value="MobileNetV2"
                 )
-                epochs = gr.Number(label="Épocas", value=3, minimum=1, maximum=20)
+                epochs = gr.Number(label="Épocas", value=3, minimum=1, maximum=10)
                 lr = gr.Number(label="Learning Rate", value=0.001, minimum=0.0001, maximum=0.1)
-
-            train_button = gr.Button("🚀 Iniciar Treinamento", variant="primary", size="lg")
-            train_output = gr.Textbox(label="📋 Status do Treinamento", lines=10)
-
-            # Conectar eventos
-            prepare_button.click(prepare_data, inputs=[batch_size, resize_input], outputs=prepare_output)
-            train_button.click(start_training, inputs=[model_name, epochs, lr], outputs=train_output)
-
+            
+            train_button = gr.Button("🚀 Treinar", variant="primary")
+            train_output = gr.Textbox(label="Status do Treinamento", lines=10)
+            
         with gr.Tab("4️⃣ Avaliação"):
-            gr.Markdown("### 📊 Avaliar Desempenho do Modelo")
-
-            with gr.Row():
-                eval_button = gr.Button("📊 Avaliar Modelo", variant="primary")
-                cm_button = gr.Button("📈 Matriz de Confusão", variant="secondary")
-
-            eval_output = gr.Textbox(label="📋 Relatório de Avaliação", lines=15)
-            cm_output = gr.Image(label="📈 Matriz de Confusão")
-
-            # Conectar eventos
-            eval_button.click(evaluate_model, outputs=eval_output)
-            cm_button.click(show_confusion_matrix, outputs=cm_output)
-
+            eval_button = gr.Button("📊 Avaliar", variant="primary")
+            eval_output = gr.Textbox(label="Relatório", lines=15)
+            
         with gr.Tab("5️⃣ Predição"):
-            gr.Markdown("### 🔮 Fazer Predições em Novas Imagens")
-
             predict_images_input = gr.File(
-                label="Upload de Imagens para Predição",
-                file_count="multiple",
-                type="filepath",
+                label="Imagens para Predição", 
+                file_count="multiple", 
                 file_types=["image"]
             )
-            predict_button = gr.Button("🔮 Predizer", variant="primary", size="lg")
-            predict_output = gr.Textbox(label="📋 Resultados das Predições", lines=10)
-
-            # Conectar eventos
-            predict_button.click(predict_images, inputs=predict_images_input, outputs=predict_output)
-
+            predict_button = gr.Button("🔮 Predizer", variant="primary")
+            predict_output = gr.Textbox(label="Resultados", lines=10)
+        
+        # Conectar eventos
+        setup_button.click(
+            fn=setup_classes,
+            inputs=[num_classes_input],
+            outputs=[setup_output, class_selector]
+        )
+        
+        set_labels_button.click(
+            fn=set_class_labels,
+            inputs=[label0, label1, label2, label3, label4],
+            outputs=[labels_output, class_selector]
+        )
+        
+        upload_button.click(
+            fn=upload_images,
+            inputs=[class_selector, images_upload],
+            outputs=[upload_output]
+        )
+        
+        prepare_button.click(
+            fn=prepare_data,
+            inputs=[batch_size],
+            outputs=[prepare_output]
+        )
+        
+        train_button.click(
+            fn=start_training,
+            inputs=[model_name, epochs, lr],
+            outputs=[train_output]
+        )
+        
+        eval_button.click(
+            fn=evaluate_model,
+            outputs=[eval_output]
+        )
+        
+        predict_button.click(
+            fn=predict_images,
+            inputs=[predict_images_input],
+            outputs=[predict_output]
+        )
+    
     return demo
 
-# ============================================================================
-# EXECUÇÃO PRINCIPAL
-# ============================================================================
-
 if __name__ == "__main__":
-    print("🎯 Criando interface...")
     demo = create_interface()
-
-    print("🚀 Iniciando aplicação...")
-    demo.launch()
-
-print("✅ Sistema pronto para uso!")
+    demo.launch(server_name="0.0.0.0", server_port=7860)