Update app.py

app.py

@@ -1,60 +1,64 @@
 # install (if needed)
-# pip install tensorflow
+#pip install tensorflow
 
 import tensorflow as tf
+
+
+# load SavedModel dir or .h5
+#model = tf.keras.models.load_model("stone_artifact_classifier.h5")        # use compile=False if custom objects
+# or: model = tf.keras.models.load_model("my_model.h5", compile=False)
+
 import gradio as gr
 from PIL import Image
 import numpy as np
+import tensorflow as tf
 
 # Load your model (handle any custom layers if needed)
-model = tf.keras.models.load_model(
-    "stone_artifact_classifier.h5",
-    compile=False,
-    custom_objects={
-        "RandomHeight": tf.keras.layers.RandomHeight,
-        "RandomWidth": tf.keras.layers.RandomWidth,
-    },
-)
+# If you have custom layers like RandomHeight, include them in custom_objects:
+model = tf.keras.models.load_model("stone_artifact_classifier.h5", compile=False, custom_objects={"RandomHeight": tf.keras.layers.RandomHeight, "RandomWidth": tf.keras.layers.RandomWidth})
+#model = tf.keras.models.load_model("stone_artifact_classifier.h5", compile=False)
+
+
+
+
+
 
 # Replace with your classes
 class_names = ["arrowheads", "axehammerstones", "scrapers"]
 
-IMG_SIZE = (224, 224)  # change this to your model's expected size
-
 from tensorflow.keras.applications.efficientnet import preprocess_input
 
 
 def predict(img):
-    """Predict function for Gradio — only highest class shows bar"""
+    """Predict function for Gradio"""
 
     img = img.convert("RGB").resize(IMG_SIZE)
     x = np.array(img)
     x = preprocess_input(x)
     x = np.expand_dims(x, 0)
+    preds = model.predict(x)[0]
+    # top = np.argmax(preds)
+    return {class_names[i]: float(preds[i]) for i in range(len(preds))}
+
+
+
+
+
+
+
+
 
-    preds = model.predict(x)[0]  # probabilities
 
-    # ✅ find highest prediction
-    max_idx = int(np.argmax(preds))
 
-    # ✅ zero out all other classes
-    output = {}
-    for i in range(len(preds)):
-        if i == max_idx:
-            output[class_names[i]] = float(preds[i])
-        else:
-            output[class_names[i]] = 0.0
 
-    return output
 
 
 # Gradio interface
 iface = gr.Interface(
-    fn=predict,
     inputs=gr.Image(type="pil"),
     outputs=gr.Label(num_top_classes=len(class_names)),
     title="Stone Artifact Classifier",
-    description="Upload an image of an artifact and the model will classify it.",
+    description="Upload an image of an artifact and the model will classify it."
 )
 
-iface.launch()
+iface.launch()