Subindo arquivos

README.md

@@ -1,13 +1,20 @@
 ---
-title: Image Classifier Interactive
-emoji: 🖼
-colorFrom: purple
-colorTo: red
+title: interactive-image-classifier
+emoji: 🤖
+colorFrom: blue
+colorTo: green
 sdk: gradio
-sdk_version: 4.26.0
+sdk_version: "4.12.0"
 app_file: app.py
 pinned: false
-license: ecl-2.0
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Image Enhancer
+
+Upload an image (.jpg, .png) per class, follow the interactive process for image classification, train, evaluate, predict and export
+
+## Versão teste 1 (16/05)
+
+- Ramon Mayor Martins
+- E-mail: [[email protected]](mailto:[email protected])
+

app.py

@@ -1,146 +1,360 @@
 import gradio as gr
+import os
+import shutil
+from sklearn.metrics import classification_report, confusion_matrix
+import matplotlib.pyplot as plt
+import seaborn as sns
 import numpy as np
-import random
-from diffusers import DiffusionPipeline
+import io
 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 joblib  #  .pkl
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-if torch.cuda.is_available():
-    torch.cuda.max_memory_allocated(device=device)
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
-    pipe.enable_xformers_memory_efficient_attention()
-    pipe = pipe.to(device)
-else: 
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)
-    pipe = pipe.to(device)
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
-
-def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
-
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-        
-    generator = torch.Generator().manual_seed(seed)
-    
-    image = pipe(
-        prompt = prompt, 
-        negative_prompt = negative_prompt,
-        guidance_scale = guidance_scale, 
-        num_inference_steps = num_inference_steps, 
-        width = width, 
-        height = height,
-        generator = generator
-    ).images[0] 
-    
-    return image
-
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
-]
-
-css="""
-#col-container {
-    margin: 0 auto;
-    max-width: 520px;
+#
+model_dict = {
+    'AlexNet': models.alexnet,
+    'ResNet18': models.resnet18,
+    'ResNet34': models.resnet34,
+    'ResNet50': models.resnet50,
+    'MobileNetV2': models.mobilenet_v2
 }
-"""
-
-if torch.cuda.is_available():
-    power_device = "GPU"
-else:
-    power_device = "CPU"
-
-with gr.Blocks(css=css) as demo:
-    
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(f"""
-        # Text-to-Image Gradio Template
-        Currently running on {power_device}.
-        """)
-        
-        with gr.Row():
-            
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
-            
-            run_button = gr.Button("Run", scale=0)
-        
-        result = gr.Image(label="Result", show_label=False)
-
-        with gr.Accordion("Advanced Settings", open=False):
-            
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
-            )
-            
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
-            )
-            
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-            
-            with gr.Row():
-                
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=512,
-                )
-                
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=512,
-                )
-            
-            with gr.Row():
-                
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=0.0,
-                )
-                
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=12,
-                    step=1,
-                    value=2,
-                )
-        
-        gr.Examples(
-            examples = examples,
-            inputs = [prompt]
-        )
-
-    run_button.click(
-        fn = infer,
-        inputs = [prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
-        outputs = [result]
-    )
-
-demo.queue().launch()
+
+#
+model = None
+train_loader = None
+val_loader = None
+test_loader = None
+dataset_path = 'dataset'
+class_dirs = []
+test_dataset_path = 'test_dataset'
+test_class_dirs = []
+num_classes = 2  # 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# 
+def setup_classes(num_classes_value):
+    global class_dirs, dataset_path, num_classes
+
+    num_classes = int(num_classes_value)  # 
+    # 
+    if os.path.exists(dataset_path):
+        shutil.rmtree(dataset_path)
+    os.makedirs(dataset_path)
+
+    # 
+    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 upload_images(class_id, images):
+    class_dir = class_dirs[int(class_id)]
+    for image in images:
+        shutil.copy(image, class_dir)
+    return f"Imagens salvas na classe {class_id}."
+
+# 
+def prepare_data(batch_size=32, resize=(224, 224)):
+    global train_loader, val_loader, test_loader, num_classes
+
+    # 
+    transform = transforms.Compose([
+        transforms.Resize(resize),
+        transforms.ToTensor(),
+    ])
+
+    dataset = datasets.ImageFolder(dataset_path, transform=transform)
+
+    if len(dataset.classes) != num_classes:
+        return f"Erro: Número de classes detectadas ({len(dataset.classes)}) não corresponde ao número esperado ({num_classes}). Verifique suas imagens."
+
+    # 
+    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])
+
+    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
+    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
+    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
+
+    return "Preparação dos dados concluída com sucesso."
+
+# 
+def start_training(model_name, epochs, lr):
+    global model, train_loader, val_loader, device
+
+    if train_loader is None or val_loader is None:
+        return "Erro: Dados não preparados."
+
+    model = model_dict[model_name](pretrained=True)
+    model.fc = nn.Linear(model.fc.in_features, num_classes)
+    model = model.to(device)
+
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.Adam(model.parameters(), lr=float(lr))
+
+    for epoch in range(int(epochs)):
+        model.train()
+        running_loss = 0.0
+        for inputs, labels in 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()
+
+        print(f"Epoch {epoch+1}/{epochs}, Loss: {running_loss/len(train_loader)}")
+
+    torch.save(model.state_dict(), 'modelo.pth')
+    return f"Treinamento concluído com sucesso. Modelo salvo."
+
+# 
+def evaluate_model(loader):
+    global model, device, num_classes
+
+    if model is None:
+        return "Erro: Modelo não treinado."
+
+    if loader is None:
+        return "Erro: Conjunto de dados de teste não está preparado."
+
+    model.eval()
+    all_preds = []
+    all_labels = []
+    try:
+        with torch.no_grad():
+            for inputs, labels in 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())
+
+        report = classification_report(all_labels, all_preds, labels=list(range(num_classes)), target_names=[f"class_{i}" for i in range(num_classes)], zero_division=0)
+        return report
+    except Exception as e:
+        return f"Erro durante a avaliação: {str(e)}"
+
+# 
+def show_confusion_matrix(loader):
+    global model, device, num_classes
+
+    if model is None:
+        return "Erro: Modelo não treinado."
+
+    model.eval()
+    all_preds = []
+    all_labels = []
+    with torch.no_grad():
+        for inputs, labels in 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=list(range(num_classes)))
+
+    plt.figure(figsize=(6, 4.8))
+    sns.heatmap(cm, annot=True, fmt="d", cmap="Blues", xticklabels=[f"class_{i}" for i in range(num_classes)], yticklabels=[f"class_{i}" for i in range(num_classes)])
+    plt.xlabel('Predictions')
+    plt.ylabel('Actuals')
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png')
+    plt.close()
+    buf.seek(0)
+    return Image.open(buf)
+
+# 
+def predict_images(images):
+    global model, device, num_classes
+
+    if model is None:
+        return "Erro: Modelo não treinado."
+
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+    ])
+
+    model.eval()
+    results = []
+
+    for image in images:
+        try:
+            img = transform(Image.open(image)).unsqueeze(0).to(device)
+            with torch.no_grad():
+                outputs = model(img)
+                _, preds = torch.max(outputs, 1)
+                predicted_class = preds.item()
+                results.append(f"Imagem {os.path.basename(image)} - Classe prevista: class_{predicted_class}")
+        except Exception as e:
+            results.append(f"Erro ao processar a imagem {image}: {str(e)}")
+
+    return results
+
+# 
+def export_model(format):
+    global model
+
+    if model is None:
+        return "Erro: Modelo não treinado."
+
+    file_path = f"modelo_exportado.{format}"
+    if format == "pth":
+        torch.save(model.state_dict(), file_path)
+    elif format == "onnx":
+        try:
+            dummy_input = torch.randn(1, 3, 224, 224).to(device)
+            torch.onnx.export(model, dummy_input, file_path, export_params=True, opset_version=10, input_names=['input'], output_names=['output'])
+        except Exception as e:
+            return f"Erro ao exportar para ONNX: {str(e)}"
+    elif format == "pkl":
+        joblib.dump(model, file_path)
+    else:
+        return f"Formato {format} não suportado."
+
+    return f"Modelo exportado com sucesso para {file_path}"
+
+# 
+def setup_test_classes():
+    global test_class_dirs, test_dataset_path
+
+    if os.path.exists(test_dataset_path):
+        shutil.rmtree(test_dataset_path)
+    os.makedirs(test_dataset_path)
+
+    # 
+    test_class_dirs = [os.path.join(test_dataset_path, f'class_{i}') for i in range(num_classes)]
+    for class_dir in test_class_dirs:
+        os.makedirs(class_dir)
+
+    return f"Criados {num_classes} diretórios para classes de teste."
+
+# 
+def upload_test_images(class_id, images):
+    class_dir = test_class_dirs[int(class_id)]
+    for image in images:
+        shutil.copy(image, class_dir)
+    return f"Imagens de teste salvas na classe {class_id}."
+
+# 
+def prepare_test_data(batch_size=32, resize=(224, 224)):
+    global test_loader, num_classes
+
+    transform = transforms.Compose([
+        transforms.Resize(resize),
+        transforms.ToTensor(),
+    ])
+
+    test_dataset = datasets.ImageFolder(test_dataset_path, transform=transform)
+
+    if len(test_dataset.classes) != num_classes:
+        return f"Erro: Número de classes detectadas ({len(test_dataset.classes)}) não corresponde ao número esperado ({num_classes}). Verifique suas imagens."
+
+    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
+
+    return "Preparação dos dados de teste concluída com sucesso."
+
+# 
+def main():
+    with gr.Blocks() as demo:
+        gr.Markdown("# Image Classification Training")
+
+        with gr.Tab("Configurar Classes"):
+            num_classes_input = gr.Number(label="Número de Classes", value=2, precision=0)
+            setup_button = gr.Button("Configurar Classes")
+            setup_output = gr.Textbox()
+            setup_button.click(setup_classes, inputs=num_classes_input, outputs=setup_output)
+
+        with gr.Tab("Upload de Imagens"):
+            upload_inputs = []
+            for i in range(num_classes):
+                with gr.Column():
+                    gr.Markdown(f"### Classe {i}")
+                    class_id = gr.Number(label=f"ID da Classe {i}", value=i, precision=0)
+                    images = gr.File(label="Upload de Imagens", file_count="multiple", type="filepath")
+                    upload_button = gr.Button("Upload")
+                    upload_output = gr.Textbox()
+
+                    upload_inputs.append((class_id, images, upload_button, upload_output))
+                    upload_button.click(upload_images, inputs=[class_id, images], outputs=upload_output)
+
+        with gr.Tab("Preparação de Dados"):
+            batch_size = gr.Number(label="Tamanho do Batch", value=32)
+            resize = gr.Textbox(label="Resize (Ex: 224,224)", value="224,224")
+            prepare_button = gr.Button("Preparar Dados")
+            prepare_output = gr.Textbox()
+            prepare_button.click(lambda batch_size, resize: prepare_data(batch_size=batch_size, resize=tuple(map(int, resize.split(',')))), inputs=[batch_size, resize], outputs=prepare_output)
+
+        with gr.Tab("Treinamento"):
+            model_name = gr.Dropdown(label="Modelo", choices=list(model_dict.keys()))
+            epochs = gr.Number(label="Épocas", value=30)
+            lr = gr.Number(label="Taxa de Aprendizado", value=0.001)
+            train_button = gr.Button("Iniciar Treinamento")
+            train_output = gr.Textbox()
+            train_button.click(start_training, inputs=[model_name, epochs, lr], outputs=train_output)
+
+        with gr.Tab("Avaliação do Modelo"):
+            eval_button = gr.Button("Avaliar Modelo")
+            eval_output = gr.Textbox()
+            eval_button.click(lambda: evaluate_model(test_loader), outputs=eval_output)
+
+            cm_button = gr.Button("Mostrar Matriz de Confusão")
+            cm_output = gr.Image()
+            cm_button.click(lambda: show_confusion_matrix(test_loader), outputs=cm_output)
+
+        with gr.Tab("Predição e Avaliação"):
+            predict_images_input = gr.File(label="Upload de Imagens para Predição", file_count="multiple", type="filepath")
+            predict_button = gr.Button("Predizer")
+            predict_output = gr.Textbox()
+            predict_button.click(predict_images, inputs=predict_images_input, outputs=predict_output)
+
+            gr.Markdown("### Upload de Imagens de Teste")
+            setup_test_button = gr.Button("Configurar Diretórios de Teste")
+            setup_test_output = gr.Textbox()
+            setup_test_button.click(setup_test_classes, outputs=setup_test_output)
+
+            upload_test_inputs = []
+            for i in range(num_classes):
+                with gr.Column():
+                    gr.Markdown(f"### Classe de Teste {i}")
+                    test_class_id = gr.Number(label=f"ID da Classe {i}", value=i, precision=0)
+                    test_images = gr.File(label="Upload de Imagens de Teste", file_count="multiple", type="filepath")
+                    upload_test_button = gr.Button("Upload Imagens de Teste")
+                    upload_test_output = gr.Textbox()
+
+                    upload_test_inputs.append((test_class_id, test_images, upload_test_button, upload_test_output))
+                    upload_test_button.click(upload_test_images, inputs=[test_class_id, test_images], outputs=upload_test_output)
+
+            prepare_test_button = gr.Button("Preparar Dados de Teste")
+            prepare_test_output = gr.Textbox()
+            prepare_test_button.click(lambda batch_size, resize: prepare_test_data(batch_size=batch_size, resize=tuple(map(int, resize.split(',')))), inputs=[batch_size, resize], outputs=prepare_test_output)
+
+            eval_test_button = gr.Button("Avaliar Conjunto de Teste")
+            eval_test_output = gr.Textbox()
+            eval_test_button.click(lambda: evaluate_model(test_loader), outputs=eval_test_output)
+
+            cm_test_button = gr.Button("Mostrar Matriz de Confusão do Conjunto de Teste")
+            cm_test_output = gr.Image()
+            cm_test_button.click(lambda: show_confusion_matrix(test_loader), outputs=cm_test_output)
+
+        with gr.Tab("Exportação"):
+            export_format = gr.Radio(label="Formato", choices=["pth", "onnx", "pkl"])
+            export_button = gr.Button("Exportar Modelo")
+            export_output = gr.Textbox()
+            export_button.click(export_model, inputs=export_format, outputs=export_output)
+
+    demo.launch()
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -1,6 +1,12 @@
-accelerate
-diffusers
-invisible_watermark
-torch
-transformers
-xformers
+torch==1.11.0
+torchvision==0.12.0
+scikit-learn==0.24.2
+matplotlib==3.4.2
+seaborn==0.11.1
+numpy==1.21.0
+Pillow==8.2.0
+gradio==4.12.0
+joblib==1.0.1
+onnx==1.10.1
+onnx-tf==1.8.0
+tensorflow==2.16.0