Update app.py

app.py

@@ -3,39 +3,49 @@ import tensorflow as tf
 import numpy as np
 import gradio as gr
 import io
+import json
 
 # Load the model
 model_path = 'final_teath_classifier.h5'
 model = tf.keras.models.load_model(model_path)
+
+# Define preprocessing function
+def preprocess_image(image):
+    # Resize the image to match input size
+    #image = image.resize((256, 256))
+    # Convert image to array and preprocess input
+    image = tf.keras.preprocessing.image.img_to_array(image)
+    # Add batch dimension
+    img_array = np.expand_dims(image, axis=0)
+    return img_array
+
 # Define prediction function
 def predict_image(image):
+    # Save the image to a file-like object
     image_bytes = io.BytesIO()
     image.save(image_bytes, format="JPEG")
-
     # Load the image from the file-like object
     image = tf.keras.preprocessing.image.load_img(image_bytes, target_size=(256, 256))
-    img_array = np.array(image) / 255.0
-    image = np.expand_dims(img_array, axis=0)
-
-    # Make a prediction
-    prediction = model.predict(image)
-
-    # Get the probability of being 'Good'
-    probability_good = prediction[0][0]  # Assuming it's a binary classification
-
-    # Define the prediction result
-    result = {
-        "prediction": probability_good
-    }
-
-    return result
+    img_array = preprocess_image(image)
+    outputs = model.predict(img_array)
+    predictions = tf.nn.softmax(outputs.logits, axis=-1)
+    predicted_class = np.argmax(predictions)
+    if predicted_class == 0:
+        predict_label = "Clean"
+    else:
+        predict_label = "Carries"
+    confidence = float(np.max(predictions))
+    prediction_dict = {"prediction": predict_label, "confidence": confidence}
+    return prediction_dict
 
 # Create the interface
-input_interface = gr.Image(type="pil")
+input_interface = gr.Image(type = "pil")
+output_interface = "json"
+
 iface = gr.Interface(
     fn=predict_image,
     inputs=input_interface,
-    outputs="json")
+    outputs=gr.Textbox("output"))
 
 # Launch the interface
 iface.launch(share=True)