Update app.py

app.py

@@ -3,13 +3,14 @@ import torch
 from transformers import ViTFeatureExtractor, ViTForImageClassification
 from PIL import Image
 import numpy as np
+import rembg
 
-# Définir le modèle et le feature extractor
+# Define the model and feature extractor
 model_name ="KhadijaAsehnoune12/ViTOrangeLeafDiseaseClassifier"
 model = ViTForImageClassification.from_pretrained(model_name, num_labels=10, ignore_mismatched_sizes=True)
 feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
 
-# Définir la mappage des labels
+# Define the label mapping
 id2label = {
     "0": "Aleurocanthus spiniferus",
     "1": "Chancre citrique",
@@ -23,34 +24,52 @@ id2label = {
     "9": "Trou de balle"
 }
 
-
+def remove_background(image):
+    # Convert the image to RGBA
+    image = image.convert("RGBA")
+    
+    # Remove the background
+    image_np = np.array(image)
+    output_np = rembg.remove(image_np)
+    
+    # Create a white background image
+    white_bg = Image.new("RGBA", image.size, "WHITE")
+    
+    # Composite the original image over the white background
+    output_image = Image.alpha_composite(white_bg, Image.fromarray(output_np))
+    
+    # Convert back to RGB
+    output_image = output_image.convert("RGB")
+    
+    return output_image
 
 def predict(image):
-   
+    # Remove the background
+    image = remove_background(image)
     
-    # Prétraiter l'image
+    # Preprocess the image
     inputs = feature_extractor(images=image, return_tensors="pt")
     
-    # Passage en avant dans le modèle
+    # Forward pass through the model
     outputs = model(**inputs)
     
-    # Obtenir les logits
+    # Get the logits
     logits = outputs.logits
     
-    # Calculer les scores de confiance avec softmax
+    # Calculate confidence scores with softmax
     probs = torch.nn.functional.softmax(logits, dim=-1)[0]
     
-    # Obtenir l'indice de la classe la plus probable
+    # Get the index of the most probable class
     predicted_class_idx = probs.argmax().item()
     
-    # Obtenir le label et le score de confiance de la classe la plus probable
+    # Get the label and confidence score of the most probable class
     predicted_label = id2label[str(predicted_class_idx)]
-    confidence_score = probs[predicted_class_idx].item() * 100  # Multiplie par 100 pour obtenir un pourcentage
+    confidence_score = probs[predicted_class_idx].item() * 100  # Multiply by 100 to get a percentage
     
-    # Retourner le label et le score de confiance
+    # Return the label and confidence score
     return f"{predicted_label}: {confidence_score:.2f}%"
 
-# Créer l'interface Gradio
+# Create the Gradio interface
 image = gr.Image(type="pil") 
 label = gr.Textbox(label="Prediction")
 
@@ -59,4 +78,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)