import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  # Disable GPU usage
import gradio as gr
import tensorflow as tf
from tensorflow.keras.preprocessing import image
import numpy as np
from transformers import pipeline

# Load image classification model from Hugging Face
model = pipeline("image-classification", model="icputrd/gelderman_decomposition_classification/head/xception_070523")


def classify_image(img):
    # Preprocess the image
    img = img.resize((299, 299))  # Resize to the input shape of Xception
    img_array = image.img_to_array(img)  # Convert image to array
    img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension
    img_array /= 255.0  # Normalize the image

    # Make prediction
    predictions = model.predict(img_array)
    predicted_class = np.argmax(predictions, axis=-1)[0]  # Get the predicted class index
    
    return str(predicted_class)  # Return the class index as a string

# Gradio interface for image classification
demo = gr.Interface(fn=classify_image, inputs=gr.inputs.Image(type="pil"), outputs="label")
demo.launch()