app.py

# ============================================================================
# SISTEMA DE CLASSIFICAÇÃO DE IMAGENS - HUGGING FACE SPACE
# ============================================================================

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.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("✅ 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
# ============================================================================

# Modelos disponíveis
model_dict = {
    'AlexNet': models.alexnet,
    '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
# ============================================================================

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"

def set_class_labels(*labels):
    """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)}"

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

def prepare_data(batch_size=32, resize_str="224,224"):
    """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(',')))

        # Transformações
        transform = transforms.Compose([
            transforms.Resize(resize),
            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)

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

        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()):
    """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)

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

        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):
                inputs, labels = inputs.to(device), labels.to(device)

                optimizer.zero_grad()
                outputs = 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_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.")

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

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

        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
# ============================================================================

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!
        
        ⚠️ **Nota**: Este é um ambiente temporário. Os dados e modelos serão perdidos ao reiniciar.
        """)

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

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

            prepare_button = gr.Button("⚙️ Preparar Dados", variant="primary")
            prepare_output = gr.Textbox(label="📋 Status da Preparação", lines=3)

            gr.Markdown("### 🚀 Treinamento")

            with gr.Row():
                model_name = gr.Dropdown(
                    label="Modelo",
                    choices=list(model_dict.keys()),
                    value="MobileNetV2"
                )
                epochs = gr.Number(label="Épocas", value=3, minimum=1, maximum=20)
                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)

        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)

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

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