Update app.py

app.py

@@ -1,20 +1,15 @@
 import json
-from PIL import Image
 import numpy as np
-import tensorflow as tf
-from transformers import TFAutoModelForSequenceClassification, AutoTokenizer
-from tensorflow.keras.models import load_model
-import ipywidgets as widgets
-from IPython.display import display
+import gradio as gr
 
+# Load the model
 model_path = 'final_teath_classifier.h5'
-
 model = tf.keras.models.load_model(model_path)
 
-# Load the model from Hugging Face model hub
-
-def preprocess_image(image: Image.Image) -> np.ndarray:
+# Define preprocessing function
+def preprocess_image(image):
     # Resize the image to match input size
+    image = Image.fromarray(image)
     image = image.resize((256, 256))
     # Convert image to array and preprocess input
     img_array = np.array(image) / 255.0
@@ -22,13 +17,9 @@ def preprocess_image(image: Image.Image) -> np.ndarray:
     img_array = np.expand_dims(img_array, axis=0)
     return img_array
 
-def predict_image(image_path):
-    img = Image.open(image_path)
-    # Preprocess the image
-    img_array = preprocess_image(img)
-    # Convert image array to string using base64 encoding (for text-based models)
-    #inputs = tokenizer.encode(img_array, return_tensors="tf")
-    # Make prediction
+# Define prediction function
+def predict_image(image):
+    img_array = preprocess_image(image)
     outputs = model(img_array)
     predictions = tf.nn.softmax(outputs.logits, axis=-1)
     predicted_class = np.argmax(predictions)
@@ -36,31 +27,13 @@ def predict_image(image_path):
         predict_label = "Clean"
     else:
         predict_label = "Carries"
+    return {"prediction": predict_label, "confidence": float(np.max(predictions))}
 
-    return predict_label, predictions.numpy().flatten()
-
-# Create a file uploader widget
-uploader = widgets.FileUpload(accept="image/*", multiple=False)
-
-# Display the file uploader widget
-display(uploader)
+# Create the interface
+input_interface = gr.inputs.Image(shape=(256, 256))
+output_interface = gr.outputs.Label(num_top_classes=2)
 
-# Define a callback function to handle the uploaded image
-def on_upload(change):
-    # Get the uploaded image file
-    image_file = list(uploader.value.values())[0]["content"]
-    # Save the image to a temporary file
-    with open("temp_image.jpg", "wb") as f:
-        f.write(image_file)
-    # Get predictions for the uploaded image
-    predict_label, logits = predict_image("temp_image.jpg")
-    # Create a JSON object with the predictions
-    predictions_json = {
-        "predicted_class": predict_label,
-        "evaluations": [f"{logit*100:.4f}%" for logit in logits]
-    }
-    # Print the JSON object
-    print(json.dumps(predictions_json, indent=4))
+iface = gr.Interface(fn=predict_image, inputs=input_interface, outputs=output_interface)
 
-# Set the callback function for when a file is uploaded
-uploader.observe(on_upload, names="value")
+# Launch the interface
+iface.launch()