Update app.py

app.py

@@ -3,16 +3,42 @@ from torchvision import transforms
 from PIL import Image
 import streamlit as st
 import json
+from torchvision.models import efficientnet_b7, EfficientNet_B7_Weights
+
 
 
 # Charger les noms des classes
 with open("class_names.json", "r") as f:
     class_names = json.load(f)
+    
+# Charger le modèle avec des poids pré-entraînés
+weights = EfficientNet_B7_Weights.DEFAULT
+base_model = efficientnet_b7(weights=weights)
+
+# Adapter le modèle pour la classification
+class CustomEfficientNet(nn.Module):
+    def __init__(self, base_model, num_classes):
+        super(CustomEfficientNet, self).__init__()
+        self.base = nn.Sequential(*list(base_model.children())[:-2])
+        self.global_avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.fc1 = nn.Linear(2560, 512)
+        self.relu = nn.ReLU()
+        self.fc2 = nn.Linear(512, num_classes)
+
+    def forward(self, x):
+        x = self.base(x)
+        x = self.global_avg_pool(x)
+        x = x.view(x.size(0), -1)
+        x = self.relu(self.fc1(x))
+        x = self.fc2(x)
+        return x
+
+# Définir le modèle final
+num_classes = 2
+model = CustomEfficientNet(base_model, num_classes).to("cuda" if torch.cuda.is_available() else "cpu")
+model.load_state_dict(torch.load("efficientnet_b7_best.pth",weights_only=False))
+model.eval()  # Passer le modèle en mode évaluation
 
-# Charger le modèle
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = torch.load("efficientnet_b7_best.pth", map_location=device)
-model.eval()  # Mode évaluation
 
 # Définir la taille de l'image
 image_size = (224, 224)