from transformers import AutoImageProcessor, AutoModelForImageClassification
import gradio as gr

def classifyArchitecture(img):
    # Call the model
    image_processor = AutoImageProcessor.from_pretrained("hanslab37/architectural_styles_classifier")
    classifier = AutoModelForImageClassification.from_pretrained("hanslab37/architectural_styles_classifier")

    # Transform our image and pass it through the model
    inputs = image_processor(img, return_tensors="pt")
    output = classifier(**inputs)

    # Predicted Class probabilities
    proba = output.logits.softmax(1)
    preds = proba.argmax(1)
    
    result = f'Predicted Architectural Styles: {classifier.config.id2label[preds.item()]} (Confidence: {proba[0][preds.item()].item()})'

    return result

page = gr.Interface(fn=classifyArchitecture, inputs="image", outputs="text", title="Find out the Architecture Style",)
page.launch()