import gradio as gr
import tensorflow as tf
from tensorflow.keras.preprocessing import image
import numpy as np
from PIL import Image
from keras import layers
from transformers import TFAutoModelForImageClassification
from transformers import AutoImageProcessor

# Load model#'model = tf.keras.models.load_model("xception-head")

# Replace with your Hugging Face model repository name
model_name = "icputrd/Inception-V3-Human-Bodypart-Classifier"

# Load the pre-trained TensorFlow model from Hugging Face
model = TFAutoModelForImageClassification.from_pretrained(model_name)

# Load the associated image processor (for preprocessing input images)
image_processor = AutoImageProcessor.from_pretrained(model_name)


# Define the labels for your classification
class_labels = ['arm', 'hand', 'foot',  'legs','fullbody','head','backside', 'torso', 'stake', 'plastic'] # Replace with your actual class names

def classify_image(img):
    # Preprocess the image to fit the model input shape
    img = img.resize((299, 299))  # Xception takes 299x299 input size
    img = np.array(img) / 255.0   # Normalize the image
    img = np.expand_dims(img, axis=0)

    # Make prediction
    predictions = model.predict(img)
    predicted_class = np.argmax(predictions, axis=1)[0]
    confidence = np.max(predictions)
    return {class_labels[i]: float(predictions[0][i]) for i in range(len(class_labels))}, confidence

# Example images (local paths or URLs)
#example_images = [
    #'examples/fresh.jpg',  # Replace with actual local file paths or URLs
#]

# Gradio interface
demo = gr.Interface(
    fn=classify_image,
    title="Human Bodypart Image Classification",
    description = "Predict the bodypart of huma. This is a demo of our human bodypart image <a href=\"https://huggingface.co/icputrd/Inception-V3-Human-Bodypart-Classifier">classifier</a>.",
    inputs=gr.Image(type="pil"),
    outputs=[gr.Label(num_top_classes=len(class_labels)), gr.Number()],
    live=True,
)

if __name__ == "__main__":
    demo.launch()