Update app.py

app.py

@@ -5,24 +5,22 @@ import numpy as np
 
 loaded_model = tf.keras.models.load_model('kidney2.h5')
 
+label_names = {
+    "1": "Cyst",
+    "2": "Normal",
+    "3": "Stone",
+    "4": "Tumor"
+}
+
 def classify_kidney_image(img):
     resize = tf.image.resize(img, (224, 224))
     gray = tfio.experimental.color.bgr_to_rgb(resize)
-    yhat = loaded_model.predict(np.expand_dims(gray / 255, 0))
-    label_names = {
-        "1": "Cyst",
-        "2": "Normal",
-        "3": "Stone",
-        "4": "Tumor"
-    }
-    classes_x = np.argmax(yhat, axis=1)
-    a = classes_x[0]
-    input_value = a + 1
-    input_str = str(input_value)
-    predicted_label = label_names[input_str]
-    q, w, e, r = yhat[0]
-
-    return {'Cyst': str(q), 'Normal': str(w), 'Stone': str(e), 'Tumor': str(r)}
+    normalized_img = gray / 255.0
+    yhat = loaded_model.predict(np.expand_dims(normalized_img, 0))
+    class_index = np.argmax(yhat, axis=1)[0]
+    predicted_label = label_names[str(class_index + 1)]
+    probabilities = {label_names[str(i + 1)]: str(prob) for i, prob in enumerate(yhat[0])}
+    return predicted_label, probabilities
 
 image = gr.inputs.Image(shape=(224, 224))
 label = gr.outputs.Label()