import gradio as gr
from transformers import AutoFeatureExtractor, AutoModelForImageClassification

# Load the feature extractor and model directly
extractor = AutoFeatureExtractor.from_pretrained("ALM-AHME/beit-large-patch16-224-finetuned-BreastCancer-Classification-BreakHis-AH-60-20-20")
model = AutoModelForImageClassification.from_pretrained("ALM-AHME/beit-large-patch16-224-finetuned-BreastCancer-Classification-BreakHis-AH-60-20-20")

# Define the prediction function using the loaded model
def classify_image(image):
    # Preprocess the image and get the features
    inputs = extractor(images=image, return_tensors="pt")
    # Make the prediction using the model
    outputs = model(**inputs)
    logits = outputs.logits
    # Get the predicted label and confidence
    predicted_label = logits.argmax(dim=1).item()
    confidence = logits.softmax(dim=1).max().item()

    # Map predicted label to "benigno" or "maligno"
    class_names = ["benigno", "maligno"]
    predicted_class = class_names[predicted_label]

    return {"prediction": predicted_class, "confidence": confidence}

# Define the Gradio interface
iface = gr.Interface(
    fn=classify_image,
    inputs=gr.inputs.Image(),
    outputs="json",
    title="Classificação de Imagens de Câncer de Mama",
    description="Este aplicativo classifica imagens de câncer de mama em diferentes classes.",
    article="Este modelo é uma versão fine-tuned do microsoft/beit-large-patch16-224 no dataset imagefolder. Alcançou os seguintes resultados no conjunto de avaliação: Loss: 0.0275, Accuracy: 0.9939.",
)

# Launch the Gradio interface
if __name__ == "__main__":
    iface.launch()