# install (if needed)
#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)
# If you have custom layers like RandomHeight, include them in custom_objects:
model = tf.keras.models.load_model("stone_artifact_classifier_with_brokenshards.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"""
    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))}

# 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 like projectile points, hammer stones or scraper rocks and the model will classify it."
)

iface.launch()