import gradio as gr
from tensorflow import keras
import numpy as np
from huggingface_hub import HfApi
import h5py
from io import BytesIO

# Authenticate and read the custom model from Hugging Face Spaces
hf_api = HfApi()
model_url = hf_api.presigned_url('dhhd255', 'idk_test', filename='best_model.h5', token='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')
r = requests.get(model_url)
model_file = h5py.File(BytesIO(r.content), 'r')

# Load your custom model
model = keras.models.load_model(model_file)

def image_classifier(inp):
    # Preprocess the input image
    inp = np.array(inp)
    inp = inp / 255.0
    inp = np.expand_dims(inp, axis=0)
    
    # Use your custom model for inference
    predictions = model.predict(inp)
    
    # Process the predictions and return the result
    result = {}
    for i, prediction in enumerate(predictions[0]):
        label = f'Label {i+1}'
        result[label] = prediction
    
    return result

demo = gr.Interface(fn=image_classifier, inputs='image', outputs='label')
demo.launch()