from transformers import pipeline
import gradio as gr

# Load a lightweight pre-trained model without specifying cache_dir
model = pipeline("image-classification", model="facebook/deit-tiny-patch16-224")

# Function to classify an image
def classify_image(image):
    predictions = model(image)
    # Format predictions as {label: confidence}
    return {pred["label"]: round(pred["score"], 4) for pred in predictions}

# Gradio interface
interface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(),
    title="Image Classifier Test",
    description="Upload an image to classify."
)

# Launch the app
if __name__ == "__main__":
    interface.launch(server_name="0.0.0.0")