Update app.py

app.py

@@ -104,7 +104,7 @@ history = model.fit(
     train_ds,
     validation_data=val_ds,
     epochs=epochs,
-    callbacks=[early_stopping]
+    
 )
 
 # Evaluate the model on validation data
@@ -131,13 +131,14 @@ plt.legend()
 plt.title('Training and Validation Accuracy')
 plt.show()
 
-# Prediction function
 def predict_image(img):
     img = np.array(img)
     img_resized = tf.image.resize(img, (img_height, img_width))
     img_4d = tf.expand_dims(img_resized, axis=0)
-    prediction = model.predict(img_4d)[0]
-    return {class_names[i]: float(prediction[i]) for i in range(num_classes)}
+    logits = model.predict(img_4d)[0]
+    probabilities = tf.nn.softmax(logits)
+    return {class_names[i]: float(probabilities[i]) for i in range(num_classes)}
+
 
 # Interface
 image = gr.Image()