app

app.py

@@ -1,5 +1,8 @@
 import streamlit as st
 import joblib
+import torch
+import torch.nn as nn
+from torchvision.models import resnet18
 from PIL import Image
 import torchvision.transforms as transforms
 from huggingface_hub import hf_hub_download
@@ -16,22 +19,34 @@ def load_model():
     label_encoder = joblib.load(label_encoder_path)
     return model, label_encoder
 
-# Função para pré-processar a imagem
-def preprocess_image(image):
+# Função para carregar a CNN (ResNet18) para extração de características
+@st.cache_resource
+def load_cnn():
+    # Carregar a ResNet18 pré-treinada e remover a última camada
+    cnn_model = resnet18(pretrained=True)
+    cnn_model = nn.Sequential(*list(cnn_model.children())[:-1])  # Remove a última camada
+    cnn_model.eval()  # Colocar o modelo em modo de avaliação
+    return cnn_model
+
+# Função para pré-processar a imagem e extrair características
+def extract_features(image, cnn_model):
     transform = transforms.Compose([
         transforms.Resize((256, 256)),
         transforms.ToTensor(),
-        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # Normalização para a ResNet18
     ])
-    image = transform(image).flatten().numpy().reshape(1, -1)
-    return image
+    image = transform(image).unsqueeze(0)  # Adicionar dimensão do batch
+    with torch.no_grad():  # Desativar cálculo de gradientes
+        features = cnn_model(image)  # Extrair características
+    return features.flatten().numpy().reshape(1, -1)  # Achatar e converter para numpy array
 
 # Interface da aplicação
 st.title("Classificador de Animais")
 st.write("Envie uma imagem de um animal (gato, cachorro ou pássaro) para classificação.")
 
-# Carregar o modelo e o LabelEncoder
+# Carregar o modelo, o LabelEncoder e a CNN
 model, label_encoder = load_model()
+cnn_model = load_cnn()
 
 # Upload da imagem
 uploaded_file = st.file_uploader("Escolha uma imagem...", type=["jpg", "jpeg", "png"])
@@ -39,11 +54,11 @@ if uploaded_file is not None:
     image = Image.open(uploaded_file)
     st.image(image, caption="Imagem enviada", use_column_width=True)
 
-    # Pré-processar a imagem
-    image_processed = preprocess_image(image)
+    # Extrair características da imagem usando a CNN
+    image_features = extract_features(image, cnn_model)
 
     # Fazer a previsão
-    prediction = model.predict(image_processed)
+    prediction = model.predict(image_features)
     predicted_class = label_encoder.inverse_transform(prediction)[0]
 
     # Exibir o resultado