Spaces:
Runtime error
Runtime error
| import os | |
| import shutil | |
| import gradio as gr | |
| import torch | |
| import torch.nn as nn | |
| import torch.optim as optim | |
| from torchvision import datasets, transforms, models | |
| from torch.utils.data import DataLoader, random_split | |
| from PIL import Image | |
| import matplotlib | |
| matplotlib.use('Agg') | |
| import matplotlib.pyplot as plt | |
| import seaborn as sns | |
| import numpy as np | |
| from sklearn.metrics import classification_report, confusion_matrix | |
| import tempfile | |
| import warnings | |
| warnings.filterwarnings("ignore") | |
| print("🖥️ Iniciando sistema...") | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| print(f"Device: {device}") | |
| # Modelos disponíveis | |
| MODELS = { | |
| 'ResNet18': models.resnet18, | |
| 'ResNet34': models.resnet34, | |
| 'ResNet50': models.resnet50, | |
| 'MobileNetV2': models.mobilenet_v2 | |
| } | |
| # Estado global | |
| 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 = ['classe_0', 'classe_1'] | |
| self.num_classes = 2 | |
| self.image_queue = [] | |
| state = AppState() | |
| def setup_classes(num_classes_value): | |
| """Configura número de classes""" | |
| try: | |
| state.num_classes = max(2, min(5, int(num_classes_value))) | |
| state.dataset_path = tempfile.mkdtemp() | |
| state.class_labels = [f'classe_{i}' for i in range(state.num_classes)] | |
| # Criar diretórios | |
| state.class_dirs = [] | |
| for i in range(state.num_classes): | |
| class_dir = os.path.join(state.dataset_path, f'classe_{i}') | |
| os.makedirs(class_dir, exist_ok=True) | |
| state.class_dirs.append(class_dir) | |
| return f"✅ Sistema configurado para {state.num_classes} classes" | |
| except Exception as e: | |
| return f"❌ Erro: {str(e)}" | |
| def set_class_labels(labels_text): | |
| """Define rótulos das classes""" | |
| try: | |
| labels = [label.strip() for label in labels_text.split(',')] | |
| if len(labels) != state.num_classes: | |
| return f"❌ Forneça {state.num_classes} rótulos separados por vírgula" | |
| state.class_labels = labels | |
| return f"✅ Rótulos definidos: {', '.join(state.class_labels)}" | |
| except Exception as e: | |
| return f"❌ Erro: {str(e)}" | |
| def add_images_to_queue(images): | |
| """Adiciona múltiplas imagens à fila""" | |
| if not images: | |
| return "❌ Nenhuma imagem selecionada", len(state.image_queue) | |
| count = 0 | |
| for image_file in images: | |
| try: | |
| if image_file is not None: | |
| # Carregar imagem | |
| img = Image.open(image_file.name).convert('RGB') | |
| state.image_queue.append(img) | |
| count += 1 | |
| except Exception as e: | |
| print(f"Erro processando imagem: {e}") | |
| return f"✅ {count} imagens adicionadas. Total na fila: {len(state.image_queue)}", len(state.image_queue) | |
| def save_queue_to_class(class_id): | |
| """Salva fila de imagens para uma classe""" | |
| try: | |
| if not state.image_queue: | |
| return "❌ Nenhuma imagem na fila" | |
| if not state.class_dirs: | |
| return "❌ Configure as classes primeiro" | |
| class_idx = max(0, min(int(class_id), len(state.class_dirs) - 1)) | |
| class_dir = state.class_dirs[class_idx] | |
| count = 0 | |
| for i, image in enumerate(state.image_queue): | |
| try: | |
| import time | |
| filename = f"img_{int(time.time())}_{i}.jpg" | |
| filepath = os.path.join(class_dir, filename) | |
| image.save(filepath) | |
| count += 1 | |
| except Exception as e: | |
| print(f"Erro salvando imagem {i}: {e}") | |
| state.image_queue = [] # Limpar fila | |
| class_name = state.class_labels[class_idx] | |
| return f"✅ {count} imagens salvas em '{class_name}'" | |
| except Exception as e: | |
| return f"❌ Erro: {str(e)}" | |
| def clear_queue(): | |
| """Limpa a fila""" | |
| state.image_queue = [] | |
| return "✅ Fila limpa", 0 | |
| def prepare_data(batch_size): | |
| """Prepara dados""" | |
| try: | |
| if not state.dataset_path: | |
| return "❌ Configure as classes primeiro" | |
| transform = transforms.Compose([ | |
| transforms.Resize((224, 224)), | |
| transforms.ToTensor(), | |
| transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) | |
| ]) | |
| dataset = datasets.ImageFolder(state.dataset_path, transform=transform) | |
| if len(dataset) < 6: | |
| return f"❌ Poucas imagens ({len(dataset)}). Mínimo: 6" | |
| # Divisão: 70% treino, 20% val, 10% teste | |
| 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) | |
| ) | |
| batch_size = max(1, min(int(batch_size), 32)) | |
| state.train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) | |
| state.val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False) | |
| state.test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False) | |
| return f"✅ Dados preparados:\n• Treino: {train_size}\n• Validação: {val_size}\n• Teste: {test_size}\n• Batch size: {batch_size}" | |
| except Exception as e: | |
| return f"❌ Erro: {str(e)}" | |
| def train_model(model_name, epochs, lr): | |
| """Treina modelo""" | |
| try: | |
| if state.train_loader is None: | |
| return "❌ Prepare os dados primeiro" | |
| # Carregar modelo | |
| state.model = MODELS[model_name](pretrained=True) | |
| # Adaptar última camada | |
| if hasattr(state.model, 'fc'): | |
| state.model.fc = nn.Linear(state.model.fc.in_features, state.num_classes) | |
| elif hasattr(state.model, 'classifier'): | |
| if isinstance(state.model.classifier, nn.Sequential): | |
| state.model.classifier[-1] = nn.Linear(state.model.classifier[-1].in_features, state.num_classes) | |
| state.model = state.model.to(device) | |
| criterion = nn.CrossEntropyLoss() | |
| optimizer = optim.Adam(state.model.parameters(), lr=float(lr)) | |
| results = [f"🚀 Treinando {model_name}"] | |
| state.model.train() | |
| for epoch in range(int(epochs)): | |
| running_loss = 0.0 | |
| correct = 0 | |
| total = 0 | |
| for inputs, labels in state.train_loader: | |
| inputs, labels = inputs.to(device), labels.to(device) | |
| optimizer.zero_grad() | |
| outputs = state.model(inputs) | |
| loss = criterion(outputs, labels) | |
| loss.backward() | |
| optimizer.step() | |
| running_loss += loss.item() | |
| _, predicted = torch.max(outputs, 1) | |
| total += labels.size(0) | |
| correct += (predicted == labels).sum().item() | |
| epoch_loss = running_loss / len(state.train_loader) | |
| epoch_acc = 100. * correct / total | |
| 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: {str(e)}" | |
| def evaluate_model(): | |
| """Avalia modelo""" | |
| try: | |
| if state.model is None or state.test_loader is None: | |
| return "❌ Modelo/dados não disponíveis" | |
| state.model.eval() | |
| all_preds = [] | |
| all_labels = [] | |
| with torch.no_grad(): | |
| for inputs, labels in state.test_loader: | |
| inputs, labels = inputs.to(device), labels.to(device) | |
| outputs = state.model(inputs) | |
| _, preds = torch.max(outputs, 1) | |
| all_preds.extend(preds.cpu().numpy()) | |
| all_labels.extend(labels.cpu().numpy()) | |
| report = classification_report(all_labels, all_preds, target_names=state.class_labels, zero_division=0) | |
| return f"📊 RELATÓRIO DE AVALIAÇÃO:\n\n{report}" | |
| except Exception as e: | |
| return f"❌ Erro: {str(e)}" | |
| def generate_confusion_matrix(): | |
| """Gera matriz de confusão""" | |
| try: | |
| if state.model is None or state.test_loader is None: | |
| return None | |
| state.model.eval() | |
| all_preds = [] | |
| all_labels = [] | |
| with torch.no_grad(): | |
| for inputs, labels in state.test_loader: | |
| inputs, labels = inputs.to(device), labels.to(device) | |
| outputs = state.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) | |
| plt.figure(figsize=(8, 6)) | |
| sns.heatmap(cm, annot=True, fmt="d", cmap="Blues", | |
| xticklabels=state.class_labels, | |
| yticklabels=state.class_labels) | |
| plt.xlabel('Predições') | |
| plt.ylabel('Valores Reais') | |
| plt.title('Matriz de Confusão') | |
| plt.tight_layout() | |
| temp_path = tempfile.NamedTemporaryFile(suffix='.png', delete=False).name | |
| plt.savefig(temp_path, dpi=150, bbox_inches='tight') | |
| plt.close() | |
| return temp_path | |
| except Exception as e: | |
| return None | |
| def predict_image(image): | |
| """Prediz imagem""" | |
| try: | |
| if state.model is None: | |
| return "❌ Treine o modelo primeiro" | |
| if image is None: | |
| return "❌ Selecione uma imagem" | |
| transform = transforms.Compose([ | |
| transforms.Resize((224, 224)), | |
| transforms.ToTensor(), | |
| transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) | |
| ]) | |
| img_tensor = transform(image).unsqueeze(0).to(device) | |
| state.model.eval() | |
| with torch.no_grad(): | |
| outputs = state.model(img_tensor) | |
| probs = torch.nn.functional.softmax(outputs[0], dim=0) | |
| _, predicted = torch.max(outputs, 1) | |
| class_id = predicted.item() | |
| confidence = probs[class_id].item() * 100 | |
| class_name = state.class_labels[class_id] | |
| return f"🎯 Predição: {class_name}\n📊 Confiança: {confidence:.2f}%" | |
| except Exception as e: | |
| return f"❌ Erro: {str(e)}" | |
| # Interface usando Gradio 3.x (sintaxe correta) | |
| def create_interface(): | |
| # Interface com abas usando Gradio 3.x | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# 🖼️ Sistema de Classificação de Imagens Completo") | |
| gr.Markdown("**Versão estável sem bugs - Funcionalidade completa mantida**") | |
| with gr.Tab("1️⃣ Configuração"): | |
| gr.Markdown("### 🎯 Configurar Classes") | |
| num_classes_input = gr.Number(value=2, label="Número de Classes (2-5)") | |
| setup_btn = gr.Button("🔧 Configurar Classes", variant="primary") | |
| setup_output = gr.Textbox(label="Status da Configuração") | |
| gr.Markdown("### 🏷️ Definir Rótulos") | |
| labels_input = gr.Textbox(value="gato,cachorro", label="Rótulos (separados por vírgula)") | |
| labels_btn = gr.Button("🏷️ Definir Rótulos") | |
| labels_output = gr.Textbox(label="Status dos Rótulos") | |
| # Conectar eventos | |
| setup_btn.click(setup_classes, inputs=[num_classes_input], outputs=[setup_output]) | |
| labels_btn.click(set_class_labels, inputs=[labels_input], outputs=[labels_output]) | |
| with gr.Tab("2️⃣ Upload de Imagens"): | |
| gr.Markdown("### 📤 Upload Múltiplo via Fila") | |
| images_upload = gr.File(file_count="multiple", label="Selecionar Múltiplas Imagens", file_types=["image"]) | |
| add_btn = gr.Button("➕ Adicionar à Fila") | |
| with gr.Row(): | |
| queue_output = gr.Textbox(label="Status da Fila") | |
| queue_count_output = gr.Number(label="Total na Fila", value=0) | |
| gr.Markdown("### 💾 Salvar por Classe") | |
| with gr.Row(): | |
| class_id_input = gr.Number(value=0, label="Classe de Destino (0, 1, 2...)") | |
| save_btn = gr.Button("💾 Salvar Fila na Classe", variant="primary") | |
| clear_btn = gr.Button("🗑️ Limpar Fila") | |
| save_output = gr.Textbox(label="Status do Upload") | |
| # Conectar eventos | |
| add_btn.click(add_images_to_queue, inputs=[images_upload], outputs=[queue_output, queue_count_output]) | |
| save_btn.click(save_queue_to_class, inputs=[class_id_input], outputs=[save_output]) | |
| clear_btn.click(clear_queue, outputs=[queue_output, queue_count_output]) | |
| with gr.Tab("3️⃣ Preparação e Treinamento"): | |
| gr.Markdown("### ⚙️ Preparar Dados") | |
| batch_size_input = gr.Number(value=8, label="Batch Size") | |
| prepare_btn = gr.Button("⚙️ Preparar Dados", variant="primary") | |
| prepare_output = gr.Textbox(label="Status da Preparação", lines=4) | |
| gr.Markdown("### 🚀 Configurar e Treinar Modelo") | |
| with gr.Row(): | |
| model_input = gr.Dropdown(choices=list(MODELS.keys()), value="MobileNetV2", label="Modelo") | |
| epochs_input = gr.Number(value=5, label="Épocas") | |
| lr_input = gr.Number(value=0.001, label="Learning Rate") | |
| train_btn = gr.Button("🚀 Iniciar Treinamento", variant="primary") | |
| train_output = gr.Textbox(label="Status do Treinamento", lines=8) | |
| # Conectar eventos | |
| prepare_btn.click(prepare_data, inputs=[batch_size_input], outputs=[prepare_output]) | |
| train_btn.click(train_model, inputs=[model_input, epochs_input, lr_input], outputs=[train_output]) | |
| with gr.Tab("4️⃣ Avaliação do Modelo"): | |
| gr.Markdown("### 📊 Avaliar Desempenho") | |
| with gr.Row(): | |
| eval_btn = gr.Button("📊 Avaliar Modelo", variant="primary") | |
| matrix_btn = gr.Button("📈 Gerar Matriz de Confusão") | |
| eval_output = gr.Textbox(label="Relatório de Avaliação", lines=12) | |
| matrix_output = gr.Image(label="Matriz de Confusão") | |
| # Conectar eventos | |
| eval_btn.click(evaluate_model, outputs=[eval_output]) | |
| matrix_btn.click(generate_confusion_matrix, outputs=[matrix_output]) | |
| with gr.Tab("5️⃣ Predição"): | |
| gr.Markdown("### 🔮 Predizer Novas Imagens") | |
| predict_image_input = gr.Image(type="pil", label="Imagem para Predição") | |
| predict_btn = gr.Button("🔮 Fazer Predição", variant="primary") | |
| predict_output = gr.Textbox(label="Resultado da Predição", lines=3) | |
| # Conectar eventos | |
| predict_btn.click(predict_image, inputs=[predict_image_input], outputs=[predict_output]) | |
| # Informações adicionais | |
| with gr.Tab("ℹ️ Informações"): | |
| gr.Markdown(""" | |
| ## 📋 Como Usar Este Sistema | |
| ### 1️⃣ **Configuração Inicial** | |
| - Defina o número de classes (2-5) | |
| - Configure rótulos personalizados separados por vírgula | |
| ### 2️⃣ **Upload de Imagens** | |
| - Selecione múltiplas imagens | |
| - Adicione à fila | |
| - Escolha a classe de destino (0, 1, 2...) | |
| - Salve a fila na classe escolhida | |
| - Repita para todas as classes | |
| ### 3️⃣ **Treinamento** | |
| - Configure batch size (recomendado: 8-16) | |
| - Prepare os dados | |
| - Escolha modelo (MobileNetV2 = mais rápido) | |
| - Configure épocas (recomendado: 3-10) | |
| - Inicie o treinamento | |
| ### 4️⃣ **Avaliação** | |
| - Avalie o modelo para ver métricas | |
| - Gere matriz de confusão para análise visual | |
| ### 5️⃣ **Predição** | |
| - Teste com novas imagens | |
| - Veja predições com níveis de confiança | |
| ## 🎯 **Dicas para Melhores Resultados** | |
| - Use pelo menos 10-20 imagens por classe | |
| - Imagens bem balanceadas entre classes | |
| - Imagens claras e bem iluminadas | |
| - Varie poses, ângulos e ambientes | |
| ## 🔧 **Modelos Disponíveis** | |
| - **MobileNetV2**: Rápido, ideal para prototipagem | |
| - **ResNet18**: Bom equilíbrio velocidade/precisão | |
| - **ResNet34/50**: Maior precisão, mais lento | |
| """) | |
| return demo | |
| if __name__ == "__main__": | |
| print("🎯 Criando interface...") | |
| demo = create_interface() | |
| print("🚀 Iniciando aplicação...") | |
| demo.launch(server_name="0.0.0.0", server_port=7860) |