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) |