Update app.py

app.py

@@ -7,8 +7,8 @@ import json
 
 app = Flask(__name__)
 
-# Load the TensorFlow SavedModel
-model = tf.keras.models.load_model('./plant_disease_detection_saved_model')
+# Load the TensorFlow model
+model = tf.keras.layers.TFSMLayer('./plant_disease_detection_saved_model', call_endpoint='serving_default')
 
 # Load categories
 with open('./categories.json') as f:
@@ -31,12 +31,11 @@ def predict():
     image_array = preprocess_image(image)
 
     # Make prediction
-    predictions = model.predict(image_array)
+    predictions = model(image_array)
     predicted_class = np.argmax(predictions, axis=1)[0]
 
     # Map to category names
-    predicted_label = [key for key, value in categories.items() if value == predicted_class]
-    predicted_label = predicted_label[0] if predicted_label else 'Unknown'
+    predicted_label = categories.get(str(predicted_class), 'Unknown')
 
     return jsonify({'class': predicted_label, 'confidence': float(predictions[0][predicted_class])})