Subindo arquivos7

app.py

@@ -12,9 +12,9 @@ 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
+import joblib  # Para salvar como .pkl
 
-#
+# Modelos para seleção
 model_dict = {
     'AlexNet': models.alexnet,
     'ResNet18': models.resnet18,
@@ -23,7 +23,7 @@ model_dict = {
     'MobileNetV2': models.mobilenet_v2
 }
 
-#
+# Variáveis globais
 model = None
 train_loader = None
 val_loader = None
@@ -32,38 +32,38 @@ dataset_path = 'dataset'
 class_dirs = []
 test_dataset_path = 'test_dataset'
 test_class_dirs = []
-num_classes = 2  # 
+num_classes = 2  # Valor padrão para o número de classes
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# 
+# Função para configurar as classes
 def setup_classes(num_classes_value):
     global class_dirs, dataset_path, num_classes
 
-    num_classes = int(num_classes_value)  # 
-    # 
+    num_classes = int(num_classes_value)  # Atualizar a variável global num_classes
+    # Limpar diretório antigo se existir
     if os.path.exists(dataset_path):
         shutil.rmtree(dataset_path)
     os.makedirs(dataset_path)
 
-    # 
+    # 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."
 
-# 
+# Função para upload de imagens
 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}."
 
-# 
+# Função para preparação dos dados
 def prepare_data(batch_size=32, resize=(224, 224)):
     global train_loader, val_loader, test_loader, num_classes
 
-    # 
+    # Transformações para os dados de treinamento e validação
     transform = transforms.Compose([
         transforms.Resize(resize),
         transforms.ToTensor(),
@@ -74,7 +74,7 @@ def prepare_data(batch_size=32, resize=(224, 224)):
     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."
 
-    # 
+    # Divisão do conjunto de dados em treinamento, validação e teste
     train_size = int(0.7 * len(dataset))
     val_size = int(0.2 * len(dataset))
     test_size = len(dataset) - train_size - val_size
@@ -86,7 +86,7 @@ def prepare_data(batch_size=32, resize=(224, 224)):
 
     return "Preparação dos dados concluída com sucesso."
 
-# 
+# Função para iniciar o treinamento
 def start_training(model_name, epochs, lr):
     global model, train_loader, val_loader, device
 
@@ -117,7 +117,7 @@ def start_training(model_name, epochs, lr):
     torch.save(model.state_dict(), 'modelo.pth')
     return f"Treinamento concluído com sucesso. Modelo salvo."
 
-# 
+# Função para avaliação do modelo com conjunto de teste
 def evaluate_model(loader):
     global model, device, num_classes
 
@@ -144,7 +144,7 @@ def evaluate_model(loader):
     except Exception as e:
         return f"Erro durante a avaliação: {str(e)}"
 
-# 
+# Função para mostrar a matriz de confusão
 def show_confusion_matrix(loader):
     global model, device, num_classes
 
@@ -174,7 +174,7 @@ def show_confusion_matrix(loader):
     buf.seek(0)
     return Image.open(buf)
 
-# 
+# Função para predição de imagens desconhecidas
 def predict_images(images):
     global model, device, num_classes
 
@@ -202,7 +202,7 @@ def predict_images(images):
 
     return results
 
-# 
+# Função para exportar o modelo
 def export_model(format):
     global model
 
@@ -225,7 +225,7 @@ def export_model(format):
 
     return f"Modelo exportado com sucesso para {file_path}"
 
-# 
+# Função para configurar diretórios de teste
 def setup_test_classes():
     global test_class_dirs, test_dataset_path
 
@@ -233,21 +233,21 @@ def setup_test_classes():
         shutil.rmtree(test_dataset_path)
     os.makedirs(test_dataset_path)
 
-    # 
+    # Criar diretórios para cada classe
     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."
 
-# 
+# Função para upload de imagens 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}."
 
-# 
+# Função para preparar dados de teste
 def prepare_test_data(batch_size=32, resize=(224, 224)):
     global test_loader, num_classes
 
@@ -265,7 +265,7 @@ def prepare_test_data(batch_size=32, resize=(224, 224)):
 
     return "Preparação dos dados de teste concluída com sucesso."
 
-# 
+# Interface Gradio
 def main():
     with gr.Blocks() as demo:
         gr.Markdown("# Image Classification Training")

requirements.txt

@@ -9,8 +9,10 @@ joblib==1.1.0
 onnx==1.10.2
 onnx-tf==1.9.0
 tensorflow==2.9.1
+tensorflow-addons
 tensorflowjs==3.20.0
 numpy==1.21.6
 Cython==0.29.24
 
 
+