Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import torch
 from transformers import ViTFeatureExtractor, ViTForImageClassification
 from PIL import Image
+import numpy as np
 
 # Définir le modèle et le feature extractor
 model_name ="KhadijaAsehnoune12/ViTOrangeLeafDiseaseClassifier"
@@ -22,7 +23,28 @@ id2label = {
     "9": "Trou de balle"
 }
 
+def remove_background(image):
+    # Convert image to numpy array
+    img_array = np.array(image)
+    
+    # Calculate the threshold to separate background from foreground
+    threshold = np.mean(img_array) + 2 * np.std(img_array)
+    
+    # Create a mask to separate background from foreground
+    mask = img_array > threshold
+    
+    # Replace background with white
+    img_array[mask] = 255
+    
+    # Convert back to PIL image
+    image = Image.fromarray(img_array)
+    
+    return image
+
 def predict(image):
+    # Remove background and replace with white
+    image = remove_background(image)
+    
     # Prétraiter l'image
     inputs = feature_extractor(images=image, return_tensors="pt")
     
@@ -54,4 +76,4 @@ gr.Interface(fn=predict,
              outputs=label,
              title="Classification des maladies des agrumes",
              description="Téléchargez une image d'une feuille d'agrume pour classer sa maladie. Le modèle est entraîné sur les maladies suivantes : Aleurocanthus spiniferus, Chancre citrique, Cochenille blanche, Dépérissement des agrumes, Feuille saine, Jaunissement des feuilles, Maladie de l'oïdium, Maladie du dragon jaune, Mineuse des agrumes, Trou de balle.",
-             examples=["maladie_du_dragon_jaune.jpg", "critique.jpg","feuille_saine.jpg"]).launch(share=True)
+             examples=["maladie_du_dragon_jaune.jpg", "critique.jpg","feuille_saine.jpg"]).launch(share=True)