Update app.py

app.py

@@ -1,10 +1,11 @@
 import streamlit as st
 from PIL import Image
 import numpy as np
-from transformers import AutoFeatureExtractor, AutoModelForImageClassification
+import tensorflow as tf
+from transformers import AutoFeatureExtractor, TFAutoModelForImageClassification
 
 def predict_disease(image_file):
-    """Predicts the disease of a plant from an image.
+    """Predicts the disease of a plant from an image using TensorFlow.
     
     Args:
         image_file: The uploaded image file.
@@ -12,24 +13,27 @@ def predict_disease(image_file):
     Returns:
         A string representing the predicted disease.
     """
-    # Load the model and feature extractor
-    model_name = "ombhojane/healthyPlantsModel"
-    feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
-    model = AutoModelForImageClassification.from_pretrained(model_name)
-    
-    # Process the image
-    image = Image.open(image_file).convert("RGB")
-    inputs = feature_extractor(images=image, return_tensors="pt")
-    
-    # Make prediction
-    outputs = model(**inputs)
-    logits = outputs.logits
-    predicted_class_idx = logits.argmax(-1).item()
-    
-    # Get the label
-    predicted_label = model.config.id2label[predicted_class_idx]
-    
-    return predicted_label
+    try:
+        # Load the model and feature extractor
+        model_name = "ombhojane/healthyPlantsModel"
+        feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
+        model = TFAutoModelForImageClassification.from_pretrained(model_name)
+        
+        # Process the image
+        image = Image.open(image_file).convert("RGB")
+        inputs = feature_extractor(images=image, return_tensors="tf")
+        
+        # Make prediction
+        outputs = model(**inputs)
+        logits = outputs.logits
+        predicted_class_idx = tf.argmax(logits, axis=-1).numpy()[0]
+        
+        # Get the label
+        predicted_label = model.config.id2label[predicted_class_idx]
+        
+        return predicted_label
+    except Exception as e:
+        return f"Error: {str(e)}"
 
 def main():
     """Creates the Streamlit app."""
@@ -50,7 +54,10 @@ def main():
                 disease = predict_disease(image_file)
             
             # Display the prediction
-            st.success(f"Predicted disease: {disease}")
+            if disease.startswith("Error"):
+                st.error(disease)
+            else:
+                st.success(f"Predicted disease: {disease}")
 
 if __name__ == "__main__":
     main()