import gradio as gr
from transformers import pipeline

# Load the models using pipeline
image_model = pipeline("image-classification", model="dima806/deepfake_vs_real_image_detection")

# Define the prediction function
def predict(image):
    result = image_model(image)
    print("Raw prediction result:", result)  # Debugging statement
        # Convert the result to the expected format
    output = {item['label']: item['score'] for item in result}
    print("Formatted prediction result:", output)  # Debugging statement
    return output

# Create Gradio interface
with gr.Blocks() as iface:
    image_input = gr.Image(type="filepath", label="Upload Image File", visible=False)
    output = gr.Label()

    submit_button = gr.Button("Submit")
    submit_button.click(fn=predict, inputs=[audio_input, image_input, model_choice], outputs=output)
    
    iface.launch()