app.py

import gradio as gr
from transformers import AutoFeatureExtractor, AutoModelForImageClassification

# Carregar o extrator de features e o modelo diretamente
extractor = AutoFeatureExtractor.from_pretrained("Devarshi/Brain_Tumor_Classification")
model = AutoModelForImageClassification.from_pretrained("Devarshi/Brain_Tumor_Classification")

# Definir a função de previsão usando o modelo carregado
def classificar_imagem(imagem):
    # Pré-processar a imagem e obter as features
    inputs = extractor(images=imagem, return_tensors="pt")
    # Fazer a previsão usando o modelo
    outputs = model(**inputs)
    logits = outputs.logits
    # Obter a classe prevista e a confiança da previsão
    classe_prevista = logits.argmax(dim=1).item()
    confianca = logits.softmax(dim=1).max().item()

    # Mapear a classe prevista para os nomes corretos
    nomes_classes = ["glioma_tumor", "meningioma_tumor",  "no_tumor", "pituitary_tumor"]
    classe_prevista_texto = nomes_classes[classe_prevista]

    return {"previsao": classe_prevista_texto, "confianca": confianca}

# Definir a interface Gradio
iface = gr.Interface(
    fn=classificar_imagem,
    inputs=gr.inputs.Image(),
    outputs="json",
    title="Classificação de Imagens de Tumor Cerebral",
    description="Este aplicativo classifica imagens de tumores cerebrais em diferentes classes.",
)

# Iniciar a interface Gradio
if __name__ == "__main__":
    iface.launch()