Update app.py

app.py

@@ -1,9 +1,10 @@
 import gradio as gr
 import torch
 import torch.nn as nn
-from torchvision import models, transforms
+from torchvision import transforms, models
 from PIL import Image
 
+# Define the model class
 class CustomEfficientNet(nn.Module):
     def __init__(self, num_classes, num_layers, neurons_per_layer):
         super(CustomEfficientNet, self).__init__()
@@ -28,6 +29,7 @@ class CustomEfficientNet(nn.Module):
         x = self.custom_classifier(x)
         return x
 
+# Function to create and load the model
 def create_model(num_classes, num_layers, neurons_per_layer):
     model = CustomEfficientNet(num_classes, num_layers, neurons_per_layer)
     return model
@@ -62,6 +64,7 @@ class_names = ['Coeur 1', 'Coeur 10', 'Coeur 2', 'Coeur 3', 'Coeur 4', 'Coeur 5'
                'Trefle Roi', 'Trefle Valet', 'carreau 1', 'carreau 10', 'carreau 2', 'carreau 3', 'carreau 4', 'carreau 5',
                'carreau 6', 'carreau 7', 'carreau 8', 'carreau 9', 'carreau Dame', 'carreau Roi', 'carreau Valet']
 
+# Define the prediction function
 def predict(image):
     image = transform(image).unsqueeze(0)
     with torch.no_grad():
@@ -69,12 +72,21 @@ def predict(image):
         _, predicted = torch.max(outputs, 1)
         return class_names[predicted[0]]
 
+# Example images
+examples = [
+    ['trefledame.jpg'],
+    ['coeurroi.jpg'],
+    ['coeur3.jpg']
+]
+
 # Create the Gradio interface
 iface = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),
     outputs="label",
-    description="Upload an image to classify"
+    description="Upload an image to classify",
+    examples=examples
 )
 
+# Launch the interface
 iface.launch()