Update app.py

app.py

@@ -1,38 +1,44 @@
-import numpy as np
+import math
+
 import gradio as gr
 import tensorflow as tf
 
-models = [{"name": "my_model_2.h5", "size": 512}, {"name": "my_model.h5", "size": 224}]
+configs = [
+    {
+        "model": "my_model_2.h5", "size": 512
+    },
+    {
+        "model": "my_model.h5", "size": 224
+    },
+]
+
+config = configs[0]
 
-def classify_image(image, model_name):
-    model_config = next(m for m in models if m["name"] == model_name)
-    model = tf.keras.models.load_model(model_name)
-    input_image = np.expand_dims(image, axis=0)
-    prediction = model.predict(input_image).flatten()
+new_model = tf.keras.models.load_model(config["model"])
+
+def classify_image(inp):
+    inp = inp.reshape((-1, config["size"], config["size"], 3))
+    prediction = new_model.predict(inp).flatten()
+    print(prediction)
     if len(prediction) > 1:
-        probability = 100 * np.exp(prediction[0]) / (np.exp(prediction[0]) + np.exp(prediction[1]))
+        probability = 100 * math.exp(prediction[0]) / (math.exp(prediction[0]) + math.exp(prediction[1]))
     else:
-        probability = round(100. / (1 + np.exp(-prediction[0])), 2)
+        probability = round(100. / (1 + math.exp(-prediction[0])), 2)
     if probability > 45:
-        label = "Glaucoma"
-    elif probability > 25:
-        label = "Unclear"
-    else:
-        label = "Not glaucoma"
-    return label, probability
+        return "Glaucoma", probability
+    if probability > 25:
+        return "Unclear", probability
+    return "Not glaucoma", probability
 
-inputs = [
-    gr.inputs.Image(shape=(224, 224), label="Eye image"),
-    gr.inputs.Dropdown(choices=[m["name"] for m in models], label="Model"),
-]
-outputs = [
-    gr.outputs.Textbox(label="Predicted label"),
-    gr.outputs.Textbox(label="Probability of glaucoma (0-100)"),
-]
-
-examples = [
-    [np.zeros((224, 224, 3)), "my_model.h5"],
-    [np.ones((224, 224, 3)) * 255, "my_model_2.h5"]
-]
 
-gr.Interface(classify_image, inputs, outputs, examples=examples, title="Glaucoma Classification").launch() 
+gr.Interface(
+    fn=classify_image, 
+    inputs=gr.inputs.Image(shape=(config["size"], config["size"])),
+    outputs=[
+    gr.outputs.Textbox(label="Label"),
+    gr.outputs.Textbox(label="Glaucoma probability (0 - 100)"),
+    ],
+    examples=["001.jpg", "002.jpg", "225.jpg"],
+    flagging_options=["Correct label", "Incorrect label"],
+    allow_flagging="manual",
+).launch()